nativelink-store: report a string key's uncompressed size from CompressionStore::has()

nativelink-store/src/compression_store.rs

@@ -20,6 +20,7 @@ use async_trait::async_trait;
 use byteorder::{ByteOrder, LittleEndian};
 use bytes::{Buf, BufMut, BytesMut};
 use futures::future::FutureExt;
+use futures::stream::{FuturesUnordered, StreamExt};
 use lz4_flex::block::{compress_into, decompress_into, get_maximum_output_size};
 use nativelink_config::stores::CompressionSpec;
 use nativelink_error::{Code, Error, ResultExt, error_if, make_err};
@@ -33,6 +34,7 @@ use nativelink_util::store_trait::{
     RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo,
 };
 use serde::{Deserialize, Serialize};
+use tracing::warn;
 use wincode::{SchemaRead, SchemaWrite};
 
 use crate::cas_utils::is_zero_digest;
@@ -113,6 +115,11 @@ pub const CHUNK_FRAME_TYPE: u8 = 0;
 /// Number representing the footer.
 pub const FOOTER_FRAME_TYPE: u8 = 1;
 
+/// Length of the last three footer fields, `uncompressed_data_sz (u64)`,
+/// `block_size (u32)` and `version (u8)`. Every field before them has a
+/// variable length, these do not, so they always end the stream.
+const FOOTER_TAIL_SIZE: u64 = (size_of::<u64>() + size_of::<u32>() + size_of::<u8>()) as u64;
+
 /// This is a partial mirror of `nativelink_config::stores::Lz4Config`.
 /// We cannot use that natively here because it could cause our
 /// serialized format to change if we added more configs.
@@ -313,6 +320,85 @@ impl CompressionStore {
         }
         Ok(())
     }
+
+    /// Turns the sizes the inner store reported for `keys`, which are the
+    /// lengths of the compressed streams, into the sizes `get()` returns.
+    ///
+    /// A digest carries its size, so that costs nothing. It is the content's
+    /// size only when the digest names the content, as in a CAS: an action
+    /// cache keys an entry by the action's digest, so a compression store
+    /// fronting one reports the action's size (`GrpcStore` reports `u64::MAX`
+    /// for those). A string key carries no size, so the footer's is read.
+    async fn uncompressed_sizes(
+        &self,
+        keys: &[StoreKey<'_>],
+        results: &mut [Option<u64>],
+    ) -> Result<(), Error> {
+        let mut footer_reads = FuturesUnordered::new();
+        for (index, (key, result)) in keys.iter().zip(results.iter_mut()).enumerate() {
+            let Some(stored_size) = *result else {
+                continue;
+            };
+            match key {
+                StoreKey::Digest(digest) => *result = Some(digest.size_bytes()),
+                StoreKey::Str(_) => footer_reads.push(async move {
+                    (index, self.read_uncompressed_size(key, stored_size).await)
+                }),
+            }
+        }
+        while let Some((index, size)) = footer_reads.next().await {
+            results[index] = size?;
+        }
+        Ok(())
+    }
+
+    /// Reads `uncompressed_data_sz` from the footer of the stream under `key`,
+    /// which the inner store reported as `stored_size` bytes long. The field
+    /// sits at a fixed distance from the end of the stream, so this is one
+    /// ranged read of `FOOTER_TAIL_SIZE` bytes, not a read of the blocks.
+    async fn read_uncompressed_size(
+        &self,
+        key: &StoreKey<'_>,
+        stored_size: u64,
+    ) -> Result<Option<u64>, Error> {
+        let tail = match self
+            .inner_store
+            .get_part_unchunked(
+                key.borrow(),
+                stored_size.saturating_sub(FOOTER_TAIL_SIZE),
+                Some(FOOTER_TAIL_SIZE),
+            )
+            .await
+        {
+            Ok(tail) => tail,
+            // Removed since the inner store's `has()`, so it is missing now.
+            Err(err) if err.code == Code::NotFound => return Ok(None),
+            Err(err) => return Err(err).err_tip(|| "Failed to read footer in compression store"),
+        };
+        if tail.len() as u64 == FOOTER_TAIL_SIZE {
+            let uncompressed_data_size = LittleEndian::read_u64(&tail[..8]);
+            let block_size = LittleEndian::read_u32(&tail[8..12]);
+            let version = tail[12];
+            if version == CURRENT_STREAM_FORMAT_VERSION
+                && block_size != 0
+                && block_size <= self.config.max_decode_block_size
+            {
+                return Ok(Some(uncompressed_data_size));
+            }
+        }
+        // The footer is not where the reported length puts it, so the inner
+        // store reported something other than the exact length of the stream
+        // (a filesystem store reports its block-rounded footprint). Finding the
+        // footer without that length means reading the whole stream; report
+        // the stored length instead, which still says the object exists.
+        warn!(
+            %key,
+            stored_size,
+            "Compression store could not locate the footer at the inner store's \
+             reported length; has() reports the compressed length instead"
+        );
+        Ok(Some(stored_size))
+    }
 }
 
 #[async_trait]
