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, 61 insertions, 0 deletions

diff --git a/tvix/store/src/chunkservice/memory.rs b/tvix/store/src/chunkservice/memory.rs
new file mode 100644
index 000000000000..1ae8b9130565
--- /dev/null
+++ b/tvix/store/src/chunkservice/memory.rs
@@ -0,0 +1,61 @@
+use data_encoding::BASE64;
+use std::{
+    collections::HashMap,
+    sync::{Arc, RwLock},
+};
+use tracing::instrument;
+
+use crate::Error;
+
+use super::ChunkService;
+
+#[derive(Clone)]
+pub struct MemoryChunkService {
+    db: Arc<RwLock<HashMap<Vec<u8>, Vec<u8>>>>,
+}
+
+impl MemoryChunkService {
+    pub fn new() -> Self {
+        let db = Arc::new(RwLock::new(HashMap::default()));
+
+        Self { db }
+    }
+}
+
+impl ChunkService for MemoryChunkService {
+    #[instrument(skip(self, digest), fields(chunk.digest=BASE64.encode(digest)))]
+    fn has(&self, digest: &[u8]) -> Result<bool, Error> {
+        let db = self.db.read().unwrap();
+        Ok(db.get(digest).is_some())
+    }
+
+    #[instrument(skip(self), fields(chunk.digest=BASE64.encode(digest)))]
+    fn get(&self, digest: &[u8]) -> Result<Option<Vec<u8>>, Error> {
+        let db = self.db.read().unwrap();
+        match db.get(digest) {
+            None => Ok(None),
+            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(data.clone()))
+            }
+        }
+    }
+
+    #[instrument(skip(self, data))]
+    fn put(&self, data: Vec<u8>) -> Result<Vec<u8>, Error> {
+        let digest = blake3::hash(&data).as_bytes().to_vec();
+
+        let mut db = self.db.write().unwrap();
+        db.insert(digest.clone(), data);
+
+        Ok(digest)
+    }
+}