use std::io::{self, Error, ErrorKind, Write};
pub struct MultiLineWriter<W: Write> {
inner: Option<W>,
buf: Vec<u8>,
last_newline: usize,
panicked: bool,
}
impl<W: Write> MultiLineWriter<W> {
pub fn new(inner: W) -> MultiLineWriter<W> {
MultiLineWriter::with_capacity(4096, inner)
}
pub fn with_capacity(capacity: usize, inner: W) -> MultiLineWriter<W> {
MultiLineWriter {
inner: Some(inner),
buf: Vec::with_capacity(capacity),
panicked: false,
last_newline: 0,
}
}
pub fn get_ref(&self) -> &W {
self.inner.as_ref().unwrap()
}
pub fn get_mut(&mut self) -> &mut W {
self.inner.as_mut().unwrap()
}
fn flush_buf(&mut self, flush_entire_buffer: bool) -> io::Result<()> {
let mut written = 0;
let len = if flush_entire_buffer {
self.buf.len()
} else {
(self.last_newline + 1).min(self.buf.len())
};
let mut ret = Ok(());
while written < len {
self.panicked = true;
let r = self.inner.as_mut().unwrap().write(&self.buf[written..len]);
self.panicked = false;
match r {
Ok(0) => {
ret = Err(Error::new(
ErrorKind::WriteZero,
"failed to write the buffered data",
));
break;
}
Ok(n) => written += n,
Err(ref e) if e.kind() == io::ErrorKind::Interrupted => {}
Err(e) => {
ret = Err(e);
break;
}
}
}
if written > 0 {
self.buf.drain(..written);
}
if flush_entire_buffer {
self.last_newline = 0;
} else if written > self.last_newline {
self.last_newline = 0
} else {
self.last_newline -= written;
}
ret
}
}
impl<W: Write> Write for MultiLineWriter<W> {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
if self.buf.len() + buf.len() > self.buf.capacity() {
self.flush_buf(false)?;
}
if buf.len() >= self.buf.capacity() {
self.panicked = true;
let r = self.get_mut().write(buf);
self.panicked = false;
r
} else {
if let Some(i) = memchr::memrchr(b'\n', buf) {
self.last_newline = self.buf.len() + i;
};
self.buf.write(buf)
}
}
fn flush(&mut self) -> io::Result<()> {
self.flush_buf(true).and_then(|()| self.get_mut().flush())
}
}
impl<W: Write> Drop for MultiLineWriter<W> {
fn drop(&mut self) {
if self.inner.is_some() && !self.panicked {
let _r = self.flush_buf(true);
}
}
}
#[cfg(test)]
mod tests {
use super::MultiLineWriter;
use std::io::Write;
const CAPACITY: usize = 16;
#[test]
fn an_oversized_record_on_a_fresh_writer_does_not_panic() {
let mut writer = MultiLineWriter::with_capacity(CAPACITY, Vec::new());
let record = vec![b'a'; CAPACITY * 4];
let written = writer
.write(&record)
.expect("a record larger than the capacity is written straight through");
assert_eq!(
written,
record.len(),
"the oversized record must be handed to the inner writer whole"
);
assert_eq!(
writer.get_ref(),
&record,
"the inner writer must receive exactly the oversized record"
);
}
#[test]
fn an_oversized_record_after_a_flush_does_not_panic() {
let mut writer = MultiLineWriter::with_capacity(CAPACITY, Vec::new());
writer.write_all(b"x\n").expect("a short line is buffered");
writer.flush().expect("the flush empties the buffer");
assert_eq!(
writer.get_ref(),
b"x\n",
"the short line reached the inner writer"
);
let record = vec![b'b'; CAPACITY * 4];
let written = writer
.write(&record)
.expect("a record larger than the capacity is written straight through");
assert_eq!(written, record.len());
assert_eq!(
&writer.get_ref()[2..],
&record[..],
"the oversized record must follow the already-flushed line"
);
}
#[test]
fn a_partial_flush_still_stops_at_the_last_newline() {
let mut writer = MultiLineWriter::with_capacity(CAPACITY, Vec::new());
writer
.write_all(b"aaaa\nbb")
.expect("buffered below capacity");
assert!(
writer.get_ref().is_empty(),
"nothing reaches the inner writer before a flush is triggered"
);
writer
.write_all(b"cccccccccc")
.expect("triggers a partial flush");
assert_eq!(
writer.get_ref(),
b"aaaa\n",
"a partial flush writes up to and including the last newline"
);
writer.flush().expect("the final flush drains the tail");
assert_eq!(
writer.get_ref(),
b"aaaa\nbbcccccccccc",
"the buffered tail must survive the partial flush"
);
}
}