snix_castore/blobstore/
blob_reader.rs

1use std::{
2    io::{self, Cursor, SeekFrom},
3    task::Poll,
4};
5
6use bytes::Buf;
7use futures::{StreamExt, TryStreamExt, ready, stream::BoxStream};
8use pin_project_lite::pin_project;
9use tokio::io::AsyncBufRead;
10use tracing::{instrument, warn};
11
12use crate::{
13    blobstore::BlobMeta,
14    chunkstore::{Chunk, ChunkStore, EMPTY_CHUNK},
15};
16
17pin_project! {
18/// A reader for one Blob.
19///
20/// Supports single-chunked blobs as well as blobs composed of multiple chunks.
21///
22/// In the latter case, it takes an existing [BlobMeta] with a specific chunking
23/// provided upfront.
24///
25/// Any fallback to other chunkings, needs to be happen elsewhere.
26// pin_project_lite doesnt't seem to support docstrings fully
27#[allow(missing_docs)]
28#[project = BlobReaderProj]
29    pub enum BlobReader<CS> {
30        /// Blob is composed of a single chunk.
31        /// This allows skipping asking for a [BlobMeta].
32        /// For very small blobs the callsite might ask a [ChunkStore] directly.
33        SingleChunk {
34            // A cursor over the single chunk.
35            #[pin] cursor: Cursor<Chunk>,
36        },
37        /// The blob is composed of multiple chunks,
38        /// and we need to query a [ChunkStore] to read data.
39        ChunkedBlob {
40            // Chunking information for this blob
41            blob_meta: BlobMeta,
42
43            // A [ChunkStore] to read chunks from
44            chunk_store: CS,
45
46            // The configured fetch concurrency
47            fetch_concurrency: usize,
48
49            // The position of the reader in the entire blob
50            pos: u64,
51
52            // The current chunk.
53            current_chunk: Cursor<Chunk>,
54
55            // A stream providing the remaining bytes from pos + current_chunk.remaining() till the end.
56            #[pin] stream: BoxStream<'static, std::io::Result<Cursor<Chunk>>>,
57        },
58    }
59}
60
61impl<CS> BlobReader<CS>
62where
63    CS: ChunkStore + Clone + 'static,
64{
65    /// Initialize a [BlobReader] with the [Chunk] of a single-chunked blob
66    pub fn from_single_chunk(chunk: Chunk) -> Self {
67        Self::SingleChunk {
68            cursor: Cursor::new(chunk),
69        }
70    }
71
72    /// Initialize a [BlobReader] with a fixed [BlobMeta],
73    /// a reference to a [ChunkStore] and a fetch concurrency.
74    pub fn from_blob_meta(
75        blob_meta: BlobMeta,
76        chunk_store: CS,
77        fetch_concurrency: usize,
78    ) -> BlobReader<CS> {
79        let stream = blob_meta
80            .bytes_stream(0, None, fetch_concurrency, chunk_store.clone())
81            .boxed();
82
83        Self::ChunkedBlob {
84            blob_meta,
85            chunk_store,
86            fetch_concurrency,
87            pos: 0,
88            current_chunk: Cursor::new(EMPTY_CHUNK),
89            stream,
90        }
91    }
92
93    /// The total length of the blob.
94    /// Length of the single chunk or, for chunked blobs, the length as
95    /// specified in the [BlobMeta].
96    fn blob_len(&self) -> u64 {
97        match self {
98            BlobReader::SingleChunk { cursor } => cursor.get_ref().as_ref().len() as u64,
99            BlobReader::ChunkedBlob { blob_meta, .. } => blob_meta.blob_len(),
100        }
101    }
102
103    /// Get the current position, from the start of the blob.
104    fn position(&self) -> u64 {
105        match self {
106            BlobReader::SingleChunk { cursor } => cursor.position(),
107            BlobReader::ChunkedBlob { pos, .. } => *pos,
108        }
109    }
110}
111
112impl<CS> tokio::io::AsyncRead for BlobReader<CS>
113where
114    CS: ChunkStore,
115{
116    fn poll_read(
117        self: std::pin::Pin<&mut Self>,
118        cx: &mut std::task::Context<'_>,
119        buf: &mut tokio::io::ReadBuf<'_>,
120    ) -> Poll<io::Result<()>> {
121        match self.project() {
122            BlobReaderProj::SingleChunk { cursor } => cursor.poll_read(cx, buf),
123            BlobReaderProj::ChunkedBlob {
124                pos,
125                current_chunk,
126                mut stream,
127                ..
128            } => {
129                // If the buffer we write to is full, we can't do anything
130                if buf.remaining() == 0 {
131                    return Poll::Ready(Ok(()));
132                }
133
134                // If we have no more bytes to read in current_chunk, poll for the next one
135                if current_chunk.remaining() == 0 {
136                    if let Some(res) = ready!(stream.poll_next_unpin(cx)) {
137                        *current_chunk = res?;
138                        debug_assert!(
139                            current_chunk.remaining() > 0,
140                            "fetched chunk should have some data"
141                        );
142                    } else {
143                        // end of stream, EOF
144                        return Poll::Ready(Ok(()));
145                    }
146                }
147
148                let b = &current_chunk.get_ref().as_ref()[current_chunk.position() as usize..];
149                let to_fill = std::cmp::min(buf.remaining(), b.len());
150                let dst = buf.initialize_unfilled_to(to_fill);
151                dst.copy_from_slice(&b[..to_fill]);
152
153                buf.advance(to_fill);
154                current_chunk.advance(to_fill);
155
156                *pos = pos.checked_add(to_fill as u64).ok_or(std::io::Error::new(
157                    std::io::ErrorKind::OutOfMemory,
158                    "position > u64::MAX bytes",
159                ))?;
160
161                Poll::Ready(Ok(()))
162            }
163        }
164    }
165}
166
167impl<CS> AsyncBufRead for BlobReader<CS>
168where
169    CS: ChunkStore,
170{
171    fn poll_fill_buf(
172        self: std::pin::Pin<&mut Self>,
173        cx: &mut std::task::Context<'_>,
174    ) -> Poll<io::Result<&[u8]>> {
175        match self.project() {
176            BlobReaderProj::SingleChunk { cursor } => cursor.poll_fill_buf(cx),
177            BlobReaderProj::ChunkedBlob {
178                current_chunk,
179                mut stream,
180                ..
181            } => {
182                if current_chunk.remaining() == 0 {
183                    if let Some(res) = ready!(stream.poll_next_unpin(cx)) {
184                        *current_chunk = res?;
185                        debug_assert!(
186                            current_chunk.remaining() > 0,
187                            "fetched chunk should have some data"
188                        );
189                    } else {
190                        // end of stream, EOF
191                        return Poll::Ready(Ok(&[]));
192                    }
193                }
194                let p = current_chunk.position();
195
196                Poll::Ready(Ok(&current_chunk.get_ref().as_ref()[(p as usize)..]))
197            }
198        }
199    }
200
201    fn consume(self: std::pin::Pin<&mut Self>, amt: usize) {
202        match self.project() {
203            BlobReaderProj::SingleChunk { cursor } => cursor.consume(amt),
204            BlobReaderProj::ChunkedBlob {
205                current_chunk, pos, ..
206            } => {
207                current_chunk.advance(amt);
208                *pos = pos
209                    .checked_add(amt as u64)
210                    .expect("consume would increase pos > u64::MAX bytes");
211            }
212        }
213    }
214}
215impl<CS> tokio::io::AsyncSeek for BlobReader<CS>
216where
217    CS: ChunkStore + Clone + 'static,
218{
219    fn start_seek(self: std::pin::Pin<&mut Self>, position: SeekFrom) -> io::Result<()> {
220        // calculate the new position
221        let new_pos = calc_position(self.position(), self.blob_len(), position)?;
222        if new_pos == self.position() {
223            return Ok(());
224        }
225
226        match self.project() {
227            BlobReaderProj::SingleChunk { cursor } => cursor.start_seek(position),
228            BlobReaderProj::ChunkedBlob {
229                blob_meta,
230                chunk_store,
231                fetch_concurrency,
232                pos,
233                current_chunk,
234                mut stream,
235            } => {
236                // if the new position is still covered by our current buffer, we can simply update our position in there.
237                // else, empty current_chunk, and construct a new stream from the new position.
238                if !seek_in_current_chunk(*pos, new_pos, current_chunk)? {
239                    *current_chunk = Cursor::new(EMPTY_CHUNK);
240                    stream.set(
241                        blob_meta
242                            .bytes_stream(new_pos, None, *fetch_concurrency, chunk_store.clone())
243                            .boxed(),
244                    );
245                }
246
247                *pos = new_pos;
248
249                Ok(())
250            }
251        }
252    }
253
254    fn poll_complete(
255        self: std::pin::Pin<&mut Self>,
256        _cx: &mut std::task::Context<'_>,
257    ) -> Poll<io::Result<u64>> {
258        Poll::Ready(Ok(self.position()))
259    }
260}
261
262/// Limited reads from a blob (starting from offset, up to size)
263pub trait LimitBlobReader {
264    /// Copies up to `size` bytes starting from `offset` into the passed writer.
265    ///
266    /// The (sync) writer is assumed to be sufficiently non-blocking,
267    /// so it can still be invoked from an async context.
268    ///
269    /// Returns the number of bytes actually written.
270    fn copy_limited(
271        &self,
272        offset: u64,
273        size: u64,
274        w: &mut dyn std::io::Write,
275    ) -> impl std::future::Future<Output = io::Result<u64>>;
276}
277
278impl<CS> LimitBlobReader for BlobReader<CS>
279where
280    CS: ChunkStore + Clone + 'static,
281{
282    #[instrument(skip_all, fields(blob.len = self.blob_len(), %offset, %size))]
283    fn copy_limited(
284        &self,
285        offset: u64,
286        size: u64,
287        w: &mut dyn std::io::Write,
288    ) -> impl std::future::Future<Output = io::Result<u64>> {
289        async move {
290            if offset >= self.blob_len() {
291                warn!("offset past blob_len");
292                return Ok(0);
293            }
294            match self {
295                BlobReader::SingleChunk { cursor } => {
296                    // We don't want to mutate the AsyncRead[Buf] state, so create our own cursor.
297                    let mut rd = Cursor::new(cursor.get_ref().to_owned());
298                    rd.set_position(offset);
299                    let mut rd = std::io::Read::take(rd, size);
300                    std::io::copy(&mut rd, w)
301                }
302                BlobReader::ChunkedBlob {
303                    blob_meta,
304                    chunk_store,
305                    fetch_concurrency,
306                    ..
307                } => {
308                    // construct a stream for the requested offset and size
309                    let mut stream = blob_meta.bytes_stream(
310                        offset,
311                        Some(size),
312                        *fetch_concurrency,
313                        chunk_store.clone(),
314                    );
315
316                    let mut bytes_written = 0;
317
318                    while let Some(mut cursor) = stream.try_next().await? {
319                        bytes_written += std::io::copy(&mut cursor, w)?;
320                    }
321
322                    Ok(bytes_written)
323                }
324            }
325        }
326    }
327}
328
329impl<CS> super::Reader for BlobReader<CS> where CS: ChunkStore + Clone + 'static {}
330
331/// For a given blob length and current position, returns the position that seek_from would seek to.
332fn calc_position(cur_pos: u64, blob_len: u64, seek_from: SeekFrom) -> std::io::Result<u64> {
333    let new_pos = match seek_from {
334        SeekFrom::Start(p) => p,
335        SeekFrom::End(p) => blob_len.checked_sub_signed(p).ok_or(std::io::Error::new(
336            std::io::ErrorKind::InvalidInput,
337            "tried to seek before beginning of blob",
338        ))?,
339        SeekFrom::Current(p) => cur_pos.checked_add_signed(p).ok_or(std::io::Error::new(
340            std::io::ErrorKind::UnexpectedEof,
341            "tried to seek way past end of blob",
342        ))?,
343    };
344
345    if new_pos > blob_len {
346        Err(std::io::Error::new(
347            std::io::ErrorKind::UnexpectedEof,
348            "tried to seek past end of blob",
349        ))
350    } else {
351        Ok(new_pos)
352    }
353}
354
355/// Seeks the cursor to the new blob position if contained in the chunk.
356///
357/// Returns true if the new position is contained in the current chunk.
358/// It is permissible to seek to the end of the chunk, (with no more bytes remaining)
359fn seek_in_current_chunk(
360    cur_blob_pos: u64,
361    new_blob_pos: u64,
362    current_chunk: &mut Cursor<impl AsRef<[u8]>>,
363) -> io::Result<bool> {
364    // Determine current_chunks blob offsets
365    let chunk_len = current_chunk.get_ref().as_ref().len();
366    let chunk_start_offset = cur_blob_pos - current_chunk.position();
367    let chunk_end_pos =
368        chunk_start_offset
369            .checked_add(chunk_len as u64)
370            .ok_or(std::io::Error::new(
371                std::io::ErrorKind::UnexpectedEof,
372                "chunk_end > u64::MAX",
373            ))?;
374
375    if (chunk_start_offset..=chunk_end_pos).contains(&new_blob_pos) {
376        let new_chunk_pos = new_blob_pos - chunk_start_offset;
377        debug_assert!(
378            new_chunk_pos <= chunk_len as u64,
379            "Snix bug: tried seeking past chunk boundary"
380        );
381        current_chunk.set_position(new_chunk_pos);
382        Ok(true)
383    } else {
384        Ok(false)
385    }
386}
387
388#[cfg(test)]
389mod test {
390    use mockall::predicate;
391    use std::{
392        io::{Cursor, SeekFrom},
393        sync::{Arc, LazyLock},
394    };
395    use tokio::io::{AsyncReadExt, AsyncSeekExt};
396
397    use crate::{
398        blobstore::BlobMeta,
399        chunkstore::{Chunk, ChunkStore, MockChunkStore, memory::MemoryChunkStore},
400    };
401
402    use super::{BlobReader, seek_in_current_chunk};
403    use proptest::prelude::*;
404    proptest! {
405        #[test]
406        fn seek_correctly_seeks(
407            // chunk data
408            chunk_data in prop::collection::vec(any::<u8>(), 0..256),
409            // position inside the chunk
410            chunk_pos in 0..=256usize,
411            // the chunk must not necessarily be at the start
412            chunk_start_offset in 0..256u64,
413            // where we want to seek to
414            seek_relative in -128..128i64,
415        ) {
416            // clamp chunk_pos to not fall outside chunk_data.
417            let chunk_pos = chunk_data.len().min(chunk_pos) as u64;
418
419            let cur_blob_pos = chunk_start_offset.checked_add(chunk_pos).expect("to not overflow");
420
421            let chunk_end_pos = chunk_start_offset.checked_add(chunk_data.len() as u64).expect("to not overflow");
422            let mut current_chunk = {
423                let mut c = Cursor::new(chunk_data);
424                c.set_position(chunk_pos);
425                c
426            };
427
428            // bail out early if we're trying to seek to a negative position
429            let new_blob_pos = match cur_blob_pos.checked_add_signed(seek_relative) {
430                None => return Ok(()),
431                Some(v) => v,
432            };
433            let success = seek_in_current_chunk(cur_blob_pos, new_blob_pos, &mut current_chunk).expect("to not error");
434
435            let target_is_inside_chunk = new_blob_pos >= chunk_start_offset && new_blob_pos <= chunk_end_pos;
436            if target_is_inside_chunk {
437                assert!(success, "seek should have been successful");
438                let new_chunk_pos = chunk_pos.checked_add_signed(seek_relative).expect("to not over/underflow");
439                assert_eq!(new_chunk_pos, current_chunk.position(), "seek should be to the new position");
440            } else {
441                assert!(!success, "seek should be unsuccessful, from {cur_blob_pos} to {new_blob_pos}, {chunk_start_offset}..={chunk_end_pos}");
442                assert_eq!(chunk_pos, current_chunk.position(), "no seek should have happened")
443            }
444        }
445    }
446
447    #[tokio::test]
448    async fn single_chunk() {
449        let mut rd = BlobReader::<MemoryChunkStore>::from_single_chunk(b"abc"[..].into());
450
451        {
452            let mut buf = Vec::new();
453
454            let bytes_read = tokio::io::copy(&mut rd, &mut buf)
455                .await
456                .expect("to succeed");
457            assert_eq!(3, bytes_read, "expect bytes_read to match");
458            assert_eq!(b"abc"[..], buf[..], "expect data to match");
459        }
460
461        // seek to pos 1
462        rd.seek(SeekFrom::Current(-2))
463            .await
464            .expect("seek to succeed");
465
466        // should read "bc". we try the bufread part
467        {
468            let mut buf = Vec::new();
469            let bytes_read = tokio::io::copy_buf(&mut rd, &mut buf)
470                .await
471                .expect("to succeed");
472            assert_eq!(2, bytes_read, "expect bytes_read to match");
473            assert_eq!(b"bc"[..], buf[..], "expect data to match");
474        }
475    }
476
477    const CHUNK_1: Chunk = Chunk::from_static(b"ab");
478    const CHUNK_2: Chunk = Chunk::from_static(b"c");
479    static BLOB_1_META: LazyLock<BlobMeta> = LazyLock::new(|| {
480        BlobMeta::from_digests_and_sizes([
481            (CHUNK_1.digest(), CHUNK_1.len() as u64),
482            (CHUNK_2.digest(), CHUNK_2.len() as u64),
483        ])
484    });
485
486    #[tokio::test]
487    async fn multiple_chunks() {
488        let chunk_store = MemoryChunkStore::default();
489        chunk_store
490            .put(CHUNK_1.to_owned())
491            .await
492            .expect("to succeed");
493        chunk_store
494            .put(CHUNK_2.to_owned())
495            .await
496            .expect("to succeed");
497
498        let mut rd = BlobReader::from_blob_meta(BLOB_1_META.to_owned(), chunk_store, 10);
499        {
500            let mut buf = Vec::new();
501            tokio::io::copy(&mut rd, &mut buf)
502                .await
503                .expect("to succeed");
504            assert_eq!(b"abc"[..].to_vec(), buf.as_slice(), "data to match");
505        }
506
507        let pos = rd.seek(SeekFrom::Start(1)).await.expect("seek to succeed");
508        assert_eq!(1, pos, "position to be correct");
509
510        {
511            let mut buf = Vec::new();
512            tokio::io::copy_buf(&mut rd, &mut buf)
513                .await
514                .expect("to succeed");
515            assert_eq!(b"bc"[..].to_vec(), buf.as_slice(), "data to match");
516        }
517
518        // use AsyncBufRead a bit, then seek relatively. This ensures consume() updates the internal position tracking.
519        {
520            // seek to the front
521            let pos = rd.seek(SeekFrom::Start(0)).await.expect("seek to succeed");
522            assert_eq!(0, pos, "position to be correct");
523
524            // use tokio::io::copy_buf to read to the end
525            {
526                let mut buf = Vec::new();
527                tokio::io::copy_buf(&mut rd, &mut buf)
528                    .await
529                    .expect("to succeed");
530                assert_eq!(b"abc"[..].to_vec(), buf.as_slice(), "data to match");
531            }
532
533            // seek relatively
534            let new_pos = rd
535                .seek(SeekFrom::Current(-1))
536                .await
537                .expect("seek to succeed");
538            assert_eq!(2, new_pos, "expect position to be correct");
539
540            // read to the end again
541            {
542                let mut buf = Vec::new();
543                tokio::io::copy_buf(&mut rd, &mut buf)
544                    .await
545                    .expect("to succeed");
546                assert_eq!(b"c"[..].to_vec(), buf.as_slice(), "data to match");
547            }
548        }
549
550        // seek to the front
551        let pos = rd.seek(SeekFrom::Start(0)).await.expect("seek to succeed");
552        assert_eq!(0, pos, "position to be correct");
553    }
554
555    #[tokio::test]
556    async fn test_seek_same_chunk() {
557        // This chunk store will only ever respond with CHUNK_1 once.
558        let mut chunk_store = MockChunkStore::new();
559        chunk_store
560            .expect_get()
561            .with(predicate::eq(CHUNK_1.digest()))
562            .return_once(|_| Ok(Some(CHUNK_1.to_owned())));
563
564        let chunk_store = Arc::new(chunk_store);
565
566        // construct BlobReader
567        let mut rd = BlobReader::from_blob_meta(
568            BLOB_1_META.to_owned(),
569            chunk_store,
570            // we explicitly set the concurrency to 1, so BLOB2 will only get fetched if would poll the stream a second time
571            // (which we don't).
572            1,
573        );
574
575        let first = rd.read_u8().await.expect("to read first byte");
576        assert_eq!(b'a', first, "expect first byte to match");
577
578        // seek backwards to start, read again
579        rd.seek(SeekFrom::Start(0)).await.expect("seek to succeed");
580        let first = rd.read_u8().await.expect("to read first byte");
581        assert_eq!(b'a', first, "expect first byte to match");
582
583        // seek to the end of this chunk
584        rd.seek(SeekFrom::Start(2)).await.expect("seek to succeed");
585        // we now don't read, so we won't poll the stream and cause the chunkservice to panic
586        // seek back to the middle of the first chunk
587        rd.seek(SeekFrom::Start(1)).await.expect("seek to succeed");
588        // and read a bit more
589        let second = rd.read_u8().await.expect("to read second byte");
590        assert_eq!(b'b', second, "expect second byte to match");
591    }
592}
snix_castore/blobstore/blob_reader.rs

snix_castore/blobstore/
blob_reader.rs
