tempfile/
spooled.rs

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

tempfile/
spooled.rs
