From 64baa70a60b9905c53ada1660cb8fe00513e6906 Mon Sep 17 00:00:00 2001
From: Wiktor Drewniak <wiktor.drewniak@rainlabs.pl>
Date: Sun, 6 Oct 2024 22:19:43 +0200
Subject: [PATCH] Handle empty file request on dedup store

---
 nativelink-store/src/dedup_store.rs        |  7 +++++
 nativelink-store/tests/dedup_store_test.rs | 32 ++++++++++++++++++++++
 2 files changed, 39 insertions(+)

diff --git a/nativelink-store/src/dedup_store.rs b/nativelink-store/src/dedup_store.rs
index ebacad27d..dcd14f92d 100644
--- a/nativelink-store/src/dedup_store.rs
+++ b/nativelink-store/src/dedup_store.rs
@@ -32,6 +32,8 @@ use tokio_util::codec::FramedRead;
 use tokio_util::io::StreamReader;
 use tracing::{event, Level};
 
+use crate::cas_utils::is_zero_digest;
+
 // NOTE: If these change update the comments in `stores.rs` to reflect
 // the new defaults.
 const DEFAULT_MIN_SIZE: u64 = 64 * 1024;
@@ -159,6 +161,11 @@ impl StoreDriver for DedupStore {
             .iter()
             .zip(results.iter_mut())
             .map(|(key, result)| async move {
+                if is_zero_digest(key.borrow()) {
+                    *result = Some(0);
+                    return Ok(());
+                }
+
                 match self.has(key.borrow()).await {
                     Ok(maybe_size) => {
                         *result = maybe_size;
diff --git a/nativelink-store/tests/dedup_store_test.rs b/nativelink-store/tests/dedup_store_test.rs
index a19c6a8ec..275cad521 100644
--- a/nativelink-store/tests/dedup_store_test.rs
+++ b/nativelink-store/tests/dedup_store_test.rs
@@ -14,6 +14,7 @@
 
 use nativelink_error::{Code, Error, ResultExt};
 use nativelink_macro::nativelink_test;
+use nativelink_store::cas_utils::ZERO_BYTE_DIGESTS;
 use nativelink_store::dedup_store::DedupStore;
 use nativelink_store::memory_store::MemoryStore;
 use nativelink_util::common::DigestInfo;
@@ -393,3 +394,34 @@ async fn has_with_no_existing_index_returns_none_test() -> Result<(), Error> {
     }
     Ok(())
 }
+
+/// Ensure that when we run a `.has()` on a dedup store to check for empty blobs it will
+/// properly return Some(0).
+#[nativelink_test]
+async fn has_with_zero_digest_returns_some_test() -> Result<(), Error> {
+    let index_store = MemoryStore::new(&nativelink_config::stores::MemoryStore::default());
+    let content_store = MemoryStore::new(&nativelink_config::stores::MemoryStore {
+        eviction_policy: Some(nativelink_config::stores::EvictionPolicy {
+            ..Default::default()
+        }),
+    });
+
+    let store = DedupStore::new(
+        &make_default_config(),
+        Store::new(index_store.clone()),
+        Store::new(content_store.clone()),
+    );
+
+    let digest = ZERO_BYTE_DIGESTS[0];
+
+    {
+        let size_info = store.has(digest).await.err_tip(|| "Failed to run .has")?;
+        assert_eq!(
+            size_info,
+            Some(0),
+            "Expected Sone(0) to be returned, got {:?}",
+            size_info
+        );
+    }
+    Ok(())
+}