disk_core/
stream_pipe.rs

1//! Copy a file between two filesystems without holding it in memory.
2//!
3//! The source's `write_file_to` runs on a scoped thread and feeds a bounded
4//! channel of chunks; the consumer (usually the target's `create_file`) reads
5//! the other end. Peak memory is a few chunks, whatever the file's size. A
6//! producer error reaches the consumer as an `io::Error` rather than a short
7//! read, so a failed source never yields a silently truncated copy.
8
9use std::io::{self, Read, Write};
10use std::sync::mpsc::{sync_channel, Receiver, SyncSender};
11
12use crate::entry::FileEntry;
13use crate::filesystem::{Filesystem, FilesystemError};
14
15const CHUNK: usize = 1 << 20;
16const DEPTH: usize = 4;
17
18type Msg = Result<Vec<u8>, String>;
19
20struct ChanWriter {
21    tx: SyncSender<Msg>,
22    buf: Vec<u8>,
23}
24
25impl Write for ChanWriter {
26    fn write(&mut self, data: &[u8]) -> io::Result<usize> {
27        let room = CHUNK - self.buf.len();
28        let n = data.len().min(room);
29        self.buf.extend_from_slice(&data[..n]);
30        if self.buf.len() == CHUNK {
31            self.flush()?;
32        }
33        Ok(n)
34    }
35
36    fn flush(&mut self) -> io::Result<()> {
37        if self.buf.is_empty() {
38            return Ok(());
39        }
40        let chunk = std::mem::replace(&mut self.buf, Vec::with_capacity(CHUNK));
41        self.tx
42            .send(Ok(chunk))
43            .map_err(|_| io::Error::new(io::ErrorKind::BrokenPipe, "consumer stopped reading"))
44    }
45}
46
47struct ChanReader {
48    rx: Receiver<Msg>,
49    cur: Vec<u8>,
50    pos: usize,
51}
52
53impl Read for ChanReader {
54    fn read(&mut self, out: &mut [u8]) -> io::Result<usize> {
55        while self.pos == self.cur.len() {
56            match self.rx.recv() {
57                Ok(Ok(chunk)) => {
58                    self.cur = chunk;
59                    self.pos = 0;
60                }
61                Ok(Err(e)) => return Err(crate::io_other(e)),
62                Err(_) => return Ok(0),
63            }
64        }
65        let n = out.len().min(self.cur.len() - self.pos);
66        out[..n].copy_from_slice(&self.cur[self.pos..self.pos + n]);
67        self.pos += n;
68        Ok(n)
69    }
70}
71
72/// Run `consume` over a reader that streams `entry`'s bytes out of `source`.
73pub fn with_file_reader<S, T>(
74    source: &mut S,
75    entry: &FileEntry,
76    consume: impl FnOnce(&mut dyn Read) -> Result<T, FilesystemError>,
77) -> Result<T, FilesystemError>
78where
79    S: Filesystem + Send + ?Sized,
80{
81    let (tx, rx) = sync_channel::<Msg>(DEPTH);
82    std::thread::scope(|scope| {
83        let producer = scope.spawn(move || {
84            let mut w = ChanWriter {
85                tx: tx.clone(),
86                buf: Vec::with_capacity(CHUNK),
87            };
88            let r = source.write_file_to(entry, &mut w).and_then(|n| {
89                w.flush()?;
90                Ok(n)
91            });
92            if let Err(e) = &r {
93                let _ = tx.send(Err(e.to_string()));
94            }
95            r
96        });
97        let mut reader = ChanReader {
98            rx,
99            cur: Vec::new(),
100            pos: 0,
101        };
102        let out = consume(&mut reader);
103        // Unblock a producer still waiting on a full channel if the consumer stopped early.
104        drop(reader);
105        let produced = producer
106            .join()
107            .map_err(|_| FilesystemError::Io(crate::io_other("source reader panicked")))?;
108        match (out, produced) {
109            (Ok(v), Ok(_)) => Ok(v),
110            (Err(e), _) => Err(e),
111            (Ok(_), Err(e)) => Err(e),
112        }
113    })
114}
115
116#[cfg(test)]
117mod tests {
118    use super::*;
119
120    /// One file at `/big`, served by the default chunked `write_file_to`.
121    struct OneFile(Vec<u8>);
122
123    impl Filesystem for OneFile {
124        fn root(&mut self) -> Result<FileEntry, FilesystemError> {
125            Ok(FileEntry::root())
126        }
127        fn list_directory(&mut self, _: &FileEntry) -> Result<Vec<FileEntry>, FilesystemError> {
128            Ok(vec![self.entry()])
129        }
130        fn read_file(&mut self, _: &FileEntry, max: usize) -> Result<Vec<u8>, FilesystemError> {
131            Ok(self.0[..self.0.len().min(max)].to_vec())
132        }
133        fn write_file_to(
134            &mut self,
135            _: &FileEntry,
136            w: &mut dyn Write,
137        ) -> Result<u64, FilesystemError> {
138            for c in self.0.chunks(64 * 1024) {
139                w.write_all(c)?;
140            }
141            Ok(self.0.len() as u64)
142        }
143        fn volume_label(&self) -> Option<&str> {
144            None
145        }
146        fn fs_type(&self) -> &str {
147            "test"
148        }
149        fn total_size(&self) -> u64 {
150            0
151        }
152        fn used_size(&self) -> u64 {
153            0
154        }
155    }
156
157    impl OneFile {
158        fn entry(&self) -> FileEntry {
159            FileEntry::new_file("big".into(), "/big".into(), self.0.len() as u64, 1)
160        }
161    }
162
163    #[test]
164    fn streams_a_file_larger_than_the_channel() {
165        let body: Vec<u8> = (0..(CHUNK * DEPTH * 3 + 17))
166            .map(|i| (i % 251) as u8)
167            .collect();
168        let mut fs = OneFile(body.clone());
169        let entry = fs.entry();
170        let got = with_file_reader(&mut fs, &entry, |r| {
171            let mut v = Vec::new();
172            r.read_to_end(&mut v)?;
173            Ok(v)
174        })
175        .unwrap();
176        assert_eq!(got, body);
177    }
178
179    #[test]
180    fn a_consumer_that_stops_early_does_not_hang() {
181        let mut fs = OneFile(vec![7u8; CHUNK * DEPTH * 4]);
182        let entry = fs.entry();
183        let got = with_file_reader(&mut fs, &entry, |r| {
184            let mut one = [0u8; 10];
185            r.read_exact(&mut one)?;
186            Ok(one[0])
187        });
188        assert!(got.is_err(), "the producer's broken pipe surfaces");
189    }
190}
disk_core/stream_pipe.rs

disk_core/
stream_pipe.rs
