feat(tvix/store): add chunkservice r/5893

This adds the simpler ChunkService trait, and an implementation for it
using sled, and one using a HashMap.

Change-Id: Icb0fdc41b37b44e9e9e4f548d0f4acae1d83b71e
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8086
Reviewed-by: raitobezarius <tvl@lahfa.xyz>
Tested-by: BuildkiteCI

1 files changed, 63 insertions, 0 deletions

diff --git a/tvix/store/src/chunkservice/sled.rs b/tvix/store/src/chunkservice/sled.rs
new file mode 100644
index 000000000000..06194ddc5943
--- /dev/null
+++ b/tvix/store/src/chunkservice/sled.rs
@@ -0,0 +1,63 @@
+use std::path::PathBuf;
+
+use data_encoding::BASE64;
+use tracing::instrument;
+
+use crate::Error;
+
+use super::ChunkService;
+
+#[derive(Clone)]
+pub struct SledChunkService {
+    db: sled::Db,
+}
+
+impl SledChunkService {
+    pub fn new(p: PathBuf) -> Result<Self, sled::Error> {
+        let config = sled::Config::default().use_compression(true).path(p);
+        let db = config.open()?;
+
+        Ok(Self { db })
+    }
+}
+
+impl ChunkService for SledChunkService {
+    #[instrument(name = "SledChunkService::has", skip(self, digest), fields(chunk.digest=BASE64.encode(digest)))]
+    fn has(&self, digest: &[u8]) -> Result<bool, Error> {
+        match self.db.get(digest) {
+            Ok(None) => Ok(false),
+            Ok(Some(_)) => Ok(true),
+            Err(e) => Err(Error::StorageError(e.to_string())),
+        }
+    }
+
+    #[instrument(name = "SledChunkService::get", skip(self), fields(chunk.digest=BASE64.encode(digest)))]
+    fn get(&self, digest: &[u8]) -> Result<Option<Vec<u8>>, Error> {
+        match self.db.get(digest) {
+            Ok(None) => Ok(None),
+            Ok(Some(data)) => {
+                // calculate the hash to verify this is really what we expect
+                let actual_digest = blake3::hash(&data).as_bytes().to_vec();
+                if actual_digest != digest {
+                    return Err(Error::StorageError(format!(
+                        "invalid hash encountered when reading chunk, expected {}, got {}",
+                        BASE64.encode(digest),
+                        BASE64.encode(&actual_digest),
+                    )));
+                }
+                Ok(Some(Vec::from(&*data)))
+            }
+            Err(e) => Err(Error::StorageError(e.to_string())),
+        }
+    }
+
+    #[instrument(name = "SledChunkService::put", skip(self, data))]
+    fn put(&self, data: Vec<u8>) -> Result<Vec<u8>, Error> {
+        let digest = blake3::hash(&data).as_bytes().to_vec();
+        let result = self.db.insert(&digest, data);
+        if let Err(e) = result {
+            return Err(Error::StorageError(e.to_string()));
+        }
+        Ok(digest)
+    }
+}