1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
|
use crate::file::tempfile;
use std::fs::File;
use std::io::{self, Cursor, Read, Seek, SeekFrom, Write};
/// A wrapper for the two states of a `SpooledTempFile`.
#[derive(Debug)]
pub enum SpooledData {
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: SpooledData,
}
/// 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
///
/// ```
/// 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 {
#[must_use]
pub fn new(max_size: usize) -> SpooledTempFile {
SpooledTempFile {
max_size,
inner: SpooledData::InMemory(Cursor::new(Vec::new())),
}
}
/// Returns true if the file has been rolled over to disk.
#[must_use]
pub fn is_rolled(&self) -> bool {
match self.inner {
SpooledData::InMemory(_) => false,
SpooledData::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 SpooledData::InMemory(cursor) = &mut self.inner {
file.write_all(cursor.get_ref())?;
file.seek(SeekFrom::Start(cursor.position()))?;
}
self.inner = SpooledData::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 &mut self.inner {
SpooledData::InMemory(cursor) => {
cursor.get_mut().resize(size as usize, 0);
Ok(())
}
SpooledData::OnDisk(file) => file.set_len(size),
}
}
/// Consumes and returns the inner `SpooledData` type.
#[must_use]
pub fn into_inner(self) -> SpooledData {
self.inner
}
}
impl Read for SpooledTempFile {
fn read(&mut self, buf: &mut [u8]) -> io::Result<usize> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.read(buf),
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
if matches! {
&self.inner, SpooledData::InMemory(cursor)
if cursor.position() as usize + buf.len() > self.max_size
} {
self.roll()?;
}
// write the bytes
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.write(buf),
SpooledData::OnDisk(file) => file.write(buf),
}
}
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 {
SpooledData::InMemory(cursor) => cursor.flush(),
SpooledData::OnDisk(file) => file.flush(),
}
}
}
impl Seek for SpooledTempFile {
fn seek(&mut self, pos: SeekFrom) -> io::Result<u64> {
match &mut self.inner {
SpooledData::InMemory(cursor) => cursor.seek(pos),
SpooledData::OnDisk(file) => file.seek(pos),
}
}
}
|
