From e958cb02517dc69e4fe489a91e4b39c7d1d583ac Mon Sep 17 00:00:00 2001
From: Florian Klink <flokli@flokli.de>
Date: Sun, 14 Apr 2024 17:04:25 +0300
Subject: [PATCH] feat(tvix/castore/blobs/object_store): chunks() method for
 small blobs

We previously returned Ok(None) when being asked for more granular
chunking info, signalling the blob does not exist at all.

This is however incorrect, we should return an empty Vec instead, as
documented in the trait.

Change-Id: I83ecc2027e0767134c7598792c2ee6d964853c66
Reviewed-on: https://cl.tvl.fyi/c/depot/+/11439
Tested-by: BuildkiteCI
Reviewed-by: Connor Brewster <cbrewster@hey.com>
Autosubmit: flokli <flokli@flokli.de>
---
 tvix/castore/src/blobservice/object_store.rs | 41 ++++++++++++++++++--
 1 file changed, 37 insertions(+), 4 deletions(-)

diff --git a/tvix/castore/src/blobservice/object_store.rs b/tvix/castore/src/blobservice/object_store.rs
index 0a9bd85e9..04fb9360b 100644
--- a/tvix/castore/src/blobservice/object_store.rs
+++ b/tvix/castore/src/blobservice/object_store.rs
@@ -218,18 +218,51 @@ impl BlobService for ObjectStoreBlobService {
 
     #[instrument(skip_all, err, fields(blob.digest=%digest))]
     async fn chunks(&self, digest: &B3Digest) -> io::Result<Option<Vec<ChunkMeta>>> {
-        let p = derive_blob_path(&self.base_path, digest);
-
-        match self.object_store.get(&p).await {
+        match self
+            .object_store
+            .get(&derive_blob_path(&self.base_path, digest))
+            .await
+        {
             Ok(get_result) => {
                 // fetch the data at the blob path
                 let blob_data = get_result.bytes().await?;
                 // parse into StatBlobResponse
                 let stat_blob_response: StatBlobResponse = StatBlobResponse::decode(blob_data)?;
 
+                debug!(
+                    chunk.count = stat_blob_response.chunks.len(),
+                    blob.size = stat_blob_response
+                        .chunks
+                        .iter()
+                        .map(|x| x.size)
+                        .sum::<u64>(),
+                    "found more granular chunks"
+                );
+
                 Ok(Some(stat_blob_response.chunks))
             }
-            Err(object_store::Error::NotFound { .. }) => Ok(None),
+            Err(object_store::Error::NotFound { .. }) => {
+                // If there's only a chunk, we must return the empty vec here, rather than None.
+                match self
+                    .object_store
+                    .head(&derive_chunk_path(&self.base_path, digest))
+                    .await
+                {
+                    Ok(_) => {
+                        // present, but no more chunks available
+                        debug!("found a single chunk");
+                        Ok(Some(vec![]))
+                    }
+                    Err(object_store::Error::NotFound { .. }) => {
+                        // Neither blob nor single chunk found
+                        debug!("not found");
+                        Ok(None)
+                    }
+                    // error checking for chunk
+                    Err(e) => Err(e.into()),
+                }
+            }
+            // error checking for blob
             Err(err) => Err(err.into()),
         }
     }