Repository: Stebalien/tempfile
Pull request #68 (2018-09-27 06:49:49)

spooled temporary file
For your consideration. Similar to python's `tempfile.SpooledTemporaryFile`: https://docs.python.org/3/library/tempfile.html#tempfile.SpooledTemporaryFile


Patch:
diff --git a/src/lib.rs b/src/lib.rs
index c5e21d2..7d9c674 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -115,9 +115,11 @@ use std::{env, io};
 mod dir;
 mod file;
 mod util;
+mod spooled;
 
 pub use dir::{tempdir, tempdir_in, TempDir};
 pub use file::{tempfile, tempfile_in, NamedTempFile, PersistError, TempPath};
+pub use spooled::{spooled_tempfile, SpooledTempFile};
 
 /// Create a new temporary file or directory with custom parameters.
 #[derive(Debug, Clone, Eq, PartialEq)]
diff --git a/src/spooled.rs b/src/spooled.rs
new file mode 100644
index 0000000..b046c13
--- /dev/null
+++ b/src/spooled.rs
@@ -0,0 +1,155 @@
+use std::fs::File;
+use std::io::{self, Read, Write, Seek, SeekFrom, Cursor};
+use file::tempfile;
+use std::mem::drop;
+
+#[derive(Debug)]
+enum SpooledInner {
+    InMemory(Cursor<Vec<u8>>),
+    OnDisk(File),
+}
+
+/// An object that behaves like a regular temporary file, but keeps data in
+/// memory until it reaches a configured size, at which point the data is
+/// written to a temporary file on disk, and further operations use the file
+/// on disk.
+#[derive(Debug)]
+pub struct SpooledTempFile {
+    max_size: usize,
+    inner: SpooledInner,
+}
+
+/// Create a new spooled temporary file.
+///
+/// # Security
+///
+/// This variant is secure/reliable in the presence of a pathological temporary
+/// file cleaner.
+///
+/// # Resource Leaking
+///
+/// The temporary file will be automatically removed by the OS when the last
+/// handle to it is closed. This doesn't rely on Rust destructors being run, so
+/// will (almost) never fail to clean up the temporary file.
+///
+/// # Examples
+///
+/// ```
+/// # extern crate tempfile;
+/// use tempfile::spooled_tempfile;
+/// use std::io::{self, Write};
+///
+/// # fn main() {
+/// #     if let Err(_) = run() {
+/// #         ::std::process::exit(1);
+/// #     }
+/// # }
+/// # fn run() -> Result<(), io::Error> {
+/// let mut file = spooled_tempfile(15);
+///
+/// writeln!(file, "short line")?;
+/// assert!(!file.is_rolled());
+///
+/// // as a result of this write call, the size of the data will exceed
+/// // `max_size` (15), so it will be written to a temporary file on disk,
+/// // and the in-memory buffer will be dropped
+/// writeln!(file, "marvin gardens")?;
+/// assert!(file.is_rolled());
+///
+/// # Ok(())
+/// # }
+/// ```
+#[inline]
+pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile {
+    SpooledTempFile::new(max_size)
+}
+
+impl SpooledTempFile {
+    pub fn new(max_size: usize) -> SpooledTempFile {
+        SpooledTempFile {
+            max_size: max_size,
+            inner: SpooledInner::InMemory(Cursor::new(Vec::new())),
+        }
+    }
+
+    /// Returns true if the file has been rolled over to disk.
+    pub fn is_rolled(&self) -> bool {
+        match self.inner {
+            SpooledInner::InMemory(_) => false,
+            SpooledInner::OnDisk(_) => true,
+        }
+    }
+
+    /// Rolls over to a file on disk, regardless of current size. Does nothing
+    /// if already rolled over.
+    pub fn roll(&mut self) -> io::Result<()> {
+        if !self.is_rolled() {
+            let mut file = tempfile()?;
+            if let SpooledInner::InMemory(ref mut cursor) = self.inner {
+                file.write(cursor.get_ref())?;
+                file.seek(SeekFrom::Start(cursor.position()))?;
+                drop(cursor);
+            }
+            self.inner = SpooledInner::OnDisk(file);
+        }
+        Ok(())
+    }
+
+    pub fn set_len(&mut self, size: u64) -> Result<(), io::Error> {
+        if size as usize > self.max_size {
+            self.roll()?; // does nothing if already rolled over
+        }
+        match self.inner {
+            SpooledInner::InMemory(ref mut cursor) => {
+                cursor.get_mut().resize(size as usize, 0);
+                Ok(())
+            },
+            SpooledInner::OnDisk(ref mut file) => file.set_len(size),
+        }
+    }
+}
+
+impl Read for SpooledTempFile {
+    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
+        match self.inner {
+            SpooledInner::InMemory(ref mut cursor) => cursor.read(buf),
+            SpooledInner::OnDisk(ref mut file) => file.read(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 SpooledInner::InMemory(ref mut cursor) = self.inner {
+            rolling = cursor.position() as usize + buf.len() > self.max_size;
+        }
+        if rolling {
+            self.roll()?;
+        }
+
+        // write the bytes
+        match self.inner {
+            SpooledInner::InMemory(ref mut cursor) => cursor.write(buf),
+            SpooledInner::OnDisk(ref mut file) => file.write(buf),
+        }
+    }
+
+    #[inline]
+    fn flush(&mut self) -> io::Result<()> {
+        match self.inner {
+            SpooledInner::InMemory(ref mut cursor) => cursor.flush(),
+            SpooledInner::OnDisk(ref mut file) => file.flush(),
+        }
+    }
+}
+
+impl Seek for SpooledTempFile {
+    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
+        match self.inner {
+            SpooledInner::InMemory(ref mut cursor) => cursor.seek(pos),
+            SpooledInner::OnDisk(ref mut file) => file.seek(pos),
+        }
+    }
+}