@@ -330,16 +416,10 @@ impl StoreDriver for CompressionStore {
         self.inner_has_with_results(digests, results).await?;
         // The inner store holds the framed, compressed bytes, so the sizes it
         // reports are the compressed ones. Callers size fast-tier writes and
-        // buffers from this answer, so for a digest key report the length the
-        // digest declares, which is what a read of this store delivers. A
-        // string key declares no length and keeps the stored size; reading
-        // the uncompressed one back from the frame header is a follow-up.
-        for (key, result) in digests.iter().zip(results.iter_mut()) {
-            if let (StoreKey::Digest(digest), Some(_)) = (key.borrow(), result.as_ref()) {
-                *result = Some(digest.size_bytes());
-            }
-        }
-        Ok(())
+        // buffers from this answer, so report the length a read of this store
+        // delivers: the digest's for a digest key, and for a string key, which
+        // declares no length, the one recorded in the stream's footer.
+        self.uncompressed_sizes(digests, results).await
     }
 
     async fn update(

nativelink-store/tests/compression_store_test.rs

@@ -671,19 +671,18 @@ async fn has_zero_digests_skips_inner_store() -> Result<(), Error> {
 #[nativelink_test]
 async fn has_mixed_digests_preserves_result_order() -> Result<(), Error> {
     let (store, inner_store) = make_recording_compression_store()?;
+    let first = DigestInfo::try_new(VALID_HASH, 1)?;
+    let second = DigestInfo::try_new(VALID_HASH, 2)?;
     let keys = [
         StoreKey::from(ZERO_BYTE_DIGESTS[0]),
-        StoreKey::new_str("first"),
+        StoreKey::from(first),
         StoreKey::from(ZERO_BYTE_DIGESTS[1]),
-        StoreKey::new_str("second"),
+        StoreKey::from(second),
     ];
     let mut results = [None; 4];
 
     store.has_with_results(&keys, &mut results).await?;
-    let expected_keys = [
-        StoreKey::new_str("first").into_owned(),
-        StoreKey::new_str("second").into_owned(),
-    ];
+    let expected_keys = [StoreKey::from(first), StoreKey::from(second)];
     assert_eq!(
         (
             results,
@@ -692,7 +691,7 @@ async fn has_mixed_digests_preserves_result_order() -> Result<(), Error> {
             inner_store.has_keys.lock().unwrap().as_slice(),
         ),
         (
-            [Some(0), Some(100), Some(0), Some(101)],
+            [Some(0), Some(1), Some(0), Some(2)],
             1,
             2,
             expected_keys.as_slice(),
@@ -704,22 +703,103 @@ async fn has_mixed_digests_preserves_result_order() -> Result<(), Error> {
 #[nativelink_test]
 async fn has_nonzero_digests_delegates_single_batch() -> Result<(), Error> {
     let (store, inner_store) = make_recording_compression_store()?;
-    let keys = [StoreKey::new_str("first"), StoreKey::new_str("second")];
+    let first = DigestInfo::try_new(VALID_HASH, 1)?;
+    let second = DigestInfo::try_new(VALID_HASH, 2)?;
+    let keys = [StoreKey::from(first), StoreKey::from(second)];
     let mut results = [None; 2];
 
     store.has_with_results(&keys, &mut results).await?;
-    let expected_keys = [
-        StoreKey::new_str("first").into_owned(),
-        StoreKey::new_str("second").into_owned(),
-    ];
+    // The inner store reports 100 and 101, standing in for compressed lengths;
+    // `has()` reports each digest's own size instead.
     assert_eq!(
         (
             results,
             inner_store.has_call_count.load(Ordering::Relaxed),
             inner_store.has_key_count.load(Ordering::Relaxed),
             inner_store.has_keys.lock().unwrap().as_slice(),
         ),
-        ([Some(100), Some(101)], 1, 2, expected_keys.as_slice())
+        ([Some(1), Some(2)], 1, 2, keys.as_slice())
+    );
+    Ok(())
+}
+
+/// Text lz4 shrinks to a fraction of its size, spanning two blocks so the
+/// footer carries an index entry.
+fn compressible_data() -> Bytes {
+    let data = Bytes::from(b"the uncompressed size, not the stored one. ".repeat(2048));
+    assert!(data.len() > DEFAULT_BLOCK_SIZE as usize);
+    data
+}
+
+fn make_compression_store_over(inner_store: &Arc<MemoryStore>) -> Result<Store, Error> {
+    Ok(Store::new(CompressionStore::new(
+        &CompressionSpec {
+            backend: StoreSpec::Memory(MemorySpec::default()),
+            compression_algorithm: nativelink_config::stores::CompressionAlgorithm::Lz4(
+                nativelink_config::stores::Lz4Config::default(),
+            ),
+        },
+        Store::new(inner_store.clone()),
+    )?))
+}
+
+// `has()` used to forward the inner store's size, which is the length of the
+// compressed stream. A fast_slow store filled its fast tier with that as an
+// `ExactSize`, and a filesystem fast tier refuses a write whose length differs
+// from it, so every populate through a compression slow tier failed.
+#[nativelink_test]
+async fn has_reports_uncompressed_size_of_digest() -> Result<(), Error> {
+    let inner_store = MemoryStore::new(&MemorySpec::default());
+    let store = make_compression_store_over(&inner_store)?;
+    let data = compressible_data();
+    let digest = DigestInfo::try_new(VALID_HASH, data.len())?;
+    store.update_oneshot(digest, data.clone()).await?;
+
+    let stored_size = inner_store.has(digest).await?.err_tip(|| "Not stored")?;
+    assert!(
+        stored_size < data.len() as u64 / 4,
+        "Expected {stored_size} to be well below {}",
+        data.len()
+    );
+    let missing = DigestInfo::try_new(VALID_HASH, data.len() + 1)?;
+    assert_eq!(
+        store
+            .has_many(&[ZERO_BYTE_DIGESTS[0].into(), digest.into(), missing.into()])
+            .await?,
+        vec![Some(0), Some(data.len() as u64), None]
+    );
+    Ok(())
+}
+
+// A string key carries no size of its own, so `has()` reads the uncompressed
+// size out of the footer at the end of the stream.
+#[nativelink_test]
+async fn has_reports_uncompressed_size_of_string_key() -> Result<(), Error> {
+    let inner_store = MemoryStore::new(&MemorySpec::default());
+    let store = make_compression_store_over(&inner_store)?;
+    let data = compressible_data();
+    let key = StoreKey::new_str("compressible");
+    let empty_key = StoreKey::new_str("empty");
+    store.update_oneshot(key.borrow(), data.clone()).await?;
+    store
+        .update_oneshot(empty_key.borrow(), Bytes::new())
+        .await?;
+
+    let stored_size = inner_store
+        .has(key.borrow())
+        .await?
+        .err_tip(|| "Not stored")?;
+    assert!(
+        stored_size < data.len() as u64 / 4,
+        "Expected {stored_size} to be well below {}",
+        data.len()
+    );
+    assert_ne!(inner_store.has(empty_key.borrow()).await?, Some(0));
+    assert_eq!(
+        store
+            .has_many(&[key, empty_key, StoreKey::new_str("missing")])
+            .await?,
+        vec![Some(data.len() as u64), Some(0), None]
     );
     Ok(())
 }
@@ -966,7 +1046,7 @@ async fn compressed_slow_tier_refills_filesystem_with_the_uncompressed_size() ->
 }
 
 #[nativelink_test]
-async fn has_preserves_missing_digests_and_stored_string_key_sizes() -> Result<(), Error> {
+async fn has_preserves_missing_digests_and_reports_string_key_sizes() -> Result<(), Error> {
     let backend = Store::new(MemoryStore::new(&MemorySpec::default()));
     let store = CompressionStore::new(
         &CompressionSpec {
@@ -983,8 +1063,9 @@ async fn has_preserves_missing_digests_and_stored_string_key_sizes() -> Result<(
         .update_oneshot(digest, vec![b'a'; 4096].into())
         .await?;
     store.update_oneshot(named.borrow(), "tiny".into()).await?;
-    let named_stored_size = backend.has(named.borrow()).await?;
-    assert_ne!(named_stored_size, Some(4));
+    // The stored, framed length is not the blob's; `has()` reads the blob's
+    // from the footer.
+    assert_ne!(backend.has(named.borrow()).await?, Some(4));
     let keys = [
         StoreKey::from(DigestInfo::try_new(VALID_HASH, 78)?),
         StoreKey::from(digest),
@@ -993,6 +1074,6 @@ async fn has_preserves_missing_digests_and_stored_string_key_sizes() -> Result<(
     ];
     let mut results = [None; 4];
     store.has_with_results(&keys, &mut results).await?;
-    assert_eq!(results, [None, Some(4096), Some(0), named_stored_size]);
+    assert_eq!(results, [None, Some(4096), Some(0), Some(4)]);
     Ok(())
 }