From d60edab03d3b49a51db8e506b611d70a0b35d90d Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 21 Oct 2026 17:31:25 +0000 Subject: [PATCH] nativelink-store: report a string key's uncompressed size from CompressionStore::has() a1e6efb55d (#2960) made CompressionStore::has() report a digest key's declared size instead of the compressed stream's length, and left string keys reporting the stored length as a follow-up. This is that follow-up. A FastSlowStore declares the slow tier's has() size to its fast tier, and since 086289a063 (#2974) the populate itself rejects a fill whose length differs from that size. A string key read through a compression slow tier therefore still fails every populate on main: Expected 514 upload bytes, received 88064 Slow store ended before the cache fill was complete A string key has no size of its own, so the size comes from the footer. Its last three fields (uncompressed_data_sz u64, block_size u32, version u8) have fixed widths and always end the stream, so one ranged get_part of 13 bytes at the inner store's reported length finds it without reading the blocks; the version and block size are checked. If the object went away since the inner has(), the key reports missing. If the footer is not at the reported length, the inner store reported something other than the stream's exact length (FilesystemStore reports its block-rounded footprint), and locating the footer would mean reading the whole stream, so has() logs a warning and reports the stored length as before; reporting None instead would call a present object missing. String keys rarely reach a compression store, since CAS traffic is digest keyed, so the extra read stays off the FindMissingBlobs path. Tests: - has_reports_uncompressed_size_of_string_key: the footer path, plus an empty object (0) and a missing key. - populate_from_compression_slow_store_into_filesystem_fast_store: end to end for a digest key and a string key. On main the string key fails as above. - has_reports_uncompressed_size_of_digest: the digest path #2960 fixed, over a two-block payload and has_many with a zero and a missing digest. - has_preserves_missing_digests_and_stored_string_key_sizes (from #2960) pinned the stored length for a string key; it now expects the blob's length and is renamed to match. - has_mixed_digests_preserves_result_order and has_nonzero_digests_delegates_single_batch used string keys against a recording inner store that reports sizes it cannot serve. They now use digest keys. Co-Authored-By: Claude Opus 5.5 --- nativelink-store/src/compression_store.rs | 100 +++++++++++++-- .../tests/compression_store_test.rs | 115 +++++++++++++++--- .../tests/fast_slow_store_test.rs | 80 +++++++++++- 3 files changed, 267 insertions(+), 28 deletions(-) diff --git a/nativelink-store/src/compression_store.rs b/nativelink-store/src/compression_store.rs index 7f95636a2b4..903740f3e3d 100644 --- a/nativelink-store/src/compression_store.rs +++ b/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::() + size_of::() + size_of::()) 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], + ) -> 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, 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( diff --git a/nativelink-store/tests/compression_store_test.rs b/nativelink-store/tests/compression_store_test.rs index d6672b817a3..78267dd086e 100644 --- a/nativelink-store/tests/compression_store_test.rs +++ b/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,14 +703,14 @@ 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, @@ -719,7 +718,88 @@ async fn has_nonzero_digests_delegates_single_batch() -> Result<(), Error> { 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) -> Result { + 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(()) } diff --git a/nativelink-store/tests/fast_slow_store_test.rs b/nativelink-store/tests/fast_slow_store_test.rs index dbbac9ee502..48d7f614465 100644 --- a/nativelink-store/tests/fast_slow_store_test.rs +++ b/nativelink-store/tests/fast_slow_store_test.rs @@ -21,11 +21,13 @@ use async_trait::async_trait; use bytes::Bytes; use futures::future::join_all; use nativelink_config::stores::{ - FastSlowSpec, FilesystemSpec, MemorySpec, NoopSpec, StoreDirection, StoreSpec, + CompressionAlgorithm, CompressionSpec, FastSlowSpec, FilesystemSpec, Lz4Config, MemorySpec, + NoopSpec, StoreDirection, StoreSpec, }; use nativelink_error::{Code, Error, ErrorContext, ResultExt, make_err}; use nativelink_macro::nativelink_test; use nativelink_metric::MetricsComponent; +use nativelink_store::compression_store::CompressionStore; use nativelink_store::fast_slow_store::FastSlowStore; use nativelink_store::filesystem_store::{FileEntry, FileEntryImpl, FilesystemStore}; use nativelink_store::memory_store::MemoryStore; @@ -2411,3 +2413,79 @@ async fn incomplete_slow_read_never_populates_memory_or_filesystem() -> Result<( } Ok(()) } + +/// A filesystem store, which refuses a write whose length is not its declared +/// `ExactSize`, with one-byte blocks so `has()` reports the stored length. +async fn make_filesystem_fast_store() -> Result<(Arc, FilesystemSpec), Error> { + let spec = FilesystemSpec { + content_path: make_temp_path("content_path"), + temp_path: make_temp_path("temp_path"), + block_size: 1, + ..Default::default() + }; + Ok((FilesystemStore::::new(&spec).await?, spec)) +} + +/// Reads `key` through a `FastSlowStore` whose fast tier is `fast` and whose +/// slow tier alone holds `data`, then checks the fast tier was filled with it. +async fn check_populates_filesystem_fast_store( + fast: Arc, + fast_spec: FilesystemSpec, + slow: Store, + key: StoreKey<'_>, + data: &Bytes, +) -> Result<(), Error> { + let fast = Store::new(fast); + let store = Store::new(FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Filesystem(fast_spec), + slow: StoreSpec::Memory(MemorySpec::default()), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + fast.clone(), + slow.clone(), + )); + slow.update_oneshot(key.borrow(), data.clone()).await?; + assert_eq!(fast.has(key.borrow()).await?, None); + + assert_eq!( + store.get_part_unchunked(key.borrow(), 0, None).await?, + data, + "{key}" + ); + assert_eq!( + fast.has(key.borrow()).await?, + Some(data.len() as u64), + "{key}" + ); + assert_eq!(fast.get_part_unchunked(key.borrow(), 0, None).await?, data); + Ok(()) +} + +// A compression store's `has()` used to report the length of the compressed +// stream, and the populate declared that to the fast tier as an `ExactSize`. +// Once the filesystem store refused a write whose length disagrees with its +// `ExactSize`, every read through a compression slow tier failed with +// "refusing to store a partial blob", so a worker could fetch no input. +#[nativelink_test] +async fn populate_from_compression_slow_store_into_filesystem_fast_store() -> Result<(), Error> { + let data = Bytes::from(b"the uncompressed size, not the stored one. ".repeat(2048)); + for key in [ + StoreKey::from(DigestInfo::try_new(VALID_HASH, data.len())?), + StoreKey::new_str("compressible"), + ] { + let compression_spec = CompressionSpec { + backend: StoreSpec::Memory(MemorySpec::default()), + compression_algorithm: CompressionAlgorithm::Lz4(Lz4Config::default()), + }; + let slow = Store::new(CompressionStore::new( + &compression_spec, + Store::new(MemoryStore::new(&MemorySpec::default())), + )?); + let (fast, fast_spec) = make_filesystem_fast_store().await?; + check_populates_filesystem_fast_store(fast, fast_spec, slow, key, &data).await?; + } + Ok(()) +}