Test patch:
diff --git a/tests/spooled.rs b/tests/spooled.rs
new file mode 100644
index 0000000..fa388aa
--- /dev/null
+++ b/tests/spooled.rs
@@ -0,0 +1,304 @@
+extern crate tempfile;
+
+use std::io::{Write, Seek, Read, SeekFrom};
+
+use tempfile::{spooled_tempfile, SpooledTempFile};
+
+#[test]
+fn test_automatic_rollover() {
+    let mut t = spooled_tempfile(10);
+    let mut buf = Vec::new();
+
+    assert!(!t.is_rolled());
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 0);
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 0);
+    assert_eq!(buf.as_slice(), b"");
+    buf.clear();
+
+    assert_eq!(t.write(b"abcde").unwrap(), 5);
+
+    assert!(!t.is_rolled());
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 5);
+    assert_eq!(buf.as_slice(), b"abcde");
+
+    assert_eq!(t.write(b"fghijklmno").unwrap(), 10);
+
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 15);
+    assert!(t.is_rolled());
+}
+
+#[test]
+fn test_explicit_rollover() {
+    let mut t = SpooledTempFile::new(100);
+    assert_eq!(t.write(b"abcdefghijklmnopqrstuvwxyz").unwrap(), 26);
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 26);
+    assert!(!t.is_rolled());
+
+    // roll over explicitly
+    assert!(t.roll().is_ok());
+    assert!(t.is_rolled());
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 26);
+
+    let mut buf = Vec::new();
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 0);
+    assert_eq!(buf.as_slice(), b"");
+    buf.clear();
+
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 26);
+    assert_eq!(buf.as_slice(), b"abcdefghijklmnopqrstuvwxyz");
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 26);
+}
+
+// called by test_seek_{buffer, file}
+// assumes t is empty and offset is 0 to start
+fn test_seek(t: &mut SpooledTempFile) {
+    assert_eq!(t.write(b"abcdefghijklmnopqrstuvwxyz").unwrap(), 26);
+
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 26); // tell()
+    assert_eq!(t.seek(SeekFrom::Current(-1)).unwrap(), 25);
+    assert_eq!(t.seek(SeekFrom::Current(1)).unwrap(), 26);
+    assert_eq!(t.seek(SeekFrom::Current(1)).unwrap(), 27);
+    assert_eq!(t.seek(SeekFrom::Current(-27)).unwrap(), 0);
+    assert!(t.seek(SeekFrom::Current(-1)).is_err());
+    assert!(t.seek(SeekFrom::Current(-1245)).is_err());
+
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    assert_eq!(t.seek(SeekFrom::Start(1)).unwrap(), 1);
+    assert_eq!(t.seek(SeekFrom::Start(26)).unwrap(), 26);
+    assert_eq!(t.seek(SeekFrom::Start(27)).unwrap(), 27);
+    // // these are build errors
+    // assert!(t.seek(SeekFrom::Start(-1)).is_err());
+    // assert!(t.seek(SeekFrom::Start(-1000)).is_err());
+
+    assert_eq!(t.seek(SeekFrom::End(0)).unwrap(), 26);
+    assert_eq!(t.seek(SeekFrom::End(-1)).unwrap(), 25);
+    assert_eq!(t.seek(SeekFrom::End(-26)).unwrap(), 0);
+    assert!(t.seek(SeekFrom::End(-27)).is_err());
+    assert!(t.seek(SeekFrom::End(-99)).is_err());
+    assert_eq!(t.seek(SeekFrom::End(1)).unwrap(), 27);
+    assert_eq!(t.seek(SeekFrom::End(1)).unwrap(), 27);
+}
+
+#[test]
+fn test_seek_buffer() {
+    let mut t = spooled_tempfile(100);
+    test_seek(&mut t);
+}
+
+#[test]
+fn test_seek_file() {
+    let mut t = SpooledTempFile::new(10);
+    test_seek(&mut t);
+}
+
+fn test_seek_read(t: &mut SpooledTempFile) {
+    assert_eq!(t.write(b"abcdefghijklmnopqrstuvwxyz").unwrap(), 26);
+
+    let mut buf = Vec::new();
+
+    // we're at the end
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 0);
+    assert_eq!(buf.as_slice(), b"");
+    buf.clear();
+
+    // seek to start, read whole thing
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 26);
+    assert_eq!(buf.as_slice(), b"abcdefghijklmnopqrstuvwxyz");
+    buf.clear();
+
+    // now we're at the end again
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 26); // tell()
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 0);
+    assert_eq!(buf.as_slice(), b"");
+    buf.clear();
+
+    // seek to somewhere in the middle, read a bit
+    assert_eq!(t.seek(SeekFrom::Start(5)).unwrap(), 5);
+    let mut buf = [0; 5];
+    assert!(t.read_exact(&mut buf).is_ok());
+    assert_eq!(buf, *b"fghij");
+
+    // read again from current spot
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 10); // tell()
+    assert!(t.read_exact(&mut buf).is_ok());
+    assert_eq!(buf, *b"klmno");
+
+    let mut buf = [0; 15];
+    // partial read
+    assert_eq!(t.read(&mut buf).unwrap(), 11);
+    assert_eq!(buf[0..11], *b"pqrstuvwxyz");
+
+    // try to read off the end: UnexpectedEof
+    assert!(t.read_exact(&mut buf).is_err());
+}
+
+#[test]
+fn test_seek_read_buffer() {
+    let mut t = spooled_tempfile(100);
+    test_seek_read(&mut t);
+}
+
+#[test]
+fn test_seek_read_file() {
+    let mut t = SpooledTempFile::new(10);
+    test_seek_read(&mut t);
+}
+
+fn test_overwrite_middle(t: &mut SpooledTempFile) {
+    assert_eq!(t.write(b"abcdefghijklmnopqrstuvwxyz").unwrap(), 26);
+
+    assert_eq!(t.seek(SeekFrom::Start(10)).unwrap(), 10);
+    assert_eq!(t.write(b"0123456789").unwrap(), 10);
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+
+    let mut buf = Vec::new();
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 26);
+    assert_eq!(buf.as_slice(), b"abcdefghij0123456789uvwxyz");
+}
+
+#[test]
+fn test_overwrite_middle_of_buffer() {
+    let mut t = spooled_tempfile(100);
+    test_overwrite_middle(&mut t);
+}
+
+#[test]
+fn test_overwrite_middle_of_file() {
+    let mut t = SpooledTempFile::new(10);
+    test_overwrite_middle(&mut t);
+}
+
+#[test]
+fn test_overwrite_and_extend_buffer() {
+    let mut t = spooled_tempfile(100);
+    assert_eq!(t.write(b"abcdefghijklmnopqrstuvwxyz").unwrap(), 26);
+    assert_eq!(t.seek(SeekFrom::End(-5)).unwrap(), 21);
+    assert_eq!(t.write(b"0123456789").unwrap(), 10);
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    let mut buf = Vec::new();
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 31);
+    assert_eq!(buf.as_slice(), b"abcdefghijklmnopqrstu0123456789");
+    assert!(!t.is_rolled());
+}
+
+#[test]
+fn test_overwrite_and_extend_rollover() {
+    let mut t = SpooledTempFile::new(20);
+    assert_eq!(t.write(b"abcdefghijklmno").unwrap(), 15);
+    assert!(!t.is_rolled());
+    assert_eq!(t.seek(SeekFrom::End(-5)).unwrap(), 10);
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 10); // tell()
+    assert!(!t.is_rolled());
+    assert_eq!(t.write(b"0123456789)!@#$%^&*(").unwrap(), 20);
+    assert!(t.is_rolled());
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 30); // tell()
+    let mut buf = Vec::new();
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 30);
+    assert_eq!(buf.as_slice(), b"abcdefghij0123456789)!@#$%^&*(");
+}
+
+fn test_sparse(t: &mut SpooledTempFile) {
+    assert_eq!(t.write(b"abcde").unwrap(), 5);
+    assert_eq!(t.seek(SeekFrom::Current(5)).unwrap(), 10);
+    assert_eq!(t.write(b"klmno").unwrap(), 5);
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    let mut buf = Vec::new();
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 15);
+    assert_eq!(buf.as_slice(), b"abcde\0\0\0\0\0klmno");
+}
+
+#[test]
+fn test_sparse_buffer() {
+    let mut t = spooled_tempfile(100);
+    test_sparse(&mut t);
+}
+
+#[test]
+fn test_sparse_file() {
+    let mut t = SpooledTempFile::new(1);
+    test_sparse(&mut t);
+}
+
+#[test]
+fn test_sparse_write_rollover() {
+    let mut t = spooled_tempfile(10);
+    assert_eq!(t.write(b"abcde").unwrap(), 5);
+    assert!(!t.is_rolled());
+    assert_eq!(t.seek(SeekFrom::Current(5)).unwrap(), 10);
+    assert!(!t.is_rolled());
+    assert_eq!(t.write(b"klmno").unwrap(), 5);
+    assert!(t.is_rolled());
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    let mut buf = Vec::new();
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 15);
+    assert_eq!(buf.as_slice(), b"abcde\0\0\0\0\0klmno");
+}
+
+fn test_set_len(t: &mut SpooledTempFile) {
+    let mut buf: Vec<u8> = Vec::new();
+
+    assert_eq!(t.write(b"abcdefghijklmnopqrstuvwxyz").unwrap(), 26);
+
+    // truncate to 10 bytes
+    assert!(t.set_len(10).is_ok());
+
+    // position should not have moved
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 26); // tell()
+
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 0);
+    assert_eq!(buf.as_slice(), b"");
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 26); // tell()
+    buf.clear();
+
+    // read whole thing
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 10);
+    assert_eq!(buf.as_slice(), b"abcdefghij");
+    buf.clear();
+
+    // set_len to expand beyond the end
+    assert!(t.set_len(40).is_ok());
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 10); // tell()
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 40);
+    assert_eq!(buf.as_slice(), &b"abcdefghij\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0"[..]);
+}
+
+#[test]
+fn test_set_len_buffer() {
+    let mut t = spooled_tempfile(100);
+    test_set_len(&mut t);
+}
+
+#[test]
+fn test_set_len_file() {
+    let mut t = spooled_tempfile(100);
+    test_set_len(&mut t);
+}
+
+#[test]
+fn test_set_len_rollover() {
+    let mut buf: Vec<u8> = Vec::new();
+
+    let mut t = spooled_tempfile(10);
+    assert_eq!(t.write(b"abcde").unwrap(), 5);
+    assert!(!t.is_rolled());
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 5); // tell()
+
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 5);
+    assert_eq!(buf.as_slice(), b"abcde");
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 5); // tell()
+    buf.clear();
+
+    assert!(t.set_len(20).is_ok());
+    assert!(t.is_rolled());
+    assert_eq!(t.seek(SeekFrom::Current(0)).unwrap(), 5); // tell()
+    assert_eq!(t.seek(SeekFrom::Start(0)).unwrap(), 0);
+    assert_eq!(t.read_to_end(&mut buf).unwrap(), 20);
+    assert_eq!(buf.as_slice(), b"abcde\0\0\0\0\0\0\0\0\0\0\0\0\0\0\0");
+}