feat(tvix/castore/blobs/object_store): chunks() method for small blobs r/7937

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>

1 files changed, 37 insertions, 4 deletions

diff --git a/tvix/castore/src/blobservice/object_store.rs b/tvix/castore/src/blobservice/object_store.rs
index 0a9bd85e98e5..04fb9360b427 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()),
         }
     }