Skip to main content

snix_castore/chunkstore/
object_store.rs

1//! Object-store backend, also supporting the local filesystem.
2
3use super::Chunk;
4use crate::{
5    B3Digest,
6    chunkstore::ChunkStore,
7    composition::{CompositionContext, ServiceBuilder},
8};
9use data_encoding::HEXLOWER;
10use object_store::{ObjectStoreExt, PutPayload, path::Path};
11use std::collections::HashMap;
12use std::{io::Cursor, sync::Arc};
13use tokio::io;
14use tonic::async_trait;
15use tracing::{Level, instrument, warn};
16
17/// Stores Chunks in any backend supported by the [object_store] crate.
18///
19/// Each [Chunk] is stored as a zstd-compressed object
20/// at `{base_path}/chunks/b3/{HEXLOWER(digest[0..2])}/{HEXLOWER(digest)}`.
21pub struct ObjectStoreChunkStore {
22    instance_name: String,
23    object_store: Arc<dyn object_store::ObjectStore>,
24    base_path: Path,
25}
26
27impl ObjectStoreChunkStore {
28    /// Constructs a new [ObjectStoreChunkStore] from already existing
29    /// `instance_name`, `object_store` and `base_name`.
30    pub fn from_parts(
31        instance_name: String,
32        object_store: Arc<dyn object_store::ObjectStore>,
33        base_path: Path,
34    ) -> Self {
35        Self {
36            instance_name,
37            object_store,
38            base_path,
39        }
40    }
41}
42
43#[instrument(level=Level::TRACE, skip_all,fields(base_path=%base_path,chunk.digest=%digest))]
44fn derive_chunk_path(base_path: &Path, digest: &B3Digest) -> Path {
45    base_path
46        .clone()
47        .join("chunks")
48        .join("b3")
49        .join(HEXLOWER.encode(&digest[..2]))
50        .join(HEXLOWER.encode(&digest[..]))
51}
52
53#[async_trait]
54impl ChunkStore for ObjectStoreChunkStore {
55    #[instrument(skip_all, err, fields(chunk.digest=%digest, instance_name=%self.instance_name))]
56    async fn has(&self, digest: &B3Digest) -> Result<bool, super::Error> {
57        let p = derive_chunk_path(&self.base_path, digest);
58        match self.object_store.head(&p).await {
59            Ok(_) => Ok(true),
60            Err(object_store::Error::NotFound { .. }) => Ok(false),
61            Err(err) => Err(Error::from(err))?,
62        }
63    }
64
65    #[instrument(skip_all, err, fields(chunk.digest=%digest, instance_name=%self.instance_name))]
66    async fn get(&self, digest: &B3Digest) -> Result<Option<Chunk>, super::Error> {
67        let p = derive_chunk_path(&self.base_path, digest);
68        match self.object_store.get(&p).await {
69            Ok(get_result) => {
70                let compressed_bytes = get_result.bytes().await.map_err(|err| {
71                    warn!(%err,"failed to read payload from object_store");
72                    Error::ObjectStore(err)
73                })?;
74                // FUTUREWORK: use zstd::bulk to prevent decompression bombs
75                let chunk = Chunk::from(
76                    zstd::stream::decode_all(Cursor::new(compressed_bytes))
77                        .map_err(Error::ChunkDecompress)?,
78                );
79
80                let actual_digest = chunk.digest();
81                if &actual_digest != digest {
82                    Err(Error::IncorrectDigest {
83                        expected: digest.to_owned(),
84                        actual: actual_digest,
85                    })?
86                }
87
88                Ok(Some(chunk))
89            }
90            Err(object_store::Error::NotFound { .. }) => Ok(None),
91            Err(err) => Err(super::Error(Box::new(err)))?,
92        }
93    }
94
95    #[instrument(skip_all, err, fields(instance_name=%self.instance_name))]
96    async fn put(&self, chunk: Chunk) -> Result<B3Digest, super::Error> {
97        let digest: B3Digest = chunk.digest();
98        let p = derive_chunk_path(&self.base_path, &digest);
99        let compressed_chunk =
100            zstd::stream::encode_all(Cursor::new(chunk.as_ref()), zstd::DEFAULT_COMPRESSION_LEVEL)
101                .map_err(Error::ChunkCompress)?;
102
103        self.object_store
104            .put(&p, PutPayload::from_bytes(compressed_chunk.into()))
105            .await
106            .map_err(Error::ObjectStore)?;
107
108        Ok(digest)
109    }
110}
111
112/// Configuration for an [ObjectStoreChunkStore].
113#[derive(serde::Deserialize)]
114#[serde(deny_unknown_fields)]
115pub struct ObjectStoreChunkStoreConfig {
116    object_store_url: url::Url,
117    #[serde(default)]
118    object_store_options: HashMap<String, String>,
119}
120
121impl TryFrom<url::Url> for ObjectStoreChunkStoreConfig {
122    type Error = Box<dyn std::error::Error + Send + Sync>;
123    fn try_from(url: url::Url) -> Result<Self, Self::Error> {
124        Ok(ObjectStoreChunkStoreConfig {
125            object_store_url: crate::object_store::trim_objectstore_prefix(&url)?,
126            object_store_options: url
127                .query_pairs()
128                .into_iter()
129                .map(|(k, v)| (k.to_string(), v.to_string()))
130                .collect(),
131        })
132    }
133}
134
135#[async_trait]
136impl ServiceBuilder for ObjectStoreChunkStoreConfig {
137    type Output = dyn ChunkStore;
138    async fn build<'a>(
139        &'a self,
140        instance_name: &str,
141        _context: &CompositionContext,
142    ) -> Result<Arc<Self::Output>, Box<dyn std::error::Error + Send + Sync>> {
143        let (object_store, path) = crate::object_store::setup_object_store(
144            &self.object_store_url,
145            &self.object_store_options,
146        )
147        .await?;
148
149        Ok(Arc::new(ObjectStoreChunkStore::from_parts(
150            instance_name.to_string(),
151            Arc::new(object_store),
152            path,
153        )))
154    }
155}
156
157/// Error returned from [ObjectStoreChunkStore].
158#[derive(thiserror::Error, Debug)]
159enum Error {
160    #[error("reading bytes from object_store")]
161    ObjectStore(#[from] object_store::Error),
162    #[error("decompressing zstd chunk")]
163    ChunkDecompress(io::Error),
164    #[error("compressing zstd chunk")]
165    ChunkCompress(io::Error),
166    #[error("unexpected digest, expected: {expected}, got: {actual}")]
167    IncorrectDigest {
168        expected: B3Digest,
169        actual: B3Digest,
170    },
171}
172
173impl From<Error> for super::Error {
174    fn from(value: Error) -> Self {
175        super::Error(Box::new(value))
176    }
177}