Add ability to read partial store

TraceMachina · Jan 3, 2021 · 0b304cc · 0b304cc
1 parent 0ed309c
commit 0b304cc
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 4 deletions.
diff --git a/cas/store/memory_store.rs b/cas/store/memory_store.rs
@@ -52,16 +52,21 @@ impl StoreTrait for MemoryStore {
         Ok(())
     }
 
-    async fn get(
+    async fn get_part(
         &self,
         digest: &DigestInfo,
         writer: &mut (dyn AsyncWrite + Send + Unpin),
+        offset: usize,
+        length: Option<usize>,
     ) -> Result<(), Error> {
         let map = self.map.lock().await;
         let value = map
             .get(&digest.packed_hash)
-            .err_tip_with_code(|_| (Code::NotFound, format!("Hash {} not found", digest.str())))?;
-        writer.write_all(value).await?;
+            .err_tip_with_code(|_| (Code::NotFound, format!("Hash {} not found", digest.str())))?
+            .as_ref();
+        let default_len = value.len() - offset;
+        let length = length.unwrap_or(default_len).min(default_len);
+        writer.write_all(&value[offset..length]).await?;
         Ok(())
     }
 }
diff --git a/cas/store/store_trait.rs b/cas/store/store_trait.rs
@@ -27,9 +27,19 @@ pub trait StoreTrait: Sync + Send {
         mut reader: Box<dyn AsyncRead + Send + Unpin + 'b>,
     ) -> Result<(), Error>;
 
-    async fn get(
+    async fn get_part(
         &self,
         digest: &DigestInfo,
         writer: &mut (dyn AsyncWrite + Send + Unpin),
+        offset: usize,
+        length: Option<usize>,
     ) -> Result<(), Error>;
+
+    async fn get(
+        &self,
+        digest: &DigestInfo,
+        writer: &mut (dyn AsyncWrite + Send + Unpin),
+    ) -> Result<(), Error> {
+        self.get_part(digest, writer, 0, None).await
+    }
 }
diff --git a/cas/store/tests/memory_store_test.rs b/cas/store/tests/memory_store_test.rs
@@ -90,6 +90,31 @@ mod memory_store_tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn read_partial() -> Result<(), Error> {
+        let store = MemoryStore::new(&StoreConfig {
+            store_type: StoreType::Memory,
+            verify_size: true,
+        });
+        const VALUE1: &str = "1234";
+        let digest = DigestInfo::try_new(&VALID_HASH1, 4).unwrap();
+        store.update(&digest, Box::new(Cursor::new(VALUE1))).await?;
+
+        let mut store_data = Vec::new();
+        store
+            .get_part(&digest, &mut Cursor::new(&mut store_data), 1, Some(3))
+            .await?;
+
+        assert_eq!(
+            VALUE1[1..3].as_bytes(),
+            store_data,
+            "Expected partial data to match, expected '{:#x?}' got: {:#x?}'",
+            VALUE1[1..3].as_bytes(),
+            store_data,
+        );
+        Ok(())
+    }
+
     #[tokio::test]
     async fn errors_with_invalid_inputs() -> Result<(), Error> {
         let mut store = MemoryStore::new(&StoreConfig {