Skip to content

Commit

Permalink
feat: implement default methods for SpooledTempFile Read/Write
Browse files Browse the repository at this point in the history
That way we can take advantage of, e.g., vectorized operations.
  • Loading branch information
Stebalien committed Apr 8, 2023
1 parent 3590dbf commit 214814f
Showing 1 changed file with 47 additions and 5 deletions.
52 changes: 47 additions & 5 deletions src/spooled.rs
Original file line number Diff line number Diff line change
Expand Up @@ -122,16 +122,43 @@ impl Read for SpooledTempFile {
SpooledData::OnDisk(file) => file.read(buf),
}
}

fn read_vectored(&mut self, bufs: &mut [io::IoSliceMut<'_>]) -> io::Result<usize> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.read_vectored(bufs),
SpooledData::OnDisk(file) => file.read_vectored(bufs),
}
}

fn read_to_end(&mut self, buf: &mut Vec<u8>) -> io::Result<usize> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.read_to_end(buf),
SpooledData::OnDisk(file) => file.read_to_end(buf),
}
}

fn read_to_string(&mut self, buf: &mut String) -> io::Result<usize> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.read_to_string(buf),
SpooledData::OnDisk(file) => file.read_to_string(buf),
}
}

fn read_exact(&mut self, buf: &mut [u8]) -> io::Result<()> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.read_exact(buf),
SpooledData::OnDisk(file) => file.read_exact(buf),
}
}
}

impl Write for SpooledTempFile {
fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
// roll over to file if necessary
let mut rolling = false;
if let SpooledData::InMemory(cursor) = &mut self.inner {
rolling = cursor.position() as usize + buf.len() > self.max_size;
}
if rolling {
if matches! {
&self.inner, SpooledData::InMemory(cursor)
if cursor.position() as usize + buf.len() > self.max_size
} {
self.roll()?;
}

Expand All @@ -142,6 +169,21 @@ impl Write for SpooledTempFile {
}
}

fn write_vectored(&mut self, bufs: &[io::IoSlice<'_>]) -> io::Result<usize> {
if matches! {
&self.inner, SpooledData::InMemory(cursor)
// Borrowed from the rust standard library.
if cursor.position() as usize + bufs.iter()
.fold(0usize, |a, b| a.saturating_add(b.len())) > self.max_size
} {
self.roll()?;
}
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.write_vectored(bufs),
SpooledData::OnDisk(file) => file.write_vectored(bufs),
}
}

#[inline]
fn flush(&mut self) -> io::Result<()> {
match &mut self.inner {
Expand Down

0 comments on commit 214814f

Please sign in to comment.