diff --git a/nativelink-store/src/compression_store.rs b/nativelink-store/src/compression_store.rs index 7f95636a2b..903740f3e3 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 d6672b817a..78267dd086 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 dbbac9ee50..48d7f61446 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(()) +}