snix_castore/chunkstore/
object_store.rs1use 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
17pub struct ObjectStoreChunkStore {
22 instance_name: String,
23 object_store: Arc<dyn object_store::ObjectStore>,
24 base_path: Path,
25}
26
27impl ObjectStoreChunkStore {
28 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 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#[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#[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}