Skip to main content

snix_castore/blobstore/
grpc.rs

1//! A gRPC client implementation.
2
3use std::sync::Arc;
4
5use tonic::{Code, async_trait};
6use tracing::instrument;
7
8use crate::{
9    B3Digest,
10    blobstore::{BlobMeta, BlobStore},
11    composition::{CompositionContext, ServiceBuilder},
12    proto::{self, GetBlobRequest, GetBlobResponse, PutBlobRequest, StatBlobRequest},
13};
14
15/// Connects to a (remote) snix-castore BlobStoreService over gRPC.
16#[derive(Clone)]
17pub struct GrpcBlobStore<T> {
18    instance_name: String,
19    /// The internal reference to a gRPC client.
20    /// Cloning it is cheap, and it internally handles concurrent requests.
21    client: proto::blob_store_service_client::BlobStoreServiceClient<T>,
22}
23
24impl<T> GrpcBlobStore<T> {
25    /// construct a [GrpcBlobStore] from a [proto::blob_store_service_client::BlobStoreServiceClient].
26    pub fn from_client(
27        instance_name: String,
28        client: proto::blob_store_service_client::BlobStoreServiceClient<T>,
29    ) -> Self {
30        Self {
31            instance_name,
32            client,
33        }
34    }
35}
36#[async_trait]
37impl<T> BlobStore for GrpcBlobStore<T>
38where
39    T: tonic::client::GrpcService<tonic::body::Body> + Send + Sync + Clone + 'static,
40    T::ResponseBody: tonic::codegen::Body<Data = tonic::codegen::Bytes> + Send + 'static,
41    <T::ResponseBody as tonic::codegen::Body>::Error: Into<tonic::codegen::StdError> + Send,
42    T::Future: Send,
43{
44    #[instrument(skip_all, err, fields(blob.digest=%digest, instance_name=%self.instance_name))]
45    async fn get(&self, digest: &B3Digest) -> Result<Option<BlobMeta>, super::Error> {
46        // We cannot easily convert from [proto::BlobMeta] to [BlobMeta] for the empty blob,
47        // and it's pointless to look it up anyways.
48        if digest == &B3Digest::EMPTY {
49            return Ok(Some(BlobMeta::from_digests_and_sizes([])));
50        }
51
52        let resp = match self
53            .client
54            .clone()
55            .get_blob(GetBlobRequest::with_digest(*digest))
56            .await
57        {
58            Ok(resp) => resp,
59            Err(err) if err.code() == Code::NotFound => return Ok(None),
60            Err(err) => Err(Error::Tonic(err))?,
61        };
62
63        let GetBlobResponse { blob_meta } = resp.into_inner();
64        let proto_blob_meta = blob_meta.ok_or(Error::BlobMetaMissing)?;
65        let blob_meta = BlobMeta::try_from(proto_blob_meta).map_err(Error::DecodeBlobMeta)?;
66
67        Ok(Some(blob_meta))
68    }
69
70    #[instrument(skip_all, err, fields(blob.digest=%digest, instance_name=%self.instance_name))]
71    async fn has(&self, digest: &B3Digest) -> Result<bool, super::Error> {
72        // We reply for the empty blob in all cases.
73        if digest == &B3Digest::EMPTY {
74            return Ok(true);
75        }
76
77        match self
78            .client
79            .clone()
80            .stat_blob(StatBlobRequest::with_digest(*digest))
81            .await
82        {
83            Ok(_) => Ok(true),
84            Err(err) if err.code() == Code::NotFound => return Ok(false),
85            Err(err) => Err(Error::Tonic(err))?,
86        }
87    }
88
89    #[instrument(skip_all, err, fields(blob.digest=%digest, instance_name=%self.instance_name))]
90    async fn put(&self, digest: &B3Digest, blob_meta: BlobMeta) -> Result<(), super::Error> {
91        match self
92            .client
93            .clone()
94            .put_blob(PutBlobRequest {
95                digest: (*digest).into(),
96                blob_meta: Some(blob_meta.into()),
97            })
98            .await
99        {
100            Ok(_) => Ok(()),
101            Err(err) => Err(Error::Tonic(err))?,
102        }
103    }
104}
105
106/// Configuration for a [GrpcBlobStore].
107#[derive(serde::Deserialize, Debug)]
108#[serde(deny_unknown_fields)]
109pub struct GRPCBlobStoreConfig {
110    url: url::Url,
111}
112
113impl TryFrom<url::Url> for GRPCBlobStoreConfig {
114    type Error = Box<dyn std::error::Error + Send + Sync>;
115    fn try_from(url: url::Url) -> Result<Self, Self::Error> {
116        Ok(GRPCBlobStoreConfig { url })
117    }
118}
119
120#[async_trait]
121impl ServiceBuilder for GRPCBlobStoreConfig {
122    type Output = dyn BlobStore;
123    async fn build<'a>(
124        &'a self,
125        instance_name: &str,
126        _context: &CompositionContext,
127    ) -> Result<Arc<Self::Output>, Box<dyn std::error::Error + Send + Sync>> {
128        let client = proto::blob_store_service_client::BlobStoreServiceClient::with_interceptor(
129            crate::tonic::channel_from_url(&self.url).await?,
130            snix_tracing::propagate::tonic::send_trace,
131        );
132        Ok(Arc::new(GrpcBlobStore::from_client(
133            instance_name.to_string(),
134            client,
135        )))
136    }
137}
138
139#[derive(thiserror::Error, Debug)]
140enum Error {
141    #[error(transparent)]
142    Tonic(tonic::Status),
143    #[error("no blob_meta in GetBlobResponse")]
144    BlobMetaMissing,
145    #[error("failed to decode BlobMeta")]
146    DecodeBlobMeta(#[from] crate::blobstore::blob_meta::DecodeError),
147}
148
149impl From<Error> for super::Error {
150    fn from(value: Error) -> Self {
151        super::Error(Box::new(value))
152    }
153}