snix_castore/blobstore/blob_meta/
bytes_stream.rs

1use futures::StreamExt;
2use std::io::{self, Cursor};
3
4use crate::{
5    B3Digest,
6    blobstore::BlobMeta,
7    chunkstore::{Chunk, ChunkStore},
8};
9
10impl BlobMeta {
11    /// Returns a stream of bytes of actual blob bytes.
12    ///
13    /// Both a start_offset and a size limit to read can be specified.
14    ///
15    /// `start_offset` must be inside the range of the blob.
16    /// If `limit_size` is set, the last element in the stream will contain a chunk with truncated data.
17    /// If `limit_size` would point past the end of the blob, it's clamped to the end.
18    ///
19    /// Chunks not part of that range are skipped entirely.
20    /// The returned Cursor into the first chunk is advanced to the desired start_offset.
21    ///
22    /// Queries the passed ChunkStore, with configurable concurrency.
23    pub fn bytes_stream(
24        &self,
25        start_offset: u64,
26        limit_size: Option<u64>,
27        fetch_concurrency: usize,
28        chunk_store: impl ChunkStore + Clone + 'static,
29    ) -> impl futures::Stream<Item = std::io::Result<Cursor<Chunk>>> + 'static {
30        // Calculate the range of chunks we want to iterate over
31
32        // The index of the chunk to start with, and the number of bytes to skip from the front of it.
33        let (start_idx, skip_from_first) = self.chunk_idx_from(start_offset);
34
35        // specifies the position in the blob that is just outside the range
36        let end = if let Some(size) = limit_size {
37            start_offset
38                .checked_add(size)
39                .expect("offset + size overlows u64::MAX")
40                .min(self.blob_len())
41        } else {
42            self.blob_len()
43        };
44
45        // Calculate the last chunk idx that holds our data,
46        // and how many bytes from there we want to include.
47        let (end_idx, len_last) = self.chunk_idx_until(end);
48
49        let chunk_metas = self.chunk_metas.clone();
50
51        // produce a stream of byte chunks, if the range is not empty
52        let chunk_idxs = (start_offset < end)
53            .then_some(start_idx..=end_idx)
54            .into_iter()
55            .flatten();
56
57        tokio_stream::iter(chunk_idxs)
58            .map(move |i| {
59                let start_pos = if i == start_idx { skip_from_first } else { 0 };
60
61                // get a cursor over the chunk, seeked to start_pos
62                fetch_chunk_and_seek(
63                    chunk_metas[i].digest,
64                    start_pos,
65                    (i == end_idx).then_some(len_last),
66                    chunk_store.clone(),
67                )
68            })
69            .buffered(fetch_concurrency)
70    }
71}
72
73/// Fetches a Chunk, wraps it in a Cursor and seeks to the specified position, optionally truncating to a size.
74async fn fetch_chunk_and_seek(
75    chunk_digest: B3Digest,
76    start_pos: u64,
77    truncate_size: Option<u64>,
78    chunk_store: impl ChunkStore,
79) -> io::Result<Cursor<Chunk>> {
80    let chunk = chunk_store
81        .get(&chunk_digest)
82        .await
83        .map_err(io::Error::other)?
84        .ok_or_else(|| {
85            io::Error::new(
86                io::ErrorKind::BrokenPipe,
87                format!("chunk {0} not found", chunk_digest),
88            )
89        })?;
90
91    let chunk_len = chunk.len();
92
93    let chunk = if let Some(s) = truncate_size {
94        debug_assert!(
95            chunk_len as u64 >= s,
96            "Snix bug: truncate_size is larger than chunk_len"
97        );
98        Chunk::from(bytes::Bytes::from(chunk).slice(0..s as usize))
99    } else {
100        chunk
101    };
102
103    let mut cursor = Cursor::new(chunk);
104
105    // Set the chunk position
106    if start_pos > 0 {
107        debug_assert!(
108            chunk_len as u64 > start_pos,
109            "Snix bug: chunk size is smaller than bytes to skip"
110        );
111        cursor.set_position(start_pos);
112    }
113
114    Ok::<_, io::Error>(cursor)
115}
116
117#[cfg(test)]
118mod test {
119    use std::{
120        io::Cursor,
121        sync::{Arc, LazyLock},
122    };
123
124    use futures::TryStreamExt;
125    use mockall::predicate;
126
127    use crate::{
128        blobstore::BlobMeta,
129        chunkstore::{Chunk, MockChunkStore},
130    };
131
132    const CHUNK_1: Chunk = Chunk::from_static(b"ab");
133    const CHUNK_2: Chunk = Chunk::from_static(b"c");
134    const CHUNK_3: Chunk = Chunk::from_static(b"de");
135    static BLOB_1_META: LazyLock<BlobMeta> = LazyLock::new(|| {
136        BlobMeta::from_digests_and_sizes([
137            (CHUNK_1.digest(), CHUNK_1.len() as u64),
138            (CHUNK_2.digest(), CHUNK_2.len() as u64),
139        ])
140    });
141    static BLOB_2_META: LazyLock<BlobMeta> = LazyLock::new(|| {
142        BlobMeta::from_digests_and_sizes([
143            (CHUNK_1.digest(), CHUNK_1.len() as u64),
144            (CHUNK_2.digest(), CHUNK_2.len() as u64),
145            (CHUNK_3.digest(), CHUNK_3.len() as u64),
146        ])
147    });
148
149    async fn collect_chunks(
150        s: impl futures::Stream<Item = std::io::Result<Cursor<Chunk>>>,
151    ) -> Vec<Vec<u8>> {
152        s.map_ok(|c| c.get_ref().as_ref()[c.position() as usize..].to_vec())
153            .try_collect()
154            .await
155            .expect("to not fail")
156    }
157
158    /// Sets offset to the beginning of the second chunk, ensures the first chunk is not fetched.
159    #[tokio::test]
160    async fn chunked_get_skip() {
161        let mut chunk_store = MockChunkStore::new();
162        chunk_store
163            .expect_get()
164            .with(predicate::eq(CHUNK_2.digest()))
165            .return_once(|_| Ok(Some(CHUNK_2.to_owned())));
166
167        let chunks =
168            collect_chunks((*BLOB_1_META).bytes_stream(2, None, 10, Arc::new(chunk_store))).await;
169
170        assert_eq!(vec![b"c".to_vec()], chunks);
171    }
172
173    /// Skip the first byte in the first chunk
174    #[tokio::test]
175    async fn chunked_skip_one_byte() {
176        let mut chunk_store = MockChunkStore::new();
177        chunk_store
178            .expect_get()
179            .returning(|digest| {
180                if *digest == CHUNK_1.digest() {
181                    Ok(Some(CHUNK_1.to_owned()))
182                } else if *digest == CHUNK_2.digest() {
183                    Ok(Some(CHUNK_2.to_owned()))
184                } else {
185                    panic!("called with unexpected digest")
186                }
187            })
188            .times(2);
189
190        let chunks =
191            collect_chunks((*BLOB_1_META).bytes_stream(1, None, 10, Arc::new(chunk_store))).await;
192        assert_eq!(vec![b"b".to_vec(), b"c".to_vec()], chunks, "data to match");
193    }
194
195    /// Skip to the end
196    #[tokio::test]
197    async fn chunked_skip_end() {
198        let mut chunk_store = MockChunkStore::new();
199        chunk_store.expect_get().never();
200
201        let chunks =
202            collect_chunks((*BLOB_1_META).bytes_stream(3, None, 10, Arc::new(chunk_store))).await;
203        assert_eq!(Vec::<Vec<u8>>::new(), chunks, "data to match");
204    }
205
206    #[tokio::test]
207    async fn read_empty() {
208        let mut chunk_store = MockChunkStore::new();
209        chunk_store.expect_get().never();
210
211        let blob_meta = BlobMeta::from_digests_and_sizes([]);
212        let chunks =
213            collect_chunks(blob_meta.bytes_stream(0, None, 10, Arc::new(chunk_store))).await;
214
215        assert_eq!(Vec::<Vec<u8>>::new(), chunks, "data to match");
216    }
217
218    #[tokio::test]
219    async fn read_right_at_end_empty() {
220        let mut chunk_store = MockChunkStore::new();
221        chunk_store.expect_get().never();
222        let chunk_store = Arc::new(chunk_store);
223
224        for (start_offset, size) in [(3, None), (3, Some(1))] {
225            let chunks = collect_chunks(BLOB_1_META.bytes_stream(
226                start_offset,
227                size,
228                10,
229                Arc::new(chunk_store.clone()),
230            ))
231            .await;
232
233            assert_eq!(0, chunks.len(), "expect 0 chunks");
234        }
235    }
236
237    #[tokio::test]
238    #[should_panic]
239    async fn offset_past_end_panic() {
240        let mut chunk_store = MockChunkStore::new();
241        chunk_store.expect_get().never();
242        let mut stream = BLOB_1_META.bytes_stream(4, None, 10, Arc::new(chunk_store));
243        let _ = stream.try_next().await;
244    }
245
246    #[tokio::test]
247    async fn size_past_end_clamp() {
248        let mut chunk_store = MockChunkStore::new();
249        chunk_store
250            .expect_get()
251            .returning(|digest| {
252                if *digest == CHUNK_1.digest() {
253                    Ok(Some(CHUNK_1.to_owned()))
254                } else if *digest == CHUNK_2.digest() {
255                    Ok(Some(CHUNK_2.to_owned()))
256                } else {
257                    panic!("called with unexpected digest")
258                }
259            })
260            .times(2);
261        let chunks =
262            collect_chunks(BLOB_1_META.bytes_stream(1, Some(20), 10, Arc::new(chunk_store))).await;
263        let data = chunks.iter().fold(vec![], |mut i, e| {
264            i.extend_from_slice(e);
265            i
266        });
267        assert_eq!(b"bc", &data[..], "expect data to match")
268    }
269
270    #[tokio::test]
271    async fn size_actually_limits() {
272        let mut chunk_store = MockChunkStore::new();
273        chunk_store
274            .expect_get()
275            .with(predicate::eq(CHUNK_1.digest()))
276            .return_once(|_| Ok(Some(CHUNK_1.to_owned())));
277        let chunks =
278            collect_chunks(BLOB_1_META.bytes_stream(0, Some(1), 10, Arc::new(chunk_store))).await;
279        let data = chunks.iter().fold(vec![], |mut i, e| {
280            i.extend_from_slice(e);
281            i
282        });
283        assert_eq!(b"a", &data[..], "expect data to match")
284    }
285
286    use proptest::prelude::*;
287    proptest! {
288        /// Reads various subslices from BLOB_2 (5 bytes) and ensures they come back correctly.
289        #[test]
290        fn test_start_end_offsets(
291            // start
292            start_offset in 0usize..5,
293            // end
294            size in 0usize..5
295        ) {
296            let blob_2_data = b"abcde";
297            // skip too large sizes
298            if start_offset + size > blob_2_data.len() {
299                return Ok(());
300            }
301
302            let mut chunk_store = MockChunkStore::new();
303            chunk_store
304                .expect_get()
305                .returning(|digest| {
306                    if *digest == CHUNK_1.digest() {
307                        Ok(Some(CHUNK_1.to_owned()))
308                    } else if *digest == CHUNK_2.digest() {
309                        Ok(Some(CHUNK_2.to_owned()))
310                    } else if *digest == CHUNK_3.digest() {
311                        Ok(Some(CHUNK_3.to_owned()))
312                    } else {
313                        panic!("called with unexpected digest")
314                    }
315                });
316
317            let tokio_runtime = tokio::runtime::Runtime::new().expect("to construct runtime");
318            tokio_runtime.block_on(async {
319                let stream = BLOB_2_META.bytes_stream(start_offset as u64, Some(size as u64), 10, Arc::new(chunk_store));
320                let mut reader = tokio_util::io::StreamReader::new(stream);
321
322                let mut buf = Vec::new();
323                use tokio::io::AsyncReadExt;
324                reader.read_to_end(&mut buf).await.expect("read to succeed");
325                assert_eq!(size, buf.len(), "expect size to match");
326                assert_eq!(&&blob_2_data[start_offset..start_offset + size], &buf.as_slice(), "expect data to be equal");
327            });
328        }
329    }
330}
snix_castore/blobstore/blob_meta/bytes_stream.rs

snix_castore/blobstore/blob_meta/
bytes_stream.rs
