snix_castore/blobstore/
blob_writer.rs

1use std::{pin::pin, task::Poll};
2
3use futures::TryStreamExt;
4use pin_project_lite::pin_project;
5use tokio::{
6    io::{SimplexStream, WriteHalf},
7    task::JoinHandle,
8};
9use tokio_util::io::InspectReader;
10
11use crate::{
12    B3Digest,
13    blobstore::BlobMeta,
14    chunkstore::{Chunk, ChunkStore},
15};
16
17pin_project! {
18/// Writes bytes to a [ChunkStore], returning the [BlobMeta] and [B3Digest] on close.
19pub struct BlobWriter{
20    // This is where we write received bytes into during a poll_write.
21    // The background task keeps the chunker and takes care of uploading.
22    writer: WriteHalf<SimplexStream>,
23
24    // Represents the completion of the upload task, which can either fail if we
25    // couldn't upload a chunk, or succeed (if it reaches EOF, which happens
26    // when we close our writer).
27    // In that case, we return the [BlobMeta] and [B3Digest] of the blob payload.
28    upload_task: JoinHandle<std::io::Result<(BlobMeta, B3Digest)>>,
29}
30}
31
32impl BlobWriter {
33    /// Constructs a new [BlobWriter] writing to the given [ChunkStore].
34    pub fn new(
35        chunk_store: impl ChunkStore + 'static,
36        avg_chunk_size: u32,
37        upload_concurrency: usize,
38    ) -> Self {
39        let (reader, writer) = tokio::io::simplex(1024 * 1024 * 10);
40
41        let upload_task = tokio::task::spawn(async move {
42            let mut blob_hasher = blake3::Hasher::new();
43            let reader = InspectReader::new(reader, |b| {
44                blob_hasher.update(b);
45            });
46
47            let mut chunker = fastcdc::v2020::AsyncStreamCDC::new(
48                reader,
49                avg_chunk_size / 2,
50                avg_chunk_size,
51                avg_chunk_size * 2,
52            );
53
54            let chunk_digests_sizes: Vec<(B3Digest, u64)> = chunker
55                .as_stream()
56                .map_ok(|chunk| {
57                    let chunk = Chunk::from(chunk.data);
58                    let chunk_len = chunk.len();
59                    // we only want to move a &ChunkStore.
60                    let chunk_store = &chunk_store;
61                    async move {
62                        let chunk_digest = chunk_store.put(chunk).await.map_err(|_| {
63                            std::io::Error::other("Failed to upload chunk".to_owned())
64                        })?;
65
66                        // on an upload error, this will cause rx to be dropped, which we will see when trying to enqueue work.
67                        Ok((chunk_digest, chunk_len as u64))
68                    }
69                })
70                .try_buffered(upload_concurrency)
71                .try_collect()
72                .await?;
73
74            Ok((
75                BlobMeta::from_digests_and_sizes(chunk_digests_sizes),
76                blob_hasher.finalize().into(),
77            ))
78        });
79
80        BlobWriter {
81            writer,
82            upload_task,
83        }
84    }
85
86    /// Get the resulting [BlobMeta] and [B3Digest] once all data has been written.
87    ///
88    /// This will consume self, drop the internal writer and wait for the
89    /// background task to finish uploading.
90    /// It is the responsibility of the caller to persist to the BlobStore.
91    pub async fn close(self) -> std::io::Result<(BlobMeta, B3Digest)> {
92        let Self {
93            mut writer,
94            upload_task,
95        } = self;
96
97        // shutdown the writer, to ensure the EOF gets propagated.
98        use tokio::io::AsyncWriteExt;
99        writer.shutdown().await?;
100
101        // drop the writer, wait for the results to come back.
102        drop(writer);
103
104        upload_task.await.expect("poisoned")
105    }
106}
107
108impl tokio::io::AsyncWrite for BlobWriter {
109    fn poll_write(
110        self: std::pin::Pin<&mut Self>,
111        cx: &mut std::task::Context<'_>,
112        buf: &[u8],
113    ) -> Poll<std::io::Result<usize>> {
114        // // Writes to the writer, which maintains its own internal buffer.
115        pin!(self.project().writer).poll_write(cx, buf)
116    }
117
118    fn poll_flush(
119        self: std::pin::Pin<&mut Self>,
120        cx: &mut std::task::Context<'_>,
121    ) -> Poll<std::io::Result<()>> {
122        pin!(self.project().writer).poll_flush(cx)
123    }
124
125    fn poll_shutdown(
126        self: std::pin::Pin<&mut Self>,
127        cx: &mut std::task::Context<'_>,
128    ) -> Poll<std::io::Result<()>> {
129        pin!(self.project().writer).poll_shutdown(cx)
130    }
131}
132
133#[cfg(test)]
134mod test {
135    use tokio::io::AsyncWriteExt;
136
137    use super::BlobWriter;
138    use crate::blobstore::BlobMeta;
139    use crate::{
140        B3Digest,
141        chunkstore::{self, ChunkStore},
142    };
143    use std::time::Duration;
144
145    #[tokio::test]
146    async fn test_write_scenario() {
147        let mut test_builder = tokio_test::io::Builder::new();
148        test_builder.read("test".as_bytes());
149        test_builder.wait(Duration::from_nanos(1));
150        test_builder.read("123".as_bytes());
151
152        let exp_digest: B3Digest = blake3::hash(b"test123").into();
153
154        let chunk_store = chunkstore::memory::MemoryChunkStore::default();
155        let mut blob_writer = BlobWriter::new(chunk_store.clone(), 4096, 10);
156
157        let written = tokio::io::copy(&mut test_builder.build(), &mut blob_writer)
158            .await
159            .expect("should write");
160        assert_eq!(7, written);
161
162        blob_writer.flush().await.expect("should flush");
163
164        let (blob_meta, blob_digest) = blob_writer.close().await.expect("should close");
165        assert_eq!(7, blob_meta.blob_len(), "BlobMeta blob len must match");
166        assert_eq!(exp_digest, blob_digest, "blob digest must match",);
167
168        let exp_blob_meta = BlobMeta::from_digests_and_sizes([(exp_digest, 7)]);
169        assert_eq!(exp_blob_meta, blob_meta, "expected blob meta to be correct");
170
171        assert_eq!(
172            &b"test123",
173            &chunk_store
174                .get(&exp_digest)
175                .await
176                .expect("ChunkStore to not fail")
177                .expect("to have the chunk")
178                .as_ref(),
179            "the returned data to be correct"
180        )
181    }
182
183    #[tokio::test]
184    async fn test_write_empty() {
185        let mut chunk_store = chunkstore::MockChunkStore::new();
186        // put shall never be called
187        chunk_store.expect_put().never();
188
189        let mut blob_writer = BlobWriter::new(chunk_store, 4096, 10);
190        let written = tokio::io::copy(&mut tokio::io::empty(), &mut blob_writer)
191            .await
192            .expect("copy to succeed");
193        assert_eq!(0, written, "expected 0 bytes to be written");
194
195        let (blob_meta, blob_digest) = blob_writer.close().await.expect("should close");
196
197        let exp_digest: B3Digest = blake3::hash(b"").into();
198        assert_eq!(exp_digest, blob_digest, "blob digest must match",);
199
200        let exp_blob_meta = BlobMeta::from_digests_and_sizes([]);
201        assert_eq!(exp_blob_meta, blob_meta, "expected blob meta to be correct");
202    }
203}
snix_castore/blobstore/blob_writer.rs

snix_castore/blobstore/
blob_writer.rs
