diff --git a/.github/styles/config/vocabularies/TraceMachina/accept.txt b/.github/styles/config/vocabularies/TraceMachina/accept.txt index b3da2fbc2..dc29391f6 100644 --- a/.github/styles/config/vocabularies/TraceMachina/accept.txt +++ b/.github/styles/config/vocabularies/TraceMachina/accept.txt @@ -14,6 +14,9 @@ Colab composable CPUs [Dd]eduplication +[Dd]emotes? +[Dd]emoted +[Dd]emotion eviction_policy ELB Eskandar diff --git a/nativelink-service/BUILD.bazel b/nativelink-service/BUILD.bazel index 43aa6d168..89daebd5c 100644 --- a/nativelink-service/BUILD.bazel +++ b/nativelink-service/BUILD.bazel @@ -66,6 +66,7 @@ rust_test_suite( "tests/bytestream_server_test.rs", "tests/capabilities_server_test.rs", "tests/cas_server_test.rs", + "tests/cas_splice_eviction_test.rs", "tests/execution_server_test.rs", "tests/fastcdc_conformance_test.rs", "tests/fetch_server_test.rs", diff --git a/nativelink-service/src/cas_server.rs b/nativelink-service/src/cas_server.rs index 953c2bc6e..389a9c45d 100644 --- a/nativelink-service/src/cas_server.rs +++ b/nativelink-service/src/cas_server.rs @@ -45,7 +45,7 @@ use nativelink_util::buf_channel::make_buf_channel_pair; use nativelink_util::common::DigestInfo; use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc, make_ctx_for_hash_func}; use nativelink_util::spawn_blocking; -use nativelink_util::store_trait::{Store, StoreLike, UploadSizeInfo}; +use nativelink_util::store_trait::{Store, StoreKey, StoreLike, UploadSizeInfo}; use opentelemetry::context::FutureExt; use prost::Message; use tokio_util::io::StreamReader; @@ -969,6 +969,10 @@ impl CasServer { blob_digest.size_bytes() )); } + let chunk_keys: Vec = chunk_digests + .iter() + .map(|chunk_digest| StoreKey::Digest(*chunk_digest)) + .collect(); // One round of existence checks: the chunks (which also touches // them, best-effort extending their lifetimes), the blob, and the @@ -994,6 +998,13 @@ impl CasServer { self.chunking_metrics .splice_already_exists .fetch_add(1, Ordering::Relaxed); + // The existence check above promoted the chunks in the eviction + // order; undo that — the assembled blob exists, so the chunks + // remain reproducible dedup hints (mirrors the demotion in the + // assembly path below). + if let Err(err) = store.demote_keys(&chunk_keys).await { + debug!(?err, "Failed to demote chunk keys after no-op splice"); + } return Ok(Response::new(SpliceBlobResponse { blob_digest: Some(blob_digest.into()), })); @@ -1017,6 +1028,7 @@ impl CasServer { let verification_failed_ref = &verification_failed; let (tx, rx) = make_buf_channel_pair(); let send_store = store.clone(); + let demote_store = store.clone(); // `tx` is moved into the future so that an early error return drops // it without an EOF, which aborts the in-flight store update instead // of leaving it waiting for more data. @@ -1075,6 +1087,19 @@ impl CasServer { .join(" / ") )); } + // Every chunk's content is now consumed (hashed and handed to + // the store update), so the chunks are once again reproducible + // dedup hints. Demote them to the front of the eviction order + // BEFORE the EOF releases the assembled blob for commit: the + // commit's own eviction pass must prefer evicting chunks over + // unrelated primary blobs. Without this, chunked uploads store + // every large blob twice (freshly-promoted chunks + assembled + // blob) and can push still-referenced small blobs out of a + // size-capped store. Best-effort: a demote failure must not + // fail the splice. + if let Err(err) = demote_store.demote_keys(&chunk_keys).await { + debug!(?err, "Failed to demote chunk keys in splice_blob"); + } tx.send_eof() .err_tip(|| "Failed to send EOF in splice_blob")?; Ok::<(), Error>(()) diff --git a/nativelink-service/tests/cas_splice_eviction_test.rs b/nativelink-service/tests/cas_splice_eviction_test.rs new file mode 100644 index 000000000..078940a06 --- /dev/null +++ b/nativelink-service/tests/cas_splice_eviction_test.rs @@ -0,0 +1,178 @@ +// Copyright 2024 The NativeLink Authors. All rights reserved. +// +// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// See LICENSE file for details +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +use std::sync::Arc; + +use bytes::Bytes; +use nativelink_config::cas_server::{CasChunkingConfig, CasStoreConfig, WithInstanceName}; +use nativelink_config::stores::{EvictionPolicy, MemorySpec, StoreSpec}; +use nativelink_error::Error; +use nativelink_macro::nativelink_test; +use nativelink_proto::build::bazel::remote::execution::v2::content_addressable_storage_server::ContentAddressableStorage; +use nativelink_proto::build::bazel::remote::execution::v2::{ + Digest, SpliceBlobRequest, chunking_function, digest_function, +}; +use nativelink_service::cas_server::CasServer; +use nativelink_service::wire_compression::RemoteCacheCompressionInstances; +use nativelink_store::default_store_factory::store_factory; +use nativelink_store::store_manager::StoreManager; +use nativelink_util::common::DigestInfo; +use nativelink_util::digest_hasher::{DigestHasher, DigestHasherFunc}; +use nativelink_util::store_trait::{Store, StoreLike}; +use tonic::Request; + +const INSTANCE_NAME: &str = "foo_instance_name"; +const CAS_MAX_BYTES: usize = 20 * 1024 * 1024; +const CHUNK_LEN: usize = 512 * 1024; +const NUM_CHUNKS: usize = 24; // 12MiB of chunks; chunks + assembled blob > cap. +const SMALL_BLOB: &[u8] = b"tiny early-build action input that queued work still references"; + +async fn make_capped_store_manager() -> Result, Error> { + let store_manager = Arc::new(StoreManager::new()); + store_manager.add_store( + "main_cas", + store_factory( + &StoreSpec::Memory(MemorySpec { + eviction_policy: Some(EvictionPolicy { + max_bytes: CAS_MAX_BYTES, + ..Default::default() + }), + }), + &store_manager, + None, + ) + .await?, + )?; + store_manager.add_store( + "chunk_index", + store_factory( + &StoreSpec::Memory(MemorySpec::default()), + &store_manager, + None, + ) + .await?, + )?; + Ok(store_manager) +} + +fn make_chunking_cas_server(store_manager: &StoreManager) -> Result { + CasServer::new( + &[WithInstanceName { + instance_name: INSTANCE_NAME.to_string(), + config: CasStoreConfig { + cas_store: "main_cas".to_string(), + experimental_chunking: Some(CasChunkingConfig { + index_store: Some("chunk_index".to_string()), + avg_chunk_size_bytes: 0, + max_chunk_count: 0, + }), + }, + }], + store_manager, + &RemoteCacheCompressionInstances::default(), + ) +} + +fn sha256_digest_info(data: &[u8]) -> DigestInfo { + let mut hasher = DigestHasherFunc::Sha256.hasher(); + hasher.update(data); + hasher.finalize_digest() +} + +// Regression test for a production failure: a splice stores its bytes twice +// (chunk blobs + the assembled blob) in the same size-capped CAS, so the +// assembled blob's insert-time eviction pass could push out unrelated, +// still-referenced blobs (a build-without-the-bytes client cannot re-upload +// an evicted intermediate output). With post-consumption chunk demotion, +// the eviction pass must reclaim space from the reproducible chunks +// instead. +#[nativelink_test] +async fn splice_evicts_demoted_chunks_not_unrelated_blobs() +-> Result<(), Box> { + let store_manager = make_capped_store_manager().await?; + let cas_server = make_chunking_cas_server(&store_manager)?; + let store: Store = store_manager.get_store("main_cas").unwrap(); + + // 1. A small unrelated blob, uploaded early (oldest LRU entry). + let small_digest = sha256_digest_info(SMALL_BLOB); + store + .update_oneshot(small_digest, Bytes::from_static(SMALL_BLOB)) + .await?; + + // 2. Chunks of a large blob (the CDC client's BatchUpdateBlobs phase). + let mut blob_hasher = DigestHasherFunc::Sha256.hasher(); + let mut chunk_digests: Vec = Vec::with_capacity(NUM_CHUNKS); + let mut chunk_infos: Vec = Vec::with_capacity(NUM_CHUNKS); + let mut assembled = Vec::with_capacity(NUM_CHUNKS * CHUNK_LEN); + let mut state = 0x0005_DEEC_E66D_u64; + for _ in 0..NUM_CHUNKS { + let mut data = vec![0u8; CHUNK_LEN]; + for word in data.chunks_mut(8) { + state = state + .wrapping_mul(6_364_136_223_846_793_005) + .wrapping_add(1_442_695_040_888_963_407); + let bytes = state.to_le_bytes(); + let len = word.len(); + word.copy_from_slice(&bytes[..len]); + } + blob_hasher.update(&data); + assembled.extend_from_slice(&data); + let chunk_info = sha256_digest_info(&data); + store.update_oneshot(chunk_info, Bytes::from(data)).await?; + chunk_digests.push(chunk_info.into()); + chunk_infos.push(chunk_info); + } + let blob_digest_info = blob_hasher.finalize_digest(); + + // 3. Splice. Chunks + assembled blob exceed the cap, so the eviction + // pass must run during the assembled blob's commit. + let response = cas_server + .splice_blob(Request::new(SpliceBlobRequest { + instance_name: INSTANCE_NAME.to_string(), + blob_digest: Some(blob_digest_info.into()), + chunk_digests, + digest_function: digest_function::Value::Sha256.into(), + chunking_function: chunking_function::Value::FastCdc2020.into(), + })) + .await? + .into_inner(); + let expected_blob_digest: Digest = blob_digest_info.into(); + assert_eq!(response.blob_digest.as_ref(), Some(&expected_blob_digest)); + + // The assembled blob must be present and byte-identical. + let stored = store.get_part_unchunked(blob_digest_info, 0, None).await?; + assert_eq!(stored, assembled, "assembled blob corrupt after splice"); + + // The unrelated small blob must have survived the eviction pass. + assert!( + store.has(small_digest).await?.is_some(), + "unrelated small blob was evicted by the splice's own commit; \ + chunk demotion failed" + ); + + // Sanity: eviction pressure was real — some chunks must be gone, + // otherwise this test is not exercising the eviction pass at all. + let mut surviving_chunks = 0usize; + for chunk_info in &chunk_infos { + if store.has(*chunk_info).await?.is_some() { + surviving_chunks += 1; + } + } + assert!( + surviving_chunks < NUM_CHUNKS, + "expected the eviction pass to reclaim some chunks \ + (cap {CAS_MAX_BYTES} vs ~24MiB stored), got {surviving_chunks}/{NUM_CHUNKS} alive" + ); + Ok(()) +} diff --git a/nativelink-store/src/fast_slow_store.rs b/nativelink-store/src/fast_slow_store.rs index a951b60da..8f13a9938 100644 --- a/nativelink-store/src/fast_slow_store.rs +++ b/nativelink-store/src/fast_slow_store.rs @@ -436,6 +436,16 @@ impl StoreDriver for FastSlowStore { Ok(()) } + async fn demote_keys(self: Pin<&Self>, keys: &[StoreKey<'_>]) -> Result<(), Error> { + // Eviction-ordering hint: forward to both tiers; each may have its + // own size cap. + try_join!( + self.fast_store.demote_keys(keys), + self.slow_store.demote_keys(keys) + )?; + Ok(()) + } + async fn has_with_results( self: Pin<&Self>, key: &[StoreKey<'_>], diff --git a/nativelink-store/src/filesystem_store.rs b/nativelink-store/src/filesystem_store.rs index 8654254ec..47c786f4e 100644 --- a/nativelink-store/src/filesystem_store.rs +++ b/nativelink-store/src/filesystem_store.rs @@ -1740,6 +1740,15 @@ impl StoreDriver for FilesystemStore { Ok(()) } + async fn demote_keys(self: Pin<&Self>, keys: &[StoreKey<'_>]) -> Result<(), Error> { + let own_keys = keys + .iter() + .map(|sk| sk.borrow().into_owned()) + .collect::>(); + self.evicting_map.demote(own_keys.iter()); + Ok(()) + } + async fn has_with_results( self: Pin<&Self>, keys: &[StoreKey<'_>], diff --git a/nativelink-store/src/memory_store.rs b/nativelink-store/src/memory_store.rs index 2a71227bc..ea64ac795 100644 --- a/nativelink-store/src/memory_store.rs +++ b/nativelink-store/src/memory_store.rs @@ -103,6 +103,15 @@ impl StoreDriver for MemoryStore { Ok(()) } + async fn demote_keys(self: Pin<&Self>, keys: &[StoreKey<'_>]) -> Result<(), Error> { + let own_keys = keys + .iter() + .map(|sk| sk.borrow().into_owned()) + .collect::>(); + self.evicting_map.demote(own_keys.iter()); + Ok(()) + } + async fn has_with_results( self: Pin<&Self>, keys: &[StoreKey<'_>], diff --git a/nativelink-util/src/evicting_map.rs b/nativelink-util/src/evicting_map.rs index 3339a6597..40fe2d46b 100644 --- a/nativelink-util/src/evicting_map.rs +++ b/nativelink-util/src/evicting_map.rs @@ -490,6 +490,23 @@ where while callbacks.next().await.is_some() {} } + /// Moves the entries for `keys` to the least-recently-used end of the + /// eviction order without firing callbacks, changing sizes, or touching + /// timestamps. Absent keys are ignored. Intended for reproducible + /// derived data (e.g. CDC chunk blobs after a successful `SpliceBlob`) + /// that must be first in line for eviction so it can never out-compete + /// primary blobs under size pressure. + pub fn demote(&self, keys: It) + where + It: IntoIterator, + R: Borrow, + { + let mut state = self.state.lock(); + for key in keys { + state.lru.demote(key.borrow()); + } + } + /// Fires the registered remove callbacks for `key` without touching the map. /// A store uses this to invalidate downstream listeners (e.g. an /// `ExistenceCacheStore`) when a write is rejected outright and never diff --git a/nativelink-util/src/store_trait.rs b/nativelink-util/src/store_trait.rs index 5fb650e2b..96adb54eb 100644 --- a/nativelink-util/src/store_trait.rs +++ b/nativelink-util/src/store_trait.rs @@ -485,6 +485,19 @@ pub trait StoreLike: Send + Sync + Sized + Unpin + 'static { .has_with_results(digests, results) } + /// Hints that `keys` are first-in-line eviction candidates. See + /// [`StoreDriver::demote_keys`] for details. + #[inline] + fn demote_keys<'a>( + &'a self, + keys: &'a [StoreKey<'a>], + ) -> impl Future> + Send + 'a { + if keys.is_empty() { + return future::ready(Ok(())).boxed(); + } + self.as_store_driver_pin().demote_keys(keys) + } + /// List all the keys in the store that are within the given range. /// `handler` is called for each key in the range. If `handler` returns /// false, the listing is stopped. @@ -650,6 +663,17 @@ pub trait StoreDriver: results: &mut [Option], ) -> Result<(), Error>; + /// Hints that `keys` hold reproducible derived data (e.g. CDC chunk + /// blobs after a successful `SpliceBlob`) that should be first in line + /// for eviction. Implementations with an eviction order move the + /// entries to the least-recently-used position; the default is a + /// no-op. This never deletes data and never fires removal callbacks — + /// it is purely an eviction-ordering hint, so wrapper stores that do + /// not forward it simply leave the ordering unchanged. + async fn demote_keys(self: Pin<&Self>, _keys: &[StoreKey<'_>]) -> Result<(), Error> { + Ok(()) + } + /// See: [`StoreLike::list`] for details. async fn list( self: Pin<&Self>, diff --git a/nativelink-util/tests/evicting_map_test.rs b/nativelink-util/tests/evicting_map_test.rs index 769000dea..95035b5a7 100644 --- a/nativelink-util/tests/evicting_map_test.rs +++ b/nativelink-util/tests/evicting_map_test.rs @@ -877,3 +877,66 @@ async fn snapshot_display_info() -> Result<(), Error> { ); Ok(()) } + +#[nativelink_test] +async fn demote_moves_entry_to_eviction_front() -> Result<(), Error> { + let evicting_map = EvictingMap::::new( + &EvictionPolicy { + max_count: 3, + max_seconds: 0, + max_bytes: 0, + evict_bytes: 0, + }, + MockInstantWrapped::default(), + ); + let digest_a = DigestInfo::try_new(HASH1, 0)?; + let digest_b = DigestInfo::try_new(HASH2, 0)?; + let digest_c = DigestInfo::try_new(HASH3, 0)?; + let digest_d = DigestInfo::try_new(HASH4, 0)?; + evicting_map.insert(digest_a, Bytes::new().into()).await; + evicting_map.insert(digest_b, Bytes::new().into()).await; + evicting_map.insert(digest_c, Bytes::new().into()).await; + + // Promote A (order oldest-first is now B, C, A), then demote C so it + // becomes the eviction candidate despite being recently inserted. + evicting_map.get(&digest_a).await; + evicting_map.demote([&digest_c]); + + // Inserting D must now evict the demoted C, not B. + evicting_map.insert(digest_d, Bytes::new().into()).await; + + let keys = [digest_a, digest_b, digest_c, digest_d]; + let mut results = [None, None, None, None]; + evicting_map + .sizes_for_keys(keys.iter(), &mut results, true /* peek */) + .await; + assert_eq!(results[0], Some(0), "Expected A to survive"); + assert_eq!(results[1], Some(0), "Expected B to survive"); + assert_eq!(results[2], None, "Expected demoted C to be evicted"); + assert_eq!(results[3], Some(0), "Expected D to survive"); + Ok(()) +} + +#[nativelink_test] +async fn demote_absent_key_is_noop() -> Result<(), Error> { + let evicting_map = EvictingMap::::new( + &EvictionPolicy::default(), + MockInstantWrapped::default(), + ); + let digest_a = DigestInfo::try_new(HASH1, 0)?; + let absent = DigestInfo::try_new(HASH2, 0)?; + evicting_map.insert(digest_a, Bytes::new().into()).await; + + evicting_map.demote([&absent]); + + let mut results = [None]; + evicting_map + .sizes_for_keys([&digest_a], &mut results, true /* peek */) + .await; + assert_eq!( + results[0], + Some(0), + "Expected A untouched by absent-key demote" + ); + Ok(()) +}