sozu_command_lib/
writer.rs1use std::io::{self, Error, ErrorKind, Write};
2
3pub struct MultiLineWriter<W: Write> {
5 inner: Option<W>,
6 buf: Vec<u8>,
7 last_newline: usize,
8 panicked: bool,
9}
10
11impl<W: Write> MultiLineWriter<W> {
12 pub fn new(inner: W) -> MultiLineWriter<W> {
13 MultiLineWriter::with_capacity(4096, inner)
14 }
15
16 pub fn with_capacity(capacity: usize, inner: W) -> MultiLineWriter<W> {
17 MultiLineWriter {
19 inner: Some(inner),
20 buf: Vec::with_capacity(capacity),
21 panicked: false,
22 last_newline: 0,
23 }
24 }
25
26 pub fn get_ref(&self) -> &W {
27 self.inner.as_ref().unwrap()
28 }
29
30 pub fn get_mut(&mut self) -> &mut W {
31 self.inner.as_mut().unwrap()
32 }
33
34 fn flush_buf(&mut self, flush_entire_buffer: bool) -> io::Result<()> {
35 let mut written = 0;
36 let len = if flush_entire_buffer {
37 self.buf.len()
38 } else {
39 self.last_newline + 1
40 };
41
42 let mut ret = Ok(());
43 while written < len {
44 self.panicked = true;
45 let r = self.inner.as_mut().unwrap().write(&self.buf[written..len]);
46 self.panicked = false;
47
48 match r {
49 Ok(0) => {
50 ret = Err(Error::new(
51 ErrorKind::WriteZero,
52 "failed to write the buffered data",
53 ));
54 break;
55 }
56 Ok(n) => written += n,
57 Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
58 Err(e) => {
59 ret = Err(e);
60 break;
61 }
62 }
63 }
64 if written > 0 {
66 self.buf.drain(..written);
68 }
69
70 if flush_entire_buffer {
71 self.last_newline = 0;
72 } else if written > self.last_newline {
73 self.last_newline = 0
74 } else {
75 self.last_newline -= written;
76 }
77
78 ret
79 }
80}
81
82impl<W: Write> Write for MultiLineWriter<W> {
83 fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
84 if self.buf.len() + buf.len() > self.buf.capacity() {
85 self.flush_buf(false)?;
86 }
87 if buf.len() >= self.buf.capacity() {
88 self.panicked = true;
89 let r = self.get_mut().write(buf);
90 self.panicked = false;
91 r
92 } else {
93 if let Some(i) = memchr::memrchr(b'\n', buf) {
94 self.last_newline = self.buf.len() + i;
95 };
96
97 self.buf.write(buf)
98 }
99 }
100
101 fn flush(&mut self) -> io::Result<()> {
102 self.flush_buf(true).and_then(|()| self.get_mut().flush())
103 }
104}
105
106impl<W: Write> Drop for MultiLineWriter<W> {
107 fn drop(&mut self) {
108 if self.inner.is_some() && !self.panicked {
109 let _r = self.flush_buf(true);
111 }
112 }
113}
114
115