sozu_command_lib/
writer.rs

1use std::io::{self, Error, ErrorKind, Write};
2
3/// A multiline writer used for logging
4pub 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 { inner: BufWriter::with_capacity(capacity, inner), capacity, last_newline: usize, need_flush: false }
18        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        //println!("buf len {} last newline {}, written {}", self.buf.len(), self.last_newline, written);
65        if written > 0 {
66            //println!("FLUSHED: {}", ::std::str::from_utf8(&self.buf[..written]).unwrap());
67            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            // dtors should not panic, so we ignore a failed flush
110            let _r = self.flush_buf(true);
111        }
112    }
113}
114
115/*
116impl<W: Write> Write for MultiLineWriter<W> {
117  fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
118    let buffer_len = self.inner.buffer().len();
119
120    let i = match memchr::memrchr(b'\n', buf) {
121      Some(i) => i,
122      None => buf.len(),
123    };
124
125    if buffer_len + i > self.capacity {
126      self.inner.flush()?;
127      self.last_newline = 0;
128    }
129
130
131
132
133  }
134}
135*/
136
137/*
138impl<W: Write> Drop for BufWriter<W> {
139  fn drop(&mut self) {
140    if self.inner.is_some() && !self.panicked {
141      // dtors should not panic, so we ignore a failed flush
142      let _r = self.flush_buf();
143    }
144  }
145}
146*/