tempfile/
spooled.rs

Help
1use crate::file::tempfile;
2use std::fs::File;
3use std::io::{self, Cursor, Read, Seek, SeekFrom, Write};
4
5#[derive(Debug)]
6enum SpooledInner {
7    InMemory(Cursor<Vec<u8>>),
8    OnDisk(File),
9}
10
11/// An object that behaves like a regular temporary file, but keeps data in
12/// memory until it reaches a configured size, at which point the data is
13/// written to a temporary file on disk, and further operations use the file
14/// on disk.
15#[derive(Debug)]
16pub struct SpooledTempFile {
17    max_size: usize,
18    inner: SpooledInner,
19}
20
21/// Create a new spooled temporary file.
22///
23/// # Security
24///
25/// This variant is secure/reliable in the presence of a pathological temporary
26/// file cleaner.
27///
28/// # Resource Leaking
29///
30/// The temporary file will be automatically removed by the OS when the last
31/// handle to it is closed. This doesn't rely on Rust destructors being run, so
32/// will (almost) never fail to clean up the temporary file.
33///
34/// # Examples
35///
36/// ```
37/// use tempfile::spooled_tempfile;
38/// use std::io::{self, Write};
39///
40/// # fn main() {
41/// #     if let Err(_) = run() {
42/// #         ::std::process::exit(1);
43/// #     }
44/// # }
45/// # fn run() -> Result<(), io::Error> {
46/// let mut file = spooled_tempfile(15);
47///
48/// writeln!(file, "short line")?;
49/// assert!(!file.is_rolled());
50///
51/// // as a result of this write call, the size of the data will exceed
52/// // `max_size` (15), so it will be written to a temporary file on disk,
53/// // and the in-memory buffer will be dropped
54/// writeln!(file, "marvin gardens")?;
55/// assert!(file.is_rolled());
56///
57/// # Ok(())
58/// # }
59/// ```
60#[inline]
61pub fn spooled_tempfile(max_size: usize) -> SpooledTempFile {
62    SpooledTempFile::new(max_size)
63}
64
65impl SpooledTempFile {
66    pub fn new(max_size: usize) -> SpooledTempFile {
67        SpooledTempFile {
68            max_size: max_size,
69            inner: SpooledInner::InMemory(Cursor::new(Vec::new())),
70        }
71    }
72
73    /// Returns true if the file has been rolled over to disk.
74    pub fn is_rolled(&self) -> bool {
75        match self.inner {
76            SpooledInner::InMemory(_) => false,
77            SpooledInner::OnDisk(_) => true,
78        }
79    }
80
81    /// Rolls over to a file on disk, regardless of current size. Does nothing
82    /// if already rolled over.
83    pub fn roll(&mut self) -> io::Result<()> {
84        if !self.is_rolled() {
85            let mut file = tempfile()?;
86            if let SpooledInner::InMemory(ref mut cursor) = self.inner {
87                file.write_all(cursor.get_ref())?;
88                file.seek(SeekFrom::Start(cursor.position()))?;
89            }
90            self.inner = SpooledInner::OnDisk(file);
91        }
92        Ok(())
93    }
94
95    pub fn set_len(&mut self, size: u64) -> Result<(), io::Error> {
96        if size as usize > self.max_size {
97            self.roll()?; // does nothing if already rolled over
98        }
99        match self.inner {
100            SpooledInner::InMemory(ref mut cursor) => {
101                cursor.get_mut().resize(size as usize, 0);
102                Ok(())
103            }
104            SpooledInner::OnDisk(ref mut file) => file.set_len(size),
105        }
106    }
107}
108
109impl Read for SpooledTempFile {
110    fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
111        match self.inner {
112            SpooledInner::InMemory(ref mut cursor) => cursor.read(buf),
113            SpooledInner::OnDisk(ref mut file) => file.read(buf),
114        }
115    }
116}
117
118impl Write for SpooledTempFile {
119    fn write(&mut self, buf: &[u8]) -> io::Result<usize> {
120        // roll over to file if necessary
121        let mut rolling = false;
122        if let SpooledInner::InMemory(ref mut cursor) = self.inner {
123            rolling = cursor.position() as usize + buf.len() > self.max_size;
124        }
125        if rolling {
126            self.roll()?;
127        }
128
129        // write the bytes
130        match self.inner {
131            SpooledInner::InMemory(ref mut cursor) => cursor.write(buf),
132            SpooledInner::OnDisk(ref mut file) => file.write(buf),
133        }
134    }
135
136    #[inline]
137    fn flush(&mut self) -> io::Result<()> {
138        match self.inner {
139            SpooledInner::InMemory(ref mut cursor) => cursor.flush(),
140            SpooledInner::OnDisk(ref mut file) => file.flush(),
141        }
142    }
143}
144
145impl Seek for SpooledTempFile {
146    fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
147        match self.inner {
148            SpooledInner::InMemory(ref mut cursor) => cursor.seek(pos),
149            SpooledInner::OnDisk(ref mut file) => file.seek(pos),
150        }
151    }
152}
tempfile/spooled.rs

tempfile/
spooled.rs
