1 | use super::buf_writer::BufWriter; |
2 | use futures_core::ready; |
3 | use futures_core::task::{Context, Poll}; |
4 | use futures_io::AsyncWrite; |
5 | use futures_io::IoSlice; |
6 | use pin_project_lite::pin_project; |
7 | use std::io; |
8 | use std::pin::Pin; |
9 | |
10 | pin_project! { |
11 | /// Wrap a writer, like [`BufWriter`] does, but prioritizes buffering lines |
12 | /// |
13 | /// This was written based on `std::io::LineWriter` which goes into further details |
14 | /// explaining the code. |
15 | /// |
16 | /// Buffering is actually done using `BufWriter`. This class will leverage `BufWriter` |
17 | /// to write on-each-line. |
18 | #[derive(Debug)] |
19 | pub struct LineWriter<W: AsyncWrite> { |
20 | #[pin] |
21 | buf_writer: BufWriter<W>, |
22 | } |
23 | } |
24 | |
25 | impl<W: AsyncWrite> LineWriter<W> { |
26 | /// Create a new `LineWriter` with default buffer capacity. The default is currently 1KB |
27 | /// which was taken from `std::io::LineWriter` |
28 | pub fn new(inner: W) -> LineWriter<W> { |
29 | LineWriter::with_capacity(1024, inner) |
30 | } |
31 | |
32 | /// Creates a new `LineWriter` with the specified buffer capacity. |
33 | pub fn with_capacity(capacity: usize, inner: W) -> LineWriter<W> { |
34 | LineWriter { buf_writer: BufWriter::with_capacity(capacity, inner) } |
35 | } |
36 | |
37 | /// Flush `buf_writer` if last char is "new line" |
38 | fn flush_if_completed_line(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
39 | let this = self.project(); |
40 | match this.buf_writer.buffer().last().copied() { |
41 | Some(b' \n' ) => this.buf_writer.flush_buf(cx), |
42 | _ => Poll::Ready(Ok(())), |
43 | } |
44 | } |
45 | |
46 | /// Returns a reference to `buf_writer`'s internally buffered data. |
47 | pub fn buffer(&self) -> &[u8] { |
48 | self.buf_writer.buffer() |
49 | } |
50 | |
51 | /// Acquires a reference to the underlying sink or stream that this combinator is |
52 | /// pulling from. |
53 | pub fn get_ref(&self) -> &W { |
54 | self.buf_writer.get_ref() |
55 | } |
56 | } |
57 | |
58 | impl<W: AsyncWrite> AsyncWrite for LineWriter<W> { |
59 | fn poll_write( |
60 | mut self: Pin<&mut Self>, |
61 | cx: &mut Context<'_>, |
62 | buf: &[u8], |
63 | ) -> Poll<io::Result<usize>> { |
64 | let mut this = self.as_mut().project(); |
65 | let newline_index = match memchr::memrchr(b' \n' , buf) { |
66 | None => { |
67 | ready!(self.as_mut().flush_if_completed_line(cx)?); |
68 | return self.project().buf_writer.poll_write(cx, buf); |
69 | } |
70 | Some(newline_index) => newline_index + 1, |
71 | }; |
72 | |
73 | ready!(this.buf_writer.as_mut().poll_flush(cx)?); |
74 | |
75 | let lines = &buf[..newline_index]; |
76 | |
77 | let flushed = { ready!(this.buf_writer.as_mut().inner_poll_write(cx, lines))? }; |
78 | |
79 | if flushed == 0 { |
80 | return Poll::Ready(Ok(0)); |
81 | } |
82 | |
83 | let tail = if flushed >= newline_index { |
84 | &buf[flushed..] |
85 | } else if newline_index - flushed <= this.buf_writer.capacity() { |
86 | &buf[flushed..newline_index] |
87 | } else { |
88 | let scan_area = &buf[flushed..]; |
89 | let scan_area = &scan_area[..this.buf_writer.capacity()]; |
90 | match memchr::memrchr(b' \n' , scan_area) { |
91 | Some(newline_index) => &scan_area[..newline_index + 1], |
92 | None => scan_area, |
93 | } |
94 | }; |
95 | |
96 | let buffered = this.buf_writer.as_mut().write_to_buf(tail); |
97 | Poll::Ready(Ok(flushed + buffered)) |
98 | } |
99 | |
100 | fn poll_write_vectored( |
101 | mut self: Pin<&mut Self>, |
102 | cx: &mut Context<'_>, |
103 | bufs: &[IoSlice<'_>], |
104 | ) -> Poll<io::Result<usize>> { |
105 | let mut this = self.as_mut().project(); |
106 | // `is_write_vectored()` is handled in original code, but not in this crate |
107 | // see https://github.com/rust-lang/rust/issues/70436 |
108 | |
109 | let last_newline_buf_idx = bufs |
110 | .iter() |
111 | .enumerate() |
112 | .rev() |
113 | .find_map(|(i, buf)| memchr::memchr(b' \n' , buf).map(|_| i)); |
114 | let last_newline_buf_idx = match last_newline_buf_idx { |
115 | None => { |
116 | ready!(self.as_mut().flush_if_completed_line(cx)?); |
117 | return self.project().buf_writer.poll_write_vectored(cx, bufs); |
118 | } |
119 | Some(i) => i, |
120 | }; |
121 | |
122 | ready!(this.buf_writer.as_mut().poll_flush(cx)?); |
123 | |
124 | let (lines, tail) = bufs.split_at(last_newline_buf_idx + 1); |
125 | |
126 | let flushed = { ready!(this.buf_writer.as_mut().inner_poll_write_vectored(cx, lines))? }; |
127 | if flushed == 0 { |
128 | return Poll::Ready(Ok(0)); |
129 | } |
130 | |
131 | let lines_len = lines.iter().map(|buf| buf.len()).sum(); |
132 | if flushed < lines_len { |
133 | return Poll::Ready(Ok(flushed)); |
134 | } |
135 | |
136 | let buffered: usize = tail |
137 | .iter() |
138 | .filter(|buf| !buf.is_empty()) |
139 | .map(|buf| this.buf_writer.as_mut().write_to_buf(buf)) |
140 | .take_while(|&n| n > 0) |
141 | .sum(); |
142 | |
143 | Poll::Ready(Ok(flushed + buffered)) |
144 | } |
145 | |
146 | /// Forward to `buf_writer` 's `BufWriter::poll_flush()` |
147 | fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
148 | self.as_mut().project().buf_writer.poll_flush(cx) |
149 | } |
150 | |
151 | /// Forward to `buf_writer` 's `BufWriter::poll_close()` |
152 | fn poll_close(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<()>> { |
153 | self.as_mut().project().buf_writer.poll_close(cx) |
154 | } |
155 | } |
156 | |