refactor(tvix/store/chunksvc): use [u8; 32] instead of Vec<u8> r/6015

Change-Id: Ie2b94aa5d69ff2c61fb77e13ae844f81f6270273
Reviewed-on: https://cl.tvl.fyi/c/depot/+/8314
Tested-by: BuildkiteCI
Reviewed-by: tazjin <tazjin@tvl.su>
Autosubmit: flokli <flokli@flokli.de>

4 files changed, 20 insertions, 20 deletions

diff --git a/tvix/store/src/chunkservice/memory.rs b/tvix/store/src/chunkservice/memory.rs
index d9a5ead4de4a..55e646ea7714 100644
--- a/tvix/store/src/chunkservice/memory.rs
+++ b/tvix/store/src/chunkservice/memory.rs
@@ -11,18 +11,18 @@ use super::ChunkService;
 
 #[derive(Clone, Default)]
 pub struct MemoryChunkService {
-    db: Arc<RwLock<HashMap<Vec<u8>, Vec<u8>>>>,
+    db: Arc<RwLock<HashMap<[u8; 32], Vec<u8>>>>,
 }
 
 impl ChunkService for MemoryChunkService {
     #[instrument(skip(self, digest), fields(chunk.digest=BASE64.encode(digest)))]
-    fn has(&self, digest: &[u8]) -> Result<bool, Error> {
+    fn has(&self, digest: &[u8; 32]) -> 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> {
+    fn get(&self, digest: &[u8; 32]) -> Result<Option<Vec<u8>>, Error> {
         let db = self.db.read().unwrap();
         match db.get(digest) {
             None => Ok(None),
@@ -42,12 +42,12 @@ impl ChunkService for MemoryChunkService {
     }
 
     #[instrument(skip(self, data))]
-    fn put(&self, data: Vec<u8>) -> Result<Vec<u8>, Error> {
-        let digest = blake3::hash(&data).as_bytes().to_vec();
+    fn put(&self, data: Vec<u8>) -> Result<[u8; 32], Error> {
+        let digest = blake3::hash(&data);
 
         let mut db = self.db.write().unwrap();
-        db.insert(digest.clone(), data);
+        db.insert(digest.as_bytes().clone(), data);
 
-        Ok(digest)
+        Ok(digest.as_bytes().clone())
     }
 }
diff --git a/tvix/store/src/chunkservice/mod.rs b/tvix/store/src/chunkservice/mod.rs
index 50365caccf4a..faf0a88f151a 100644
--- a/tvix/store/src/chunkservice/mod.rs
+++ b/tvix/store/src/chunkservice/mod.rs
@@ -17,12 +17,12 @@ pub use self::util::upload_chunk;
 /// chunking information.
 pub trait ChunkService {
     /// check if the service has a chunk, given by its digest.
-    fn has(&self, digest: &[u8]) -> Result<bool, Error>;
+    fn has(&self, digest: &[u8; 32]) -> Result<bool, Error>;
 
     /// retrieve a chunk by its digest. Implementations MUST validate the digest
     /// matches.
-    fn get(&self, digest: &[u8]) -> Result<Option<Vec<u8>>, Error>;
+    fn get(&self, digest: &[u8; 32]) -> Result<Option<Vec<u8>>, Error>;
 
     /// insert a chunk. returns the digest of the chunk, or an error.
-    fn put(&self, data: Vec<u8>) -> Result<Vec<u8>, Error>;
+    fn put(&self, data: Vec<u8>) -> Result<[u8; 32], Error>;
 }
diff --git a/tvix/store/src/chunkservice/sled.rs b/tvix/store/src/chunkservice/sled.rs
index 5cb469b2663d..43c98813a240 100644
--- a/tvix/store/src/chunkservice/sled.rs
+++ b/tvix/store/src/chunkservice/sled.rs
@@ -30,7 +30,7 @@ impl SledChunkService {
 
 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> {
+    fn has(&self, digest: &[u8; 32]) -> Result<bool, Error> {
         match self.db.get(digest) {
             Ok(None) => Ok(false),
             Ok(Some(_)) => Ok(true),
@@ -39,7 +39,7 @@ impl ChunkService for SledChunkService {
     }
 
     #[instrument(name = "SledChunkService::get", skip(self), fields(chunk.digest=BASE64.encode(digest)))]
-    fn get(&self, digest: &[u8]) -> Result<Option<Vec<u8>>, Error> {
+    fn get(&self, digest: &[u8; 32]) -> Result<Option<Vec<u8>>, Error> {
         match self.db.get(digest) {
             Ok(None) => Ok(None),
             Ok(Some(data)) => {
@@ -59,12 +59,12 @@ impl ChunkService for SledChunkService {
     }
 
     #[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);
+    fn put(&self, data: Vec<u8>) -> Result<[u8; 32], Error> {
+        let digest = blake3::hash(&data);
+        let result = self.db.insert(&digest.as_bytes(), data);
         if let Err(e) = result {
             return Err(Error::StorageError(e.to_string()));
         }
-        Ok(digest)
+        Ok(digest.as_bytes().clone())
     }
 }
diff --git a/tvix/store/src/chunkservice/util.rs b/tvix/store/src/chunkservice/util.rs
index 761a58a646a3..1121a791fc94 100644
--- a/tvix/store/src/chunkservice/util.rs
+++ b/tvix/store/src/chunkservice/util.rs
@@ -9,7 +9,7 @@ use super::ChunkService;
 pub fn upload_chunk<CS: ChunkService>(
     chunk_service: &CS,
     chunk_data: Vec<u8>,
-) -> Result<Vec<u8>, Error> {
+) -> Result<[u8; 32], Error> {
     let mut hasher = blake3::Hasher::new();
     update_hasher(&mut hasher, &chunk_data);
     let digest = hasher.finalize();
@@ -19,9 +19,9 @@ pub fn upload_chunk<CS: ChunkService>(
     }
     let digest_resp = chunk_service.put(chunk_data)?;
 
-    assert_eq!(digest_resp, digest.as_bytes());
+    assert_eq!(&digest_resp, digest.as_bytes());
 
-    Ok(digest.as_bytes().to_vec())
+    Ok(digest.as_bytes().clone())
 }
 
 /// reads through a reader, writes chunks to a [ChunkService] and returns a
@@ -56,7 +56,7 @@ pub fn read_all_and_chunk<CS: ChunkService, R: Read>(
         let chunk_digest = upload_chunk(chunk_service, chunk.data)?;
 
         blob_meta.chunks.push(proto::blob_meta::ChunkMeta {
-            digest: chunk_digest,
+            digest: chunk_digest.to_vec(),
             size: chunk_len,
         });
     }