diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index b8e827492..9a2c04691 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1003,6 +1003,23 @@ pub struct FastSlowSpec { /// Default: disabled (0) #[serde(default, deserialize_with = "convert_data_size_with_shellexpand")] pub bypass_dedup_threshold_bytes: u64, + + /// Treat a hit in the `fast` store as proof that a blob exists when + /// answering existence checks (`FindMissingBlobs`, action result + /// completeness checks). The `fast` store is queried first and only the + /// keys it does not hold go to the `slow` store. + /// + /// When `false`, existence checks consult the `slow` store only. A blob + /// held solely by the `fast` tier then reads as missing, so the client + /// re-uploads it and the blob reaches durable storage. Setting this to + /// `true` gives up that guarantee in exchange for far fewer `slow` store + /// requests: if the `fast` tier evicts a blob that never reached the + /// `slow` tier, the blob is gone. Only enable this when the `fast` tier + /// does not evict, or when another path makes the `slow` tier + /// authoritative. + /// Default: false + #[serde(default, deserialize_with = "convert_boolean_with_shellexpand")] + pub trust_fast_store_for_has: bool, } #[derive(Serialize, Deserialize, Debug, Default, Clone, Copy)] diff --git a/nativelink-store/src/fast_slow_store.rs b/nativelink-store/src/fast_slow_store.rs index a6e9d28b1..4126cf28d 100644 --- a/nativelink-store/src/fast_slow_store.rs +++ b/nativelink-store/src/fast_slow_store.rs @@ -64,6 +64,8 @@ pub struct FastSlowStore { slow_direction: StoreDirection, /// See [`FastSlowSpec::bypass_dedup_threshold_bytes`]. bypass_dedup_threshold_bytes: u64, + /// See [`FastSlowSpec::trust_fast_store_for_has`]. + trust_fast_store_for_has: bool, weak_self: Weak, #[metric] metrics: FastSlowStoreMetrics, @@ -166,6 +168,7 @@ impl FastSlowStore { slow_direction: spec.slow_direction, // 0 (default) disables the bypass entirely (always dedup). bypass_dedup_threshold_bytes: spec.bypass_dedup_threshold_bytes, + trust_fast_store_for_has: spec.trust_fast_store_for_has, weak_self: weak_self.clone(), metrics: FastSlowStoreMetrics::default(), populating_digests: Mutex::new(HashMap::new()), @@ -206,6 +209,34 @@ impl FastSlowStore { .is_some_and(|size| size >= self.bypass_dedup_threshold_bytes) } + /// Asks the slow store only about the keys still unresolved in + /// `results` and writes its answers back into the matching slots. Used + /// by `has_with_results` when a fast-store hit is trusted as existence. + async fn has_in_slow_store_for_misses( + &self, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + let (miss_indexes, miss_keys): (Vec, Vec>) = keys + .iter() + .zip(results.iter()) + .enumerate() + .filter(|(_, (_, result))| result.is_none()) + .map(|(i, (key, _))| (i, key.borrow())) + .unzip(); + if miss_keys.is_empty() { + return Ok(()); + } + let mut slow_results = vec![None; miss_keys.len()]; + self.slow_store + .has_with_results(&miss_keys, &mut slow_results) + .await?; + for (i, size) in miss_indexes.into_iter().zip(slow_results) { + results[i] = size; + } + Ok(()) + } + pub const fn fast_store(&self) -> &Store { &self.fast_store } @@ -479,8 +510,18 @@ impl StoreDriver for FastSlowStore { return self.fast_store.has_with_results(key, results).await; } - // Check with the slow store first. - self.slow_store.has_with_results(key, results).await?; + if self.trust_fast_store_for_has { + self.fast_store.has_with_results(key, results).await?; + self.has_in_slow_store_for_misses(key, results).await?; + } else { + // NOTE: By default we intentionally *NEVER* check the fast store, + // this is to ensure that we re-upload data to the slow store if + // it only exists in the fast store. This does not affect workers + // as they do not check existence through `has` and instead go + // direct to loading the data which bypasses the check and will + // load from the fast store if it does not exist in the slow store. + self.slow_store.has_with_results(key, results).await?; + } // Check for any in-flight requests to the slow store next. let mut in_flight_futs = FuturesUnordered::new(); @@ -502,13 +543,6 @@ impl StoreDriver for FastSlowStore { results[i] = size; } - // NOTE: We intentionally *NEVER* check the fast store, this is to - // ensure that we re-upload data to the slow store if it only exists - // in the fast store. This does not affect workers as they do not - // check existence through `has` and instead go direct to loading the - // data which bypasses the check and will load from the fast store if - // it does not exist in the slow store. - Ok(()) } diff --git a/nativelink-store/tests/fast_slow_store_test.rs b/nativelink-store/tests/fast_slow_store_test.rs index b6ee7ee5a..231e2c4be 100644 --- a/nativelink-store/tests/fast_slow_store_test.rs +++ b/nativelink-store/tests/fast_slow_store_test.rs @@ -54,6 +54,7 @@ fn make_stores_direction( fast_direction, slow_direction, bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, fast_store.clone(), slow_store.clone(), @@ -362,6 +363,7 @@ async fn drop_on_eof_completes_store_futures() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, fast_store, slow_store, @@ -406,6 +408,7 @@ async fn ignore_value_in_fast_store() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, fast_store.clone(), slow_store, @@ -432,6 +435,7 @@ async fn has_checks_fast_store_when_noop() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }; let fast_slow_store = Arc::new(FastSlowStore::new( &fast_slow_store_config, @@ -673,6 +677,7 @@ fn make_stores_with_lazy_slow() -> (Store, Store, Store) { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, fast_store.clone(), slow_store.clone(), @@ -821,6 +826,7 @@ fn make_fast_slow_with_instrumented_slow( fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes, + trust_fast_store_for_has: false, }, fast, Store::new(slow.clone()), @@ -1069,6 +1075,7 @@ async fn has_sees_in_flight_slow_writes() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, fast, Store::new(slow.clone()), @@ -1229,6 +1236,7 @@ async fn has_does_not_consult_fast_store_when_slow_store_hits() -> Result<(), Er fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, fast, slow.clone(), @@ -1348,6 +1356,7 @@ async fn dropping_update_future_cleans_up_in_flight_entry() -> Result<(), Error> fast_direction: StoreDirection::ReadOnly, slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, fast, Store::new(slow.clone()), @@ -1515,6 +1524,7 @@ async fn has_with_results_handles_mixed_key_sources() -> Result<(), Error> { fast_direction: StoreDirection::ReadOnly, slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, fast.clone(), Store::new(slow.clone()), @@ -1685,6 +1695,7 @@ async fn huge_blob_bypasses_dedup_and_skips_populate() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: THRESHOLD, + trust_fast_store_for_has: false, }, fast_store.clone(), slow_store, @@ -1750,6 +1761,7 @@ async fn small_blob_still_dedups_and_populates() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: THRESHOLD, + trust_fast_store_for_has: false, }, fast_store.clone(), slow_store, @@ -1813,6 +1825,7 @@ async fn bypass_threshold_is_inclusive_at_exact_size() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: SIZE as u64, + trust_fast_store_for_has: false, }, fast_store.clone(), slow_store, @@ -1936,6 +1949,7 @@ fn make_stores_with_stale_fast( fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, fast_store, slow_store.clone(), @@ -2000,3 +2014,270 @@ async fn get_part_propagates_not_found_after_partial_fast_read() -> Result<(), E Ok(()) } + +/// Memory-backed store that counts `has_with_results` calls so a test can +/// prove which tier an existence check reached. +#[derive(MetricsComponent)] +struct HasCountingStore { + inner: Arc, + has_calls: AtomicUsize, +} + +#[async_trait] +impl StoreDriver for HasCountingStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + self.has_calls.fetch_add(1, Ordering::SeqCst); + Pin::new(self.inner.as_ref()) + .has_with_results(keys, results) + .await + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + Pin::new(self.inner.as_ref()) + .update(key, reader, size_info) + .await + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + Pin::new(self.inner.as_ref()) + .get_part(key, writer, offset, length) + .await + } + + fn inner_store(&self, _key: Option) -> &'_ dyn StoreDriver { + self + } + + fn as_any(&self) -> &(dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } + + fn register_remove_callback(self: Arc, _callback: RemoveCallback) -> Result<(), Error> { + Ok(()) + } +} + +default_health_status_indicator!(HasCountingStore); + +fn make_has_counting_stores( + trust_fast_store_for_has: bool, +) -> (Store, Arc, Arc) { + let fast = Arc::new(HasCountingStore { + inner: MemoryStore::new(&MemorySpec::default()), + has_calls: AtomicUsize::new(0), + }); + let slow = Arc::new(HasCountingStore { + inner: MemoryStore::new(&MemorySpec::default()), + has_calls: AtomicUsize::new(0), + }); + let fast_slow = Store::new(FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Memory(MemorySpec::default()), + slow: StoreSpec::Memory(MemorySpec::default()), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has, + }, + Store::new(fast.clone()), + Store::new(slow.clone()), + )); + (fast_slow, fast, slow) +} + +/// Default polarity: a blob held only by the fast store reads as missing and +/// the fast store is never asked, so the client re-uploads it to the slow +/// tier. +#[nativelink_test] +async fn has_by_default_checks_slow_store_only() -> Result<(), Error> { + let (fast_slow, fast, slow) = make_has_counting_stores(false); + let data = make_random_data(100); + let digest = DigestInfo::try_new(VALID_HASH, data.len()).unwrap(); + Store::new(fast.clone()) + .update_oneshot(digest, data.into()) + .await?; + + assert_eq!( + fast_slow.has(digest).await?, + None, + "Fast-only blob must read as missing by default" + ); + assert_eq!( + fast.has_calls.load(Ordering::SeqCst), + 0, + "Fast store must not be consulted by default" + ); + assert_eq!( + slow.has_calls.load(Ordering::SeqCst), + 1, + "Slow store must be consulted exactly once" + ); + Ok(()) +} + +/// With `trust_fast_store_for_has`, a batch that the fast store fully +/// answers never reaches the slow store. +#[nativelink_test] +async fn has_trusting_fast_store_hit_skips_slow_store() -> Result<(), Error> { + let (fast_slow, fast, slow) = make_has_counting_stores(true); + let data = make_random_data(100); + let digest = DigestInfo::try_new(VALID_HASH, data.len()).unwrap(); + Store::new(fast.clone()) + .update_oneshot(digest, data.clone().into()) + .await?; + + assert_eq!( + fast_slow.has(digest).await?, + Some(data.len() as u64), + "Fast hit must report the blob size" + ); + assert_eq!( + fast.has_calls.load(Ordering::SeqCst), + 1, + "Fast store must be consulted exactly once" + ); + assert_eq!( + slow.has_calls.load(Ordering::SeqCst), + 0, + "Slow store must not be consulted when every key hits the fast store" + ); + Ok(()) +} + +/// With `trust_fast_store_for_has`, only the keys the fast store misses go +/// to the slow store, and the slow store's answers land in the right slots +/// while fast hits and true misses keep their values. +#[nativelink_test] +async fn has_trusting_fast_store_miss_falls_through_to_slow_store() -> Result<(), Error> { + let (fast_slow, fast, slow) = make_has_counting_stores(true); + let fast_data = make_random_data(100); + let slow_data = make_random_data(200); + let fast_digest = DigestInfo::try_new(VALID_HASH, fast_data.len()).unwrap(); + let slow_digest = DigestInfo::try_new(VALID_HASH_B, slow_data.len()).unwrap(); + let missing_digest = DigestInfo::try_new(VALID_HASH_C, 300).unwrap(); + Store::new(fast.clone()) + .update_oneshot(fast_digest, fast_data.clone().into()) + .await?; + Store::new(slow.clone()) + .update_oneshot(slow_digest, slow_data.clone().into()) + .await?; + + let results = fast_slow + .has_many(&[ + fast_digest.into(), + slow_digest.into(), + missing_digest.into(), + ]) + .await?; + assert_eq!( + results, + vec![ + Some(fast_data.len() as u64), + Some(slow_data.len() as u64), + None + ], + "Fast hit kept, slow hit filled in, miss stays None" + ); + assert_eq!( + fast.has_calls.load(Ordering::SeqCst), + 1, + "Fast store must be consulted exactly once" + ); + assert_eq!( + slow.has_calls.load(Ordering::SeqCst), + 1, + "Slow store must be consulted once, for the fast misses only" + ); + Ok(()) +} + +/// With `trust_fast_store_for_has`, a key the fast store misses must still +/// be merged against in-flight slow writes, so a concurrent `has()` waits for +/// and sees the write instead of returning a false miss. +#[nativelink_test] +async fn has_trusting_fast_store_still_sees_in_flight_slow_writes() -> Result<(), Error> { + let (gate_tx, gate_rx) = tokio::sync::oneshot::channel::<()>(); + let (started_tx, started_rx) = tokio::sync::oneshot::channel::<()>(); + let slow = Arc::new(GatedSlowStore2 { + gate: Mutex::new(Some(gate_rx)), + started_tx: Mutex::new(Some(started_tx)), + }); + // A NoopStore fast tier with `ReadOnly` direction holds nothing, so the + // trusted fast lookup misses and the in-flight map is the only thing + // that can make the concurrent `has()` report the blob. + let fast_slow = Arc::new(FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Noop(NoopSpec::default()), + slow: StoreSpec::Memory(MemorySpec::default()), + fast_direction: StoreDirection::ReadOnly, + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: true, + }, + Store::new(NoopStore::new()), + Store::new(slow.clone()), + )); + + let data = make_random_data(64); + let digest = DigestInfo::try_new(VALID_HASH, data.len()).unwrap(); + + let writer_store = fast_slow.clone(); + let writer_data = data.clone(); + let writer = tokio::spawn(async move { + writer_store + .update_oneshot(digest, writer_data.into()) + .await + }); + started_rx + .await + .map_err(|e| make_err!(Code::Internal, "started signal lost: {e:?}"))?; + + let observer_store = fast_slow.clone(); + let mut observer = tokio::spawn(async move { observer_store.has(digest).await }); + + // The observer must block on the in-flight write, not resolve early. + tokio::select! { + _ = &mut observer => panic!("Observer resolved before writer completed"), + () = tokio::time::sleep(Duration::from_millis(10)) => {} + } + + gate_tx + .send(()) + .map_err(|()| make_err!(Code::Internal, "Failed to release slow-store gate"))?; + writer + .await + .map_err(|e| make_err!(Code::Internal, "writer join error: {e:?}"))??; + + let has_result = observer + .await + .map_err(|e| make_err!(Code::Internal, "observer join error: {e:?}"))??; + assert_eq!( + has_result, + Some(data.len() as u64), + "Concurrent has() must wait for and see the in-flight slow write" + ); + Ok(()) +} diff --git a/nativelink-worker/src/directory_cache.rs b/nativelink-worker/src/directory_cache.rs index 259e74810..f97db2184 100644 --- a/nativelink-worker/src/directory_cache.rs +++ b/nativelink-worker/src/directory_cache.rs @@ -1502,6 +1502,7 @@ mod tests { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, Store::new(fast_store), Store::new(slow_store.clone()), diff --git a/nativelink-worker/tests/directory_cache_stress_test.rs b/nativelink-worker/tests/directory_cache_stress_test.rs index 8358f73b7..f7e0b2a73 100644 --- a/nativelink-worker/tests/directory_cache_stress_test.rs +++ b/nativelink-worker/tests/directory_cache_stress_test.rs @@ -370,6 +370,7 @@ async fn stress_concurrent_construction() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, Store::new(fast_store), Store::new(delay_store), diff --git a/nativelink-worker/tests/directory_cache_test.rs b/nativelink-worker/tests/directory_cache_test.rs index b26025a55..e76f92cdb 100644 --- a/nativelink-worker/tests/directory_cache_test.rs +++ b/nativelink-worker/tests/directory_cache_test.rs @@ -79,6 +79,7 @@ async fn make_cas_store(slow_store: Arc) -> Arc { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, Store::new(fast_store), Store::new(slow_store), @@ -453,6 +454,7 @@ async fn cold_construct_fetches_concurrently() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, Store::new(fast_store), Store::new(tracking_store), @@ -1196,6 +1198,7 @@ async fn construction_failure_does_not_poison_digest() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, Store::new(fast_store), Store::new(flaky), @@ -1618,6 +1621,7 @@ async fn get_tree_prefetch_follows_server_pagination() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, Store::new(fast_store), Store::new(grpc_store), diff --git a/nativelink-worker/tests/local_worker_test.rs b/nativelink-worker/tests/local_worker_test.rs index 48ade6a53..5ba635d04 100644 --- a/nativelink-worker/tests/local_worker_test.rs +++ b/nativelink-worker/tests/local_worker_test.rs @@ -503,6 +503,7 @@ async fn new_local_worker_creates_work_directory_test() -> Result<(), Error> { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, Store::new( ::new(&FilesystemSpec { @@ -545,6 +546,7 @@ async fn new_local_worker_removes_work_directory_before_start_test() -> Result<( fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, Store::new( ::new(&FilesystemSpec { diff --git a/nativelink-worker/tests/running_actions_manager_test.rs b/nativelink-worker/tests/running_actions_manager_test.rs index 6d18a5ea7..7ce3e9df7 100644 --- a/nativelink-worker/tests/running_actions_manager_test.rs +++ b/nativelink-worker/tests/running_actions_manager_test.rs @@ -137,6 +137,7 @@ mod tests { fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + trust_fast_store_for_has: false, }, Store::new(fast_store.clone()), Store::new(slow_store.clone()),