snix_castore/blob_engine/chunking/
writer.rs

1use std::pin::pin;
2
3use pin_project_lite::pin_project;
4use tokio::io::AsyncWrite;
5use tonic::async_trait;
6use tracing::Span;
7
8use crate::{B3Digest, blobstore::BlobStore};
9
10pin_project! {
11/// Wrapper to implementing [crate::blobstore::Writer] for
12/// [crate::blobstore::BlobWriter] (persisting [BlobMeta] in a [BlobStore])
13pub(super) struct BlobWriterWrapper<BS>{
14    inner: Option<(crate::blobstore::BlobWriter, BS)>
15}
16}
17
18impl<BS> BlobWriterWrapper<BS> {
19    pub fn new(blob_writer: crate::blobstore::BlobWriter, blob_service: BS) -> Self {
20        Self {
21            inner: Some((blob_writer, blob_service)),
22        }
23    }
24}
25
26impl<BS> AsyncWrite for BlobWriterWrapper<BS>
27where
28    BS: BlobStore,
29{
30    fn poll_write(
31        self: std::pin::Pin<&mut Self>,
32        cx: &mut std::task::Context<'_>,
33        buf: &[u8],
34    ) -> std::task::Poll<std::io::Result<usize>> {
35        let (writer, _) = self.project().inner.as_mut().ok_or(std::io::Error::new(
36            std::io::ErrorKind::BrokenPipe,
37            "writer closed",
38        ))?;
39
40        pin!(writer).poll_write(cx, buf)
41    }
42
43    fn poll_flush(
44        self: std::pin::Pin<&mut Self>,
45        cx: &mut std::task::Context<'_>,
46    ) -> std::task::Poll<std::io::Result<()>> {
47        let (writer, _) = self.project().inner.as_mut().ok_or(std::io::Error::new(
48            std::io::ErrorKind::BrokenPipe,
49            "writer closed",
50        ))?;
51
52        pin!(writer).poll_flush(cx)
53    }
54
55    fn poll_shutdown(
56        self: std::pin::Pin<&mut Self>,
57        cx: &mut std::task::Context<'_>,
58    ) -> std::task::Poll<std::io::Result<()>> {
59        let (writer, _) = self.project().inner.as_mut().ok_or(std::io::Error::new(
60            std::io::ErrorKind::BrokenPipe,
61            "writer closed",
62        ))?;
63
64        pin!(writer).poll_shutdown(cx)
65    }
66}
67
68#[async_trait]
69impl<BS: BlobStore> crate::blobstore::Writer for BlobWriterWrapper<BS> {
70    #[tracing::instrument(skip_all, err, fields(blob.digest=tracing::field::Empty))]
71    async fn close(&mut self) -> std::io::Result<B3Digest> {
72        let (writer, blob_store) = self.inner.take().ok_or(std::io::Error::new(
73            std::io::ErrorKind::BrokenPipe,
74            "writer closed",
75        ))?;
76
77        let (blob_meta, blob_digest) = writer.close().await?;
78
79        Span::current().record("blob.digest", blob_digest.to_string());
80
81        if !blob_store
82            .has(&blob_digest)
83            .await
84            .map_err(std::io::Error::other)?
85        {
86            blob_store
87                .put(&blob_digest, blob_meta)
88                .await
89                .map_err(std::io::Error::other)?;
90        } else {
91            tracing::debug!("blob already exists, skipping upload")
92        }
93
94        return Ok(blob_digest);
95    }
96}
snix_castore/blob_engine/chunking/writer.rs

snix_castore/blob_engine/chunking/
writer.rs
