1 | use crate::io::{self, BufWriter, IoSlice, Write}; |
2 | use crate::sys_common::memchr; |
3 | |
4 | /// Private helper struct for implementing the line-buffered writing logic. |
5 | /// This shim temporarily wraps a BufWriter, and uses its internals to |
6 | /// implement a line-buffered writer (specifically by using the internal |
7 | /// methods like write_to_buf and flush_buf). In this way, a more |
8 | /// efficient abstraction can be created than one that only had access to |
9 | /// `write` and `flush`, without needlessly duplicating a lot of the |
10 | /// implementation details of BufWriter. This also allows existing |
11 | /// `BufWriters` to be temporarily given line-buffering logic; this is what |
12 | /// enables Stdout to be alternately in line-buffered or block-buffered mode. |
13 | #[derive (Debug)] |
14 | pub struct LineWriterShim<'a, W: ?Sized + Write> { |
15 | buffer: &'a mut BufWriter<W>, |
16 | } |
17 | |
18 | impl<'a, W: ?Sized + Write> LineWriterShim<'a, W> { |
19 | pub fn new(buffer: &'a mut BufWriter<W>) -> Self { |
20 | Self { buffer } |
21 | } |
22 | |
23 | /// Get a reference to the inner writer (that is, the writer |
24 | /// wrapped by the BufWriter). |
25 | fn inner(&self) -> &W { |
26 | self.buffer.get_ref() |
27 | } |
28 | |
29 | /// Get a mutable reference to the inner writer (that is, the writer |
30 | /// wrapped by the BufWriter). Be careful with this writer, as writes to |
31 | /// it will bypass the buffer. |
32 | fn inner_mut(&mut self) -> &mut W { |
33 | self.buffer.get_mut() |
34 | } |
35 | |
36 | /// Get the content currently buffered in self.buffer |
37 | fn buffered(&self) -> &[u8] { |
38 | self.buffer.buffer() |
39 | } |
40 | |
41 | /// Flush the buffer iff the last byte is a newline (indicating that an |
42 | /// earlier write only succeeded partially, and we want to retry flushing |
43 | /// the buffered line before continuing with a subsequent write) |
44 | fn flush_if_completed_line(&mut self) -> io::Result<()> { |
45 | match self.buffered().last().copied() { |
46 | Some(b' \n' ) => self.buffer.flush_buf(), |
47 | _ => Ok(()), |
48 | } |
49 | } |
50 | } |
51 | |
52 | impl<'a, W: ?Sized + Write> Write for LineWriterShim<'a, W> { |
53 | /// Write some data into this BufReader with line buffering. This means |
54 | /// that, if any newlines are present in the data, the data up to the last |
55 | /// newline is sent directly to the underlying writer, and data after it |
56 | /// is buffered. Returns the number of bytes written. |
57 | /// |
58 | /// This function operates on a "best effort basis"; in keeping with the |
59 | /// convention of `Write::write`, it makes at most one attempt to write |
60 | /// new data to the underlying writer. If that write only reports a partial |
61 | /// success, the remaining data will be buffered. |
62 | /// |
63 | /// Because this function attempts to send completed lines to the underlying |
64 | /// writer, it will also flush the existing buffer if it ends with a |
65 | /// newline, even if the incoming data does not contain any newlines. |
66 | fn write(&mut self, buf: &[u8]) -> io::Result<usize> { |
67 | let newline_idx = match memchr::memrchr(b' \n' , buf) { |
68 | // If there are no new newlines (that is, if this write is less than |
69 | // one line), just do a regular buffered write (which may flush if |
70 | // we exceed the inner buffer's size) |
71 | None => { |
72 | self.flush_if_completed_line()?; |
73 | return self.buffer.write(buf); |
74 | } |
75 | // Otherwise, arrange for the lines to be written directly to the |
76 | // inner writer. |
77 | Some(newline_idx) => newline_idx + 1, |
78 | }; |
79 | |
80 | // Flush existing content to prepare for our write. We have to do this |
81 | // before attempting to write `buf` in order to maintain consistency; |
82 | // if we add `buf` to the buffer then try to flush it all at once, |
83 | // we're obligated to return Ok(), which would mean suppressing any |
84 | // errors that occur during flush. |
85 | self.buffer.flush_buf()?; |
86 | |
87 | // This is what we're going to try to write directly to the inner |
88 | // writer. The rest will be buffered, if nothing goes wrong. |
89 | let lines = &buf[..newline_idx]; |
90 | |
91 | // Write `lines` directly to the inner writer. In keeping with the |
92 | // `write` convention, make at most one attempt to add new (unbuffered) |
93 | // data. Because this write doesn't touch the BufWriter state directly, |
94 | // and the buffer is known to be empty, we don't need to worry about |
95 | // self.buffer.panicked here. |
96 | let flushed = self.inner_mut().write(lines)?; |
97 | |
98 | // If buffer returns Ok(0), propagate that to the caller without |
99 | // doing additional buffering; otherwise we're just guaranteeing |
100 | // an "ErrorKind::WriteZero" later. |
101 | if flushed == 0 { |
102 | return Ok(0); |
103 | } |
104 | |
105 | // Now that the write has succeeded, buffer the rest (or as much of |
106 | // the rest as possible). If there were any unwritten newlines, we |
107 | // only buffer out to the last unwritten newline that fits in the |
108 | // buffer; this helps prevent flushing partial lines on subsequent |
109 | // calls to LineWriterShim::write. |
110 | |
111 | // Handle the cases in order of most-common to least-common, under |
112 | // the presumption that most writes succeed in totality, and that most |
113 | // writes are smaller than the buffer. |
114 | // - Is this a partial line (ie, no newlines left in the unwritten tail) |
115 | // - If not, does the data out to the last unwritten newline fit in |
116 | // the buffer? |
117 | // - If not, scan for the last newline that *does* fit in the buffer |
118 | let tail = if flushed >= newline_idx { |
119 | &buf[flushed..] |
120 | } else if newline_idx - flushed <= self.buffer.capacity() { |
121 | &buf[flushed..newline_idx] |
122 | } else { |
123 | let scan_area = &buf[flushed..]; |
124 | let scan_area = &scan_area[..self.buffer.capacity()]; |
125 | match memchr::memrchr(b' \n' , scan_area) { |
126 | Some(newline_idx) => &scan_area[..newline_idx + 1], |
127 | None => scan_area, |
128 | } |
129 | }; |
130 | |
131 | let buffered = self.buffer.write_to_buf(tail); |
132 | Ok(flushed + buffered) |
133 | } |
134 | |
135 | fn flush(&mut self) -> io::Result<()> { |
136 | self.buffer.flush() |
137 | } |
138 | |
139 | /// Write some vectored data into this BufReader with line buffering. This |
140 | /// means that, if any newlines are present in the data, the data up to |
141 | /// and including the buffer containing the last newline is sent directly |
142 | /// to the inner writer, and the data after it is buffered. Returns the |
143 | /// number of bytes written. |
144 | /// |
145 | /// This function operates on a "best effort basis"; in keeping with the |
146 | /// convention of `Write::write`, it makes at most one attempt to write |
147 | /// new data to the underlying writer. |
148 | /// |
149 | /// Because this function attempts to send completed lines to the underlying |
150 | /// writer, it will also flush the existing buffer if it contains any |
151 | /// newlines. |
152 | /// |
153 | /// Because sorting through an array of `IoSlice` can be a bit convoluted, |
154 | /// This method differs from write in the following ways: |
155 | /// |
156 | /// - It attempts to write the full content of all the buffers up to and |
157 | /// including the one containing the last newline. This means that it |
158 | /// may attempt to write a partial line, that buffer has data past the |
159 | /// newline. |
160 | /// - If the write only reports partial success, it does not attempt to |
161 | /// find the precise location of the written bytes and buffer the rest. |
162 | /// |
163 | /// If the underlying vector doesn't support vectored writing, we instead |
164 | /// simply write the first non-empty buffer with `write`. This way, we |
165 | /// get the benefits of more granular partial-line handling without losing |
166 | /// anything in efficiency |
167 | fn write_vectored(&mut self, bufs: &[IoSlice<'_>]) -> io::Result<usize> { |
168 | // If there's no specialized behavior for write_vectored, just use |
169 | // write. This has the benefit of more granular partial-line handling. |
170 | if !self.is_write_vectored() { |
171 | return match bufs.iter().find(|buf| !buf.is_empty()) { |
172 | Some(buf) => self.write(buf), |
173 | None => Ok(0), |
174 | }; |
175 | } |
176 | |
177 | // Find the buffer containing the last newline |
178 | let last_newline_buf_idx = bufs |
179 | .iter() |
180 | .enumerate() |
181 | .rev() |
182 | .find_map(|(i, buf)| memchr::memchr(b' \n' , buf).map(|_| i)); |
183 | |
184 | // If there are no new newlines (that is, if this write is less than |
185 | // one line), just do a regular buffered write |
186 | let last_newline_buf_idx = match last_newline_buf_idx { |
187 | // No newlines; just do a normal buffered write |
188 | None => { |
189 | self.flush_if_completed_line()?; |
190 | return self.buffer.write_vectored(bufs); |
191 | } |
192 | Some(i) => i, |
193 | }; |
194 | |
195 | // Flush existing content to prepare for our write |
196 | self.buffer.flush_buf()?; |
197 | |
198 | // This is what we're going to try to write directly to the inner |
199 | // writer. The rest will be buffered, if nothing goes wrong. |
200 | let (lines, tail) = bufs.split_at(last_newline_buf_idx + 1); |
201 | |
202 | // Write `lines` directly to the inner writer. In keeping with the |
203 | // `write` convention, make at most one attempt to add new (unbuffered) |
204 | // data. Because this write doesn't touch the BufWriter state directly, |
205 | // and the buffer is known to be empty, we don't need to worry about |
206 | // self.panicked here. |
207 | let flushed = self.inner_mut().write_vectored(lines)?; |
208 | |
209 | // If inner returns Ok(0), propagate that to the caller without |
210 | // doing additional buffering; otherwise we're just guaranteeing |
211 | // an "ErrorKind::WriteZero" later. |
212 | if flushed == 0 { |
213 | return Ok(0); |
214 | } |
215 | |
216 | // Don't try to reconstruct the exact amount written; just bail |
217 | // in the event of a partial write |
218 | let lines_len = lines.iter().map(|buf| buf.len()).sum(); |
219 | if flushed < lines_len { |
220 | return Ok(flushed); |
221 | } |
222 | |
223 | // Now that the write has succeeded, buffer the rest (or as much of the |
224 | // rest as possible) |
225 | let buffered: usize = tail |
226 | .iter() |
227 | .filter(|buf| !buf.is_empty()) |
228 | .map(|buf| self.buffer.write_to_buf(buf)) |
229 | .take_while(|&n| n > 0) |
230 | .sum(); |
231 | |
232 | Ok(flushed + buffered) |
233 | } |
234 | |
235 | fn is_write_vectored(&self) -> bool { |
236 | self.inner().is_write_vectored() |
237 | } |
238 | |
239 | /// Write some data into this BufReader with line buffering. This means |
240 | /// that, if any newlines are present in the data, the data up to the last |
241 | /// newline is sent directly to the underlying writer, and data after it |
242 | /// is buffered. |
243 | /// |
244 | /// Because this function attempts to send completed lines to the underlying |
245 | /// writer, it will also flush the existing buffer if it contains any |
246 | /// newlines, even if the incoming data does not contain any newlines. |
247 | fn write_all(&mut self, buf: &[u8]) -> io::Result<()> { |
248 | match memchr::memrchr(b' \n' , buf) { |
249 | // If there are no new newlines (that is, if this write is less than |
250 | // one line), just do a regular buffered write (which may flush if |
251 | // we exceed the inner buffer's size) |
252 | None => { |
253 | self.flush_if_completed_line()?; |
254 | self.buffer.write_all(buf) |
255 | } |
256 | Some(newline_idx) => { |
257 | let (lines, tail) = buf.split_at(newline_idx + 1); |
258 | |
259 | if self.buffered().is_empty() { |
260 | self.inner_mut().write_all(lines)?; |
261 | } else { |
262 | // If there is any buffered data, we add the incoming lines |
263 | // to that buffer before flushing, which saves us at least |
264 | // one write call. We can't really do this with `write`, |
265 | // since we can't do this *and* not suppress errors *and* |
266 | // report a consistent state to the caller in a return |
267 | // value, but here in write_all it's fine. |
268 | self.buffer.write_all(lines)?; |
269 | self.buffer.flush_buf()?; |
270 | } |
271 | |
272 | self.buffer.write_all(tail) |
273 | } |
274 | } |
275 | } |
276 | } |
277 | |