nectar_primitives/file/
write_at.rs

1//! Async random-access write sink for out-of-order file reassembly.
2
3use std::io;
4use std::sync::Mutex;
5
6use futures::StreamExt;
7
8use super::error::FileError;
9use super::joiner::GenericJoiner;
10use super::mode::JoinMode;
11use crate::store::{ChunkGet, MaybeSend, MaybeSync};
12
13/// Async random-access write sink. Each leaf body is written at its absolute
14/// offset; when every offset is written the sink is whole. Not `Send`-bound so
15/// a single-threaded browser writable can implement it.
16pub trait WriteAt {
17    /// Error returned by the sink operations.
18    type Error: core::error::Error + MaybeSend + MaybeSync + 'static;
19
20    /// Write all of `buf` at `offset`.
21    fn write_at(
22        &self,
23        offset: u64,
24        buf: &[u8],
25    ) -> impl std::future::Future<Output = Result<(), Self::Error>> + MaybeSend;
26
27    /// Pre-size the sink so every in-range `write_at` lands.
28    fn set_len(
29        &self,
30        len: u64,
31    ) -> impl std::future::Future<Output = Result<(), Self::Error>> + MaybeSend;
32
33    /// Flush buffered writes. Defaults to a no-op.
34    fn flush(&self) -> impl std::future::Future<Output = Result<(), Self::Error>> + MaybeSend {
35        async { Ok(()) }
36    }
37}
38
39#[cfg(all(not(target_arch = "wasm32"), unix))]
40impl WriteAt for std::fs::File {
41    type Error = io::Error;
42
43    async fn write_at(&self, offset: u64, buf: &[u8]) -> io::Result<()> {
44        use std::os::unix::fs::FileExt;
45        self.write_all_at(buf, offset)
46    }
47
48    #[allow(clippy::use_self)]
49    async fn set_len(&self, len: u64) -> io::Result<()> {
50        std::fs::File::set_len(self, len)
51    }
52
53    async fn flush(&self) -> io::Result<()> {
54        self.sync_all()
55    }
56}
57
58#[cfg(all(not(target_arch = "wasm32"), windows))]
59impl WriteAt for std::fs::File {
60    type Error = io::Error;
61
62    async fn write_at(&self, offset: u64, buf: &[u8]) -> io::Result<()> {
63        use std::os::windows::fs::FileExt;
64        let mut written = 0;
65        while written < buf.len() {
66            let n = self.seek_write(&buf[written..], offset + written as u64)?;
67            if n == 0 {
68                return Err(io::Error::from(io::ErrorKind::WriteZero));
69            }
70            written += n;
71        }
72        Ok(())
73    }
74
75    #[allow(clippy::use_self)]
76    async fn set_len(&self, len: u64) -> io::Result<()> {
77        std::fs::File::set_len(self, len)
78    }
79
80    async fn flush(&self) -> io::Result<()> {
81        self.sync_all()
82    }
83}
84
85impl WriteAt for Mutex<Vec<u8>> {
86    type Error = io::Error;
87
88    async fn write_at(&self, offset: u64, buf: &[u8]) -> io::Result<()> {
89        let mut vec = self
90            .lock()
91            .map_err(|_| io::Error::other("sink lock poisoned"))?;
92        let end = offset as usize + buf.len();
93        if vec.len() < end {
94            vec.resize(end, 0);
95        }
96        vec[offset as usize..end].copy_from_slice(buf);
97        Ok(())
98    }
99
100    async fn set_len(&self, len: u64) -> io::Result<()> {
101        let mut vec = self
102            .lock()
103            .map_err(|_| io::Error::other("sink lock poisoned"))?;
104        vec.resize(len as usize, 0);
105        Ok(())
106    }
107}
108
109impl<T: WriteAt + ?Sized> WriteAt for &T {
110    type Error = T::Error;
111
112    fn write_at(
113        &self,
114        offset: u64,
115        buf: &[u8],
116    ) -> impl std::future::Future<Output = Result<(), Self::Error>> + MaybeSend {
117        (**self).write_at(offset, buf)
118    }
119
120    fn set_len(
121        &self,
122        len: u64,
123    ) -> impl std::future::Future<Output = Result<(), Self::Error>> + MaybeSend {
124        (**self).set_len(len)
125    }
126
127    fn flush(&self) -> impl std::future::Future<Output = Result<(), Self::Error>> + MaybeSend {
128        (**self).flush()
129    }
130}
131
132impl<G, M, const BODY_SIZE: usize> GenericJoiner<G, M, BODY_SIZE>
133where
134    G: ChunkGet<BODY_SIZE> + 'static,
135    M: JoinMode + MaybeSend + Sync,
136{
137    /// Reassemble the whole file into `sink`, writing each out-of-order leaf at
138    /// its offset. Peak memory is the in-flight width, never the file size.
139    /// Cancel-safe: dropping the returned future drops the chunk stream
140    /// (cancelling in-flight fetches) and the sink; a partially-written sink is
141    /// sparse-valid and safe to resume.
142    pub async fn download_into<S: WriteAt>(self, sink: S) -> super::error::Result<()> {
143        self.download_into_with_progress(sink, |_, _| {}).await
144    }
145
146    /// As [`download_into`](Self::download_into), invoking
147    /// `on_progress(written, total)` after each leaf lands.
148    pub async fn download_into_with_progress<S: WriteAt, F: FnMut(u64, u64)>(
149        self,
150        sink: S,
151        mut on_progress: F,
152    ) -> super::error::Result<()> {
153        let total = self.size();
154        sink.set_len(total).await.map_err(FileError::sink)?;
155        let mut stream = std::pin::pin!(self.into_offset_stream_chunked());
156        let mut written = 0u64;
157        while let Some(item) = stream.next().await {
158            let (offset, body) = item?;
159            sink.write_at(offset, &body)
160                .await
161                .map_err(FileError::sink)?;
162            written += body.len() as u64;
163            on_progress(written, total);
164        }
165        sink.flush().await.map_err(FileError::sink)?;
166        Ok(())
167    }
168}
169
170#[cfg(test)]
171mod tests {
172    use super::*;
173
174    use futures::executor::block_on;
175
176    #[test]
177    fn sparse_writes_reassemble() {
178        block_on(async {
179            let sink = Mutex::new(Vec::new());
180            sink.set_len(10).await.unwrap();
181            // High offset first, then a low offset.
182            sink.write_at(6, b"world").await.unwrap();
183            sink.write_at(0, b"hello").await.unwrap();
184            sink.write_at(5, b" ").await.unwrap();
185            assert_eq!(sink.into_inner().unwrap(), b"hello world");
186        });
187    }
188
189    #[test]
190    fn out_of_order_full_coverage() {
191        block_on(async {
192            let sink = Mutex::new(Vec::new());
193            sink.set_len(6).await.unwrap();
194            sink.write_at(4, b"23").await.unwrap();
195            sink.write_at(0, b"01").await.unwrap();
196            sink.write_at(2, b"45").await.unwrap();
197            assert_eq!(sink.into_inner().unwrap(), b"014523");
198        });
199    }
200
201    #[test]
202    fn ref_forwarding_writes_through() {
203        block_on(async {
204            let sink = Mutex::new(Vec::new());
205            let r = &sink;
206            r.set_len(4).await.unwrap();
207            r.write_at(0, b"abcd").await.unwrap();
208            assert_eq!(sink.into_inner().unwrap(), b"abcd");
209        });
210    }
211
212    #[test]
213    fn flush_default_is_noop() {
214        block_on(async {
215            let sink = Mutex::new(Vec::new());
216            sink.write_at(0, b"complete").await.unwrap();
217            sink.flush().await.unwrap();
218            assert_eq!(sink.into_inner().unwrap(), b"complete");
219        });
220    }
221
222    #[test]
223    fn set_len_presizes() {
224        block_on(async {
225            let sink = Mutex::new(Vec::new());
226            sink.set_len(8).await.unwrap();
227            assert_eq!(sink.into_inner().unwrap(), vec![0u8; 8]);
228        });
229    }
230}
231
232#[cfg(all(test, feature = "tokio"))]
233mod download_tests {
234    use super::*;
235
236    use crate::DEFAULT_BODY_SIZE;
237    use crate::chunk::AnyChunk;
238    use crate::file::{Joiner, split};
239    use futures::executor::block_on;
240    use std::collections::HashMap;
241
242    type Store = HashMap<crate::ChunkAddress, AnyChunk>;
243
244    fn split_and_store(data: &[u8]) -> (crate::ChunkAddress, Store) {
245        let (root, store) = split::<DEFAULT_BODY_SIZE>(data).unwrap();
246        (root, store.into_chunks())
247    }
248
249    fn sample(len: usize) -> Vec<u8> {
250        (0..len).map(|i| (i % 256) as u8).collect()
251    }
252
253    /// `download_into` reassembles bytes equal to `read_all` across tree shapes.
254    fn assert_download_matches(data: &[u8], width: usize) {
255        block_on(async {
256            let (root, store) = split_and_store(data);
257
258            let expected = Joiner::new(store.clone(), root)
259                .await
260                .unwrap()
261                .read_all()
262                .await
263                .unwrap();
264
265            let joiner = Joiner::new(store, root)
266                .await
267                .unwrap()
268                .with_concurrency(width);
269            let sink = Mutex::new(Vec::new());
270            joiner.download_into(&sink).await.unwrap();
271
272            assert_eq!(sink.into_inner().unwrap(), expected);
273        });
274    }
275
276    #[test]
277    fn download_into_small() {
278        assert_download_matches(b"hello world", DEFAULT_BODY_SIZE);
279    }
280
281    #[test]
282    fn download_into_exact_chunk() {
283        assert_download_matches(&sample(DEFAULT_BODY_SIZE), 8);
284    }
285
286    #[test]
287    fn download_into_multi_level() {
288        let refs_per_chunk = DEFAULT_BODY_SIZE / super::super::constants::REF_SIZE;
289        assert_download_matches(&sample(DEFAULT_BODY_SIZE * (refs_per_chunk + 1) + 17), 8);
290    }
291
292    #[test]
293    fn download_into_width_one() {
294        assert_download_matches(&sample(DEFAULT_BODY_SIZE * 5 + 7), 1);
295    }
296
297    #[test]
298    fn download_with_progress_monotonic() {
299        block_on(async {
300            let data = sample(DEFAULT_BODY_SIZE * 4 + 99);
301            let (root, store) = split_and_store(&data);
302            let joiner = Joiner::new(store, root).await.unwrap();
303            let total = joiner.size();
304
305            let sink = Mutex::new(Vec::new());
306            let mut last = 0u64;
307            let mut last_total = None;
308            joiner
309                .download_into_with_progress(&sink, |written, t| {
310                    assert!(written >= last, "written must be non-decreasing");
311                    assert!(written <= t, "written must not exceed total");
312                    last = written;
313                    last_total = Some(t);
314                })
315                .await
316                .unwrap();
317
318            assert_eq!(last, total, "final written equals size");
319            assert_eq!(last_total, Some(total));
320            assert_eq!(sink.into_inner().unwrap(), data);
321        });
322    }
323
324    #[cfg(feature = "encryption")]
325    mod encrypted {
326        use super::*;
327        use crate::chunk::encryption::EncryptedChunkRef;
328        use crate::file::{EncryptedJoiner, split_encrypted};
329
330        fn encrypted_split_and_store(data: &[u8]) -> (EncryptedChunkRef, Store) {
331            let (root_ref, store) = split_encrypted::<DEFAULT_BODY_SIZE>(data).unwrap();
332            (root_ref, store.into_chunks())
333        }
334
335        fn assert_encrypted_download_matches(data: &[u8], width: usize) {
336            block_on(async {
337                let (root_ref, store) = encrypted_split_and_store(data);
338
339                let expected = EncryptedJoiner::new(store.clone(), root_ref.clone())
340                    .await
341                    .unwrap()
342                    .read_all()
343                    .await
344                    .unwrap();
345
346                let joiner = EncryptedJoiner::new(store, root_ref)
347                    .await
348                    .unwrap()
349                    .with_concurrency(width);
350                let sink = Mutex::new(Vec::new());
351                joiner.download_into(&sink).await.unwrap();
352
353                assert_eq!(sink.into_inner().unwrap(), expected);
354            });
355        }
356
357        #[test]
358        fn encrypted_download_into_small() {
359            assert_encrypted_download_matches(b"hello world", DEFAULT_BODY_SIZE);
360        }
361
362        #[test]
363        fn encrypted_download_into_multi_level() {
364            let refs_per_chunk = DEFAULT_BODY_SIZE / super::super::super::constants::REF_SIZE;
365            assert_encrypted_download_matches(
366                &sample(DEFAULT_BODY_SIZE * (refs_per_chunk + 1) + 17),
367                8,
368            );
369        }
370
371        #[test]
372        fn encrypted_download_into_width_one() {
373            assert_encrypted_download_matches(&sample(DEFAULT_BODY_SIZE * 5 + 7), 1);
374        }
375    }
376}
nectar_primitives/file/write_at.rs

nectar_primitives/file/
write_at.rs
