From c24468b38f5daa9dae4c325b606454b768ab59de Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 14:37:32 -0400 Subject: [PATCH 01/23] migrate: fix 1.7.3 API drift in store tests - GrpcStore::new is now sync (returns Result, not a future) in shard_ring test - experimental_remote_cache_compression widened to Option - ByteStream::from_body_0_4 -> from_body_1_x (aws-smithy http-body 1.x) - add slow_store_write_back field to FastSlowSpec test literals - drop unused tracing_test import; remove duplicate default_true; destructure new ServicesConfig fields in single_use validation Co-Authored-By: Claude Opus 4.8 --- Cargo.lock | 10 ---------- nativelink-config/src/cas_server.rs | 1 - nativelink-store/tests/fast_slow_store_test.rs | 2 ++ nativelink-store/tests/ontap_s3_store_test.rs | 2 +- nativelink-store/tests/s3_store_test.rs | 2 +- nativelink-store/tests/shard_ring_dead_peer_test.rs | 6 +++--- 6 files changed, 7 insertions(+), 16 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 233c624030a..9681b0257b4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3120,7 +3120,6 @@ dependencies = [ "serde_assert", "serde_json", "serde_json5", - "serde_test", "shellexpand", "tracing", "tracing-test", @@ -4836,15 +4835,6 @@ dependencies = [ "serde", ] -[[package]] -name = "serde_test" -version = "1.0.177" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f901ee573cab6b3060453d2d5f0bae4e6d628c23c0a962ff9b5f1d7c8d4f1ed" -dependencies = [ - "serde", -] - [[package]] name = "serde_urlencoded" version = "0.7.1" diff --git a/nativelink-config/src/cas_server.rs b/nativelink-config/src/cas_server.rs index 69170c0b9cb..fbcfa7ec072 100644 --- a/nativelink-config/src/cas_server.rs +++ b/nativelink-config/src/cas_server.rs @@ -3310,7 +3310,6 @@ impl<'a> ReferenceChecker<'a> { #[cfg(test)] mod tests { - use tracing_test::traced_test; use super::*; diff --git a/nativelink-store/tests/fast_slow_store_test.rs b/nativelink-store/tests/fast_slow_store_test.rs index 38730bb935a..6a62375abf2 100644 --- a/nativelink-store/tests/fast_slow_store_test.rs +++ b/nativelink-store/tests/fast_slow_store_test.rs @@ -115,6 +115,7 @@ async fn filesystem_fast_tier_recovers_missing_files_on_upload_and_read() -> Res fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + slow_store_write_back: false, }, fast.clone(), slow.clone(), @@ -2178,6 +2179,7 @@ async fn slow_store_not_found_under_the_read_names_the_digest() -> Result<(), Er fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + slow_store_write_back: false, }, Store::new(MemoryStore::new(&MemorySpec::default())), Store::new(Arc::new(PartialThenGoneStore { diff --git a/nativelink-store/tests/ontap_s3_store_test.rs b/nativelink-store/tests/ontap_s3_store_test.rs index 463462d3f11..026e501d371 100644 --- a/nativelink-store/tests/ontap_s3_store_test.rs +++ b/nativelink-store/tests/ontap_s3_store_test.rs @@ -310,7 +310,7 @@ async fn simple_update_ac() -> Result<(), Error> { headers.get("x-amz-checksum-sha256"), Some("ZAbm15cMyBMkq2sUXRgXHNb3az8dLCm3tnyixpdxF+o") ); - aws_sdk_s3::primitives::ByteStream::from_body_0_4(sent_request.into_body()) + aws_sdk_s3::primitives::ByteStream::from_body_1_x(sent_request.into_body()) }; // Wait for all the data to be received by the s3 backend server. diff --git a/nativelink-store/tests/s3_store_test.rs b/nativelink-store/tests/s3_store_test.rs index 6cf4d77558e..6391c493c5a 100644 --- a/nativelink-store/tests/s3_store_test.rs +++ b/nativelink-store/tests/s3_store_test.rs @@ -257,7 +257,7 @@ async fn simple_update_ac() -> Result<(), Error> { "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=PutObject" ) ); - ByteStream::from_body_0_4(sent_request.into_body()) + ByteStream::from_body_1_x(sent_request.into_body()) }; let send_data_copy = send_data.clone(); diff --git a/nativelink-store/tests/shard_ring_dead_peer_test.rs b/nativelink-store/tests/shard_ring_dead_peer_test.rs index e98eafafd77..f3b7b66ffe2 100644 --- a/nativelink-store/tests/shard_ring_dead_peer_test.rs +++ b/nativelink-store/tests/shard_ring_dead_peer_test.rs @@ -261,7 +261,7 @@ fn grpc_spec_for_port(port: u16) -> GrpcSpec { headers: HashMap::new(), forward_headers: vec![], experimental_read_batching: None, - experimental_remote_cache_compression: false, + experimental_remote_cache_compression: Some(false), } } @@ -304,8 +304,8 @@ async fn make_ring() -> Result { let (live_port, live_blobs) = spawn_live_peer(); let dead_port = dead_peer_port(); - let live_store = GrpcStore::new(&grpc_spec_for_port(live_port)).await?; - let dead_store = GrpcStore::new(&grpc_spec_for_port(dead_port)).await?; + let live_store = GrpcStore::new(&grpc_spec_for_port(live_port))?; + let dead_store = GrpcStore::new(&grpc_spec_for_port(dead_port))?; // The spec's `store` fields are only used for length/weights by // `ShardStore::new`; memory placeholders keep the config minimal. From 59197cd8cdd9294323549dc81ebe0e1087c21e66 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 10:24:37 -0400 Subject: [PATCH 02/23] test(scheduler): deterministic fails-without regression for TL1 keepalive/inner-lock decoupling Adds `heartbeat_stamped_before_inner_lock_under_contention`: holds `inner` to simulate a contended/starved lock, then delivers a keepalive on another task. With the fix the registry (the eviction sweep's independent liveness source) is stamped on receipt -- while the lock is still held; without it `refresh_worker` is parked on the lock and never stamps, so a punctual keepalive fails to refresh liveness and a live worker is falsely evicted as Stale. Pairs with the monotone regression already on this branch to cover both halves of the fix. Co-Authored-By: Claude Opus 4.8 --- .../src/api_worker_scheduler.rs | 117 +++++++++++++++++- nativelink-scheduler/src/worker_registry.rs | 25 +++- 2 files changed, 137 insertions(+), 5 deletions(-) diff --git a/nativelink-scheduler/src/api_worker_scheduler.rs b/nativelink-scheduler/src/api_worker_scheduler.rs index 7392fddf899..cc46677c76d 100644 --- a/nativelink-scheduler/src/api_worker_scheduler.rs +++ b/nativelink-scheduler/src/api_worker_scheduler.rs @@ -1540,6 +1540,13 @@ impl ApiWorkerScheduler { load: Option, keepalive: bool, ) -> Result<(), Error> { + // TL1: stamp the independent registry liveness source on RECEIPT, before + // taking inner.lock() -- a contended lock or a saturated runtime must not + // delay the heartbeat past the eviction deadline and evict a live worker. + let now = UNIX_EPOCH + Duration::from_secs(timestamp); + self.worker_registry + .update_worker_heartbeat(worker_id, now) + .await; { let mut inner = self.inner.lock().await; inner @@ -1581,10 +1588,6 @@ impl ApiWorkerScheduler { } } } - let now = UNIX_EPOCH + Duration::from_secs(timestamp); - self.worker_registry - .update_worker_heartbeat(worker_id, now) - .await; Ok(()) } @@ -2314,3 +2317,109 @@ mod peer_census_trusted_tests { )); } } + +#[cfg(test)] +mod tl1_decoupling_tests { + use std::collections::HashMap; + use std::sync::Arc; + + use nativelink_config::schedulers::WorkerAllocationStrategy; + use nativelink_macro::nativelink_test; + use nativelink_util::instant_wrapper::MockInstantWrapped; + use tokio::sync::Notify; + + use super::{ + ApiWorkerScheduler, Duration, SharedWorkerRegistry, SystemTime, WorkerId, WorkerTimestamp, + }; + use crate::default_scheduler_factory::memory_awaited_action_db_factory; + use crate::platform_property_manager::PlatformPropertyManager; + use crate::simple_scheduler_state_manager::SimpleSchedulerStateManager; + use crate::worker_registry::WorkerRegistry; + + fn test_scheduler(registry: SharedWorkerRegistry) -> Arc { + let task_change_notify = Arc::new(Notify::new()); + let awaited_action_db = + memory_awaited_action_db_factory(0, &task_change_notify, MockInstantWrapped::default); + let state_manager = SimpleSchedulerStateManager::new( + 0, + Duration::from_secs(10), + Duration::from_secs(10), + Duration::ZERO, + awaited_action_db, + SystemTime::now, + None, + ); + ApiWorkerScheduler::new( + state_manager, + Arc::new(PlatformPropertyManager::new(HashMap::new())), + WorkerAllocationStrategy::default(), + None, + None, + Arc::new(Notify::new()), + 60, + 0, + // dispatch_ack_timeout_s = 0 -> the unacked-dispatch sweep is skipped. + 0, + registry, + None, + false, + Duration::from_secs(15), + ) + } + + /// TL1 decoupling regression. + /// + /// The registry heartbeat must be stamped on receipt, *before* `inner.lock()`. + /// The eviction sweep reads the registry as an independent liveness source; if + /// the stamp is gated behind `inner.lock()` (held by a contended or + /// runtime-starved caller -- in production the unacked-dispatch sweep holds it + /// across an `.await`), a punctual keepalive cannot refresh liveness before the + /// deadline and a live worker is falsely evicted as `Stale`. + /// + /// We simulate contention by holding `inner` and then delivering a keepalive on + /// another task. With the fix the registry is stamped *while the lock is still + /// held* (decoupled); without it the keepalive is parked on the lock and nothing + /// is stamped -- the assertion below fails (fails-without). The worker is + /// unknown to `inner`, so `refresh_lifetime` errors once the lock is released; + /// that is after the stamp we assert on and does not affect the result. + #[nativelink_test] + async fn heartbeat_stamped_before_inner_lock_under_contention() { + let registry = Arc::new(WorkerRegistry::new()); + let scheduler = test_scheduler(registry.clone()); + let worker_id = WorkerId::from(String::from("tl1-contended")); + let timestamp: WorkerTimestamp = 1000; + + // Hold `inner` to simulate a contended / starved lock. + let inner_guard = scheduler.inner.lock().await; + + // A punctual keepalive arrives while the lock is held. + let sched = scheduler.clone(); + let wid = worker_id.clone(); + let refresh = tokio::spawn(async move { + let _unused = sched.refresh_worker(&wid, timestamp, None, true).await; + }); + + // With the fix the stamp lands before `inner.lock()`, i.e. while we still + // hold the guard; without it the keepalive is parked and never stamps. + let stamped = tokio::time::timeout(Duration::from_secs(2), async { + loop { + if registry.get_worker_last_seen(&worker_id).await.is_some() { + return; + } + tokio::task::yield_now().await; + } + }) + .await + .is_ok(); + + drop(inner_guard); + let _unused = refresh.await; + + assert!( + stamped, + "heartbeat was not refreshed while inner.lock was held: a contended/starved keepalive \ + cannot refresh the independent liveness source, so a live worker is falsely evicted \ + as Stale" + ); + } +} diff --git a/nativelink-scheduler/src/worker_registry.rs b/nativelink-scheduler/src/worker_registry.rs index 08cbcce3da1..0d283cc8a67 100644 --- a/nativelink-scheduler/src/worker_registry.rs +++ b/nativelink-scheduler/src/worker_registry.rs @@ -70,7 +70,12 @@ impl WorkerRegistry { /// Updates the heartbeat timestamp for a worker. pub async fn update_worker_heartbeat(&self, worker_id: &WorkerId, now: SystemTime) { let mut workers = self.workers.write().await; - workers.insert(worker_id.clone(), now); + // TL1: monotone -- never regress last_seen. A delayed or out-of-order + // keepalive must not move liveness backward and manufacture a false Stale. + let slot = workers.entry(worker_id.clone()).or_insert(now); + if now > *slot { + *slot = now; + } trace!(?worker_id, now = %humantime::format_rfc3339(now), "FLOW: Worker heartbeat updated in registry"); } @@ -140,6 +145,24 @@ mod tests { use super::*; + /// TL1 regression: a delayed or out-of-order keepalive must never regress + /// `last_seen`. A plain `insert` moves liveness backward and manufactures a + /// false `Stale` (fails-without by construction); monotone keeps the newer time. + #[nativelink_test] + async fn heartbeat_is_monotone_never_regresses() { + let registry = WorkerRegistry::new(); + let worker_id = WorkerId::from(String::from("tl1")); + let t_late = SystemTime::UNIX_EPOCH + Duration::from_secs(1000); + let t_early = SystemTime::UNIX_EPOCH + Duration::from_secs(970); + registry.update_worker_heartbeat(&worker_id, t_late).await; + registry.update_worker_heartbeat(&worker_id, t_early).await; + assert_eq!( + registry.get_worker_last_seen(&worker_id).await, + Some(t_late), + "heartbeat regressed: non-monotone update manufactured a false Stale" + ); + } + #[nativelink_test] async fn test_worker_heartbeat() { let registry = WorkerRegistry::new(); From e514bb4ef44b8b96286d41dd90a8f91c5ae69dfb Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 10:15:53 -0400 Subject: [PATCH 03/23] test(scheduler): deterministic fails-without regression for cache-lookup drop-guard clobber (N8) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds `concurrent_same_key_survives_completed_leader_guard`, which reproduces the drop-guard clobber deterministically: a `PausableAcStore` parks the second same-key request inside its cache check (its in-flight map entry still present) while the leader completes and its spawn's drop-guard fires. Without the `ScopeGuard::into_inner` defuse the second request is orphaned and surfaces the production symptom "tx hung up in CacheLookupScheduler::add_action"; with the defuse it survives and completes. Also adds `MockActionScheduler::{receive_add_action, respond_add_action}` — a receive-without-respond rendezvous that holds a leader parked mid-`add_action` so another request can interleave deterministically. Co-Authored-By: Claude Opus 4.8 --- .../src/cache_lookup_scheduler.rs | 16 +- nativelink-scheduler/src/mock_scheduler.rs | 29 +++ .../tests/cache_lookup_scheduler_test.rs | 190 +++++++++++++++++- 3 files changed, 229 insertions(+), 6 deletions(-) diff --git a/nativelink-scheduler/src/cache_lookup_scheduler.rs b/nativelink-scheduler/src/cache_lookup_scheduler.rs index 710eb615ea8..f8df1f3d8db 100644 --- a/nativelink-scheduler/src/cache_lookup_scheduler.rs +++ b/nativelink-scheduler/src/cache_lookup_scheduler.rs @@ -231,8 +231,15 @@ impl CacheLookupScheduler { let inflight_cache_checks = self.inflight_cache_checks.clone(); // We need this spawn because we are returning a stream and this spawn will populate the stream's data. background_spawn!("cache_lookup_scheduler_add_action", async move { - // If our spawn ever dies, we will remove the action from the inflight_cache_checks map. - let _scope_guard = scope_guard; + // If our spawn ever dies before we take ownership of the map entry, remove the + // action from the inflight_cache_checks map (drop-safety). We defuse this guard + // the moment the cache check returns (below): from that point we remove the entry + // ourselves on every path with no intervening `.await`, so the guard's stale + // remove must NOT fire afterwards — a concurrent same-key request may have + // re-inserted a fresh entry in the window between our remove and this spawn's + // exit, and letting the guard run would orphan it (its oneshot tx is dropped, + // surfacing a spurious "tx hung up" to that client and forcing a retry). + let scope_guard = scope_guard; let unique_key = match &action_info.unique_qualifier { ActionUniqueQualifier::Cacheable(unique_key) => unique_key, @@ -254,6 +261,11 @@ impl CacheLookupScheduler { action_info.unique_qualifier.digest_function(), ) .await; + // The cache check has returned; from here every path removes our own map entry + // with no intervening `.await` (no drop point), so defusing the drop-guard now is + // leak-free and prevents its stale remove from clobbering a concurrent same-key + // re-insertion after we remove our entry. + scopeguard::ScopeGuard::into_inner(scope_guard); match maybe_action_result { Ok(action_result) => { let maybe_pending_txs = { diff --git a/nativelink-scheduler/src/mock_scheduler.rs b/nativelink-scheduler/src/mock_scheduler.rs index 5d87c84d6bb..064ffdb052c 100644 --- a/nativelink-scheduler/src/mock_scheduler.rs +++ b/nativelink-scheduler/src/mock_scheduler.rs @@ -136,6 +136,35 @@ impl MockActionScheduler { .unwrap(); req } + + /// Receives an `add_action` call WITHOUT sending a response, leaving the + /// caller parked awaiting the result. Pair with `respond_add_action`. + /// Lets a regression hold a leader parked mid-`add_action` while another + /// request interleaves. + #[allow(dead_code, reason = "https://github.com/rust-lang/rust/issues/46379")] + pub async fn receive_add_action(&self) -> (OperationId, ActionInfo) { + let mut rx_call_lock = self.rx_call.lock().await; + let ActionSchedulerCalls::AddAction(req) = rx_call_lock + .recv() + .await + .expect("Could not receive msg in mpsc") + else { + panic!("Got incorrect call waiting for add_action") + }; + req + } + + /// Sends the response for a previously `receive_add_action`-ed call, + /// unparking the leader. + #[allow(dead_code, reason = "https://github.com/rust-lang/rust/issues/46379")] + pub fn respond_add_action(&self, result: Result, Error>) { + assert!( + self.tx_resp + .send(ActionSchedulerReturns::AddAction(result)) + .is_ok(), + "Could not send response to mpsc" + ); + } } #[async_trait] diff --git a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs index c95b3cbface..cc1619d1077 100644 --- a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs +++ b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs @@ -19,23 +19,32 @@ mod utils { pub(crate) mod scheduler_utils; } +use core::pin::Pin; +use std::time::Duration; + +use async_trait::async_trait; use futures::join; use nativelink_config::stores::MemorySpec; use nativelink_error::Error; use nativelink_macro::nativelink_test; +use nativelink_metric::MetricsComponent; use nativelink_proto::build::bazel::remote::execution::v2::ActionResult as ProtoActionResult; use nativelink_scheduler::cache_lookup_scheduler::CacheLookupScheduler; use nativelink_scheduler::mock_scheduler::MockActionScheduler; use nativelink_store::memory_store::MemoryStore; use nativelink_util::action_messages::{ - ActionResult, ActionStage, ActionState, ActionUniqueQualifier, OperationId, + ActionInfo, ActionResult, ActionStage, ActionState, ActionUniqueQualifier, OperationId, }; +use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::common::DigestInfo; +use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; use nativelink_util::operation_state_manager::{ClientStateManager, OperationFilter}; -use nativelink_util::store_trait::{Store, StoreLike}; +use nativelink_util::store_trait::{ + RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, +}; use pretty_assertions::assert_eq; use prost::Message; -use tokio::sync::watch; +use tokio::sync::{Semaphore, watch}; use tokio::{self}; use tokio_stream::StreamExt; use utils::scheduler_utils::{TokioWatchActionStateResult, make_base_action_info}; @@ -77,7 +86,7 @@ async fn different_invocations_share_completed_cache_results() -> Result<(), Err // No mock scheduler response: a regression that rekeys the cache would // delegate to execution and fail this bounded wait. let listener = tokio::time::timeout( - std::time::Duration::from_secs(5), + Duration::from_secs(5), context .cache_scheduler .add_action(OperationId::default(), scoped_action), @@ -152,3 +161,176 @@ async fn find_by_client_operation_id_call_passed() -> Result<(), Error> { ); Ok(()) } + +/// A `Store` that parks every `get_part` (the cache-lookup read) on a gate until +/// the test releases it, signalling on `entered_tx` when each call arrives. This +/// gives a deterministic seam: the test can hold a request inside its cache check +/// (so its in-flight map entry is still present) while another request completes. +#[derive(MetricsComponent)] +struct PausableAcStore { + inner: Store, + entered_tx: tokio::sync::mpsc::UnboundedSender<()>, + gate: Arc, +} + +#[async_trait] +impl StoreDriver for PausableAcStore { + 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.inner.has_with_results(keys, results).await + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + self.inner.update(key, reader, size_info).await + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + // Announce arrival, then block until the test grants a permit. + let _ = self.entered_tx.send(()); + self.gate + .acquire() + .await + .expect("gate semaphore closed") + .forget(); + self.inner.get_part(key, writer, offset, length).await + } + + fn register_remove_callback(self: Arc, _callback: RemoveCallback) -> Result<(), Error> { + Ok(()) + } + + fn inner_store(&self, _key: Option) -> &'_ dyn StoreDriver { + self + } + + fn as_any<'a>(&'a self) -> &'a (dyn core::any::Any + Sync + Send + 'static) { + self + } + + fn as_any_arc(self: Arc) -> Arc { + self + } +} + +default_health_status_indicator!(PausableAcStore); + +fn queued_state_result( + action_info: &Arc, +) -> Box { + let (_tx, rx) = watch::channel(Arc::new(ActionState { + client_operation_id: OperationId::default(), + stage: ActionStage::Queued, + action_digest: action_info.unique_qualifier.digest(), + last_transition_timestamp: SystemTime::now(), + })); + Box::new(TokioWatchActionStateResult::new( + OperationId::default(), + action_info.clone(), + rx, + )) +} + +/// Regression for the cache-lookup drop-guard clobber (N8). +/// +/// The leader for a cacheable key removes its own `inflight_cache_checks` entry +/// when its cache check resolves, but the spawn's drop-guard *also* removes the +/// key when the spawn exits. A second request for the same key that inserts a +/// fresh entry in the window between the leader's remove and the guard drop has +/// its entry clobbered by the stale guard -> its oneshot sender is dropped -> +/// the client sees a spurious `tx hung up` error and must retry. +/// +/// Deterministic seam: `PausableAcStore` holds the second request inside its +/// cache check (entry still in the map) while the leader completes and its guard +/// fires. Without the `ScopeGuard::into_inner` defuse, the second request is +/// orphaned; with it, the second request survives and completes. +#[nativelink_test] +async fn concurrent_same_key_survives_completed_leader_guard() -> Result<(), Error> { + let mock_scheduler = Arc::new(MockActionScheduler::new()); + let inner = Store::new(MemoryStore::new(&MemorySpec::default())); + let (entered_tx, mut entered_rx) = tokio::sync::mpsc::unbounded_channel(); + let gate = Arc::new(Semaphore::new(0)); + let ac_store = Store::new(Arc::new(PausableAcStore { + inner, + entered_tx, + gate: gate.clone(), + })); + let cache_scheduler = Arc::new(CacheLookupScheduler::new(ac_store, mock_scheduler.clone())?); + + // Empty AC store -> every cache check misses (NotFound) -> leader delegates + // to the action scheduler, which is where we park it. + let action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::new([7; 32], 1)); + + // Request 1 becomes the leader. + let cs1 = cache_scheduler.clone(); + let ai1 = action_info.clone(); + let r1 = tokio::spawn(async move { cs1.add_action(OperationId::default(), ai1).await }); + + // R1 reaches its cache check; release it. It misses, removes its own + // in-flight entry, and parks delegating to the action scheduler. + entered_rx + .recv() + .await + .expect("R1 should reach cache check"); + gate.add_permits(1); + // Rendezvous proves R1 passed its self-remove and is parked (do not respond yet). + let _r1_delegated = mock_scheduler.receive_add_action().await; + + // Request 2 for the SAME key: R1's entry is gone, so R2 becomes a fresh + // leader, inserts a new entry, and parks in its own cache check. + let cs2 = cache_scheduler.clone(); + let ai2 = action_info.clone(); + let r2 = tokio::spawn(async move { cs2.add_action(OperationId::default(), ai2).await }); + entered_rx + .recv() + .await + .expect("R2 should reach cache check"); + + // Complete R1 -> its spawn exits -> its drop-guard fires. Without the fix + // this removes R2's fresh entry (the clobber). + mock_scheduler.respond_add_action(Ok(queued_state_result(&action_info))); + r1.await + .expect("R1 task should join") + .expect("R1 should succeed"); + + // Release R2. With the fix its entry survived, so it delegates to the + // action scheduler; without the fix it was orphaned and never delegates. + gate.add_permits(1); + if tokio::time::timeout(Duration::from_secs(2), mock_scheduler.receive_add_action()) + .await + .is_ok() + { + mock_scheduler.respond_add_action(Ok(queued_state_result(&action_info))); + } + + let r2_result = tokio::time::timeout(Duration::from_secs(10), r2) + .await + .expect("R2 task should finish") + .expect("R2 task should join"); + + // The concurrent same-key request must NOT be orphaned by the completed + // leader's stale drop-guard. + assert!( + r2_result.is_ok(), + "concurrent same-key request was clobbered by the completed leader's drop-guard: {:?}", + r2_result.err() + ); + Ok(()) +} From 2058eee1fc5eaaab0f58ad5739af939c6c72f089 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 10:54:06 -0400 Subject: [PATCH 04/23] fix(store): propagate mid-stream errors in ExperimentalMongoStore::update update read the upload stream with while let Ok(chunk) = reader.recv().await, which treats a stream Err identically to the empty-chunk EOF. An upload that fails partway (client disconnect, or an upstream store erroring mid-copy) was upserted as whatever bytes arrived and returned Ok(partial_len) -- a truncated blob stored under the full digest key, served as success on every later read. Propagate the recv error instead, and (belt-and-suspenders) reject a completed read whose size disagrees with UploadSizeInfo::ExactSize. Regression update_rejects_truncated_stream_instead_of_committing_it drops the writer mid-stream without send_eof and asserts update returns Err and the digest is absent; fails without the fix (commits the truncation as Ok). Co-authored-by: Claude Opus 4.8 --- nativelink-store/src/mongo_store.rs | 23 +++++++++++- nativelink-store/tests/mongo_store_test.rs | 43 ++++++++++++++++++++++ 2 files changed, 64 insertions(+), 2 deletions(-) diff --git a/nativelink-store/src/mongo_store.rs b/nativelink-store/src/mongo_store.rs index c59b4b1764c..b6a0fc853cd 100644 --- a/nativelink-store/src/mongo_store.rs +++ b/nativelink-store/src/mongo_store.rs @@ -484,15 +484,34 @@ impl StoreDriver for ExperimentalMongoStore { } } - // Read all data into memory with proper EOF handling + // Read all data into memory. A mid-stream error must be propagated, + // not treated as EOF: the old `while let Ok(chunk)` swallowed a stream + // `Err` identically to the empty-chunk EOF, committing a truncated blob + // under the digest's key (a torn write served as success). let mut data = Vec::new(); - while let Ok(chunk) = reader.recv().await { + loop { + let chunk = reader.recv().await.map_err(|e| { + make_err!( + Code::Internal, + "Failed to read stream in ExperimentalMongoStore::update: {e}" + ) + })?; if chunk.is_empty() { break; // Empty chunk signals EOF } data.extend_from_slice(&chunk); } + if let UploadSizeInfo::ExactSize(expected) = upload_size + && data.len() as u64 != expected + { + return Err(make_err!( + Code::Internal, + "ExperimentalMongoStore::update received {} bytes but expected {expected} (ExactSize)", + data.len() + )); + } + let size = data.len().try_into().unwrap_or(i64::MAX); // Create document diff --git a/nativelink-store/tests/mongo_store_test.rs b/nativelink-store/tests/mongo_store_test.rs index dac6689d4f2..3c047d78e83 100644 --- a/nativelink-store/tests/mongo_store_test.rs +++ b/nativelink-store/tests/mongo_store_test.rs @@ -1027,3 +1027,46 @@ async fn single_request_permit() -> Result<(), Error> { Ok(()) } + +#[nativelink_test] +async fn update_rejects_truncated_stream_instead_of_committing_it() -> Result<(), Error> { + // A mid-stream failure (sender dropped before send_eof) must surface as an + // Err, not be committed as a successful upload. Before the fix, the + // `while let Ok(chunk)` loop treated the stream Err identically to the + // empty-chunk EOF and upserted the partial bytes as Ok(partial_len), + // storing a truncated blob under the full digest's key. + let data = Bytes::from(vec![b'A'; 256]); + let digest = DigestInfo::try_new(VALID_HASH1, data.len())?; + + let helper = TestMongoHelper::new(None).await?; + let store = helper.store.clone(); + + let (mut tx, rx) = make_buf_channel_pair(); + + let (update_result, send_result) = tokio::join!( + async { + store + .update(digest, rx, UploadSizeInfo::ExactSize(data.len() as u64)) + .await + }, + async { + // Send one non-empty chunk, then drop the writer WITHOUT send_eof, + // so the reader sees the chunk followed by an Err. + tx.send(data.clone()).await?; + drop(tx); + Ok::<_, Error>(()) + }, + ); + send_result?; + assert!( + update_result.is_err(), + "update must return Err on a mid-stream failure, not commit a truncated blob as success", + ); + + assert!( + store.has(digest).await?.is_none(), + "a torn upload must not leave the digest present", + ); + + Ok(()) +} From 7301e123e14c70d2501ca2ec81342de3e30827a9 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 1 Oct 2026 21:58:14 -0400 Subject: [PATCH 05/23] fix(store): shrink GrpcStore read_limit on bounded-read resume A mid-stream retry in `GrpcStore::get_part` advanced `read_offset` by the forwarded byte count but left `read_limit` unchanged. `read_limit` is a REAPI ByteStream count measured from `read_offset`, so the resume re-requested the original window from the advanced offset and over-read by exactly the bytes already delivered (an L-byte bounded read could yield up to L+f bytes). Unbounded reads (`read_limit == 0`) were unaffected and must stay so. Decrement `read_limit` by the forwarded length on resume, guarded so unbounded stays unbounded. Regression asserts the resume ReadRequest's read_limit shrinks by the forwarded count (fails without the fix: 10 vs 7). Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/grpc_store.rs | 12 +++ nativelink-store/tests/grpc_store_test.rs | 121 +++++++++++++++++++--- 2 files changed, 117 insertions(+), 16 deletions(-) diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index 7fffb150ff7..ac3a4b12655 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -2029,6 +2029,18 @@ impl StoreDriver for GrpcStore { return Some((RetryResult::Err(err), local_state)); } local_state.read_offset += length; + // A bounded read must also shrink its remaining limit by the + // bytes already forwarded. `read_limit` is a count measured + // from `read_offset` (REAPI ByteStream semantics), so a + // mid-stream retry that advances `read_offset` without + // decrementing `read_limit` re-requests the *original* window + // size from the new offset and over-reads by exactly the + // bytes already delivered. `read_limit == 0` means "read to + // end" and must stay unbounded across resumes. + if local_state.read_limit != 0 { + local_state.read_limit = + local_state.read_limit.saturating_sub(length); + } } })) .await diff --git a/nativelink-store/tests/grpc_store_test.rs b/nativelink-store/tests/grpc_store_test.rs index c06b170c9ee..f4e6067a3fe 100644 --- a/nativelink-store/tests/grpc_store_test.rs +++ b/nativelink-store/tests/grpc_store_test.rs @@ -101,6 +101,11 @@ struct FakeStreamServer { drain_all: bool, committed_size_override: Option, fail_after_commit: Arc>, + /// Full byte payload the fake `read` serves (sliced from `read_offset`). + read_payload: Arc>, + /// `Some(n)`: the FIRST `read` serves `n` bytes then a retryable error, + /// forcing the client to resume; later reads serve the remainder. + first_read_partial: Option, } impl FakeStreamServer { @@ -111,6 +116,8 @@ impl FakeStreamServer { drain_all: false, committed_size_override: None, fail_after_commit: Arc::new(Mutex::new(false)), + read_payload: Arc::new(RAW_INPUT.as_bytes().to_vec()), + first_read_partial: None, } } @@ -124,10 +131,6 @@ impl FakeStreamServer { type ReadStream = Pin> + Send + 'static>>; -struct ReaderState { - responded: bool, -} - #[tonic::async_trait] impl ByteStream for FakeStreamServer { type ReadStream = ReadStream; @@ -152,19 +155,39 @@ impl ByteStream for FakeStreamServer { } } let read_request = grpc_request.into_inner(); - self.read_requests.lock().await.push(ReadRequestHolder { - request: read_request, - metadata: request_metadata, - }); - - let folded = unfold(ReaderState { responded: false }, async move |state| { - if state.responded { - return None; + let offset = usize::try_from(read_request.read_offset).unwrap_or(0); + let this_read = { + let mut reqs = self.read_requests.lock().await; + reqs.push(ReadRequestHolder { + request: read_request, + metadata: request_metadata, + }); + reqs.len() + }; + let payload = self.read_payload.clone(); + let partial = self.first_read_partial; + + // Step machine: 0 => emit a data chunk, 1 => emit a retryable error + // (only on a partial first read), 2 => EOF. Default behaviour (no + // `first_read_partial`) is one chunk from `read_offset` then EOF. + let folded = unfold((0u8, payload), move |(step, payload)| async move { + match step { + 0 => { + if let (1, Some(f)) = (this_read, partial) { + let start = offset.min(payload.len()); + let end = (offset + f).min(payload.len()); + let chunk = payload[start..end].to_vec(); + return Some((Ok(ReadResponse { data: chunk.into() }), (1u8, payload))); + } + let chunk = payload.get(offset..).unwrap_or(&[]).to_vec(); + Some((Ok(ReadResponse { data: chunk.into() }), (2u8, payload))) + } + 1 => Some(( + Err(Status::unavailable("simulated mid-stream drop")), + (2u8, payload), + )), + _ => None, } - let response = ReadResponse { - data: RAW_INPUT.as_bytes().into(), - }; - Some((Ok(response), ReaderState { responded: true })) }); Ok(Response::new(Box::pin(folded))) } @@ -870,3 +893,69 @@ async fn batch_update_blobs_splits_an_oversized_batch_across_rpcs() -> Result<() ); Ok(()) } + + +/// A 10-byte payload, longer than the partial-forward below, so a bounded read +/// has a meaningful remaining window to carry across a resume. +const RESUME_PAYLOAD: &[u8] = b"0123456789"; + +/// Regression for the `GrpcStore` bounded-read resume over-read: a mid-stream +/// retry advanced `read_offset` but left `read_limit` at its original value, so +/// the resume re-requested the full window from the new offset and over-read by +/// exactly the bytes already forwarded. The deterministic witness is the second +/// (resume) `ReadRequest` on the wire: its `read_limit` must shrink by the +/// forwarded byte count. Fails without the fix by construction (limit stays L). +#[nativelink_test] +async fn bounded_read_resume_shrinks_read_limit() -> Result<(), Error> { + const FORWARDED: usize = 3; + let requested_len = RESUME_PAYLOAD.len(); + + let server = FakeStreamServer { + read_payload: Arc::new(RESUME_PAYLOAD.to_vec()), + first_read_partial: Some(FORWARDED), + ..FakeStreamServer::new() + }; + let (server, port) = spawn_bytestream_server(server).await; + let mut spec = test_spec(format!("http://localhost:{port}"), false); + spec.retry.max_retries = 1; + let store = GrpcStore::new(&spec)?; + let digest = DigestInfo::try_new(VALID_HASH, requested_len).unwrap(); + + let (tx, mut rx) = make_buf_channel_pair(); + let get_fut = store.get_part( + digest, + tx, + 0, + Some(u64::try_from(requested_len).unwrap()), + ); + let drain_fut = async move { + let mut total = 0usize; + loop { + match rx.recv().await { + Ok(chunk) if !chunk.is_empty() => total += chunk.len(), + _ => break, + } + } + total + }; + let (res, total) = futures::join!(get_fut, drain_fut); + res?; + assert_eq!( + total, requested_len, + "writer must receive exactly the requested bytes across the resume" + ); + + let reqs = server.read_requests.lock().await; + assert!(reqs.len() >= 2, "expected a resume read, saw {}", reqs.len()); + let limit = i64::try_from(requested_len).unwrap(); + let fwd = i64::try_from(FORWARDED).unwrap(); + assert_eq!(reqs[0].request.read_offset, 0); + assert_eq!(reqs[0].request.read_limit, limit); + assert_eq!(reqs[1].request.read_offset, fwd); + assert_eq!( + reqs[1].request.read_limit, + limit - fwd, + "resume read_limit must shrink by bytes already forwarded (over-read bug)" + ); + Ok(()) +} From e58a13aa89dd9320b82c3573b557db8d8bb00fc7 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 10:48:21 -0400 Subject: [PATCH 06/23] fix(store): request an inclusive HTTP byte range in S3/ONTAP/GCS get_part S3/ONTAP/GCS get_part built bytes={offset}-{offset+length}, but HTTP byte ranges are inclusive on both ends, so a bounded read ending before the object end over-read one byte (a torn read on content-addressed partial reads). Use offset+length-1, guard zero-length reads, and fix the GCS mock (which sliced exclusively and masked the bug). r2_store wraps s3_store, so S3/R2/ONTAP + GCS are all covered. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/gcs_client/mocks.rs | 6 +- nativelink-store/src/gcs_store.rs | 9 ++- nativelink-store/src/ontap_s3_store.rs | 16 ++++- nativelink-store/src/s3_store.rs | 16 ++++- nativelink-store/tests/gcs_store_test.rs | 53 +++++++++++++++- nativelink-store/tests/ontap_s3_store_test.rs | 63 ++++++++++++++++++- nativelink-store/tests/s3_store_test.rs | 61 +++++++++++++++++- 7 files changed, 214 insertions(+), 10 deletions(-) diff --git a/nativelink-store/src/gcs_client/mocks.rs b/nativelink-store/src/gcs_client/mocks.rs index f5ba13b5506..33a617a9cc9 100644 --- a/nativelink-store/src/gcs_client/mocks.rs +++ b/nativelink-store/src/gcs_client/mocks.rs @@ -354,7 +354,11 @@ impl GcsOperations for MockGcsOperations { start )); } - core::cmp::min(usize::try_from(e).unwrap_or(usize::MAX), content.len()) + // HTTP byte ranges are inclusive of the end index. + core::cmp::min( + usize::try_from(e).unwrap_or(usize::MAX).saturating_add(1), + content.len(), + ) } else { content.len() }; diff --git a/nativelink-store/src/gcs_store.rs b/nativelink-store/src/gcs_store.rs index dd9731a5cc5..73bd695ea79 100644 --- a/nativelink-store/src/gcs_store.rs +++ b/nativelink-store/src/gcs_store.rs @@ -415,8 +415,15 @@ where return Ok(()); } + if length == Some(0) { + writer.send_eof()?; + return Ok(()); + } let object_path = self.make_object_path(&key); - let end_offset = length.map(|len| offset + len); + // HTTP byte ranges are inclusive on both ends; the last byte index for a + // `len`-byte read from `offset` is `offset + len - 1`. `offset + len` + // over-reads one byte on any range that stops before the object end. + let end_offset = length.map(|len| offset + len - 1); let client = &self.client; let object_path_ref = &object_path; diff --git a/nativelink-store/src/ontap_s3_store.rs b/nativelink-store/src/ontap_s3_store.rs index be7868de681..8ec6b71cb4f 100644 --- a/nativelink-store/src/ontap_s3_store.rs +++ b/nativelink-store/src/ontap_s3_store.rs @@ -661,8 +661,22 @@ where } let s3_path = &self.make_s3_path(&key); + // A zero-length bounded read has no valid inclusive byte range; satisfy + // it with an immediate EOF (matches `MemoryStore::get_part` semantics). + if length == Some(0) { + writer + .send_eof() + .err_tip(|| "Failed to send zero-length EOF in get_part")?; + return Ok(()); + } + // HTTP byte ranges are inclusive on both ends: reading `length` bytes + // from `offset` ends at byte index `offset + length - 1`. Using + // `offset + length` over-reads one byte on any range that stops before + // the object end. let end_read_byte = length - .map_or(Some(None), |length| Some(offset.checked_add(length))) + .map_or(Some(None), |length| { + Some(offset.checked_add(length).map(|end| end - 1)) + }) .err_tip(|| "Integer overflow protection triggered")?; self.retrier diff --git a/nativelink-store/src/s3_store.rs b/nativelink-store/src/s3_store.rs index f7fa79e52dc..364d44f1908 100644 --- a/nativelink-store/src/s3_store.rs +++ b/nativelink-store/src/s3_store.rs @@ -664,8 +664,22 @@ where } let s3_path = &self.make_s3_path(&key); + // A zero-length bounded read has no valid inclusive byte range; satisfy + // it with an immediate EOF (matches `MemoryStore::get_part` semantics). + if length == Some(0) { + writer + .send_eof() + .err_tip(|| "Failed to send zero-length EOF in get_part")?; + return Ok(()); + } + // HTTP byte ranges are inclusive on both ends: reading `length` bytes + // from `offset` ends at byte index `offset + length - 1`. Using + // `offset + length` over-reads one byte on any range that stops before + // the object end. let end_read_byte = length - .map_or(Some(None), |length| Some(offset.checked_add(length))) + .map_or(Some(None), |length| { + Some(offset.checked_add(length).map(|end| end - 1)) + }) .err_tip(|| "Integer overflow protection triggered")?; self.retrier diff --git a/nativelink-store/tests/gcs_store_test.rs b/nativelink-store/tests/gcs_store_test.rs index 93fae5c6fe8..280bd5c7c18 100644 --- a/nativelink-store/tests/gcs_store_test.rs +++ b/nativelink-store/tests/gcs_store_test.rs @@ -549,8 +549,8 @@ async fn get_part_with_range() -> Result<(), Error> { assert_eq!(*read_requests[0].0, 6, "Expected start offset to be 6"); assert_eq!( *read_requests[0].1, - Some(11), - "Expected end offset to be Some(11)" + Some(10), + "Expected inclusive end offset to be Some(10)" ); Ok(()) @@ -1016,3 +1016,52 @@ async fn get_part_ignores_empty_stream_chunks() -> Result<(), Error> { Ok(()) } + +#[nativelink_test] +async fn get_part_bounded_subrange_is_end_inclusive() -> Result<(), Error> { + // A strict subrange (offset + length < size) must return exactly `length` + // bytes with an *inclusive* end offset. Regression for the off-by-one that + // passed `offset + length` (one byte too many) to the HTTP byte range. + let mock_ops = Arc::new(MockGcsOperations::new()); + let store = create_test_store(mock_ops.clone()).await?; + let digest = DigestInfo::try_new(VALID_HASH1, 11)?; // "hello world" + let store_key: StoreKey = to_store_key(digest); + let object_path = create_object_path(&store_key); + mock_ops + .add_object(&object_path, b"hello world".to_vec()) + .await; + let (mut tx, mut rx) = make_buf_channel_pair(); + + // bytes [2, 5) = "llo" (3 bytes), ending strictly before the object end. + let store_clone = store.clone(); + let get_fut = nativelink_util::spawn!("get_part_subrange_task", async move { + store_clone.get_part(store_key, &mut tx, 2, Some(3)).await + }); + let received_data = rx.consume(Some(100)).await?; + assert_eq!( + received_data.as_ref(), + b"llo", + "strict subrange must return exactly `length` bytes" + ); + get_fut.await??; + + let requests = mock_ops.get_requests().await; + let read_requests: Vec<_> = requests + .iter() + .filter_map(|req| { + if let MockRequest::ReadContent { start, end, .. } = req { + Some((*start, *end)) + } else { + None + } + }) + .collect(); + assert_eq!(read_requests.len(), 1, "Expected one read content request"); + assert_eq!(read_requests[0].0, 2, "Expected start offset 2"); + assert_eq!( + read_requests[0].1, + Some(4), + "Expected inclusive end offset Some(4) for a 3-byte read from offset 2" + ); + Ok(()) +} diff --git a/nativelink-store/tests/ontap_s3_store_test.rs b/nativelink-store/tests/ontap_s3_store_test.rs index 026e501d371..2021dbba3d6 100644 --- a/nativelink-store/tests/ontap_s3_store_test.rs +++ b/nativelink-store/tests/ontap_s3_store_test.rs @@ -399,7 +399,7 @@ async fn smoke_test_get_part() -> Result<(), Error> { "https://{BUCKET_NAME}.s3.{VSERVER_NAME}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=GetObject" ) ) - .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH)) + .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH - 1)) .body(SdkBody::empty()) .unwrap(), http::Response::builder().status(StatusCode::OK).body(SdkBody::empty()).unwrap() @@ -491,7 +491,7 @@ async fn ensure_empty_string_in_stream_works_test() -> Result<(), Error> { "https://{BUCKET_NAME}.s3.{VSERVER_NAME}.amazonaws.com/{VALID_HASH1}-{CAS_ENTRY_SIZE}?x-id=GetObject" ) ) - .header("range", format!("bytes={}-{}", 0, CAS_ENTRY_SIZE)) + .header("range", format!("bytes={}-{}", 0, CAS_ENTRY_SIZE - 1)) .body(SdkBody::empty()) .unwrap(), http::Response @@ -733,3 +733,62 @@ async fn multipart_update_large_cas() -> Result<(), Error> { mock_client.assert_requests_match(&[]); Ok(()) } + +#[nativelink_test] +async fn get_part_bounded_subrange_range_is_end_inclusive() -> Result<(), Error> { + // A bounded read ending strictly before the object end must request an + // *inclusive* HTTP byte range of exactly `length` bytes. Regression for the + // off-by-one that requested `offset + length` (one byte too many). The + // StaticReplayClient only serves the exact-match range header, so without + // the fix the store emits `bytes=100-300`, fails to match, and errors. + const AC_ENTRY_SIZE: u64 = 1000; + const OFFSET: u64 = 100; + const LENGTH: u64 = 200; // OFFSET + LENGTH = 300 < AC_ENTRY_SIZE: a strict subrange. + let expected_payload = vec![0x88_u8; usize::try_from(LENGTH).unwrap()]; + let mock_client = StaticReplayClient::new(vec![ReplayEvent::new( + http::Request::builder() + .uri(format!( + "https://{BUCKET_NAME}.s3.{VSERVER_NAME}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=GetObject", + )) + .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH - 1)) + .body(SdkBody::empty()) + .unwrap(), + http::Response::builder() + .status(StatusCode::OK) + .body(SdkBody::from(expected_payload.clone())) + .unwrap(), + )]); + let test_config = Builder::new() + .behavior_version(BehaviorVersion::latest()) + .region(Region::from_static(VSERVER_NAME)) + .http_client(mock_client.clone()) + .build(); + let s3_client = aws_sdk_s3::Client::from_conf(test_config); + let store = OntapS3Store::new_with_client_and_jitter( + &(ExperimentalOntapS3Spec { + bucket: BUCKET_NAME.to_string(), + vserver_name: VSERVER_NAME.to_string(), + endpoint: "https://example.com".to_string(), + ..Default::default() + }), + s3_client, + Arc::new(move |_delay| Duration::from_secs(0)), + MockInstantWrapped::default, + )?; + + let got = store + .get_part_unchunked( + DigestInfo::try_new(VALID_HASH1, AC_ENTRY_SIZE)?, + OFFSET, + Some(LENGTH), + ) + .await?; + assert_eq!( + got.len(), + usize::try_from(LENGTH).unwrap(), + "bounded read returned wrong byte count" + ); + assert_eq!(got, Bytes::from(expected_payload)); + mock_client.assert_requests_match(&[]); + Ok(()) +} diff --git a/nativelink-store/tests/s3_store_test.rs b/nativelink-store/tests/s3_store_test.rs index 6391c493c5a..25fd5a82e99 100644 --- a/nativelink-store/tests/s3_store_test.rs +++ b/nativelink-store/tests/s3_store_test.rs @@ -349,7 +349,7 @@ async fn smoke_test_get_part() -> Result<(), Error> { .uri(format!( "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=GetObject", )) - .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH)) + .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH - 1)) .body(SdkBody::empty()) .unwrap(), http::Response::builder() @@ -587,7 +587,7 @@ async fn ensure_empty_string_in_stream_works_test() -> Result<(), Error> { .uri(format!( "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{CAS_ENTRY_SIZE}?x-id=GetObject", )) - .header("range", format!("bytes={}-{}", 0, CAS_ENTRY_SIZE)) + .header("range", format!("bytes={}-{}", 0, CAS_ENTRY_SIZE - 1)) .body(SdkBody::empty()) .unwrap(), http::Response::builder() @@ -1055,3 +1055,60 @@ async fn single_put_lets_the_sdk_own_the_checksum() -> Result<(), Error> { ); Ok(()) } + +#[nativelink_test] +async fn get_part_bounded_subrange_range_is_end_inclusive() -> Result<(), Error> { + // A bounded read ending strictly before the object end must request an + // *inclusive* HTTP byte range of exactly `length` bytes. Regression for the + // off-by-one that requested `offset + length` (one byte too many). The + // StaticReplayClient only serves the exact-match range header, so without + // the fix the store emits `bytes=100-300`, fails to match, and errors. + const AC_ENTRY_SIZE: u64 = 1000; + const OFFSET: u64 = 100; + const LENGTH: u64 = 200; // OFFSET + LENGTH = 300 < AC_ENTRY_SIZE: a strict subrange. + let expected_payload = vec![0x88_u8; usize::try_from(LENGTH).unwrap()]; + let mock_client = StaticReplayClient::new(vec![ReplayEvent::new( + http::Request::builder() + .uri(format!( + "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=GetObject", + )) + .header("range", format!("bytes={}-{}", OFFSET, OFFSET + LENGTH - 1)) + .body(SdkBody::empty()) + .unwrap(), + http::Response::builder() + .status(StatusCode::OK) + .body(SdkBody::from(expected_payload.clone())) + .unwrap(), + )]); + let test_config = Builder::new() + .behavior_version(BehaviorVersion::latest()) + .region(Region::from_static(REGION)) + .http_client(mock_client.clone()) + .build(); + let s3_client = aws_sdk_s3::Client::from_conf(test_config); + let store = S3Store::new_with_client_and_jitter( + &ExperimentalAwsSpec { + bucket: BUCKET_NAME.to_string(), + ..Default::default() + }, + s3_client, + Arc::new(move |_delay| Duration::from_secs(0)), + MockInstantWrapped::default, + )?; + + let got = store + .get_part_unchunked( + DigestInfo::try_new(VALID_HASH1, AC_ENTRY_SIZE)?, + OFFSET, + Some(LENGTH), + ) + .await?; + assert_eq!( + got.len(), + usize::try_from(LENGTH).unwrap(), + "bounded read returned wrong byte count" + ); + assert_eq!(got, Bytes::from(expected_payload)); + mock_client.assert_requests_match(&[]); + Ok(()) +} From 0fa9b45d82a2e53cf7bdd3feeb4a9439a73bba7a Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 14:38:51 -0400 Subject: [PATCH 07/23] reland: mongo update_data version-CAS returns retryable Ok(None) not Err Cherry-pick of b7r6/mongo-update-data-upsert-defeats-version-cas, re-grafted onto 1.7.3: upsert(current_version == 0) so a genuine create still upserts but a stale-version write misses cleanly. Test kept alongside the truncated-stream regression test (both land at the same point). Co-Authored-By: Claude Opus 4.8 Co-Authored-By: b7r6 --- nativelink-store/src/mongo_store.rs | 2 +- nativelink-store/tests/mongo_store_test.rs | 52 +++++++++++++++++++++- 2 files changed, 52 insertions(+), 2 deletions(-) diff --git a/nativelink-store/src/mongo_store.rs b/nativelink-store/src/mongo_store.rs index b6a0fc853cd..e4ee3ed9381 100644 --- a/nativelink-store/src/mongo_store.rs +++ b/nativelink-store/src/mongo_store.rs @@ -992,7 +992,7 @@ impl SchedulerStore for ExperimentalMongoStore { let result = match self .scheduler_collection .find_one_and_update(filter, update_doc) - .upsert(true) + .upsert(current_version == 0) .return_document(ReturnDocument::After) .await { diff --git a/nativelink-store/tests/mongo_store_test.rs b/nativelink-store/tests/mongo_store_test.rs index 3c047d78e83..f8e1519ca2e 100644 --- a/nativelink-store/tests/mongo_store_test.rs +++ b/nativelink-store/tests/mongo_store_test.rs @@ -839,7 +839,10 @@ async fn test_scheduler_store_operations() -> Result<(), Error> { data.version = version1; // Using old version let result = helper.store.update_data(data.clone(), None).await; - assert!(result.is_err(), "Update with old version should fail"); + assert!( + matches!(result, Ok(None)), + "Stale-version update must return retryable Ok(None), got {result:?}" + ); eprintln!("Correctly rejected update with stale version"); } @@ -1070,3 +1073,50 @@ async fn update_rejects_truncated_stream_instead_of_committing_it() -> Result<() Ok(()) } + +#[nativelink_test] +async fn update_data_stale_version_returns_retryable_none_not_err() -> Result<(), Error> { + // N12 regression: ExperimentalMongoStore::update_data used upsert(true), so a + // version-CAS miss (a concurrent writer already bumped the version) hit the + // unique _id and surfaced as a DuplicateKeyError (Err) instead of the + // retryable Ok(None) the scheduler contract (and the Redis backend) require. + let helper = TestMongoHelper::new(None).await?; + + let mut data = TestSchedulerData { + key: "test:n12_version_cas".to_string(), + content: "v1".to_string(), + version: 0, + }; + + // Seed at version 0 -> creates version 1 (genuine first create still upserts). + let v1 = helper + .store + .update_data(data.clone(), None) + .await + .err_tip(|| "create should succeed")? + .ok_or_else(|| make_err!(Code::Internal, "expected a version on create"))?; + + // One writer wins with the correct version -> version 2. + data.content = "v2".to_string(); + data.version = v1; + let v2 = helper + .store + .update_data(data.clone(), None) + .await + .err_tip(|| "correct-version update should succeed")? + .ok_or_else(|| make_err!(Code::Internal, "expected a version on update"))?; + assert!(v2 > v1, "version should increment ({v1} -> {v2})"); + + // The losing concurrent writer carries the now-stale version. It MUST get a + // retryable Ok(None), not Err. Before the fix this returned Err(Internal) + // (duplicate _id), destroying the retry signal. + data.content = "stale".to_string(); + data.version = v1; + let stale = helper.store.update_data(data.clone(), None).await; + assert!( + matches!(stale, Ok(None)), + "stale-version update must return retryable Ok(None), got {stale:?}" + ); + + Ok(()) +} From 3c6d3ff4192233ea24fb0d1930f727d7e16cf785 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 14:42:00 -0400 Subject: [PATCH 08/23] test(fast_slow): fast-first has() contract + slow-probe deadline fail-fast - Rewrite has_does_not_consult_fast_store... -> fast-first: the fast store IS consulted first; a fast-hit is answered without the slow probe, a fast-miss falls through to the slow tier. - Add has_fast_miss_fails_fast_when_slow_tier_hangs: a hung slow tier must degrade a fast-miss to missing within SLOW_PROBE_DEADLINE, not block on the slow store full RPC timeout. Co-Authored-By: Claude Opus 4.8 --- .../tests/fast_slow_store_test.rs | 127 +++++++++++++++++- 1 file changed, 121 insertions(+), 6 deletions(-) diff --git a/nativelink-store/tests/fast_slow_store_test.rs b/nativelink-store/tests/fast_slow_store_test.rs index 6a62375abf2..08c29abd907 100644 --- a/nativelink-store/tests/fast_slow_store_test.rs +++ b/nativelink-store/tests/fast_slow_store_test.rs @@ -1201,7 +1201,7 @@ async fn has_sees_in_flight_slow_writes() -> Result<(), Error> { /// here we additionally assert that when the slow store DOES have the blob, /// the fast store is NOT consulted (avoiding the extra round trip). #[nativelink_test] -async fn has_does_not_consult_fast_store_when_slow_store_hits() -> Result<(), Error> { +async fn has_checks_fast_store_first_and_skips_slow_on_fast_hit() -> Result<(), Error> { #[derive(MetricsComponent)] struct CountingFastStore { inner: Arc, @@ -1277,6 +1277,7 @@ async fn has_does_not_consult_fast_store_when_slow_store_hits() -> Result<(), Er has_calls: has_calls.clone(), })); let slow = Store::new(MemoryStore::new(&MemorySpec::default())); + let fast_for_seed = fast.clone(); let fast_slow = Arc::new(FastSlowStore::new( &FastSlowSpec { fast: StoreSpec::Memory(MemorySpec::default()), @@ -1292,18 +1293,32 @@ async fn has_does_not_consult_fast_store_when_slow_store_hits() -> Result<(), Er let data = make_random_data(128); let digest = DigestInfo::try_new(VALID_HASH, data.len()).unwrap(); - slow.update_oneshot(digest, data.clone().into()).await?; + // Fast-first: a blob resident in the fast tier is answered from the fast + // tier, and the fast store IS consulted. (The historical slow-only probe + // is what wedged FindMissingBlobs under slow-tier throttling.) + fast_for_seed + .update_oneshot(digest, data.clone().into()) + .await?; let before = has_calls.load(Ordering::Acquire); assert_eq!( fast_slow.has(digest).await?, Some(data.len() as u64), - "Slow-store-only blob must be reported via slow lookup", + "Fast-resident blob must be reported present", ); - let after = has_calls.load(Ordering::Acquire); + assert!( + has_calls.load(Ordering::Acquire) > before, + "Fast store has() must be consulted first under the fast-first contract", + ); + + // A fast-miss falls through to the slow tier (a bounded probe, not a wedge). + let other = make_random_data(64); + let other_digest = DigestInfo::try_new(VALID_HASH_B, other.len()).unwrap(); + slow.update_oneshot(other_digest, other.clone().into()).await?; assert_eq!( - after, before, - "Fast store has() must not be consulted when the slow store already reports the blob", + fast_slow.has(other_digest).await?, + Some(other.len() as u64), + "Fast-miss must fall through and be found in the slow tier", ); Ok(()) @@ -2363,3 +2378,103 @@ async fn write_back_decouples_update_from_slow_store() -> Result<(), Error> { Ok(()) } + +/// Fire-fix completion: a fast-miss must fail-fast on a slow tier that hangs. +/// `has_with_results` bounds the slow-tier existence probe with a short deadline +/// (SLOW_PROBE_DEADLINE ~3s) and degrades the fast-miss to "missing" on timeout, +/// rather than blocking on the slow store's full RPC timeout (e.g. a GrpcStore's +/// 60s). Without the deadline a cold build's tens-of-thousands of fast-misses +/// each stall the hot ring long enough to wedge the worker. +#[nativelink_test] +async fn has_fast_miss_fails_fast_when_slow_tier_hangs() -> Result<(), Error> { + #[derive(MetricsComponent)] + struct HangingSlowStore; + + #[async_trait] + impl StoreDriver for HangingSlowStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + _keys: &[StoreKey<'_>], + _results: &mut [Option], + ) -> Result<(), Error> { + // Simulate a throttled/unreachable slow tier whose existence probe + // never returns within the caller's patience. Far longer than + // SLOW_PROBE_DEADLINE so the bound (not this future) decides. + tokio::time::sleep(Duration::from_secs(600)).await; + Ok(()) + } + + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + _reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + Ok(0) + } + + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + _writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + Ok(()) + } + + 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!(HangingSlowStore); + + let fast = Store::new(MemoryStore::new(&MemorySpec::default())); + let slow = Store::new(Arc::new(HangingSlowStore)); + let fast_slow = Arc::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, + slow_store_write_back: false, + }, + fast, + slow, + )); + + let digest = DigestInfo::try_new(VALID_HASH, 128).unwrap(); + + // The fast tier misses; the slow probe hangs. The call must still return + // (degraded to "missing") well before the slow store's 600s hang, bounded + // by SLOW_PROBE_DEADLINE. Give generous client-side slack over the 3s bound. + let result = tokio::time::timeout(Duration::from_secs(30), fast_slow.has(digest)).await; + let has = result.expect("has() must fail-fast, not block on the hung slow tier")?; + assert_eq!( + has, None, + "a fast-miss against a hung slow tier must degrade to missing, not wedge", + ); + + Ok(()) +} From 87a16916c610bba434814d9addefa1e524ccc531 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 14:46:56 -0400 Subject: [PATCH 09/23] migrate: fix remaining 1.7.3 test drift - running_actions_manager_test: add slow_store_write_back to FastSlowSpec literals - fast_slow deadline test: HangingSlowStore must be a braces struct (MetricsComponent derive rejects unit structs) Co-Authored-By: Claude Opus 4.8 --- nativelink-store/tests/fast_slow_store_test.rs | 4 ++-- nativelink-worker/tests/running_actions_manager_test.rs | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/nativelink-store/tests/fast_slow_store_test.rs b/nativelink-store/tests/fast_slow_store_test.rs index 08c29abd907..a2ffa55cef3 100644 --- a/nativelink-store/tests/fast_slow_store_test.rs +++ b/nativelink-store/tests/fast_slow_store_test.rs @@ -2388,7 +2388,7 @@ async fn write_back_decouples_update_from_slow_store() -> Result<(), Error> { #[nativelink_test] async fn has_fast_miss_fails_fast_when_slow_tier_hangs() -> Result<(), Error> { #[derive(MetricsComponent)] - struct HangingSlowStore; + struct HangingSlowStore {} #[async_trait] impl StoreDriver for HangingSlowStore { @@ -2450,7 +2450,7 @@ async fn has_fast_miss_fails_fast_when_slow_tier_hangs() -> Result<(), Error> { default_health_status_indicator!(HangingSlowStore); let fast = Store::new(MemoryStore::new(&MemorySpec::default())); - let slow = Store::new(Arc::new(HangingSlowStore)); + let slow = Store::new(Arc::new(HangingSlowStore {})); let fast_slow = Arc::new(FastSlowStore::new( &FastSlowSpec { fast: StoreSpec::Memory(MemorySpec::default()), diff --git a/nativelink-worker/tests/running_actions_manager_test.rs b/nativelink-worker/tests/running_actions_manager_test.rs index acf033bdbd4..80f355a4c22 100644 --- a/nativelink-worker/tests/running_actions_manager_test.rs +++ b/nativelink-worker/tests/running_actions_manager_test.rs @@ -3761,6 +3761,7 @@ exit 0 fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + slow_store_write_back: false, }, Store::new(fast_store), Store::new(slow_store.clone()), @@ -4003,6 +4004,7 @@ exit 0 fast_direction: StoreDirection::default(), slow_direction: StoreDirection::default(), bypass_dedup_threshold_bytes: 0, + slow_store_write_back: false, }, Store::new(fast_store), Store::new(gated.clone()), From bb85f38c175a69cecc8f5b6625fd1a5c1b0703cd Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 14:48:14 -0400 Subject: [PATCH 10/23] migrate: fix OCI config drift in service tests - fetch_server_test: FetchConfig gained fork oci field -> oci: None - oci_bridge_integration_test: OciFetchConfig gained self_registry -> None Co-Authored-By: Claude Opus 4.8 --- nativelink-service/tests/fetch_server_test.rs | 1 + nativelink-service/tests/oci_bridge_integration_test.rs | 1 + 2 files changed, 2 insertions(+) diff --git a/nativelink-service/tests/fetch_server_test.rs b/nativelink-service/tests/fetch_server_test.rs index 048e20a56a2..27d78d6e457 100644 --- a/nativelink-service/tests/fetch_server_test.rs +++ b/nativelink-service/tests/fetch_server_test.rs @@ -128,6 +128,7 @@ async fn test_fetch_blob_corrupt_store_entry_returns_error() -> Result<(), Statu instance_name: instance_name.clone(), config: FetchConfig { fetch_store: String::from("test_fetch_store"), + oci: None, }, }], &store_manager, diff --git a/nativelink-service/tests/oci_bridge_integration_test.rs b/nativelink-service/tests/oci_bridge_integration_test.rs index e12b0a51092..2fb6aea97f0 100644 --- a/nativelink-service/tests/oci_bridge_integration_test.rs +++ b/nativelink-service/tests/oci_bridge_integration_test.rs @@ -96,6 +96,7 @@ async fn oci_bridge_projects_to_producer_hint() -> Result<(), Error> { password: None, bearer_token: None, }], + self_registry: None, }), }, }], From 22ef1fe37463fbc75b3bee1843345a0bed7eb0ab Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 14:52:19 -0400 Subject: [PATCH 11/23] test(fast_slow): update has() tests to the fast-first contract The prior tests asserted the slow-only has() behavior (fast-only blobs reported absent; a concurrent has() blocking on the gated slow write). Under fast-first a fast-resident blob is reported present and a concurrent has() may resolve as soon as the write-through populates the fast tier. Updated the inverted assertions (ignore_value_in_fast_store -> fast_store_only_value_is_reported_by_has, mixed_key_sources fast-only key, the in-flight observer tests). Tidied core::time::Duration qualification and unused mut. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/fast_slow_store.rs | 2 +- .../tests/fast_slow_store_test.rs | 64 ++++++++++--------- 2 files changed, 34 insertions(+), 32 deletions(-) diff --git a/nativelink-store/src/fast_slow_store.rs b/nativelink-store/src/fast_slow_store.rs index 9cdac7fc086..87fbee1c1bb 100644 --- a/nativelink-store/src/fast_slow_store.rs +++ b/nativelink-store/src/fast_slow_store.rs @@ -530,7 +530,7 @@ impl StoreDriver for FastSlowStore { // probes, and without a short deadline every fast-miss stalls the // hot ring long enough to wedge the worker. On timeout we treat the // blob as missing; the write-back path still carries durability. - const SLOW_PROBE_DEADLINE: core::time::Duration = core::time::Duration::from_secs(3); + const SLOW_PROBE_DEADLINE: Duration = Duration::from_secs(3); let probe = self .slow_store .has_with_results(&slow_keys, &mut slow_results); diff --git a/nativelink-store/tests/fast_slow_store_test.rs b/nativelink-store/tests/fast_slow_store_test.rs index a2ffa55cef3..c5168353d48 100644 --- a/nativelink-store/tests/fast_slow_store_test.rs +++ b/nativelink-store/tests/fast_slow_store_test.rs @@ -446,7 +446,7 @@ async fn drop_on_eof_completes_store_futures() -> Result<(), Error> { } #[nativelink_test] -async fn ignore_value_in_fast_store() -> Result<(), Error> { +async fn fast_store_only_value_is_reported_by_has() -> Result<(), Error> { let fast_store = Store::new(MemoryStore::new(&MemorySpec::default())); let slow_store = Store::new(MemoryStore::new(&MemorySpec::default())); let fast_slow_store = Arc::new(FastSlowStore::new( @@ -465,9 +465,12 @@ async fn ignore_value_in_fast_store() -> Result<(), Error> { fast_store .update_oneshot(digest, make_random_data(100).into()) .await?; - assert!( - fast_slow_store.has(digest).await?.is_none(), - "Expected data to not exist in store" + // Fast-first: a blob resident only in the fast tier is reported present. + // (The historical slow-only probe is what wedged FindMissingBlobs.) + assert_eq!( + fast_slow_store.has(digest).await?, + Some(100), + "fast-only blob must be reported present under the fast-first contract", ); Ok(()) } @@ -1158,14 +1161,13 @@ async fn has_sees_in_flight_slow_writes() -> Result<(), Error> { // has_with_results above), so the only way this can be Some is via // the in-flight map wait. let observer_store = fast_slow.clone(); - let mut observer = tokio::spawn(async move { observer_store.has(digest).await }); + let observer = tokio::spawn(async move { observer_store.has(digest).await }); - // Prove that the observer is blocked waiting for the in-flight write. - // It should not resolve before the gate is released. - tokio::select! { - _ = &mut observer => panic!("Observer resolved before writer completed"), - () = tokio::time::sleep(Duration::from_millis(10)) => {} - } + // Under the fast-first contract the observer may resolve as soon as the + // fast tier receives the blob (the write-through populates fast before the + // gated slow write completes); it is no longer required to block on the + // slow store. Give the fast write a moment to land. + tokio::time::sleep(Duration::from_millis(10)).await; // Release the writer and confirm the in-flight tracker is cleaned up. gate_tx @@ -1185,12 +1187,13 @@ async fn has_sees_in_flight_slow_writes() -> Result<(), Error> { "Concurrent has() must wait for and see in-flight slow write", ); - // After completion the fast store still has the blob, but has() should - // return None since we never fallback to checking the fast store. + // After completion the blob is resident in the fast tier, so fast-first + // has() reports it present (the historical slow-only probe returned None + // here and is exactly the wedge this contract removes). assert_eq!( fast_slow.has(digest).await?, - None, - "Post-write has() should not see the blob via fast-store fallback", + Some(data.len() as u64), + "Post-write has() must see the fast-resident blob under fast-first", ); Ok(()) @@ -1443,14 +1446,13 @@ async fn dropping_update_future_cleans_up_in_flight_entry() -> Result<(), Error> .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 }); + let observer = tokio::spawn(async move { observer_store.has(digest).await }); - // Prove that the observer is blocked waiting for the in-flight write. - // It should not resolve before the gate is released. - tokio::select! { - _ = &mut observer => panic!("Observer resolved before writer completed"), - () = tokio::time::sleep(Duration::from_millis(10)) => {} - } + // Under the fast-first contract the observer may resolve as soon as the + // fast tier receives the blob (the write-through populates fast before the + // gated slow write completes); it is no longer required to block on the + // slow store. Give the fast write a moment to land. + tokio::time::sleep(Duration::from_millis(10)).await; // Cancel the writer. The guard's Drop should remove the entry. writer.abort(); @@ -1623,7 +1625,7 @@ async fn has_with_results_handles_mixed_key_sources() -> Result<(), Error> { ]; let observer_store = fast_slow.clone(); - let mut observer = tokio::spawn(async move { + let observer = tokio::spawn(async move { let mut results: [Option; 4] = [None; 4]; observer_store .as_store_driver_pin() @@ -1632,12 +1634,11 @@ async fn has_with_results_handles_mixed_key_sources() -> Result<(), Error> { Ok::<_, Error>(results) }); - // Prove that the observer is blocked waiting for the in-flight write. - // It should not resolve before the gate is released. - tokio::select! { - _ = &mut observer => panic!("Observer resolved before writer completed"), - () = tokio::time::sleep(Duration::from_millis(10)) => {} - } + // Under the fast-first contract the observer may resolve as soon as the + // fast tier receives the blob (the write-through populates fast before the + // gated slow write completes); it is no longer required to block on the + // slow store. Give the fast write a moment to land. + tokio::time::sleep(Duration::from_millis(10)).await; // Cleanup: release the gated writer. gate_tx @@ -1654,8 +1655,9 @@ async fn has_with_results_handles_mixed_key_sources() -> Result<(), Error> { assert_eq!(results[0], Some(slow_only_size), "slow-only key"); assert_eq!(results[1], Some(in_flight_size), "in-flight key"); assert_eq!( - results[2], None, - "fast-only key should be None because we do not check fast store" + results[2], + Some(fast_only_size), + "fast-only key must be reported present under the fast-first contract", ); assert_eq!(results[3], None, "missing key must stay None"); From 82693832dfb00d23a034e3d43137259c40fe7323 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 14:56:37 -0400 Subject: [PATCH 12/23] migrate: take upstream get_part single-shot heal over fork bounded-retry Upstream 1.7.3 redesigned FilesystemStore::get_part around a single-shot read with a remove_if(Arc::ptr_eq) self-heal, and added tests that pin that contract (stale_read_does_not_remove_replacement_upload et al). The fork had an older bounded-retry loop that re-fetched the entry and read the replacement, which fails those upstream tests. Taking upstream's version; the fork's map/disk-divergence test still passes against it. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/filesystem_store.rs | 69 ++++++++---------------- 1 file changed, 23 insertions(+), 46 deletions(-) diff --git a/nativelink-store/src/filesystem_store.rs b/nativelink-store/src/filesystem_store.rs index 26eda1cc91b..70a7e69c3b8 100644 --- a/nativelink-store/src/filesystem_store.rs +++ b/nativelink-store/src/filesystem_store.rs @@ -2247,9 +2247,6 @@ impl StoreDriver for FilesystemStore { offset: u64, length: Option, ) -> Result<(), Error> { - // Bounded retries for the transient map/disk race handled in the read - // path below (declared here so it precedes any statement in the block). - const MAX_MISS_RETRIES: usize = 3; if is_zero_digest(key.borrow()) { self.has(key.borrow()) .await @@ -2260,53 +2257,33 @@ impl StoreDriver for FilesystemStore { return Ok(()); } let owned_key = key.into_owned(); + let entry = self.evicting_map.get(&owned_key).await.ok_or_else(|| { + make_err!( + Code::NotFound, + "{} not found in filesystem store here", + owned_key.as_str() + ) + })?; let read_limit = length.unwrap_or(u64::MAX); - // A `get()` hit followed by an `open()` ENOENT is ambiguous. Under - // concurrent same-key writes an evicted generation's `unref()` renames - // the shared content path away while a fresh generation is still being - // emplaced, so the entry the map just handed us can momentarily have no - // file on disk even though a present replacement is landing. Retry a - // bounded number of times, re-fetching the (newer) entry each pass, so - // a transient race self-heals instead of surfacing a spurious miss. - // Only once the miss persists past the retries do we treat it as a - // genuine map/disk divergence. - let mut miss_attempts: usize = 0; - // Keep the resolved `entry` bound for the whole stream below. Dropping - // the last `Arc` is what triggers deletion of an evicted - // (renamed-to-temp) file, and an in-flight reader must be able to - // finish streaming from its already-open fd first — so the entry must - // outlive the read loop exactly as it did before the retry was added. - let (entry, mut temp_file) = loop { - let entry = self.evicting_map.get(&owned_key).await.ok_or_else(|| { - make_err!( - Code::NotFound, - "{} not found in filesystem store here", - owned_key.as_str() - ) - })?; - match entry.read_file_part(offset, read_limit).await { - Ok(file) => break (entry, file), - Err(err) if err.code == Code::NotFound && miss_attempts < MAX_MISS_RETRIES => { - // Transient map/disk race: yield so a racing writer's - // emplace can commit, then re-fetch and re-open. - miss_attempts += 1; - tokio::task::yield_now().await; - } - Err(err) if err.code == Code::NotFound => { - // Retries exhausted: the map says the file is present but - // `open()` keeps hitting ENOENT, so this is a genuine - // map/disk divergence. Remove the stale entry — a - // fast/slow caller then re-populates from the slow store — - // and surface the miss. Recoverable, hence a warn. + let mut temp_file = match entry.read_file_part(offset, read_limit).await { + Ok(file) => file, + Err(err) => { + // If the file is not found, we need to remove it from the eviction map. + if err.code == Code::NotFound { + // Map said the file was present but `open()` hit ENOENT. + // Self-heals: we remove the stale entry below and a + // fast/slow caller re-populates from the slow store, so + // this is a recoverable warn, not a fatal error. warn!( ?err, key = ?owned_key, "Filesystem store map/disk divergence: removing entry; reader will fall through to slow store", ); - self.evicting_map.remove(&owned_key).await; - return Err(err); + self.evicting_map + .remove_if(&owned_key, |map_entry| Arc::ptr_eq(map_entry, &entry)) + .await; } - Err(err) => return Err(err), + return Err(err); } }; @@ -2331,9 +2308,9 @@ impl StoreDriver for FilesystemStore { .send_eof() .err_tip(|| "Filed to send EOF in filesystem store get_part")?; - // Hold the file entry until the body is fully streamed (see the - // binding above): releasing it earlier could let an evicted file be - // deleted out from under this in-flight read. + // Hold the file entry until the body is fully streamed: releasing the + // last Arc earlier could let an evicted file be deleted out from under + // this in-flight read. drop(entry); Ok(()) } From f1e519e76e74810f9c81dd4c9660d6c9509b0a1c Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 00:48:30 -0400 Subject: [PATCH 13/23] redis: close two subscription-lifecycle windows that silently drop events MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Two confirmed bugs left the RedisStore keyspace/pattern subscriptions recorded as active while never actually established on the live connection. Both come with a deterministic repro test that fails on the old code and passes on the fix. 1. psubscribe lost across a concurrent reconnect (StandardRedisManager). `psubscribe` used to fetch a connection, then insert the pattern into the tracked `subscriptions` set. If a concurrent `reconnect` snapshotted+replayed that set on the new connection in between, the pattern was PSUBSCRIBEd only on the old, about-to-be-demoted connection and recorded as active — so every PMessage for it was dropped until an unrelated later reconnect replayed the now-complete set. Fix: order registration against the reconnect snapshot. `psubscribe` now inserts the pattern BEFORE fetching a connection, and `reconnect` recomputes the missing delta and subscribes it while holding the `connection_manager` write lock just before publishing the new handle. A racing `psubscribe` then either lands in the replayed snapshot or blocks on the write lock and subscribes on the already-published live connection — no window where a pattern is "active" but subscribed only on a stale connection. 2. transient init error permanently disabled remove callbacks. `register_remove_callback` returned Ok(()) and deferred the CONFIG GET + PSUBSCRIBE "__key*__:*" to a background one-shot `OnceCell::get_or_try_init`. A single transient Redis error (e.g. -LOADING during a restart/failover) was only logged; the OnceCell stayed empty and nothing retried, so eviction/expiry callbacks were silently dead forever (ExistenceCacheStore never invalidated → stale `exists` hits). Fix: retry the init with capped backoff until it succeeds, and demote the advisory notify-keyspace-events CONFIG GET check to non-fatal (log misconfig, subscribe anyway) so a transient read failure can no longer permanently disable the subscription. Both repro tests fail on the unfixed code and pass on the fix (verified both directions on the build fleet). Full `cargo test -p nativelink-store` green, clippy clean, nightly rustfmt applied. Provenance: adversarial review fleet. Co-Authored-By: Claude Fable 5 --- nativelink-store/src/redis_store.rs | 366 +++++++++++++++++++-- nativelink-store/tests/redis_store_test.rs | 82 ++++- 2 files changed, 417 insertions(+), 31 deletions(-) diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index ef150ba6ccc..f52b199cd6c 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -324,12 +324,31 @@ impl RedisManager for StandardRedisManager let guard = self.subscriptions.lock(); guard.iter().map(Clone::clone).collect::>() }; - for subscription in subscriptions { - connection_manager.psubscribe(&subscription).await?; + for subscription in &subscriptions { + connection_manager.psubscribe(subscription).await?; } - // Publish the new handle under a brief exclusive lock. + let mut replayed = subscriptions.into_iter().collect::>(); + // Publish the new handle under a brief exclusive lock. Before + // publishing, replay any patterns registered since the snapshot above: + // `psubscribe` records its pattern in `self.subscriptions` *before* it + // fetches a connection, so a pattern missing from this final delta was + // inserted while we hold the write lock, and its `get_connection` will + // block here and observe the new handle — it subscribes on the live + // connection itself. Either way there is no window where a pattern is + // recorded as active but subscribed only on the old connection. { let mut guard = self.connection_manager.write().await; + let missing = { + let subs = self.subscriptions.lock(); + subs.iter() + .filter(|s| !replayed.contains(*s)) + .cloned() + .collect::>() + }; + for subscription in missing { + connection_manager.psubscribe(&subscription).await?; + replayed.insert(subscription); + } *guard = (connection_manager.clone(), new_uuid); } info!(old = %uuid, new = %new_uuid, "StandardRedisManager re-resolved the Redis master"); @@ -342,10 +361,21 @@ impl RedisManager for StandardRedisManager async fn psubscribe(&self, pattern: &str) -> Result<(), Error> { debug!(pattern, "new psubscribe"); - let mut connection = self.get_connection().await?.0; + // Record the pattern BEFORE fetching a connection. A concurrent + // `reconnect` either sees the pattern in the snapshot it replays under + // its pre-publish write lock (subscribing it on the new connection), + // or it published first — in which case `get_connection` below returns + // the new live connection and we subscribe on it ourselves. Inserting + // after fetching the connection (the previous order) left a window + // where the pattern was recorded as active but PSUBSCRIBEd only on the + // old, about-to-be-replaced connection. let new_subscription = self.subscriptions.lock().insert(String::from(pattern)); if new_subscription { - let result = connection.psubscribe(pattern).await; + let result = async { + let mut connection = self.get_connection().await?.0; + connection.psubscribe(pattern).await.map_err(Error::from) + } + .await; if result.is_err() { self.subscriptions.lock().remove(pattern); } @@ -1471,35 +1501,68 @@ where fn register_remove_callback(self: Arc, callback: RemoveCallback) -> Result<(), Error> { debug!(?callback, "New callback"); - let local_self = self.clone(); background_spawn!("remove_callback_subscribe", async move { self.remove_callbacks.lock().await.push(callback); - if let Err(err) = local_self.clone().has_remove_callback_subscribe - .get_or_try_init(|| async move { - let mut client = local_self.get_client().await?; - let cfg = redis::cmd("CONFIG").arg("GET").arg("notify-keyspace-events").to_owned().query_async::>(&mut client.connection_manager).await.map_err(|e| Error::from(e).append("Parsing notify-keyspace-events"))?; - if cfg.len() != 1 { - warn!(?cfg, "Got multiple items for CONFIG GET, expected one"); - return Err(make_input_err!("Got multiple items for CONFIG GET, expected one")); - } - let events_cfg = &cfg.first().ok_or_else(|| make_err!(Code::InvalidArgument, "Only one item"))?.1; - if events_cfg.is_empty() { - error!("notify-keyspace-events not enabled for Redis, will fail to get remove callbacks"); - } else if !events_cfg.contains('K') { - error!(notify_keyspace_events=events_cfg, "notify-keyspace-events does not contain 'K' so won't get keyspace events we need for eviction events"); - } else if !events_cfg.contains('A') { - error!(notify_keyspace_events=events_cfg, "notify-keyspace-events does not contain 'A' so we won't get eviction events"); + // The subscription init is retried until it succeeds: this + // function already told its caller Ok(()) (e.g. + // ExistenceCacheStore, which relies on eviction callbacks for + // invalidation), so a transient Redis error here (e.g. -LOADING + // during a restart/failover) must not permanently disable remove + // callbacks. Previously a one-shot `get_or_try_init` failure was + // only logged, the OnceCell stayed empty, and nothing ever + // retried, so the `__key*__:*` subscription was silently never + // established. + let mut retry_delay = Duration::from_millis(50); + loop { + let local_self = self.clone(); + let init_result = self + .has_remove_callback_subscribe + .get_or_try_init(|| async move { + // The notify-keyspace-events check is advisory: we log + // misconfiguration loudly, but a failure to *read* the + // config must not block the subscription itself. + let cfg_result = async { + let mut client = local_self.get_client().await?; + redis::cmd("CONFIG").arg("GET").arg("notify-keyspace-events").to_owned().query_async::>(&mut client.connection_manager).await.map_err(|e| Error::from(e).append("Parsing notify-keyspace-events")) + }.await; + match &cfg_result { + Ok(cfg) => { + if cfg.len() != 1 { + warn!(?cfg, "Got multiple items for CONFIG GET, expected one"); + } else if let Some((_, events_cfg)) = cfg.first() { + if events_cfg.is_empty() { + error!("notify-keyspace-events not enabled for Redis, will fail to get remove callbacks"); + } else if !events_cfg.contains('K') { + error!(notify_keyspace_events=events_cfg, "notify-keyspace-events does not contain 'K' so won't get keyspace events we need for eviction events"); + } else if !events_cfg.contains('A') { + error!(notify_keyspace_events=events_cfg, "notify-keyspace-events does not contain 'A' so we won't get eviction events"); + } + // FIXME: Redis events spec appears unreliable, so we subscribe anyways + // It should just need Ke as per https://redis.io/docs/latest/develop/pubsub/keyspace-notifications/ + // but I'm yet to get reliable eviction events out of that + info!(notify_keyspace_events=events_cfg, "Attempting to subscribe to eviction events"); + } + } + Err(err) => { + error!(?err, "Error while trying to initialise remove_callback_subscribe: could not verify notify-keyspace-events; subscribing anyway"); + } + } + local_self.connection_manager.psubscribe("__key*__:*").await?; + Ok::<(), Error>(()) + }) + .await; + match init_result { + Ok(()) => break, + Err(err) => { + error!( + ?err, + "Error while trying to initialise remove_callback_subscribe; will retry" + ); + sleep(retry_delay).await; + retry_delay = (retry_delay * 2).min(Duration::from_secs(5)); } - // FIXME: Redis events spec appears unreliable, so we subscribe anyways - // It should just need Ke as per https://redis.io/docs/latest/develop/pubsub/keyspace-notifications/ - // but I'm yet to get reliable eviction events out of that - info!(notify_keyspace_events=events_cfg, "Attempting to subscribe to eviction events"); - self.connection_manager.psubscribe("__key*__:*").await?; - Ok::<(), Error>(()) - }) - .await { - error!(?err, "Error while trying to initialise remove_callback_subscribe"); } + } }); Ok(()) } @@ -2597,3 +2660,246 @@ where )) } } + +#[cfg(test)] +// This module hand-rolls a minimal RESP3 server plus a dedicated +// current-thread runtime so the psubscribe/reconnect interleaving is exact and +// timing-independent; that legitimately needs the otherwise-disallowed tokio +// runtime/spawn primitives. +#[expect( + clippy::disallowed_methods, + reason = "deterministic in-test RESP server needs a hand-built runtime and raw spawns" +)] +#[expect( + clippy::doc_markdown, + reason = "choreography docs reference bare Redis/RESP identifiers (PSUBSCRIBE, Sha1) by name" +)] +#[expect( + trivial_casts, + reason = "explicit Pin> coercion for the connect closure's return type" +)] +mod psubscribe_reconnect_race_tests { + use core::time::Duration; + use std::collections::HashSet; + use std::sync::Arc; + + use redis::aio::{ConnectionManager, ConnectionManagerConfig}; + use tokio::io::{AsyncReadExt, AsyncWriteExt}; + use tokio::net::{TcpListener, TcpStream}; + use tokio::sync::Notify; + + use super::{RedisManager, StandardRedisManager}; + + /// Reads one CRLF-terminated line from the stream. + async fn read_line(stream: &mut TcpStream) -> Option { + let mut line = Vec::new(); + let mut byte = [0u8; 1]; + loop { + if stream.read_exact(&mut byte).await.is_err() { + return None; + } + if byte[0] == b'\n' { + break; + } + if byte[0] != b'\r' { + line.push(byte[0]); + } + } + Some(String::from_utf8_lossy(&line).into_owned()) + } + + /// Parses one RESP command (array of bulk strings) from the client. + async fn parse_command(stream: &mut TcpStream) -> Option> { + let header = read_line(stream).await?; + if !header.starts_with('*') { + return Some(vec![header]); + } + let count: usize = header[1..].parse().ok()?; + let mut parts = Vec::with_capacity(count); + for _ in 0..count { + let len_line = read_line(stream).await?; + let len: usize = len_line.strip_prefix('$')?.parse().ok()?; + let mut data = vec![0u8; len + 2]; // include trailing CRLF + stream.read_exact(&mut data).await.ok()?; + parts.push(String::from_utf8_lossy(&data[..len]).into_owned()); + } + Some(parts) + } + + struct FakeServerState { + /// Sha1 to return for SCRIPT LOAD (must match the client's own hash). + script_hash: String, + /// Patterns each connection (by accept order) was PSUBSCRIBEd to. + patterns_per_conn: parking_lot::Mutex>>, + /// Signalled when conn #1 (the reconnect connection) receives its + /// first PSUBSCRIBE (the replay of the pre-existing pattern). + replay_psubscribe_arrived: Notify, + /// The server holds conn #1's first PSUBSCRIBE reply until this is + /// notified, keeping reconnect() parked inside its replay loop. + release_replay_reply: Notify, + } + + async fn handle_conn(mut stream: TcpStream, conn_idx: usize, state: Arc) { + let mut first_psubscribe_on_this_conn = true; + while let Some(cmd) = parse_command(&mut stream).await { + let name = cmd + .first() + .map(|c| c.to_ascii_uppercase()) + .unwrap_or_default(); + match name.as_str() { + "HELLO" => { + stream + .write_all( + b"%3\r\n$6\r\nserver\r\n$5\r\nredis\r\n$7\r\nversion\r\n$5\r\n7.4.0\r\n$5\r\nproto\r\n:3\r\n", + ) + .await + .unwrap(); + } + "SCRIPT" => { + let reply = + format!("${}\r\n{}\r\n", state.script_hash.len(), state.script_hash); + stream.write_all(reply.as_bytes()).await.unwrap(); + } + "PSUBSCRIBE" => { + let pattern = cmd.get(1).cloned().unwrap_or_default(); + state.patterns_per_conn.lock()[conn_idx].push(pattern.clone()); + if conn_idx == 1 && first_psubscribe_on_this_conn { + first_psubscribe_on_this_conn = false; + // Tell the test that reconnect() is now parked awaiting + // this reply inside its replay loop, then hold the + // reply until the test releases it. + state.replay_psubscribe_arrived.notify_one(); + state.release_replay_reply.notified().await; + } + let reply = format!( + ">3\r\n$10\r\npsubscribe\r\n${}\r\n{}\r\n:1\r\n", + pattern.len(), + pattern + ); + stream.write_all(reply.as_bytes()).await.unwrap(); + } + _ => { + stream.write_all(b"+OK\r\n").await.unwrap(); + } + } + } + } + + /// Deterministic repro: a `psubscribe` that interleaves with a `reconnect` + /// whose subscription snapshot was taken before the insert lands only on + /// the old (about-to-be-replaced) connection. The new live connection is + /// never subscribed, yet the pattern is recorded as active in + /// `self.subscriptions`, so notifications are silently lost until an + /// unrelated future reconnect replays the set. + /// + /// Choreography (all real await points, no timing races): + /// 1. Manager connects (conn #0) and psubscribes "pre". + /// 2. Task B calls reconnect(): connects conn #1, snapshots + /// subscriptions = {"pre"}, and parks in the replay loop because the + /// fake server withholds conn #1's PSUBSCRIBE "pre" reply. + /// 3. The test (task A) calls psubscribe("pat"): get_connection still + /// returns conn #0 (reconnect has not published yet), the pattern is + /// inserted into `subscriptions`, and PSUBSCRIBE succeeds on conn #0. + /// 4. The server releases the reply; reconnect publishes conn #1. + /// 5. Conn #1 — now the live connection — was never subscribed to "pat". + #[test] + fn psubscribe_lost_across_reconnect() { + let runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .unwrap(); + runtime.block_on(async { + let state = Arc::new(FakeServerState { + script_hash: redis::Script::new(super::LUA_VERSION_SET_SCRIPT) + .get_hash() + .to_string(), + patterns_per_conn: parking_lot::Mutex::new(Vec::new()), + replay_psubscribe_arrived: Notify::new(), + release_replay_reply: Notify::new(), + }); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + { + let state = state.clone(); + tokio::spawn(async move { + let mut conn_idx = 0usize; + loop { + let (stream, _) = listener.accept().await.unwrap(); + state.patterns_per_conn.lock().push(Vec::new()); + tokio::spawn(handle_conn(stream, conn_idx, state.clone())); + conn_idx += 1; + } + }); + } + + // Keep the push receiver alive for the whole test. + let (push_tx, _push_rx) = tokio::sync::mpsc::unbounded_channel(); + let url = format!("redis://127.0.0.1:{port}/?protocol=resp3"); + let connect_func = Box::new(move || { + let url = url.clone(); + let push_tx = push_tx.clone(); + Box::pin(async move { + let client = redis::Client::open(url.as_str()).unwrap(); + let config = ConnectionManagerConfig::new() + .set_number_of_retries(0) + .set_connection_timeout(Some(Duration::from_secs(5))) + .set_response_timeout(Some(Duration::from_secs(5))) + .set_push_sender(push_tx); + ConnectionManager::new_with_config(client, config) + .await + .map_err(Into::into) + }) as core::pin::Pin> + }); + + let manager = Arc::new(StandardRedisManager::new(connect_func).await.unwrap()); + + // Step 1: a pre-existing subscription so reconnect's replay loop + // has an await point we can park it on. + manager.psubscribe("pre").await.unwrap(); + let (_conn, uuid0) = manager.get_connection().await.unwrap(); + + // Step 2: start reconnect; it will connect conn #1, snapshot + // subscriptions (= {"pre"}), and park replaying "pre". + let reconnect_task = { + let manager = manager.clone(); + tokio::spawn(async move { manager.reconnect(uuid0).await }) + }; + tokio::time::timeout( + Duration::from_secs(5), + state.replay_psubscribe_arrived.notified(), + ) + .await + .expect("reconnect never reached its subscription replay loop"); + + // Step 3: concurrent psubscribe while reconnect is parked + // post-snapshot, pre-publish. It succeeds — on the old connection. + manager.psubscribe("pat").await.unwrap(); + + // Step 4: let reconnect finish and publish conn #1. + state.release_replay_reply.notify_one(); + let (_new_conn, new_uuid) = reconnect_task.await.unwrap().unwrap(); + assert_ne!( + uuid0, new_uuid, + "reconnect must have published a new connection" + ); + + // The manager believes "pat" is an active subscription... + let recorded: HashSet = manager.subscriptions.lock().clone(); + assert!( + recorded.contains("pat"), + "sanity: pattern recorded as active" + ); + + // ...but the live connection (conn #1) was never subscribed to it. + let conn1_patterns = state.patterns_per_conn.lock()[1].clone(); + assert!( + conn1_patterns.contains(&"pat".to_string()), + "RACE REPRODUCED: pattern \"pat\" is recorded as an active \ + subscription but the live post-reconnect connection was only \ + subscribed to {conn1_patterns:?}; keyspace notifications for \ + \"pat\" are silently lost until the next reconnect", + ); + }); + } +} diff --git a/nativelink-store/tests/redis_store_test.rs b/nativelink-store/tests/redis_store_test.rs index b668633a778..cd898f28cde 100644 --- a/nativelink-store/tests/redis_store_test.rs +++ b/nativelink-store/tests/redis_store_test.rs @@ -24,7 +24,7 @@ use nativelink_config::stores::{ExistenceCacheSpec, RedisMode, RedisSpec, StoreS use nativelink_error::{Code, Error, ErrorContext, ResultExt, make_err}; use nativelink_macro::nativelink_test; use nativelink_redis_tester::{ - ReadOnlyRedis, SubscriptionManagerNotify, add_lua_script, add_to_response, + ReadOnlyRedis, SubscriptionManagerNotify, add_lua_script, add_to_response, add_to_response_raw, fake_redis_sentinel_master_stream, fake_redis_sentinel_stream, fake_redis_stream, make_fake_redis_with_responses, }; @@ -2717,6 +2717,86 @@ async fn callback_for_eviction_good() -> Result<(), Error> { Ok(()) } +// Repro for a fail-silent one-shot init bug: `register_remove_callback` +// returns Ok(()) immediately and performs the CONFIG GET + PSUBSCRIBE lazily +// inside a `background_spawn` via `OnceCell::get_or_try_init`. If that init +// hits a single transient Redis error (here: `-LOADING`), the error is only +// logged, the OnceCell stays uninitialized, and nothing ever retries — so the +// `__key*__:*` subscription is never established and evicted/expired keys +// never fire remove callbacks, silently and permanently. +#[nativelink_test] +async fn transient_init_error_permanently_disables_remove_callbacks() -> Result<(), Error> { + let redis_span = info_span!("redis"); + + let mut responses = add_lua_version_script(fake_redis_stream()); + // The one-time init query fails with a transient, retryable server error + // (Redis restarting/failing over). The server is otherwise healthy and + // would happily accept a PSUBSCRIBE. + add_to_response_raw( + &mut responses, + redis::cmd("CONFIG") + .arg("GET") + .arg("notify-keyspace-events"), + "-LOADING Redis is loading the dataset in memory\r\n".to_string(), + ); + add_to_response( + &mut responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + vec![Value::Nil], + ); + let redis_port = make_fake_redis_with_responses(responses) + .instrument(redis_span) + .await; + let spec = RedisSpec { + addresses: vec![format!("redis://127.0.0.1:{redis_port}/")], + mode: RedisMode::Standard, + ..Default::default() + }; + let mut raw_store = + Arc::into_inner(RedisStore::new_standard(spec).await.expect("Working spec")).unwrap(); + raw_store.replace_temp_name_generator(mock_uuid_generator); + let store = Arc::new(raw_store); + // The caller (e.g. ExistenceCacheStore at construction, its only call + // site) is told Ok(()) and believes eviction tracking is live. + store.register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + + // Deterministic sync point: wait until the background init task has hit + // the transient error and logged it. + timeout(Duration::from_secs(5), async { + loop { + if logs_contain("Error while trying to initialise remove_callback_subscribe") { + break; + } + sleep(Duration::from_millis(20)).await; + } + }) + .await + .expect("background init never ran/logged its error"); + + // Since register_remove_callback reported success, the keyspace-event + // subscription must (eventually) be established — otherwise remove + // callbacks are silently dead forever. On current code there is no retry + // path at all (the OnceCell is left empty and only another + // register_remove_callback call could re-drive it), so this never happens. + let subscribed = timeout(Duration::from_secs(2), async { + loop { + if logs_contain("new psubscribe complete pattern=\"__key*__:*\"") { + break; + } + sleep(Duration::from_millis(50)).await; + } + }) + .await + .is_ok(); + assert!( + subscribed, + "register_remove_callback returned Ok(()) but a single transient CONFIG GET \ + error permanently disabled the __key*__:* subscription; evicted/expired keys \ + will never invoke remove callbacks (stale ExistenceCacheStore entries)" + ); + Ok(()) +} + #[nativelink_test] async fn send_eviction_to_subscription_channel() -> Result<(), Error> { let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); From 6465dcc64a63c8823e84a5a10801aebb7bdd052e Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 20:36:58 -0400 Subject: [PATCH 14/23] existence_cache: drain removals queued before a cancelled last holder (C20/N6 drop residual) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The RAII PauseGuard's Drop decremented the pause count (fixing N6's permanently-wedged pause) but discarded the pending-removal vector that `release()` returns when the last holder goes away. On the evict-before-cancel ordering — evict cached A while upload B holds the sole pause, then cancel B — A's queued invalidation was dropped on the floor and the existence cache kept reporting A forever even though the inner store had evicted it. Drop now drains those queued removals. Since Drop cannot await and the removal is async, it spawns a detached background task (via a weak self-reference) to run drain_pending_removals. New regression `cancelled_update_drains_removals_queued_before_cancel` uses the evict-*before*-cancel ordering and fails on the prior Drop (removals discarded); the existing N6 test `cancelled_update_does_not_leave_callbacks_paused` (evict-after-cancel) is kept. Also closes the C13 has-path residual: inner_has_with_results now pauses remove callbacks around the inner `has` + the positive-result insert, so an eviction delivered mid-`has` lands after the insert instead of being a no-op before it. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/existence_cache_store.rs | 186 ++++- .../tests/existence_store_test.rs | 748 +++++++++++++++++- 2 files changed, 906 insertions(+), 28 deletions(-) diff --git a/nativelink-store/src/existence_cache_store.rs b/nativelink-store/src/existence_cache_store.rs index ffd2625f180..370168b7e3d 100644 --- a/nativelink-store/src/existence_cache_store.rs +++ b/nativelink-store/src/existence_cache_store.rs @@ -23,6 +23,7 @@ use futures::stream::FuturesUnordered; use nativelink_config::stores::{EvictionPolicy, ExistenceCacheSpec}; use nativelink_error::{Error, ResultExt, error_if}; use nativelink_metric::MetricsComponent; +use nativelink_util::background_spawn; use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::common::DigestInfo; use nativelink_util::evicting_map::{EvictingMap, LenEntry}; @@ -49,17 +50,108 @@ impl LenEntry for ExistenceItem { } } +/// Tracks in-flight operations that must defer remove callbacks. +/// +/// `pause_count` is the number of concurrent operations (e.g. `update()` +/// calls) currently holding the pause. While it is non-zero, remove +/// callbacks from the inner store are queued in `pending` instead of being +/// applied immediately. Only when the *last* holder releases the pause +/// (count drops to zero) is the queue drained, after every holder has +/// inserted its key into the existence cache. This guarantees that an +/// eviction notification for a key delivered during that key's own +/// operation is applied *after* the insert, so the cache cannot be left +/// claiming existence of a blob the inner store dropped. +#[derive(Debug, Default)] +struct PauseState { + pause_count: usize, + pending: Vec>, +} + +impl PauseState { + /// Release one hold. Returns the queued removals to drain iff this was the + /// last holder; otherwise returns an empty vec and leaves the queue for a + /// later holder to drain. + fn release(&mut self) -> Vec> { + self.pause_count = self + .pause_count + .checked_sub(1) + .expect("PauseState::release called without a matching pause"); + if self.pause_count == 0 { + core::mem::take(&mut self.pending) + } else { + Vec::new() + } + } +} + +/// RAII hold on the remove-callback pause. +/// +/// Dropping the guard always decrements the pause count, so a cancelled +/// operation cannot leave invalidation globally suspended (finding N6). The +/// normal path calls [`PauseGuard::resume`] instead. If the guard is dropped +/// while it is the last holder with a non-empty queue (only reachable on +/// cancellation), those removals must still be applied — otherwise the cache +/// keeps claiming a blob exists after the inner store evicted it (C20/N6 drop +/// residual). Since `Drop` cannot await and the removal is async, `Drop` spawns +/// a detached task (via the weak self-reference) to drain them. +struct PauseGuard<'a, I: InstantWrapper> { + store: Option<&'a ExistenceCacheStore>, +} + +impl PauseGuard<'_, I> { + /// Release the hold and drain queued removals if this was the last holder. + async fn resume(mut self) { + // Take the store so `Drop` becomes a no-op; we handle release here. + let Some(store) = self.store.take() else { + return; + }; + let keys = store.pause_remove_callbacks.lock().release(); + store.drain_pending_removals(keys).await; + } +} + +impl Drop for PauseGuard<'_, I> { + fn drop(&mut self) { + let Some(store) = self.store.take() else { + return; + }; + // Cancellation path: release the count so the pause mechanism is never + // left permanently engaged (finding N6). + let keys = store.pause_remove_callbacks.lock().release(); + if keys.is_empty() { + return; + } + // We were the last holder and there are removals that were queued + // *before* this operation was cancelled (C20/N6 drop residual). These + // must still be applied, otherwise the cache can keep claiming a blob + // exists after the inner store evicted it. `Drop` cannot await, and the + // removal is async, so spawn a detached task that drains them. The + // weak self-reference keeps this `'static` without extending the + // store's lifetime past its last strong owner. + let weak = store.self_ref.lock().clone(); + if let Some(store) = weak.upgrade() { + background_spawn!("existence_cache_drain_pending_removals", async move { + store.drain_pending_removals(keys).await; + }); + } + } +} + #[derive(Debug, MetricsComponent)] pub struct ExistenceCacheStore { #[metric(group = "inner_store")] inner_store: Store, existence_cache: EvictingMap, - // We need to pause them temporarily when inserting into the inner store - // as if it immediately expires them, we should only apply the remove callbacks - // afterwards. If this is None, we're not pausing; if it's Some it's the location to - // store them in temporarily - pause_remove_callbacks: Mutex>>>, + // We need to pause remove callbacks temporarily while operating on the + // inner store; see `PauseState`. + pause_remove_callbacks: Mutex, + + // A weak self-reference, set once at construction. The cancellation path in + // `PauseGuard::drop` needs an owned (`'static`) handle to the store so it + // can spawn the async drain of any removals queued before cancellation; a + // borrowed `&self` cannot outlive the `Drop` call. See `PauseGuard::drop`. + self_ref: Mutex>, } impl ExistenceCacheStore { @@ -96,9 +188,11 @@ impl RemoveItemCallback for ExistenceCacheCallback { ) -> Pin + Send + 'a>> { let cache = self.cache.upgrade(); if let Some(local_cache) = cache { - if let Some(callbacks) = local_cache.pause_remove_callbacks.lock().as_mut() { - callbacks.push(store_key.into_owned()); + let mut state = local_cache.pause_remove_callbacks.lock(); + if state.pause_count > 0 { + state.pending.push(store_key.into_owned()); } else { + drop(state); let store_key = store_key.into_owned(); return Box::pin(async move { local_cache.callback(store_key).await; @@ -122,8 +216,10 @@ impl ExistenceCacheStore { let existence_cache_store = Arc::new(Self { inner_store, existence_cache: EvictingMap::new(eviction_policy, anchor_time), - pause_remove_callbacks: Mutex::new(None), + pause_remove_callbacks: Mutex::new(PauseState::default()), + self_ref: Mutex::new(Weak::new()), }); + *existence_cache_store.self_ref.lock() = Arc::downgrade(&existence_cache_store); let other_ref = Arc::downgrade(&existence_cache_store); existence_cache_store .inner_store @@ -144,6 +240,30 @@ impl ExistenceCacheStore { self.existence_cache.remove(digest).await; } + /// Begin deferring remove callbacks for the duration of an operation and + /// return an RAII guard that releases the hold when dropped. + /// + /// The guard's `Drop` decrements the pause count unconditionally, so a + /// cancelled operation (e.g. an `update`/`get_part` future dropped mid-await + /// by a timeout) can never leave the pause count permanently elevated and + /// disable invalidation globally. On the normal path, call + /// [`PauseGuard::resume`] to release the hold *and* drain any queued + /// removals once the last holder is gone. + fn pause_remove_callbacks(&self) -> PauseGuard<'_, I> { + self.pause_remove_callbacks.lock().pause_count += 1; + PauseGuard { store: Some(self) } + } + + /// Drain the queued removals when the last holder has released. Called by + /// [`PauseGuard::resume`]; the count has already been decremented there. + async fn drain_pending_removals(&self, keys: Vec>) { + let mut callbacks: FuturesUnordered<_> = keys + .into_iter() + .map(|store_key| self.callback(store_key)) + .collect(); + while callbacks.next().await.is_some() {} + } + async fn inner_has_with_results( self: Pin<&Self>, keys: &[DigestInfo], @@ -164,12 +284,25 @@ impl ExistenceCacheStore { return Ok(()); } + // Pause remove callbacks over the inner `has` + the insert below + // (finding C13). The inner store may evict one of `not_cached_keys` + // between answering `has` positively and our insert; without the pause + // that eviction's callback would run before the insert and be a no-op, + // leaving the cache claiming existence of a blob the inner store + // dropped. Queuing it until after the insert lets the drain remove it. + let pause = self.pause_remove_callbacks(); + // Now query only the items not found in the cache. let mut inner_results = vec![None; not_cached_keys.len()]; - self.inner_store + let has_result = self + .inner_store .has_with_results(¬_cached_keys, &mut inner_results) .await - .err_tip(|| "In ExistenceCacheStore::inner_has_with_results")?; + .err_tip(|| "In ExistenceCacheStore::inner_has_with_results"); + if let Err(err) = has_result { + pause.resume().await; + return Err(err); + } // Insert found from previous query into our cache. { @@ -185,6 +318,10 @@ impl ExistenceCacheStore { drop(self.existence_cache.insert_many(inserts).await); } + // Insert is done; release the pause and drain any eviction queued + // during the inner `has` so it lands after our insert (C13). + pause.resume().await; + // Merge the results from the cache and the query. { let mut inner_results_iter = inner_results.into_iter(); @@ -251,12 +388,7 @@ impl StoreDriver for ExistenceCacheStore { .err_tip(|| "In ExistenceCacheStore::update")?; return Ok(size); } - { - let mut locked_callbacks = self.pause_remove_callbacks.lock(); - if locked_callbacks.is_none() { - locked_callbacks.replace(vec![]); - } - } + let pause = self.pause_remove_callbacks(); trace!(?digest, "Inserting into inner cache"); let result = self.inner_store.update(digest, reader, size_info).await; if let Ok(size) = &result { @@ -266,16 +398,11 @@ impl StoreDriver for ExistenceCacheStore { .insert(digest, ExistenceItem(*size)) .await; } - { - let maybe_keys = self.pause_remove_callbacks.lock().take(); - if let Some(keys) = maybe_keys { - let mut callbacks: FuturesUnordered<_> = keys - .into_iter() - .map(|store_key| self.callback(store_key)) - .collect(); - while callbacks.next().await.is_some() {} - } - } + // Note: the insert above happens *before* releasing the pause, so a + // removal for `digest` queued during the inner update is applied + // after the insert and cannot be lost. If this future is cancelled + // before here, `pause`'s Drop still releases the hold. + pause.resume().await; result } @@ -287,6 +414,12 @@ impl StoreDriver for ExistenceCacheStore { length: Option, ) -> Result<(), Error> { let digest = key.into_digest(); + // Pause remove callbacks over the read + insert: the inner store may + // evict `digest` while a read of it is still in flight (e.g. a + // filesystem store keeps streaming from an open file handle). The + // eviction is queued and replayed after our insert, so the cache + // cannot resurrect an entry the inner store dropped mid-read. + let pause = self.pause_remove_callbacks(); let result = self .inner_store .get_part(digest, writer, offset, length) @@ -297,6 +430,7 @@ impl StoreDriver for ExistenceCacheStore { .insert(digest, ExistenceItem(digest.size_bytes())) .await; } + pause.resume().await; result } diff --git a/nativelink-store/tests/existence_store_test.rs b/nativelink-store/tests/existence_store_test.rs index 5bba222568e..5d4b4dced88 100644 --- a/nativelink-store/tests/existence_store_test.rs +++ b/nativelink-store/tests/existence_store_test.rs @@ -12,20 +12,36 @@ // See the License for the specific language governing permissions and // limitations under the License. +use core::pin::Pin; +use core::sync::atomic::{AtomicBool, Ordering}; use core::time::Duration; +use std::borrow::Cow; +use std::collections::HashSet; +use std::sync::Arc; +use async_trait::async_trait; +use bytes::Bytes; use mock_instant::thread_local::MockClock; use nativelink_config::stores::{ EvictionPolicy, ExistenceCacheSpec, MemorySpec, NoopSpec, StoreSpec, }; -use nativelink_error::{Error, ResultExt}; +use nativelink_error::{Code, Error, ResultExt, make_err}; use nativelink_macro::nativelink_test; +use nativelink_metric::{ + MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, +}; use nativelink_store::existence_cache_store::ExistenceCacheStore; use nativelink_store::memory_store::MemoryStore; +use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; use nativelink_util::common::DigestInfo; +use nativelink_util::health_utils::{HealthStatus, HealthStatusIndicator}; use nativelink_util::instant_wrapper::MockInstantWrapped; -use nativelink_util::store_trait::{Store, StoreLike}; +use nativelink_util::store_trait::{ + RemoveCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, +}; +use parking_lot::Mutex; use pretty_assertions::assert_eq; +use tokio::sync::Notify; const VALID_HASH1: &str = "0123456789abcdef000000000000000000010000000000000123456789abcdef"; @@ -196,3 +212,731 @@ async fn copes_with_dropped_items() -> Result<(), Error> { Ok(()) } + +// Reproduces the pause_remove_callbacks race: the pause slot is a single +// non-refcounted Option shared by all concurrent update() calls. The first +// update to finish take()s it, ending the pause for every other in-flight +// update. A remove callback (eviction under pressure) firing inside the +// second update's inner_store.update await then runs immediately, BEFORE +// that update inserts into the existence cache, leaving a permanent stale +// exists-entry for a blob the inner store does not hold. Subsequent uploads +// of that key are drained and silently discarded. +#[nativelink_test] +async fn concurrent_update_pause_refcount_race_test() -> Result<(), Error> { + use core::pin::Pin; + use std::collections::HashMap; + use std::sync::Arc; + + use async_trait::async_trait; + use nativelink_metric::MetricsComponent; + use nativelink_util::buf_channel::{DropCloserReadHalf, DropCloserWriteHalf}; + use nativelink_util::health_utils::{HealthStatusIndicator, default_health_status_indicator}; + use nativelink_util::store_trait::{RemoveCallback, StoreDriver, StoreKey, UploadSizeInfo}; + use tokio::sync::Notify; + + const VALID_HASH2: &str = "0223456789abcdef000000000000000000010000000000000123456789abcdef"; + + type GateMap = HashMap, Arc)>; + + #[derive(MetricsComponent)] + struct GatedStore { + inner: Store, + mem: Arc, + // Captured ExistenceCacheCallback registered by ExistenceCacheStore. + callback: Mutex>, + // (entered, release) per digest. Each gate fires at most once; once + // consumed the digest is removed so a later re-upload of the same key + // passes straight through to the inner store. + gates: Mutex, + // Key whose insertion should trigger an inline eviction (i.e. the + // inner store evicts the just-inserted entry under cache pressure + // and fires the remove callback while still inside update()). + evict_key: DigestInfo, + } + + impl core::fmt::Debug for GatedStore { + fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result { + f.write_str("GatedStore") + } + } + + #[async_trait] + impl StoreDriver for GatedStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + digests: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + self.inner.has_with_results(digests, results).await + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + let digest = key.borrow().into_digest(); + let gate = self.gates.lock().remove(&digest); + let was_gated = gate.is_some(); + if let Some((entered, release)) = gate { + entered.notify_one(); + release.notified().await; + } + let result = self.inner.update(key, reader, size_info).await; + // Only the first (gated) upload of evict_key simulates the + // eviction-under-pressure; a later re-upload must persist. + if was_gated && digest == self.evict_key { + // Simulate eviction-under-pressure of the just-inserted + // entry: the evicting map removes it and fires the remove + // callback inline, inside this awaited update (see + // filesystem_store.rs "Insertion can unref this entry + // immediately under cache pressure"). + self.mem.remove_entry(digest.into()).await; + let callback = self.callback.lock().clone(); + if let Some(callback) = callback { + callback.callback(digest.into()).await; + } + } + result + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + self.inner.get_part(key, writer, offset, length).await + } + + fn inner_store(&self, _digest: 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> { + *self.callback.lock() = Some(callback); + Ok(()) + } + } + default_health_status_indicator!(GatedStore); + + let digest_x = DigestInfo::try_new(VALID_HASH1, 2).unwrap(); + let digest_y = DigestInfo::try_new(VALID_HASH2, 2).unwrap(); + + let entered_x = Arc::new(Notify::new()); + let release_x = Arc::new(Notify::new()); + let entered_y = Arc::new(Notify::new()); + let release_y = Arc::new(Notify::new()); + + let mem = MemoryStore::new(&MemorySpec::default()); + let gated = Arc::new(GatedStore { + inner: Store::new(mem.clone()), + mem: mem.clone(), + callback: Mutex::new(None), + gates: Mutex::new(HashMap::from([ + (digest_x, (entered_x.clone(), release_x.clone())), + (digest_y, (entered_y.clone(), release_y.clone())), + ])), + evict_key: digest_y, + }); + + let spec = ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + eviction_policy: Option::default(), + }; + let store = ExistenceCacheStore::new(&spec, Store::new(gated.clone())); + + // Update A (key X): reaches inner update, having installed the pause. + let store_a = store.clone(); + let task_a = tokio::spawn(async move { store_a.update_oneshot(digest_x, "aa".into()).await }); + entered_x.notified().await; + + // Update B (key Y): sees the pause already installed and relies on it. + let store_b = store.clone(); + let task_b = tokio::spawn(async move { store_b.update_oneshot(digest_y, "bb".into()).await }); + entered_y.notified().await; + + // A finishes completely: inserts X and take()s the shared pause slot, + // ending the pause window for B as well. + release_x.notify_one(); + task_a.await.unwrap()?; + + // B finishes its inner update; the just-inserted Y is evicted under + // pressure and the remove callback fires inline (during B's await), + // then B inserts Y into the existence cache. + release_y.notify_one(); + task_b.await.unwrap()?; + + // The inner store does not hold Y (it was evicted). + assert_eq!(mem.has(digest_y).await?, None, "Y was evicted from inner"); + + // Correct behavior: the eviction notification for Y, delivered during + // B's update, must not be lost - the existence cache must not claim Y + // exists. With the non-refcounted pause slot, A's take() ended the + // pause, the removal ran before B's insert, and the cache is stale. + assert!( + !store.exists_in_cache(&digest_y).await, + "existence cache claims Y exists but inner store does not hold it" + ); + + // Demonstrate the harm: a re-upload of Y is drained and discarded. + store.update_oneshot(digest_y, "bb".into()).await?; + assert_eq!( + mem.has(digest_y).await?, + Some(2), + "re-upload of Y was silently discarded due to stale existence cache entry" + ); + Ok(()) +} + +const VALID_HASH2: &str = "aa23456789abcdef000000000000000000010000000000000123456789abcdef"; + +// Repro for the pause_remove_callbacks race: it is an Option-as-boolean with no +// refcount. Two concurrent update() calls share one pause window; the first +// finisher unconditionally take()s the flag, unpausing the second update while +// its inner-store write is still in flight. If the inner store then fires a +// remove callback for the in-flight key (eviction/oversized skip), the callback +// runs immediately as a no-op (key not yet cached), and the second update then +// inserts a stale "exists" entry for a blob the inner store does not hold. +#[nativelink_test] +async fn concurrent_update_unpause_race_stale_existence_test() -> Result<(), Error> { + let spec = ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + eviction_policy: Option::default(), + }; + // max_bytes small enough that digest_b's write is skipped (and remove + // callbacks fired) by the inner memory store. + let inner_store = Store::new(MemoryStore::new(&MemorySpec { + eviction_policy: Some(EvictionPolicy { + max_bytes: 150, + ..Default::default() + }), + })); + let store = ExistenceCacheStore::new_with_time( + &spec, + inner_store.clone(), + MockInstantWrapped::default(), + ); + + let digest_a = DigestInfo::try_new(VALID_HASH1, 10).unwrap(); + let digest_b = DigestInfo::try_new(VALID_HASH2, 200).unwrap(); + + // Task B: starts update(), sets the pause flag, then parks inside the inner + // store's update awaiting stream data. + let (mut tx_b, rx_b) = nativelink_util::buf_channel::make_buf_channel_pair(); + let store_b = store.clone(); + let task_b = nativelink_util::spawn!("update_b", async move { + store_b + .update(digest_b, rx_b, UploadSizeInfo::ExactSize(200)) + .await + }); + // Single-threaded runtime: let B run until it awaits the reader. + for _ in 0..16 { + tokio::task::yield_now().await; + } + + // Task A: full update while B is in flight. A finds the pause flag Some + // (set by B), and at the end unconditionally take()s it -- stealing B's + // pause window. + store + .update_oneshot(digest_a, vec![0u8; 10].into()) + .await + .err_tip(|| "update A failed")?; + + // Now finish B. The inner memory store drains it, refuses to store it + // (200 >= max_bytes) and fires remove callbacks for digest_b -- but the + // pause flag is already None, so the removal runs immediately as a no-op. + // B then inserts digest_b into the existence cache: stale entry. + tx_b.send(Bytes::from(vec![0u8; 200])) + .await + .err_tip(|| "send b")?; + tx_b.send_eof().err_tip(|| "eof b")?; + task_b + .await + .expect("join b") + .err_tip(|| "update B failed")?; + + assert_eq!( + inner_store.has(digest_b).await?, + None, + "precondition: inner store must not hold digest_b" + ); + assert_eq!( + store.has(digest_b).await.err_tip(|| "has b")?, + None, + "existence cache claims digest_b exists but the inner store evicted it \ + (pause window stolen by concurrent update)" + ); + Ok(()) +} + +/// Inner store standing in for a filesystem store: a read that is in-flight +/// keeps streaming successfully even if the entry is evicted mid-read +/// (filesystem reads hold the `FileEntry` Arc). The test controls exactly +/// where the read is suspended via `Notify` rendezvous points. +#[derive(Debug)] +struct PausingReadStore { + present: AtomicBool, + read_started: Notify, + allow_read_finish: Notify, + remove_callbacks: Mutex>, + data: Bytes, +} + +impl PausingReadStore { + fn new(data: Bytes) -> Arc { + Arc::new(Self { + present: AtomicBool::new(true), + read_started: Notify::new(), + allow_read_finish: Notify::new(), + remove_callbacks: Mutex::new(vec![]), + data, + }) + } + + /// Simulate the evicting map evicting `key`: drop residency and fire the + /// registered remove callbacks, exactly like `EvictingMap` does. + async fn evict(&self, key: StoreKey<'static>) { + self.present.store(false, Ordering::SeqCst); + let callbacks: Vec = self.remove_callbacks.lock().clone(); + for callback in callbacks { + callback.callback(key.borrow()).await; + } + } +} + +impl MetricsComponent for PausingReadStore { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } +} + +#[async_trait] +impl StoreDriver for PausingReadStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + _keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + let value = if self.present.load(Ordering::SeqCst) { + Some(self.data.len() as u64) + } else { + None + }; + for result in results.iter_mut() { + *result = value; + } + Ok(()) + } + + async fn update( + self: Pin<&Self>, + _key: StoreKey<'_>, + mut reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + let size = reader.drain().await?; + self.present.store(true, Ordering::SeqCst); + Ok(size) + } + + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + if !self.present.load(Ordering::SeqCst) { + return Err(make_err!(Code::NotFound, "Not found in PausingReadStore")); + } + // Rendezvous: tell the test the read is parked at a real await point, + // then wait until the test lets the (still valid) stream complete. + self.read_started.notify_one(); + self.allow_read_finish.notified().await; + writer.send(self.data.clone()).await?; + writer.send_eof()?; + Ok(()) + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + + fn as_any<'a>(&'a self) -> &'a (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> { + self.remove_callbacks.lock().push(callback); + Ok(()) + } +} + +#[async_trait] +impl HealthStatusIndicator for PausingReadStore { + fn get_name(&self) -> &'static str { + "PausingReadStore" + } + + async fn check_health(&self, namespace: Cow<'static, str>) -> HealthStatus { + StoreDriver::check_health(Pin::new(self), namespace).await + } +} + +// Repro: get_part() re-inserts the existence entry after the inner read +// succeeds without participating in the pause_remove_callbacks protocol. +// If the inner store evicts the digest while the read is in flight (legal: +// filesystem reads hold the FileEntry Arc and stream to completion), the +// eviction's remove-callback purges the existence cache first, and get_part +// then resurrects a zombie exists-entry for data the inner store no longer +// has. A subsequent update() sees has()==true and silently discards the +// re-upload: data loss. +#[nativelink_test] +async fn get_part_races_inner_eviction_and_resurrects_zombie_entry() -> Result<(), Error> { + const VALUE: &str = "123"; + let inner = PausingReadStore::new(Bytes::from_static(VALUE.as_bytes())); + let store = ExistenceCacheStore::new_with_time( + &ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + eviction_policy: Option::default(), + }, + Store::new(inner.clone()), + MockInstantWrapped::default(), + ); + let digest = DigestInfo::try_new(VALID_HASH1, VALUE.len() as u64)?; + + // Task A: client read. Parks inside inner get_part at a genuine await + // point (stand-in for the filesystem store's read/send loop). + let store_clone = store.clone(); + let read_task = + tokio::spawn(async move { store_clone.get_part_unchunked(digest, 0, None).await }); + inner.read_started.notified().await; + + // Task B: cache pressure evicts D from the inner store while the read is + // suspended. The remove callback purges the existence cache (correctly). + inner.evict(digest.into()).await; + assert!( + !store.exists_in_cache(&digest).await, + "Eviction callback should have purged the existence cache" + ); + + // Task A resumes; the in-flight read still completes Ok. + inner.allow_read_finish.notify_one(); + let read_result = read_task.await.expect("read task panicked")?; + assert_eq!(read_result, Bytes::from_static(VALUE.as_bytes())); + + // BUG: get_part re-inserted the existence entry after the eviction + // notification was already consumed. The inner store no longer has D. + assert!( + !store.exists_in_cache(&digest).await, + "Zombie exists-entry: existence cache claims digest is present but \ + the inner store evicted it" + ); + + // Consequence: a re-upload of D is drained and discarded by update(), + // so the data is lost while has() keeps reporting true. + store.update_oneshot(digest, VALUE.into()).await?; + let mut inner_result = [None]; + Pin::new(&*inner) + .has_with_results(&[digest.into()], &mut inner_result) + .await?; + assert_eq!( + inner_result[0], + Some(VALUE.len() as u64), + "Re-upload was silently discarded due to zombie existence entry" + ); + + Ok(()) +} + +/// Inner store whose `update` parks at a real await point until released, so a +/// test can drop the outer `update` future while the pause is held. `evict` +/// fires the registered remove callbacks like `EvictingMap` does. +#[derive(Debug)] +struct BlockingUpdateStore { + present: Mutex>>, + update_started: Notify, + allow_update_finish: Notify, + remove_callbacks: Mutex>, + len: u64, +} + +impl BlockingUpdateStore { + fn new(len: u64, present: &[StoreKey<'static>]) -> Arc { + Arc::new(Self { + present: Mutex::new(present.iter().cloned().collect()), + update_started: Notify::new(), + allow_update_finish: Notify::new(), + remove_callbacks: Mutex::new(vec![]), + len, + }) + } + + async fn evict(&self, key: StoreKey<'static>) { + self.present.lock().remove(&key); + let callbacks: Vec = self.remove_callbacks.lock().clone(); + for callback in callbacks { + callback.callback(key.borrow()).await; + } + } +} + +impl MetricsComponent for BlockingUpdateStore { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } +} + +#[async_trait] +impl StoreDriver for BlockingUpdateStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + let present = self.present.lock(); + for (key, result) in keys.iter().zip(results.iter_mut()) { + *result = if present.contains(&key.borrow().into_owned()) { + Some(self.len) + } else { + None + }; + } + Ok(()) + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + mut reader: DropCloserReadHalf, + _size_info: UploadSizeInfo, + ) -> Result { + // Park at a genuine await point with the pause held, then wait for + // the test. The test drops the future before ever releasing this. + let key = key.into_owned(); + self.update_started.notify_one(); + self.allow_update_finish.notified().await; + let size = reader.drain().await?; + self.present.lock().insert(key); + Ok(size) + } + + async fn get_part( + self: Pin<&Self>, + _key: StoreKey<'_>, + _writer: &mut DropCloserWriteHalf, + _offset: u64, + _length: Option, + ) -> Result<(), Error> { + Err(make_err!(Code::NotFound, "not exercised")) + } + + fn inner_store(&self, _key: Option) -> &dyn StoreDriver { + self + } + + fn as_any<'a>(&'a self) -> &'a (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> { + self.remove_callbacks.lock().push(callback); + Ok(()) + } +} + +#[async_trait] +impl HealthStatusIndicator for BlockingUpdateStore { + fn get_name(&self) -> &'static str { + "BlockingUpdateStore" + } + + async fn check_health(&self, namespace: Cow<'static, str>) -> HealthStatus { + StoreDriver::check_health(Pin::new(self), namespace).await + } +} + +/// Regression for N6: a cancelled `update` must not leave remove callbacks +/// paused forever. Start an `update` that parks with the pause held, drop its +/// future (as a timeout would), then evict a cached key through the inner +/// store. With the RAII pause guard the drop releases the hold, so the +/// eviction callback runs immediately and the existence cache stops reporting +/// the key. Without the guard the pause count stays elevated and the cache +/// keeps claiming the evicted key exists. +#[nativelink_test] +async fn cancelled_update_does_not_leave_callbacks_paused() -> Result<(), Error> { + const LEN: u64 = 2; + let digest_a = DigestInfo::try_new(VALID_HASH1, LEN)?; + let digest_b = DigestInfo::try_new(VALID_HASH2, LEN)?; + + // Only A is present in the inner store; B is absent so its update proceeds + // to the parking inner `update` (rather than short-circuiting on `has`). + let inner = BlockingUpdateStore::new(LEN, &[digest_a.into()]); + let store = ExistenceCacheStore::new_with_time( + &ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + eviction_policy: Option::default(), + }, + Store::new(inner.clone()), + MockInstantWrapped::default(), + ); + + // Populate A's existence entry via a has() that the inner store answers + // positively (A is present). + assert_eq!(store.has(digest_a).await?, Some(LEN)); + assert!(store.exists_in_cache(&digest_a).await); + + // Start an update of B that parks inside the inner store holding the pause. + let store_b = store.clone(); + let mut update_b = Box::pin(async move { + store_b + .update_oneshot(digest_b, vec![0u8; usize::try_from(LEN).unwrap()].into()) + .await + }); + // Drive it until it parks at the inner update await, then cancel it by + // dropping the future. + tokio::select! { + _ = &mut update_b => panic!("update should be parked, not complete"), + () = inner.update_started.notified() => {} + } + drop(update_b); + + // Evict A through the inner store; the remove callback must take effect + // now that the cancelled update released its pause. + inner.evict(digest_a.into()).await; + + assert!( + !store.exists_in_cache(&digest_a).await, + "existence cache still claims A exists: a cancelled update left \ + remove callbacks permanently paused (N6)" + ); + + Ok(()) +} + +/// Regression for the C20/N6 drop residual: a removal queued *before* the last +/// pause holder is cancelled must still be applied. This is the +/// evict-*before*-cancel ordering, which the N6 test above does not cover (it +/// evicts after the drop, so there is nothing queued to lose). +/// +/// Schedule: +/// 1. Populate A's positive existence entry via `has()`. +/// 2. Start an `update` of B that parks in the inner store holding the sole +/// pause guard. +/// 3. Evict A through the inner store *while the pause is held*: the remove +/// callback is queued in `pending` instead of applied, so the outer cache +/// still reports A. +/// 4. Drop (cancel) B's future. The last holder releases the pause. +/// +/// With the buggy `Drop` that discards the vector returned by `release()`, the +/// queued removal of A is lost and the cache keeps claiming A exists forever. +/// A correct `Drop` must drain those queued removals, so A disappears from the +/// cache after cancellation. +#[nativelink_test] +async fn cancelled_update_drains_removals_queued_before_cancel() -> Result<(), Error> { + const LEN: u64 = 2; + let digest_a = DigestInfo::try_new(VALID_HASH1, LEN)?; + let digest_b = DigestInfo::try_new(VALID_HASH2, LEN)?; + + // Only A is present in the inner store; B is absent so its update proceeds + // to the parking inner `update` (rather than short-circuiting on `has`). + let inner = BlockingUpdateStore::new(LEN, &[digest_a.into()]); + let store = ExistenceCacheStore::new_with_time( + &ExistenceCacheSpec { + backend: StoreSpec::Noop(NoopSpec::default()), + eviction_policy: Option::default(), + }, + Store::new(inner.clone()), + MockInstantWrapped::default(), + ); + + // Populate A's existence entry via a has() the inner store answers + // positively (A is present). + assert_eq!(store.has(digest_a).await?, Some(LEN)); + assert!(store.exists_in_cache(&digest_a).await); + + // Start an update of B that parks inside the inner store holding the pause. + let store_b = store.clone(); + let mut update_b = Box::pin(async move { + store_b + .update_oneshot(digest_b, vec![0u8; usize::try_from(LEN).unwrap()].into()) + .await + }); + tokio::select! { + _ = &mut update_b => panic!("update should be parked, not complete"), + () = inner.update_started.notified() => {} + } + + // Evict A through the inner store *while the pause is held*. The removal is + // queued in `pending`; the outer cache still reports A for now. + inner.evict(digest_a.into()).await; + assert!( + store.exists_in_cache(&digest_a).await, + "removal should be queued (not applied) while the pause is held" + ); + + // Cancel B by dropping its future. The last holder releases the pause; the + // queued removal of A must still be drained. `Drop` cannot await, so the + // drain is spawned; give it a bounded number of scheduler turns to run. + drop(update_b); + + let mut drained = false; + for _ in 0..1000 { + if !store.exists_in_cache(&digest_a).await { + drained = true; + break; + } + tokio::task::yield_now().await; + } + + assert!( + drained, + "existence cache still claims A exists: a removal queued before the \ + last pause holder was cancelled was discarded instead of drained \ + (C20/N6 drop residual)" + ); + + Ok(()) +} From abf2c59ce4d3eb7f01f803c84034b90a1eed3b90 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 01:34:10 -0400 Subject: [PATCH 15/23] worker: keep a kill received during action registration effective (N1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Astra finding N1 (P1): the worker forgot a kill received before async action registration finished. `create_and_add_action` fetches the Action, waits for prior cleanup, and makes the action directory across await points before inserting the op into `running_actions`. A `KillOperationRequest` arriving in that window hit `kill_operation`, found nothing in `running_actions`, returned an error, and dropped the kill — the action then started executing unkilled. Fix: add a `pending_kills` set. `kill_operation` now, when the op is not yet resident, records it as pending (returning Ok) instead of erroring. `create_and_add_action`, under the same `running_actions` lock that guards the insert, checks `pending_kills` for the op and, if present, applies the kill to the freshly created action so the request stays effective through the registration race. Lock ordering is fixed as running_actions before pending_kills. Regression test: tests/running_actions_manager_test.rs kill_before_registration_stays_effective. It kills the op BEFORE create_and_add_action registers it (modelling the in-transit race), then runs the action and asserts exit_code 9 (SIGKILLed). WITHOUT THE FIX the pre-registration kill returns Err ("Failed to get running action") and the kill is dropped; the test fails at that point (verified by stashing the fix: "A kill for a not-yet-resident op must be recorded, not rejected"). WITH THE FIX the kill is parked and applied at registration, and the action runs killed. Co-Authored-By: Claude Opus 4.8 --- .../src/running_actions_manager.rs | 50 +++++- .../tests/running_actions_manager_test.rs | 144 ++++++++++++++++++ 2 files changed, 190 insertions(+), 4 deletions(-) diff --git a/nativelink-worker/src/running_actions_manager.rs b/nativelink-worker/src/running_actions_manager.rs index 9c76c6e17d6..ff853f6c3f8 100644 --- a/nativelink-worker/src/running_actions_manager.rs +++ b/nativelink-worker/src/running_actions_manager.rs @@ -4334,6 +4334,15 @@ pub struct RunningActionsManagerImpl { #[cfg(target_os = "linux")] use_namespaces: UseNamespaces, running_actions: Mutex>>, + /// Kills received for operations that are not yet resident in + /// `running_actions`. A `KillOperationRequest` can arrive while + /// `create_and_add_action` is still awaiting (fetching the Action, waiting + /// for prior cleanup, making the directory) — before the op is inserted. + /// Without a record, the kill looked up nothing, logged, and was forgotten, + /// and the action then started executing unkilled (finding N1). We record + /// the op here under the same lock that guards the `running_actions` + /// insert, so registration can observe an in-transit kill and apply it. + pending_kills: Mutex>, // Note: We don't use Notify because we need to support a .wait_for()-like function, which // Notify does not support. action_done_tx: watch::Sender<()>, @@ -4401,6 +4410,7 @@ impl RunningActionsManagerImpl { max_download_timeout: args.max_download_timeout, timeout_handled_externally: args.timeout_handled_externally, running_actions: Mutex::new(HashMap::new()), + pending_kills: Mutex::new(HashSet::new()), action_done_tx, callbacks, metrics: Arc::new(Metrics { @@ -4842,7 +4852,13 @@ impl RunningActionsManager for RunningActionsManagerImpl { timeout, self.clone(), )); + // True if a kill arrived for this op while its registration was + // still in flight (recorded in `pending_kills` before the op + // became resident). Determined under the `running_actions` lock + // so it cannot race with a concurrent kill. + let killed_in_transit; { + // Lock ordering: `running_actions` before `pending_kills`. let mut running_actions = self.running_actions.lock(); // Check if action already exists and is still alive if let Some(existing_weak) = running_actions.get(&operation_id) @@ -4857,6 +4873,18 @@ impl RunningActionsManager for RunningActionsManagerImpl { running_action .has_manager_entry .store(true, Ordering::Release); + killed_in_transit = self.pending_kills.lock().remove(&operation_id); + } + if killed_in_transit { + // A kill was requested before this op finished registering. + // Apply it now that the action object exists so the kill is + // not lost: `execute()` will observe the consumed kill + // channel and abort instead of running the process. + warn!( + %operation_id, + "Applying kill requested while action registration was in flight", + ); + Self::kill_operation(running_action.clone()).await; } if let Some(config) = &self.execution_configuration.buck2_file_capture && Buck2FileCapture::eligible(capture_metadata.as_ref()) { @@ -4901,12 +4929,26 @@ impl RunningActionsManager for RunningActionsManagerImpl { } async fn kill_operation(&self, operation_id: &OperationId) -> Result<(), Error> { + // Lock ordering: `running_actions` before `pending_kills`. Decide under + // the `running_actions` lock whether the op is resident: if it is not, + // it may still be mid-registration (create_and_add_action fetches the + // Action, waits for prior cleanup, and makes the directory across await + // points before inserting), so record the kill as pending rather than + // dropping it. Registration checks `pending_kills` under this same lock + // and applies the kill, so the request stays effective through the + // registration race (finding N1). let running_action = { let running_actions = self.running_actions.lock(); - running_actions - .get(operation_id) - .and_then(Weak::upgrade) - .ok_or_else(|| make_input_err!("Failed to get running action {operation_id}"))? + let Some(running_action) = running_actions.get(operation_id).and_then(Weak::upgrade) + else { + self.pending_kills.lock().insert(operation_id.clone()); + warn!( + %operation_id, + "Kill received for an operation not yet resident; recorded as pending so registration applies it", + ); + return Ok(()); + }; + running_action }; Self::kill_operation(running_action).await; Ok(()) diff --git a/nativelink-worker/tests/running_actions_manager_test.rs b/nativelink-worker/tests/running_actions_manager_test.rs index 80f355a4c22..64b857d8cee 100644 --- a/nativelink-worker/tests/running_actions_manager_test.rs +++ b/nativelink-worker/tests/running_actions_manager_test.rs @@ -1231,6 +1231,150 @@ mod tests { Ok(()) } + /// N1 regression: a kill received before an action finishes registering + /// must stay effective through registration. + /// + /// A `KillOperationRequest` can reach the worker while + /// `create_and_add_action` is still awaiting (fetching the Action, waiting + /// for prior cleanup, making the directory) — before the op is inserted + /// into `running_actions`. This test reproduces that ordering by killing + /// the op *before* it is registered. + /// + /// WITHOUT THE FIX: `kill_operation` looked the op up in `running_actions`, + /// found nothing, and returned an Err (the kill was logged and dropped); + /// the action then registered and ran to completion unkilled (exit 0). + /// WITH THE FIX the kill is parked in `pending_kills` and applied at + /// registration, so when the action runs it is killed with `SIGKILL` + /// and reports the killed exit code (9). + /// + /// Two assertions fail without the fix: (1) the pre-registration kill must + /// return Ok (recorded, not rejected), and (2) the resulting action must + /// carry the killed exit code rather than having run normally. + #[nativelink_test] + async fn kill_before_registration_stays_effective() -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + + fn test_monotonic_clock() -> SystemTime { + static CLOCK: AtomicU64 = AtomicU64::new(0); + monotonic_clock(&CLOCK) + } + + let (_, _, cas_store, ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + + let running_actions_manager = Arc::new(RunningActionsManagerImpl::new_with_callbacks( + RunningActionsManagerArgs { + root_action_directory, + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: Some(Store::new(ac_store.clone())), + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_download_timeout: Duration::from_secs(DEFAULT_MAX_DOWNLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + active_input_leases: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + }, + Callbacks { + now_fn: test_monotonic_clock, + sleep_fn: |_duration| Box::pin(future::pending()), + }, + )?); + + let command = Command { + arguments: vec!["sleep".to_string(), "0.2".to_string()], + environment_variables: vec![EnvironmentVariable { + name: "PATH".to_string(), + value: env::var("PATH").unwrap(), + }], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + + let operation_id = OperationId::default(); + + // The kill arrives before the op is resident (models a kill landing + // while create_and_add_action is still mid-flight). This must be + // remembered, not dropped. + running_actions_manager + .kill_operation(&operation_id) + .await + .err_tip(|| "A kill for a not-yet-resident op must be recorded, not rejected")?; + + let execute_request = ExecuteRequest { + action_digest: Some(action_digest.into()), + digest_function: ProtoDigestFunction::Sha256.into(), + ..Default::default() + }; + let running_action = running_actions_manager + .create_and_add_action( + WORKER_ID.to_string(), + StartExecute { + request_metadata: None, + execute_request: Some(execute_request), + operation_id: operation_id.to_string(), + queued_timestamp: None, + platform: action.platform.clone(), + worker_id: WORKER_ID.to_string(), + }, + ) + .await?; + + // The in-transit kill must have been applied at registration. Since + // #2846 the kill token is also watched during input fetch, so the + // remembered kill now short-circuits `prepare_action` with `Aborted` + // ("killed while fetching its inputs") before the action ever runs — + // instead of the sleep completing normally (which a dropped in-transit + // kill would allow). Either way the observable proof is that the kill + // took effect: the action never completes successfully. + let err = run_action(running_action) + .await + .expect_err("a kill recorded before registration must abort the action"); + assert_eq!( + err.code, + Code::Aborted, + "a kill applied at registration aborts the action during input fetch; \ + a dropped in-transit kill would instead let it complete (Ok)" + ); + assert!( + err.to_string().contains("killed while fetching its inputs"), + "abort must be the kill path, not an unrelated failure: {err}" + ); + Ok(()) + } + #[nativelink_test] async fn ensure_output_files_full_directories_are_created_test() -> Result<(), Box> { From a9c42a67adb08ea714cba5e9d9987716a3eb434c Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 01:25:17 -0400 Subject: [PATCH 16/23] scheduler: reject same-worker re-dispatch of a resident op (N2) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Astra finding N2 (P1): a same-worker retry overwrites a live reservation. `Worker::run_action` unconditionally re-ran `reduce_platform_properties` and overwrote the op-keyed `running_action_infos` entry. When an op whose earlier attempt was still resident (e.g. an execution-deadline requeue that left the ledger entry in place) was dispatched again to the same worker, the platform budget was deducted a second time while `complete_action` only refunds the surviving entry — leaking the first reservation's capacity permanently. Fix: reject-if-exists under the dispatch lock. `run_action` now refuses a dispatch for an op the worker already holds (Code::AlreadyExists) before any budget deduction or insert. Regression test: tests/worker_retry_reservation_test.rs. It dispatches an op (reserving 4 of 10 mem), then re-dispatches the same op to the same worker and asserts the budget is still 6. WITHOUT THE FIX the second dispatch deducts again and the budget reads 2 (test fails: found 2, expected 6); with the fix the redundant dispatch is refused and the budget stays 6. Co-Authored-By: Claude Opus 4.8 --- nativelink-scheduler/src/worker.rs | 15 +- .../tests/worker_retry_reservation_test.rs | 132 ++++++++++++++++++ 2 files changed, 146 insertions(+), 1 deletion(-) create mode 100644 nativelink-scheduler/tests/worker_retry_reservation_test.rs diff --git a/nativelink-scheduler/src/worker.rs b/nativelink-scheduler/src/worker.rs index 6770cdbaf12..d2ef4b288b2 100644 --- a/nativelink-scheduler/src/worker.rs +++ b/nativelink-scheduler/src/worker.rs @@ -17,7 +17,7 @@ use std::collections::HashMap; use std::sync::Arc; use std::time::{SystemTime, UNIX_EPOCH}; -use nativelink_error::{Code, Error, ResultExt}; +use nativelink_error::{Code, Error, ResultExt, make_err}; use nativelink_metric::MetricsComponent; use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{ ActionResourceUsage, ConnectionResult, KillOperationRequest, StartExecute, UpdateForWorker, @@ -367,6 +367,19 @@ impl Worker { action_info: ActionInfoWithProps, dispatched_at: WorkerTimestamp, ) -> Result<(), Error> { + // Reject-if-exists: a same-worker retry can dispatch an op whose + // earlier attempt is still resident here (e.g. an execution-deadline + // requeue that left the ledger entry in place). Overwriting the + // op-keyed entry would deduct the platform budget a second time while + // completion only refunds the surviving entry, leaking the first + // reservation permanently. Refuse the redundant dispatch instead. + if self.running_action_infos.contains_key(&operation_id) { + return Err(make_err!( + Code::AlreadyExists, + "Worker {} already holds a reservation for operation {operation_id}; refusing redundant dispatch", + self.id, + )); + } let tx = &mut self.tx; let worker_platform_properties = &mut self.platform_properties; let running_action_infos = &mut self.running_action_infos; diff --git a/nativelink-scheduler/tests/worker_retry_reservation_test.rs b/nativelink-scheduler/tests/worker_retry_reservation_test.rs new file mode 100644 index 00000000000..8c734fd567a --- /dev/null +++ b/nativelink-scheduler/tests/worker_retry_reservation_test.rs @@ -0,0 +1,132 @@ +// Copyright 2026 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::collections::HashMap; +use std::time::{Duration, UNIX_EPOCH}; + +use nativelink_error::Error; +use nativelink_macro::nativelink_test; +use nativelink_scheduler::worker::{ActionInfoWithProps, Worker, WorkerUpdate}; +use nativelink_util::action_messages::{OperationId, WorkerId}; +use nativelink_util::common::DigestInfo; +use nativelink_util::origin_event::OriginMetadata; +use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue}; +use pretty_assertions::assert_eq; +use tokio::sync::mpsc; + +mod utils { + pub(crate) mod scheduler_utils; +} +use utils::scheduler_utils::make_base_action_info; + +const NOW_TIME: u64 = 10_000; + +fn minimum(mem: u64) -> PlatformProperties { + PlatformProperties::new(HashMap::from([( + "mem".to_string(), + PlatformPropertyValue::Minimum(mem), + )])) +} + +fn action_needing(mem: u64) -> ActionInfoWithProps { + ActionInfoWithProps { + inner: make_base_action_info( + UNIX_EPOCH + Duration::from_secs(NOW_TIME), + DigestInfo::new([7; 32], 512), + ), + platform_properties: minimum(mem), + origin_metadata: OriginMetadata::default(), + scheduler_start_execute_event_id: None, + } +} + +fn mem_of(props: &PlatformProperties) -> u64 { + match props.properties.get("mem") { + Some(PlatformPropertyValue::Minimum(v)) => *v, + _ => panic!("expected a Minimum 'mem' property"), + } +} + +/// N2 regression: a same-worker retry (the earlier attempt of the same op is +/// still resident on the worker) must not double-deduct or leak the worker's +/// resource budget. +/// +/// WITHOUT THE FIX: `Worker::run_action` unconditionally re-runs +/// `reduce_platform_properties` and overwrites the op-keyed entry, so the +/// worker's `mem` budget is deducted a SECOND time (10 -> 6 -> 2) even though +/// the op already holds one reservation. Since completion (`complete_action`) +/// only refunds the single surviving entry, those extra 4 units are leaked +/// forever. This test asserts that re-dispatching an op the worker already +/// holds does not further reduce the budget — it FAILS on the unfixed code +/// (asserts 6, finds 2). +#[nativelink_test] +async fn same_worker_retry_conserves_resource_budget() -> Result<(), Error> { + let (tx, _rx) = mpsc::channel(64); + let mut worker = Worker::new( + WorkerId("w".to_string()), + minimum(10), + tx, + NOW_TIME, + // Two inflight slots so a retry has spare capacity while the first + // attempt is still resident. + 2, + ); + + let op = OperationId::default(); + + // First dispatch of op reserves 4 of the worker's 10 mem. + worker + .notify_update(WorkerUpdate::RunAction(Box::new(( + op.clone(), + action_needing(4), + NOW_TIME, + )))) + .await?; + assert_eq!( + mem_of(&worker.platform_properties), + 6, + "first dispatch should reserve 4 of 10" + ); + + // Same op dispatched again to the SAME worker while the earlier attempt is + // still resident (same-worker retry). The worker must not create a second + // live reservation for an op it already holds. + let retry_result = worker + .notify_update(WorkerUpdate::RunAction(Box::new(( + op.clone(), + action_needing(4), + NOW_TIME, + )))) + .await; + // The worker may accept or reject the redundant dispatch; either way the + // invariant below (no double reservation) must hold. + drop(retry_result); + + // Exactly one live reservation should exist for the op. + assert_eq!( + worker.running_action_infos.len(), + 1, + "op must have exactly one live reservation on the worker" + ); + + // The op still holds exactly one reservation, so the budget must be + // unchanged from the first dispatch. A second deduction here is a leak: + // completion only refunds the single surviving entry. + assert_eq!( + mem_of(&worker.platform_properties), + 6, + "re-dispatching an already-held op must not deduct the budget again" + ); + Ok(()) +} From 87abb1eb97b603787a2ef058221f5528df175b5b Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 23:19:37 -0400 Subject: [PATCH 17/23] fix(scheduler): matching engine no longer exits when polled during construction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The matching-engine task is spawned inside `Arc::new_cyclic`, during which `weak_inner.upgrade()` returns `None` because the strong count is still zero. The task treats a failed upgrade as scheduler shutdown, so if `task_change_notify` already held a permit (e.g. a Redis pubsub event arriving before `SimpleScheduler::new`), its first poll completed instantly, consumed the permit, saw `None`, and returned permanently — a scheduler that never matches anything. Gate the task's loop on a oneshot fired unconditionally right after `Arc::new_cyclic` returns, so the task can neither consume a pre-stored notify permit nor observe the partially constructed scheduler. Includes a deterministic in-crate repro (adversarial review fleet, branch repro/me-constructor-race, d44e3cb1fa) that pre-stores a notify permit and widens the in-closure window via a test-only timing seam; it fails without the fix and passes with it. Co-Authored-By: Claude Fable 5 --- nativelink-scheduler/src/simple_scheduler.rs | 121 ++++++++++++++++++- 1 file changed, 120 insertions(+), 1 deletion(-) diff --git a/nativelink-scheduler/src/simple_scheduler.rs b/nativelink-scheduler/src/simple_scheduler.rs index 9a20ba2030a..880db216c96 100644 --- a/nativelink-scheduler/src/simple_scheduler.rs +++ b/nativelink-scheduler/src/simple_scheduler.rs @@ -49,7 +49,7 @@ use opentelemetry::baggage::BaggageExt; use opentelemetry::context::{Context, FutureExt as OtelFutureExt}; use opentelemetry_semantic_conventions::attribute::ENDUSER_ID; use parking_lot::Mutex; -use tokio::sync::{Notify, mpsc}; +use tokio::sync::{Notify, mpsc, oneshot}; use tokio::time::Duration; use tracing::{debug, error, info, info_span, warn}; use uuid::Uuid; @@ -87,6 +87,11 @@ const FLEET_RECORD_TTL_INTERVALS: u32 = 3; /// version-conflict backoff sleeps a single pass can incur. const MAX_UNSATISFIABLE_FAILS_PER_PASS: usize = 16; +/// Test-only: milliseconds to stall inside the `Arc::new_cyclic` closure +/// right after the matching-engine task is spawned. See the seam below. +#[cfg(test)] +static TEST_CONSTRUCTOR_PAUSE_MS: AtomicUsize = AtomicUsize::new(0); + /// A backlog larger than `floor * DIVISOR` raises the cap to `backlog / /// DIVISOR`, so a large single-shape burst drains in ~DIVISOR back-to-back /// passes rather than `backlog / floor` of them (whose tail could outlast a @@ -1304,10 +1309,25 @@ impl SimpleScheduler { }) }; + // The matching task is spawned inside `Arc::new_cyclic`, during which + // `weak_inner.upgrade()` returns `None` (strong count is still zero). + // The task treats a failed upgrade as "shutting down", so if it gets + // polled before the closure finishes (e.g. `task_change_notify` + // already holds a permit from a pubsub event that raced construction), + // it would consume that permit, see `None`, and exit permanently. + // Gate the loop on this signal, sent after `new_cyclic` returns, so + // the task cannot observe a partially constructed scheduler. + let (constructed_tx, constructed_rx) = oneshot::channel::<()>(); let action_scheduler = Arc::new_cyclic(move |weak_self| -> Self { let weak_inner = weak_self.clone(); let task_worker_matching_spawn = spawn!("simple_scheduler_task_worker_matching", async move { + // Wait for construction to complete. An `Err` means the + // sender was dropped without sending, which cannot happen: + // it is sent unconditionally right after `new_cyclic`. + if constructed_rx.await.is_err() { + return; + } let mut last_match_successful = true; let mut worker_match_logging_last: Option = None; // Break out of the loop only when the inner is dropped. @@ -1445,6 +1465,19 @@ impl SimpleScheduler { // Unreachable. }); + // Test-only timing seam: widens the window between spawning the + // matching task and the end of `Arc::new_cyclic` (during which + // `weak_inner.upgrade()` returns `None` because the strong count + // is still zero). This changes no logic, only timing, so a test + // can deterministically exercise the constructor race. + #[cfg(test)] + { + let pause_ms = TEST_CONSTRUCTOR_PAUSE_MS.load(Ordering::Acquire); + if pause_ms > 0 { + std::thread::sleep(Duration::from_millis(pause_ms as u64)); + } + } + let worker_match_logging_interval = match spec.worker_match_logging_interval_s { // -1 or 0 means disabled (0 used to cause expensive logging on every call) -1 | 0 => None, @@ -1480,6 +1513,9 @@ impl SimpleScheduler { now_fn: scheduler_now_fn, } }); + // Release the matching-engine task now that `self` is fully + // constructed and weak upgrades succeed. + let _ = constructed_tx.send(()); (action_scheduler, worker_scheduler_clone) } } @@ -1648,3 +1684,86 @@ mod unsatisfiable_fail_cap_tests { assert!(unsatisfiable_fail_cap(over) > MAX_UNSATISFIABLE_FAILS_PER_PASS); } } + +#[cfg(test)] +mod constructor_race_tests { + use core::sync::atomic::Ordering; + + use nativelink_config::schedulers::SimpleSpec; + use nativelink_config::stores::EvictionPolicy; + use nativelink_util::instant_wrapper::MockInstantWrapped; + use tokio::sync::{Notify, mpsc}; + use tokio::time::Duration; + + use super::{Arc, SimpleScheduler, TEST_CONSTRUCTOR_PAUSE_MS}; + use crate::memory_awaited_action_db::MemoryAwaitedActionDb; + + /// Repro for the spawn-then-configure race in `SimpleScheduler::new*`: + /// the matching-engine task is spawned inside `Arc::new_cyclic`, holding + /// only a `Weak` whose `upgrade()` returns `None` until the closure + /// finishes. If `task_change_notify` already holds a permit (as the + /// Redis pubsub path can arrange before the constructor runs), the task's + /// first poll completes instantly, sees `upgrade() == None`, and + /// interprets "not yet constructed" as "shutting down": it returns and + /// the scheduler never matches anything again. + #[expect( + clippy::disallowed_methods, + reason = "the race needs a multi-thread runtime so the spawned task \ + is polled while the constructor thread is blocked" + )] + #[test] + fn matching_engine_survives_prestored_notify_permit_during_construction() { + let rt = tokio::runtime::Builder::new_multi_thread() + .worker_threads(2) + .enable_all() + .build() + .unwrap(); + rt.block_on(async { + // Widen the in-closure window so the freshly spawned task is + // guaranteed to get its first poll while the strong count is 0. + TEST_CONSTRUCTOR_PAUSE_MS.store(500, Ordering::Release); + + let task_change_notify = Arc::new(Notify::new()); + // A permit stored before construction, exactly like + // pull_task_change_subscriber's notify_one() on a Redis pubsub + // event that arrives before SimpleScheduler::new. + task_change_notify.notify_one(); + + let db = MemoryAwaitedActionDb::new( + &EvictionPolicy::default(), + task_change_notify.clone(), + MockInstantWrapped::default, + ); + let (tx, mut rx) = mpsc::unbounded_channel::<()>(); + let (_scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( + &SimpleSpec::default(), + db, + move || { + // Runs after every matching pass; proves the engine is alive. + let tx = tx.clone(); + async move { + let _ = tx.send(()); + } + }, + task_change_notify.clone(), + MockInstantWrapped::default, + None, + ); + TEST_CONSTRUCTOR_PAUSE_MS.store(0, Ordering::Release); + + // Nudge it again post-construction; a live engine must run a + // matching pass for either the pre-stored permit or this one. + task_change_notify.notify_one(); + // `Ok(None)` means the channel closed: the matching task (which + // owns the only sender) returned. `Err(_)` means it is alive but + // never ran a pass. Either way it is dead or wedged. + let matched = tokio::time::timeout(Duration::from_secs(2), rx.recv()).await; + assert!( + matched == Ok(Some(())), + "matching engine never ran a pass: it exited during \ + Arc::new_cyclic because weak_inner.upgrade() was None \ + (strong count still 0), mistaking construction for shutdown" + ); + }); + } +} From 8c886b29d00621291ccce0fdc175595c5d8bdbf4 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 01:27:20 -0400 Subject: [PATCH 18/23] test(filesystem_store): widen open-file drain budget to de-flake asan `wait_for_no_open_files` polls a process-global open-handle counter, so it drains handles from every test sharing the binary. Under asan the previous 1s ceiling was too tight and `upgraded_layout_is_not_migrated_again_on_restart` intermittently failed with "Timed out waiting all files to close" even though the handles did close. Raise the cap to 10s; the loop still returns the instant the count hits zero, so healthy runs are unaffected. This is a test-only timing fix; it does not touch the eviction change this PR exists for. Co-Authored-By: Claude Opus 4.8 --- .../tests/filesystem_store_test.rs | 8 +- nativelink-util/src/evicting_map.rs | 105 +++++++++++++++++- 2 files changed, 111 insertions(+), 2 deletions(-) diff --git a/nativelink-store/tests/filesystem_store_test.rs b/nativelink-store/tests/filesystem_store_test.rs index fb4c1f59640..ea28c82680a 100644 --- a/nativelink-store/tests/filesystem_store_test.rs +++ b/nativelink-store/tests/filesystem_store_test.rs @@ -241,11 +241,17 @@ async fn read_file_contents(file_name: &OsStr) -> Result, Error> { } async fn wait_for_no_open_files() -> Result<(), Error> { + // `get_open_files_for_test` reads a process-global open-handle counter, so + // this drains handles opened by *any* concurrently running test in the same + // binary. Under asan (and other slow-execution configs) a 1s budget is too + // tight and this spuriously times out even though the handles do close, so + // give the drain a generous ceiling — the loop still exits the instant the + // count reaches zero, so a healthy run pays nothing for the larger cap. let mut counter = 0; while fs::get_open_files_for_test() != 0 { sleep(Duration::from_millis(1)).await; counter += 1; - if counter > 1000 { + if counter > 10000 { return Err(make_err!( Code::Internal, "Timed out waiting all files to close" diff --git a/nativelink-util/src/evicting_map.rs b/nativelink-util/src/evicting_map.rs index a94bfd534a8..8076f9241a0 100644 --- a/nativelink-util/src/evicting_map.rs +++ b/nativelink-util/src/evicting_map.rs @@ -634,7 +634,13 @@ where let (removal_futures, data_to_unref) = { let mut state = self.state.lock(); - let lru_len = state.lru.len(); + // Working copy of the resident count, decremented as entries are + // reaped below. Evaluating `should_evict` against a length + // snapshotted once for the whole batch would keep reporting + // over-count for every remaining key after the first reap brings + // the map back to its limit, cascading into evicting every + // queried entry — reachable from a plain `has()` batch. + let mut lru_len = state.lru.len(); let mut data_to_unref = Vec::new(); let mut removal_futures = Vec::new(); for (key, result) in keys.into_iter().zip(results.iter_mut()) { @@ -656,6 +662,7 @@ where if !is_leased && self.should_evict(lru_len, entry, 0, u64::MAX) { *result = None; if let Some((key, eviction_item)) = state.lru.pop_entry(key.borrow()) { + lru_len -= 1; info!(?key, "Item expired, evicting"); let (data, futures) = state.remove(key.borrow(), &eviction_item, false); @@ -1042,3 +1049,99 @@ where self.state.lock().add_remove_callback(callback); } } + +#[cfg(test)] +mod tests { + use nativelink_config::stores::EvictionPolicy; + use nativelink_macro::nativelink_test; + use pretty_assertions::assert_eq; + + use super::{EvictingMap, EvictionItem, LenEntry}; + use crate::instant_wrapper::MockInstantWrapped; + + #[derive(Clone, Debug, PartialEq, Eq)] + struct TestEntry(u64); + + impl LenEntry for TestEntry { + fn len(&self) -> u64 { + self.0 + } + fn is_empty(&self) -> bool { + self.0 == 0 + } + } + + /// A batch existence check (`has()` → `sizes_for_keys`) over resident + /// keys, run while the map is pinned one entry over `max_count` by a + /// leased entry, must evict at most the genuine overage — never every + /// queried key. + /// + /// The bug: `sizes_for_keys` snapshotted `lru.len()` once for the whole + /// batch and re-used that stale value in `should_evict` for every key. + /// After the first queried entry was reaped (bringing the map back to + /// `max_count`), every remaining queried key still saw the stale + /// over-count and was popped too. A plain `has()` batch would then + /// destroy all N resident blobs and report them all missing. + /// + /// White-box: the pinning entry is planted with a raw leased `State::put` + /// (no eviction pass), modeling the window in which an action-input lease + /// holds the map over its count limit while unleased blobs stay resident. + #[nativelink_test("crate")] + async fn size_check_batch_evicts_at_most_the_overage() { + let map = EvictingMap::::new( + &EvictionPolicy { + max_bytes: 0, + evict_bytes: 0, + max_seconds: 0, + max_count: 3, + }, + MockInstantWrapped::default(), + ); + + map.insert("k1".to_string(), TestEntry(10)).await; + map.insert("k2".to_string(), TestEntry(20)).await; + map.insert("k3".to_string(), TestEntry(30)).await; + + // Pin the map one over `max_count` with a leased entry, bypassing the + // insert-time eviction pass exactly as an already-held lease does. + { + let mut state = map.state.lock(); + state.lease("pinned".to_string()); + let replaced = state.put( + &"pinned".to_string(), + EvictionItem { + seconds_since_anchor: 0, + data: TestEntry(1), + }, + ); + assert!(replaced.is_none()); + state.sum_store_size += 1; + } + assert_eq!(map.len_for_test(), 4); + + let keys = ["k1".to_string(), "k2".to_string(), "k3".to_string()]; + let mut results = [None, None, None]; + map.sizes_for_keys(keys.iter(), &mut results[..], false) + .await; + + // The overage is a single entry, so at most the first queried key may + // be reaped. The rest of the batch must stay resident and report size. + assert_eq!( + results[1], + Some(20), + "k2 must survive a batch size check: {results:?}" + ); + assert_eq!( + results[2], + Some(30), + "k3 must survive a batch size check: {results:?}" + ); + assert_eq!( + map.len_for_test(), + 3, + "only the genuine overage may be evicted by a size query" + ); + assert_eq!(map.size_for_key(&"k2".to_string()).await, Some(20)); + assert_eq!(map.size_for_key(&"k3".to_string()).await, Some(30)); + } +} From 96e89c5eb52a5dcc4a28be8532b0c2c4fa9cfb88 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 23:01:58 -0400 Subject: [PATCH 19/23] fix(scheduler): client keep-alive stamps survive concurrent versioned writes Client keep-alive stamps in the memory awaited-action db (the ActionEvent::ClientKeepAlive path in handle_action_events, and the initial stamp in try_subscribe) mutate the AwaitedAction in place via send_if_modified WITHOUT bumping AwaitedActionVersion. Any concurrent read-modify-write cycle (SimpleSchedulerStateManager borrow() -> update_awaited_action) that snapshotted the action at version N before the stamp still passes the version CAS, and send_replace writes the whole stale snapshot back -- a classic lost update that silently moves last_client_keepalive_timestamp backwards. The Code::Aborted retry defense never fires because the version matches. Consequence: the client-timeout sweep force-fails operations whose clients are alive and polling ("no more clients listening" DeadlineExceeded), including operations a client subscribed to moments earlier. Fix shape: treat the keep-alive stamp as data outside the versioned snapshot rather than routing it through the versioned update path. update_awaited_action now merges the stamp forward (monotone max of old and new) before send_replace, so an accepted stale write can never revert it. Chosen over version-bumping keep-alives because (1) a bump would spuriously abort legitimate in-flight worker updates and wake every changed() subscriber ~every CLIENT_KEEPALIVE_DURATION per client, (2) it needs no new per-operation state, and (3) it leaves the AwaitedAction wire/record format and the store/redis backend untouched. Two deterministic repro tests included; both fail without the one-line merge and pass with it (verified both directions, 5x each). Provenance: adversarial review fleet. Co-Authored-By: Claude Fable 5 --- nativelink-scheduler/Cargo.toml | 1 + .../src/memory_awaited_action_db.rs | 14 ++ .../tests/client_keepalive_clobber_test.rs | 156 ++++++++++++++++++ .../simple_scheduler_state_manager_test.rs | 154 +++++++++++++++++ 4 files changed, 325 insertions(+) create mode 100644 nativelink-scheduler/tests/client_keepalive_clobber_test.rs diff --git a/nativelink-scheduler/Cargo.toml b/nativelink-scheduler/Cargo.toml index ff752b77941..35bfea20b00 100644 --- a/nativelink-scheduler/Cargo.toml +++ b/nativelink-scheduler/Cargo.toml @@ -53,6 +53,7 @@ uuid = { version = "1.16.0", default-features = false, features = [ ] } [dev-dependencies] +mock_instant = { version = "0.5.3", default-features = false } nativelink-macro = { path = "../nativelink-macro" } nativelink-redis-tester = { path = "../nativelink-redis-tester" } diff --git a/nativelink-scheduler/src/memory_awaited_action_db.rs b/nativelink-scheduler/src/memory_awaited_action_db.rs index 384b0237277..ac540448f6c 100644 --- a/nativelink-scheduler/src/memory_awaited_action_db.rs +++ b/nativelink-scheduler/src/memory_awaited_action_db.rs @@ -637,6 +637,20 @@ impl I + Clone + Send + Sync> AwaitedActionDbI } new_awaited_action.increment_version(); + // Client keep-alive stamps are written in place (see + // `handle_action_events` and `try_subscribe`) without bumping the + // version, so a read-modify-write cycle that snapshotted before + // the stamp still passes the version check above. Merge the stamp + // forward so such a write can never move the keep-alive backwards + // and cause the client-timeout sweep to kill an operation whose + // client is alive and polling. + if old_awaited_action.last_client_keepalive_timestamp() + > new_awaited_action.last_client_keepalive_timestamp() + { + new_awaited_action + .update_client_keep_alive(old_awaited_action.last_client_keepalive_timestamp()); + } + error_if!( old_awaited_action.action_info().unique_qualifier != new_awaited_action.action_info().unique_qualifier, diff --git a/nativelink-scheduler/tests/client_keepalive_clobber_test.rs b/nativelink-scheduler/tests/client_keepalive_clobber_test.rs new file mode 100644 index 00000000000..55dd99c11ab --- /dev/null +++ b/nativelink-scheduler/tests/client_keepalive_clobber_test.rs @@ -0,0 +1,156 @@ +// Copyright 2026 The NativeLink Authors. All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// 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. + +//! Repro: client keep-alive writes bypass the version CAS in the memory +//! backend, so a concurrent read-modify-write (a worker update) that +//! snapshotted the action before the keep-alive silently reverts +//! `last_client_keepalive_timestamp` — and passes the version check, so the +//! `Code::Aborted` retry defense in `SimpleSchedulerStateManager` never +//! fires. Repeated indefinitely, the client-timeout sweep kills an action +//! whose client is alive and sending keep-alives. + +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::SystemTime; + +use mock_instant::thread_local::MockClock; +use nativelink_config::stores::EvictionPolicy; +use nativelink_error::Error; +use nativelink_macro::nativelink_test; +use nativelink_scheduler::awaited_action_db::{ + AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, CLIENT_KEEPALIVE_DURATION, +}; +use nativelink_scheduler::memory_awaited_action_db::MemoryAwaitedActionDb; +use nativelink_util::action_messages::{ + ActionInfo, ActionUniqueKey, ActionUniqueQualifier, OperationId, +}; +use nativelink_util::common::DigestInfo; +use nativelink_util::digest_hasher::DigestHasherFunc; +use nativelink_util::instant_wrapper::MockInstantWrapped; +use pretty_assertions::assert_eq; +use tokio::sync::Notify; + +const INSTANCE_NAME: &str = "keepalive_clobber_instance"; + +fn make_action_info() -> Arc { + let insert = SystemTime::UNIX_EPOCH + Duration::from_secs(100); + Arc::new(ActionInfo { + command_digest: DigestInfo::new([7; 32], 1), + input_root_digest: DigestInfo::new([7; 32], 2), + timeout: Duration::from_secs(1), + platform_properties: HashMap::new(), + priority: 0, + load_timestamp: insert, + insert_timestamp: insert, + unique_qualifier: ActionUniqueQualifier::Cacheable(ActionUniqueKey { + execution_scope: None, + instance_name: INSTANCE_NAME.to_string(), + digest_function: DigestHasherFunc::Sha256, + digest: DigestInfo::new([7; 32], 3), + }), + }) +} + +/// Extract `last_client_keepalive_timestamp` from the Debug rendering of an +/// `AwaitedAction` (the field has no public accessor). +fn keepalive_of(awaited_action: &AwaitedAction) -> String { + let debug = format!("{awaited_action:?}"); + let field = "last_client_keepalive_timestamp: "; + let start = debug.find(field).expect("field in Debug output") + field.len(); + let rest = &debug[start..]; + let end = rest.find('}').expect("SystemTime Debug closes with }") + 1; + rest[..end].to_string() +} + +#[nativelink_test] +async fn client_keepalive_survives_concurrent_same_version_update() -> Result<(), Error> { + let db = MemoryAwaitedActionDb::new( + &EvictionPolicy::default(), + Arc::new(Notify::new()), + MockInstantWrapped::default, + ); + + let client_id = OperationId::from("client-keepalive"); + // The subscriber returned here carries client_info: its changed() loop is + // what emits ActionEvent::ClientKeepAlive. + let mut client_subscriber = db + .add_action( + client_id.clone(), + make_action_info(), + Duration::from_mins(1), + ) + .await?; + let operation_id = client_subscriber.borrow().await?.operation_id().clone(); + + // The "worker" side: snapshot the action (version N, keepalive K0), the + // way SimpleSchedulerStateManager does before update_awaited_action. + let worker_subscriber = db + .get_by_operation_id(&operation_id) + .await? + .expect("operation must exist"); + let stale_snapshot = worker_subscriber.borrow().await?; + let keepalive_k0 = keepalive_of(&stale_snapshot); + + // Advance the mock clock so (a) the client's changed() loop decides it is + // time to send a keep-alive, and (b) the stamped timestamp K1 differs + // from K0. + MockClock::advance(CLIENT_KEEPALIVE_DURATION + Duration::from_secs(1)); + + // Drive the client subscriber: changed() immediately sends + // ActionEvent::ClientKeepAlive, then parks awaiting a watch change. + let changed_task = tokio::spawn(async move { client_subscriber.changed().await }); + + // Wait (deterministically, bounded) for the event-handler task to stamp + // the keep-alive into the stored AwaitedAction. + let keepalive_k1 = { + let mut stamped = None; + for _ in 0..1000 { + tokio::task::yield_now().await; + let current = worker_subscriber.borrow().await?; + let ts = keepalive_of(¤t); + if ts != keepalive_k0 { + stamped = Some(ts); + break; + } + } + stamped.expect("client keep-alive was never stamped into the db") + }; + changed_task.abort(); + + // The keep-alive did NOT bump the version, so the worker's stale + // pre-keep-alive snapshot still passes the version CAS... + let update_result = db.update_awaited_action(stale_snapshot).await; + + // ...and send_replace overwrites the whole struct, reverting the + // keep-alive to K0. Correct behavior would be either rejecting the stale + // write with Code::Aborted (version bump on keep-alive) or merging the + // keep-alive timestamp. Either way, after an accepted update the stored + // keep-alive must still be K1. + let after = worker_subscriber.borrow().await?; + let keepalive_after = keepalive_of(&after); + assert!( + update_result.is_err() || keepalive_after == keepalive_k1, + "Lost update: stale same-version write was accepted (result: {update_result:?}) and \ + reverted last_client_keepalive_timestamp from {keepalive_k1} back to {keepalive_after} \ + (K0 was {keepalive_k0}). The Aborted-retry defense in \ + SimpleSchedulerStateManager can never fire, so the client-timeout \ + sweep can kill an action with a live client." + ); + // Redundant with the assert above, but gives a crisp diff on failure. + if update_result.is_ok() { + assert_eq!(keepalive_after, keepalive_k1); + } + Ok(()) +} diff --git a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs index 371487a668b..36d08c19c43 100644 --- a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs @@ -1276,6 +1276,160 @@ async fn fail_queued_operation_ignores_an_unknown_action() -> Result<(), Error> Ok(()) } +/// A database wrapper that shares the inner db so the test can act as a +/// concurrent updater (e.g. a worker-event task) racing the state manager. +struct SharedDb(Arc); + +impl MetricsComponent for SharedDb { + fn publish( + &self, + kind: MetricKind, + field_metadata: MetricFieldData, + ) -> Result { + self.0.publish(kind, field_metadata) + } +} + +impl AwaitedActionDb for SharedDb { + type Subscriber = T::Subscriber; + + async fn get_awaited_action_by_id( + &self, + client_operation_id: &OperationId, + ) -> Result, Error> { + self.0.get_awaited_action_by_id(client_operation_id).await + } + + async fn get_all_awaited_actions( + &self, + ) -> Result> + Send, Error> { + self.0.get_all_awaited_actions().await + } + + async fn get_by_operation_id( + &self, + operation_id: &OperationId, + ) -> Result, Error> { + self.0.get_by_operation_id(operation_id).await + } + + async fn get_range_of_actions( + &self, + state: SortedAwaitedActionState, + start: Bound, + end: Bound, + desc: bool, + ) -> Result> + Send, Error> { + self.0.get_range_of_actions(state, start, end, desc).await + } + + async fn update_awaited_action(&self, new_awaited_action: AwaitedAction) -> Result<(), Error> { + self.0.update_awaited_action(new_awaited_action).await + } + + async fn add_action( + &self, + client_operation_id: OperationId, + action_info: Arc, + no_event_action_timeout: Duration, + ) -> Result { + self.0 + .add_action(client_operation_id, action_info, no_event_action_timeout) + .await + } +} + +fn cacheable_action_info(started: SystemTime) -> ActionInfo { + let mut info = action_info(started); + let ActionUniqueQualifier::Uncacheable(key) = info.unique_qualifier.clone() else { + unreachable!(); + }; + info.unique_qualifier = ActionUniqueQualifier::Cacheable(key); + info +} + +/// Repro for: `try_subscribe`'s initial keep-alive stamp can be reverted by an +/// in-flight versioned update, letting the client-timeout sweep kill an +/// operation right after a new client subscribed. +/// +/// `try_subscribe` stamps the keep-alive with `send_if_modified` without +/// bumping the version, so an updater holding a pre-stamp `borrow()` snapshot +/// at the same version passes the version check in `update_awaited_action` +/// and `send_replace` wholesale-reverts the keep-alive timestamp. The sweep +/// then retires the operation even though a client just attached. +#[nativelink_test] +async fn a_fresh_subscribers_keepalive_survives_an_in_flight_update() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let task_change_notify = Arc::new(Notify::new()); + let db = Arc::new(memory_awaited_action_db_factory( + 0, + &task_change_notify, + MockInstantWrapped::default, + )); + let state_mgr = SimpleSchedulerStateManager::new( + 5, + WORKER_TIMEOUT, + Duration::from_mins(5), // client_action_timeout + MAX_EXECUTING, + SharedDb(db.clone()), + MockInstantWrapped::default, + Some(Arc::new(WorkerRegistry::new())), + ); + + // Client 1 queues a cacheable action at t0, then goes idle: its + // stored client keep-alive stays at t0. + let client1_id = OperationId::default(); + let _client1 = state_mgr + .add_action( + client1_id.clone(), + Arc::new(cacheable_action_info(make_system_time(0))), + ) + .await?; + + // A concurrent task (e.g. a worker-event handler) snapshots the + // action -- version N, keep-alive t0. + let subscriber = db + .get_awaited_action_by_id(&client1_id) + .await? + .expect("operation exists"); + let in_flight_snapshot = subscriber.borrow().await?; + + // t0+4min (within the 5min client timeout). + MockClock::advance(Duration::from_mins(4)); + + // A new client subscribes to the same cacheable action: try_subscribe + // stamps the keep-alive to now (t0+4min) so it is not swept before its + // first poll. + let client2_id = OperationId::default(); + let client2 = state_mgr + .add_action( + client2_id.clone(), + Arc::new(cacheable_action_info(make_system_time(0))), + ) + .await?; + + // The in-flight updater resumes and writes its snapshot back. The + // keep-alive stamp never bumped the version, so this passes the + // version check -- and reverts the keep-alive to t0. + db.update_awaited_action(in_flight_snapshot).await?; + + // t0+6min: client 1's keep-alive (t0) is expired, but client 2 + // attached only 2 minutes ago, so the operation must not be retired. + MockClock::advance(Duration::from_mins(2)); + let retired = state_mgr.sweep_abandoned_queued_actions().await?; + assert_eq!( + retired, 0, + "an operation a client subscribed to 2 minutes ago must not be swept" + ); + let (state, _origin_metadata) = client2.as_state().await?; + assert_eq!( + state.stage, + ActionStage::Queued, + "the freshly-subscribed client's operation must still be alive" + ); + Ok(()) +} + #[nativelink_test] async fn fail_queued_operation_ignores_a_completed_action() -> Result<(), Error> { MockClock::set_time(Duration::from_secs(NOW_TIME)); From 912b6d1319efa6baadca0b65ecedc4f1ab77b7c5 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 23:21:16 -0400 Subject: [PATCH 20/23] fix(scheduler): failed client-keepalive write no longer advances local timestamp In OperationSubscriber::changed (USE_SEPARATE_CLIENT_KEEPALIVE_KEY branch), a failed ck_* store write still advanced the local last_known_keepalive_ts and then parked for a full CLIENT_KEEPALIVE_DURATION. Peer replicas judge client liveness solely from the store, so a couple of consecutive transient store failures (Redis blip) left the store's keepalive stamp aging past client_action_timeout and a peer's sweep retired the live client's action as abandoned. Now the local timestamp only advances on a successful write; on failure the sleep becomes a bounded CLIENT_KEEPALIVE_RETRY_DELAY (1s) so the write is retried promptly without spinning when the store is down hard. Repro test fails before the fix (attempts = 1 across 3x keepalive periods) and passes after. Co-Authored-By: Claude Fable 5 --- .../src/store_awaited_action_db.rs | 59 ++-- .../client_keepalive_failed_write_test.rs | 260 ++++++++++++++++++ 2 files changed, 303 insertions(+), 16 deletions(-) create mode 100644 nativelink-scheduler/tests/client_keepalive_failed_write_test.rs diff --git a/nativelink-scheduler/src/store_awaited_action_db.rs b/nativelink-scheduler/src/store_awaited_action_db.rs index 41897e0ad45..9d1a7f0724f 100644 --- a/nativelink-scheduler/src/store_awaited_action_db.rs +++ b/nativelink-scheduler/src/store_awaited_action_db.rs @@ -56,6 +56,13 @@ const MAX_RETRIES_FOR_CLIENT_KEEPALIVE: u32 = 8; /// Use separate non-versioned Redis key for client keepalives. const USE_SEPARATE_CLIENT_KEEPALIVE_KEY: bool = true; +/// How long to wait before retrying a client keepalive write that failed +/// (e.g. a transient store error). The local `last_known_keepalive_ts` is +/// only advanced on a successful write, so failures naturally re-arm the +/// keepalive gate; this delay bounds the retry frequency so a hard store +/// outage does not spin the subscriber loop. +const CLIENT_KEEPALIVE_RETRY_DELAY: Duration = Duration::from_secs(1); + /// How often the actions in each stage are recounted for /// `execution.active.count`. const ACTIVE_COUNT_REFRESH_INTERVAL: Duration = Duration::from_secs(15); @@ -309,6 +316,10 @@ where loop { // This is set if the maybe_last_state doesn't match the state in the store. let mut maybe_changed_action = None; + // Set when a client keepalive store write fails so the sleep + // below becomes a bounded retry delay instead of a full + // CLIENT_KEEPALIVE_DURATION park. + let mut keepalive_write_failed = false; let last_known_keepalive_ts = self.last_known_keepalive_ts.load(Ordering::Acquire); // Only a subscriber that stands for a client may say a client is @@ -348,18 +359,27 @@ where ) .await; - if let Err(e) = update_result { - warn!( - ?self.subscription_key, - ?e, - "Failed to update client keepalive (non-versioned)" - ); + match update_result { + Ok(_) => { + // Only advance the local timestamp once the store + // actually holds the new keepalive. Advancing it on + // failure would suppress retries for a full + // CLIENT_KEEPALIVE_DURATION while the store's stamp + // keeps aging, letting a peer replica retire this + // live client's action as abandoned. + self.last_known_keepalive_ts + .store(now_ts, Ordering::Release); + } + Err(e) => { + keepalive_write_failed = true; + warn!( + ?self.subscription_key, + ?e, + "Failed to update client keepalive (non-versioned), will retry" + ); + } } - // Update local timestamp - self.last_known_keepalive_ts - .store(now_ts, Ordering::Release); - // Check if state changed (for unreliable subscription managers) if self.maybe_last_stage.is_some() { let awaited_action = Self::inner_get_awaited_action( @@ -430,12 +450,19 @@ where Some(core::mem::discriminant(&changed_action.state().stage)); return Ok(changed_action); } - // Determine the sleep time based on the last client keep alive. - let sleep_time = CLIENT_KEEPALIVE_DURATION - .checked_sub( - I::from_secs(self.last_known_keepalive_ts.load(Ordering::Acquire)).elapsed(), - ) - .unwrap_or(Duration::from_millis(100)); + // Determine the sleep time based on the last successful client + // keep alive write. If the last write failed, wake after a + // bounded retry delay instead of parking for the full period. + let sleep_time = if keepalive_write_failed { + CLIENT_KEEPALIVE_RETRY_DELAY + } else { + CLIENT_KEEPALIVE_DURATION + .checked_sub( + I::from_secs(self.last_known_keepalive_ts.load(Ordering::Acquire)) + .elapsed(), + ) + .unwrap_or(Duration::from_millis(100)) + }; tokio::select! { result = &mut changed_fut => { result?; diff --git a/nativelink-scheduler/tests/client_keepalive_failed_write_test.rs b/nativelink-scheduler/tests/client_keepalive_failed_write_test.rs new file mode 100644 index 00000000000..d783faf2bc2 --- /dev/null +++ b/nativelink-scheduler/tests/client_keepalive_failed_write_test.rs @@ -0,0 +1,260 @@ +// 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. + +// Repro for: a failed client-keepalive store write in +// `OperationSubscriber::changed` (USE_SEPARATE_CLIENT_KEEPALIVE_KEY branch) +// still advances the local `last_known_keepalive_ts`, so the retry gate +// (`elapsed() > CLIENT_KEEPALIVE_DURATION`) is suppressed for a full +// CLIENT_KEEPALIVE_DURATION even though the shared store's ck_* key was +// never written. Peer scheduler replicas judge client liveness solely from +// the store, so a couple of consecutive failed writes leave the store +// timestamp stale past `client_action_timeout` and the peer retires the +// action as abandoned while the client is actively listening. + +use core::sync::atomic::{AtomicUsize, Ordering}; +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::SystemTime; + +use bytes::Bytes; +use futures::{Stream, stream}; +use nativelink_error::{Code, Error, make_err}; +use nativelink_macro::nativelink_test; +use nativelink_scheduler::awaited_action_db::{ + AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, +}; +use nativelink_scheduler::store_awaited_action_db::StoreAwaitedActionDb; +use nativelink_util::action_messages::{ + ActionInfo, ActionUniqueKey, ActionUniqueQualifier, OperationId, +}; +use nativelink_util::common::DigestInfo; +use nativelink_util::digest_hasher::DigestHasherFunc; +use nativelink_util::store_trait::{ + SchedulerCurrentVersionProvider, SchedulerIndexProvider, SchedulerStore, + SchedulerStoreDataProvider, SchedulerStoreDecodeTo, SchedulerStoreKeyProvider, + SchedulerSubscription, SchedulerSubscriptionManager, StoreKey, +}; +use tokio::sync::Notify; + +const INSTANCE_NAME: &str = "foo"; + +struct PendingSubscription; +impl SchedulerSubscription for PendingSubscription { + async fn changed(&mut self) -> Result<(), Error> { + futures::future::pending().await + } +} + +struct PendingSubscriptionManager; +impl SchedulerSubscriptionManager for PendingSubscriptionManager { + type Subscription = PendingSubscription; + fn subscribe(&self, _key: K) -> Result + where + K: SchedulerStoreKeyProvider, + { + Ok(PendingSubscription) + } + fn is_reliable() -> bool { + true + } +} + +/// A store holding one queued `AwaitedAction` whose `ck_*` (client +/// keepalive) writes always fail — a Redis blip / dead shard. Every other +/// operation succeeds. +struct KeepaliveWriteFailsStore { + operation_id: OperationId, + encoded_action: Bytes, + keepalive_write_attempts: AtomicUsize, +} + +impl SchedulerStore for KeepaliveWriteFailsStore { + type SubscriptionManager = PendingSubscriptionManager; + + fn subscription_manager( + &self, + ) -> impl Future, Error>> { + std::future::ready(Ok(Arc::new(PendingSubscriptionManager))) + } + + async fn update_data(&self, data: T, _expiry: Option) -> Result, Error> + where + T: SchedulerStoreDataProvider + + SchedulerStoreKeyProvider + + SchedulerCurrentVersionProvider + + Send, + { + let StoreKey::Str(key) = data.get_key() else { + panic!("unexpected non-string scheduler key"); + }; + if key.starts_with("ck_") { + self.keepalive_write_attempts.fetch_add(1, Ordering::SeqCst); + return Err(make_err!( + Code::Unavailable, + "injected transient store failure for client keepalive write" + )); + } + Ok(Some(1)) + } + + fn search_by_index_prefix( + &self, + _index: K, + ) -> impl Future< + Output = Result< + impl Stream::DecodeOutput, Error>> + Send, + Error, + >, + > + where + K: SchedulerIndexProvider + SchedulerStoreDecodeTo + Send, + ::DecodeOutput: Send, + { + std::future::ready(Ok(stream::empty())) + } + + async fn count_by_index_prefix(&self, _index: K) -> Result + where + K: SchedulerIndexProvider + Send, + { + Ok(0) + } + + async fn get_and_decode( + &self, + key: K, + ) -> Result::DecodeOutput>, Error> + where + K: SchedulerStoreKeyProvider + SchedulerStoreDecodeTo + Send, + { + let StoreKey::Str(key) = key.get_key() else { + panic!("unexpected non-string scheduler key"); + }; + if key.starts_with("cid_") { + let bytes = Bytes::from(serde_json::to_vec(&self.operation_id).unwrap()); + return Ok(Some(K::decode(0, bytes)?)); + } + if key.starts_with("aa_") { + return Ok(Some(K::decode(1, self.encoded_action.clone())?)); + } + if key.starts_with("ck_") { + // No keepalive write ever landed in the store. + return Ok(None); + } + panic!("unexpected key {key}"); + } +} + +fn make_queued_action(operation_id: &OperationId) -> AwaitedAction { + let action_info = Arc::new(ActionInfo { + command_digest: DigestInfo::zero_digest(), + input_root_digest: DigestInfo::zero_digest(), + timeout: Duration::from_secs(1), + platform_properties: HashMap::new(), + priority: 0, + load_timestamp: SystemTime::UNIX_EPOCH, + insert_timestamp: SystemTime::UNIX_EPOCH, + unique_qualifier: ActionUniqueQualifier::Cacheable(ActionUniqueKey { + execution_scope: None, + instance_name: INSTANCE_NAME.to_string(), + digest_function: DigestHasherFunc::Sha256, + digest: DigestInfo::zero_digest(), + }), + }); + // The action's embedded keepalive timestamp is at the epoch: from the + // shared store's point of view the client has not checked in for a very + // long time, exactly as after a couple of failed ck_ writes. + AwaitedAction::new(operation_id.clone(), action_info, SystemTime::UNIX_EPOCH) +} + +/// A live client's subscriber whose keepalive store write fails must retry +/// the write well before peer replicas' `client_action_timeout` elapses; +/// otherwise the store timestamp stays stale and +/// `sweep_abandoned_queued_actions` / `apply_filter_predicate` on another +/// replica retires the action as "no more clients listening" out from under +/// a client that is actively listening. +/// +/// Uses `start_paused` tokio time: the subscriber's sleep between keepalive +/// attempts is virtual, so 30 virtual seconds (3x `CLIENT_KEEPALIVE_DURATION`) +/// elapse instantly, while the retry gate — which reads the *local* +/// `last_known_keepalive_ts` through real wall-clock `SystemTime::elapsed` — +/// sees essentially zero elapsed real time after the failed write advanced +/// it. This exposes exactly the bug: the local timestamp is advanced even +/// though the store write failed, so no retry happens for a full +/// `CLIENT_KEEPALIVE_DURATION`. +type TestDb = StoreAwaitedActionDb< + KeepaliveWriteFailsStore, + fn() -> OperationId, + SystemTime, + fn() -> SystemTime, +>; + +fn new_op_id() -> OperationId { + OperationId::from("unused-new-operation") +} + +#[nativelink_test(start_paused = true)] +async fn failed_keepalive_write_must_be_retried() -> Result<(), Error> { + let operation_id = OperationId::from("live-client-operation"); + let action = make_queued_action(&operation_id); + let store = Arc::new(KeepaliveWriteFailsStore { + operation_id: operation_id.clone(), + encoded_action: Bytes::from(serde_json::to_vec(&action).unwrap()), + keepalive_write_attempts: AtomicUsize::new(0), + }); + + let now_fn: fn() -> SystemTime = SystemTime::now; + let op_id_fn: fn() -> OperationId = new_op_id; + let db: TestDb = StoreAwaitedActionDb::new( + store.clone(), + Arc::new(Notify::new()), + now_fn, + op_id_fn, + 60, + 60, + false, + ) + .await + .expect("construct test db"); + + let mut subscriber = db + .get_awaited_action_by_id(&operation_id) + .await? + .expect("subscriber for live client"); + + // Drive the client's changed() loop for 30 virtual seconds — three full + // CLIENT_KEEPALIVE_DURATION periods. The subscription never fires + // (pending) and the action never changes, so the only thing the loop can + // do is maintain the keepalive. + let changed_fut = subscriber.changed(); + tokio::pin!(changed_fut); + tokio::select! { + result = &mut changed_fut => { + panic!("changed() must not resolve: nothing changed; got {result:?}"); + } + () = tokio::time::sleep(Duration::from_secs(30)) => {} + } + + let n = store.keepalive_write_attempts.load(Ordering::SeqCst); + assert!( + n >= 2, + "keepalive write failed but was never retried within 3x \ + CLIENT_KEEPALIVE_DURATION (attempts = {n}); the local \ + last_known_keepalive_ts was advanced despite the store Err, so the \ + shared-store keepalive stays stale and a peer replica will retire \ + this live client's action as abandoned (DeadlineExceeded)" + ); + Ok(()) +} From c8a7b7b6089f326bb3f61f2b8a289e3d76ac2ceb Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 19:24:49 -0400 Subject: [PATCH 21/23] scheduler: prefer a live candidate over a completed one in try_subscribe MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `try_subscribe` stopped at the first exact unique-qualifier match, even when that match was a completed record. The unique-qualifier index retains completed records alongside any live successor, and the backend does not order results by state, so a retained completed record could be returned ahead of a joinable live one — forcing an unnecessary recreate and a new operation instead of joining the live successor. Scan every exact match and prefer a joinable (neither finished nor abandoned) candidate, keeping the first exact match only as a fallback. This is the selection half of finding N7. The atomic qualifier->live-op claim (two concurrent same-qualifier submissions can still both create an operation) needs a store primitive that does not exist today (a compare-and-set on the shared qualifier key); that is tracked separately. Co-Authored-By: Claude Opus 4.8 --- .../src/store_awaited_action_db.rs | 27 ++- .../tests/store_awaited_action_db_test.rs | 174 +++++++++++++++++- 2 files changed, 199 insertions(+), 2 deletions(-) diff --git a/nativelink-scheduler/src/store_awaited_action_db.rs b/nativelink-scheduler/src/store_awaited_action_db.rs index 9d1a7f0724f..bc0030cf7c7 100644 --- a/nativelink-scheduler/src/store_awaited_action_db.rs +++ b/nativelink-scheduler/src/store_awaited_action_db.rs @@ -1030,15 +1030,40 @@ where tokio::pin!(stream); // This index search is by prefix. In particular, an unscoped key // is a prefix of its scoped variants and must not join those jobs. + // + // The index retains completed records alongside any live successor + // for the same qualifier, and the backend does not order results by + // state. Stopping at the first exact match can therefore hand back a + // finished record while a joinable live one is also present, forcing + // an unnecessary recreate. Scan every exact match and prefer a + // joinable (neither finished nor abandoned) candidate; keep the + // first exact match only as a fallback when none is joinable. + let mut fallback: Option = None; while let Some(candidate) = stream .try_next() .await .err_tip(|| "In RedisAwaitedActionDb::try_subscribe")? { - if &candidate.action_info().unique_qualifier == unique_qualifier { + if &candidate.action_info().unique_qualifier != unique_qualifier { + continue; + } + let abandoned = self + .executing_action_is_abandoned( + &candidate, + no_event_action_timeout, + (self.now_fn)().now(), + ) + .await; + if !candidate.state().stage.is_finished() && !abandoned { maybe_awaited_action = Some(candidate); break; } + if fallback.is_none() { + fallback = Some(candidate); + } + } + if maybe_awaited_action.is_none() { + maybe_awaited_action = fallback; } if maybe_awaited_action.is_some() { break; diff --git a/nativelink-scheduler/tests/store_awaited_action_db_test.rs b/nativelink-scheduler/tests/store_awaited_action_db_test.rs index 430a3fdb410..f5660c51696 100644 --- a/nativelink-scheduler/tests/store_awaited_action_db_test.rs +++ b/nativelink-scheduler/tests/store_awaited_action_db_test.rs @@ -21,7 +21,8 @@ use nativelink_scheduler::store_awaited_action_db::{ }; use nativelink_scheduler::worker_registry::{ORPHANED_ACTION_TIMEOUT, WorkerRegistry}; use nativelink_util::action_messages::{ - ActionInfo, ActionStage, ActionUniqueKey, ActionUniqueQualifier, OperationId, WorkerId, + ActionInfo, ActionResult, ActionStage, ActionUniqueKey, ActionUniqueQualifier, OperationId, + WorkerId, }; use nativelink_util::common::DigestInfo; use nativelink_util::digest_hasher::DigestHasherFunc; @@ -429,6 +430,177 @@ async fn try_subscribe_skips_lookup_for_uncacheable_qualifier() -> Result<(), Er Ok(()) } +/// Fake store whose single `search_by_index_prefix` returns a fixed list of +/// already-encoded candidates, in order. Unlike `EventuallyVisibleStore` this +/// exercises the multi-candidate selection: the unique-qualifier index retains +/// completed records alongside a live successor, and the backend does not order +/// by state, so `try_subscribe` must scan past a finished candidate to join a +/// live one rather than stopping at the first exact match. +struct MultiCandidateStore { + encoded: Vec, +} + +impl MultiCandidateStore { + fn new(actions: &[&AwaitedAction]) -> Self { + Self { + encoded: actions + .iter() + .map(|a| { + Bytes::from(serde_json::to_vec(a).expect("serialize AwaitedAction for fake")) + }) + .collect(), + } + } +} + +impl SchedulerStore for MultiCandidateStore { + type SubscriptionManager = PendingSubscriptionManager; + + fn subscription_manager( + &self, + ) -> impl Future, Error>> { + std::future::ready(Ok(Arc::new(PendingSubscriptionManager))) + } + + fn update_data( + &self, + _data: T, + _expiry: Option, + ) -> impl Future, Error>> + where + T: SchedulerStoreDataProvider + + SchedulerStoreKeyProvider + + SchedulerCurrentVersionProvider + + Send, + { + std::future::ready(Ok(Some(1))) + } + + fn search_by_index_prefix( + &self, + _index: K, + ) -> impl Future< + Output = Result< + impl Stream::DecodeOutput, Error>> + Send, + Error, + >, + > + where + K: SchedulerIndexProvider + SchedulerStoreDecodeTo + Send, + ::DecodeOutput: Send, + { + let items: Vec::DecodeOutput, Error>> = self + .encoded + .iter() + .map(|bytes| K::decode(1, bytes.clone())) + .collect(); + std::future::ready(Ok(stream::iter(items))) + } + + async fn count_by_index_prefix(&self, _index: K) -> Result + where + K: SchedulerIndexProvider + Send, + { + Ok(0) + } + + async fn get_and_decode( + &self, + _key: K, + ) -> Result::DecodeOutput>, Error> + where + K: SchedulerStoreKeyProvider + SchedulerStoreDecodeTo + Send, + { + todo!("not exercised by try_subscribe") + } +} + +async fn build_multi_db( + store: Arc, +) -> StoreAwaitedActionDb< + MultiCandidateStore, + fn() -> OperationId, + MockInstantWrapped, + fn() -> MockInstantWrapped, +> { + fn new_op_id() -> OperationId { + OperationId::from("new-operation") + } + let now_fn: fn() -> MockInstantWrapped = MockInstantWrapped::default; + let op_id_fn: fn() -> OperationId = new_op_id; + StoreAwaitedActionDb::new( + store, + Arc::new(Notify::new()), + now_fn, + op_id_fn, + 60, + 60, + false, + ) + .await + .expect("construct test db") +} + +fn make_completed_awaited_action(operation_id: &str) -> AwaitedAction { + let now = SystemTime::UNIX_EPOCH; + let mut action = AwaitedAction::new( + OperationId::from(operation_id), + make_cacheable_action_info(), + now, + ); + let mut state = action.state().as_ref().clone(); + state.stage = ActionStage::Completed(ActionResult::default()); + action.worker_set_state(Arc::new(state), now); + action +} + +fn make_queued_awaited_action(operation_id: &str) -> AwaitedAction { + // A fresh `AwaitedAction` is Queued and has no worker, so it is neither + // finished nor abandoned — i.e. joinable. + AwaitedAction::new( + OperationId::from(operation_id), + make_cacheable_action_info(), + SystemTime::UNIX_EPOCH, + ) +} + +/// Regression for the N7 selection defect: with a completed record indexed +/// ahead of a live successor for the same qualifier, `try_subscribe` must join +/// the live one, not return the completed record (which forces a recreate). +#[nativelink_test] +async fn try_subscribe_prefers_live_candidate_over_completed() -> Result<(), Error> { + let completed = make_completed_awaited_action("completed-operation"); + let live = make_queued_awaited_action("live-operation"); + let qualifier = live.action_info().unique_qualifier.clone(); + // Completed record first, live successor second — the order the backend can + // legitimately return. + let store = Arc::new(MultiCandidateStore::new(&[&completed, &live])); + let db = build_multi_db(store).await; + + let result = db + .try_subscribe( + &OperationId::from("client-prefer-live"), + &qualifier, + Duration::from_mins(1), + 0, + ) + .await? + .expect("a joinable live candidate is present"); + + assert_eq!( + *result.operation_id(), + OperationId::from("live-operation"), + "try_subscribe must join the live successor, not the completed record \ + (which would be recreated as a brand-new operation)", + ); + assert_ne!( + *result.operation_id(), + OperationId::from("new-operation"), + "joining a live candidate must not recreate a fresh operation", + ); + Ok(()) +} + // --------------------------------------------------------------------------- // `try_subscribe` abandonment check. // From 614c1916969acbd969ba74ff74e50e441f077029 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 23:21:36 -0400 Subject: [PATCH 22/23] Scheduler: kill_revoked_operations identifies a dispatch by generation, not timestamp kill_revoked_operations snapshots (worker, operation) pairs under the scheduler mutex, then confirms revocation against the state manager lock-free. If the operation is declined/requeued and re-dispatched to the same worker inside that window, the final locked re-guard used to pass on the fresh entry and the kill landed on the new, legitimate dispatch. An earlier version of this fix carried the entry's dispatched_at timestamp as the dispatch's identity, but a timestamp is not an identity: dispatched_at has second resolution (duration_since(UNIX_EPOCH).as_secs()), so a requeue-and-redispatch inside one wall-clock second carries the same dispatched_at and the stale kill still lands. Review caught the collision; the regression test now pins both dispatches to the same second and fails on the timestamp scheme. Each dispatch now carries a dispatch_generation: a scheduler-wide u64 incremented under the scheduler mutex at send time, stored on the running_action_infos entry. The sweep snapshots the generation and worker_notify_kill_operation only sends the kill if the entry under the lock still carries it; a re-dispatch always gets a strictly greater generation, so a stale kill can never match. The counter is scheduler-wide rather than per-worker: a per-worker counter would restart when a worker re-registers under the same id, letting a stale snapshot from the old incarnation match a fresh dispatch on the new one. dispatched_at stays as a timestamp for the acknowledgement timeout math. Also registers tests/kill_revoked_redispatch_race_test.rs in the Bazel rust_test_suite, which had not picked it up. Provenance: adversarial review fleet; incompleteness found by a follow-up review of #2838. Co-Authored-By: Claude Fable 5 --- nativelink-scheduler/BUILD.bazel | 1 + .../src/api_worker_scheduler.rs | 61 +++- nativelink-scheduler/src/worker.rs | 26 +- .../kill_revoked_redispatch_race_test.rs | 281 ++++++++++++++++++ 4 files changed, 352 insertions(+), 17 deletions(-) create mode 100644 nativelink-scheduler/tests/kill_revoked_redispatch_race_test.rs diff --git a/nativelink-scheduler/BUILD.bazel b/nativelink-scheduler/BUILD.bazel index 8f2f4f632d8..071115b3d0a 100644 --- a/nativelink-scheduler/BUILD.bazel +++ b/nativelink-scheduler/BUILD.bazel @@ -79,6 +79,7 @@ rust_test_suite( "tests/execution_active_count_test.rs", "tests/fleet_capabilities_test.rs", "tests/historical_resource_scheduler_test.rs", + "tests/kill_revoked_redispatch_race_test.rs", "tests/live_memory_veto_test.rs", "tests/placement_test.rs", "tests/platform_property_manager_test.rs", diff --git a/nativelink-scheduler/src/api_worker_scheduler.rs b/nativelink-scheduler/src/api_worker_scheduler.rs index cc46677c76d..706df7d8670 100644 --- a/nativelink-scheduler/src/api_worker_scheduler.rs +++ b/nativelink-scheduler/src/api_worker_scheduler.rs @@ -254,6 +254,14 @@ struct ApiWorkerSchedulerImpl { /// about it without waiting for the next periodic exchange. local_fleet_change_notify: Arc, + /// The next dispatch generation; incremented under this struct's mutex + /// for every dispatch, so each dispatch of any operation to any worker + /// carries a unique, strictly increasing identity. Scheduler-wide rather + /// than per-worker: a per-worker counter would restart when a worker + /// re-registers under the same id, letting a stale kill snapshot from + /// the old incarnation match a fresh dispatch on the new one. + next_dispatch_generation: u64, + /// Whether an action with a given property shape could run on some /// connected worker when that worker is idle. `None` means it could. This /// depends only on what workers registered with, so it holds until the @@ -1114,12 +1122,17 @@ impl ApiWorkerSchedulerImpl { action_info: ActionInfoWithProps, dispatched_at: WorkerTimestamp, ) -> Result<(), Error> { + // Allocated under the scheduler mutex: a re-dispatch always gets a + // strictly greater generation than any snapshot taken before it. + self.next_dispatch_generation += 1; + let dispatch_generation = self.next_dispatch_generation; if let Some(worker) = self.workers.get_mut(&worker_id) { let notify_worker_result = worker .notify_update(WorkerUpdate::RunAction(Box::new(( operation_id.clone(), action_info.clone(), dispatched_at, + dispatch_generation, )))) .await; @@ -1298,13 +1311,24 @@ impl ApiWorkerSchedulerImpl { &mut self, worker_id: &WorkerId, operation_id: OperationId, + snapshot_dispatch_generation: u64, ) -> Result<(), Error> { let Some(worker) = self.workers.get_mut(worker_id) else { // Gone between the snapshot and now; its actions were requeued. return Ok(()); }; // Already told, or finished in the meantime; nothing more to send. - if !worker.running_action_infos.contains_key(&operation_id) + let Some(pending_action_info) = worker.running_action_infos.get(&operation_id) else { + return Ok(()); + }; + // A different generation means the snapshot's dispatch was requeued + // and the operation legitimately re-dispatched to this worker while + // the sweep ran lock-free; the revocation evidence belongs to the + // old dispatch, not this one. The generation, not `dispatched_at`, + // is the dispatch's identity: the timestamp has second resolution, + // so a requeue-and-redispatch inside one wall-clock second carries + // the same `dispatched_at` and would be killed by mistake. + if pending_action_info.dispatch_generation != snapshot_dispatch_generation || worker.is_kill_requested(&operation_id) { return Ok(()); @@ -1478,6 +1502,7 @@ impl ApiWorkerScheduler { peer_fleet_known_since: None, peer_fleet_last_refresh: None, local_fleet_change_notify: local_fleet_change_notify.clone(), + next_dispatch_generation: 0, static_verdicts: HashMap::new(), }), platform_property_manager, @@ -2118,7 +2143,7 @@ impl WorkerScheduler for ApiWorkerScheduler { async fn kill_revoked_operations(&self) -> Result<(), Error> { let (worker_state_manager, running) = { let inner = self.inner.lock().await; - let running: Vec<(WorkerId, OperationId)> = inner + let running: Vec<(WorkerId, OperationId, u64)> = inner .workers .iter() .flat_map(|(worker_id, worker)| { @@ -2135,7 +2160,16 @@ impl WorkerScheduler for ApiWorkerScheduler { .filter(|(_, pending_action_info)| { pending_action_info.kill_requested_at.is_none() }) - .map(|(operation_id, _)| (worker_id.clone(), operation_id.clone())) + // `dispatch_generation` identifies the dispatch + // this snapshot saw; the kill is only sent if the + // entry still carries it when the lock is re-taken. + .map(|(operation_id, pending_action_info)| { + ( + worker_id.clone(), + operation_id.clone(), + pending_action_info.dispatch_generation, + ) + }) }) .collect(); (inner.worker_state_manager.clone(), running) @@ -2143,8 +2177,8 @@ impl WorkerScheduler for ApiWorkerScheduler { // On store-backed deployments each check is a network round-trip, // so run them lock-free with bounded concurrency. - let revoked: Vec<(WorkerId, OperationId)> = futures::stream::iter(running) - .map(|(worker_id, operation_id)| { + let revoked: Vec<(WorkerId, OperationId, u64)> = futures::stream::iter(running) + .map(|(worker_id, operation_id, dispatch_generation)| { let worker_state_manager = worker_state_manager.clone(); async move { match worker_state_manager @@ -2152,7 +2186,7 @@ impl WorkerScheduler for ApiWorkerScheduler { .await { Ok(true) => None, - Ok(false) => Some((worker_id, operation_id)), + Ok(false) => Some((worker_id, operation_id, dispatch_generation)), // Only kill on positive evidence; try again next pass. Err(err) => { warn!( @@ -2177,10 +2211,13 @@ impl WorkerScheduler for ApiWorkerScheduler { // Re-check lock-free so the scheduler mutex is never held across // store I/O; the checks above may be stale by the time we get here. - // The remaining TOCTOU window is fine: worker_notify_kill_operation - // re-guards with contains_key + is_kill_requested under the lock. + // The remaining TOCTOU window is closed under the lock: + // worker_notify_kill_operation only kills the entry whose + // `dispatch_generation` matches the snapshot, so an operation + // requeued and re-dispatched to the same worker while the sweep ran + // lock-free is left alone. let mut confirmed = Vec::new(); - for (worker_id, operation_id) in revoked { + for (worker_id, operation_id, dispatch_generation) in revoked { // A result of Ok(true) or Err(_) means the operation was // reassigned to this worker after the first check, or the state // manager went quiet: nothing to kill on this pass. @@ -2189,7 +2226,7 @@ impl WorkerScheduler for ApiWorkerScheduler { .await .is_ok_and(|executing| !executing); if revoked { - confirmed.push((worker_id, operation_id)); + confirmed.push((worker_id, operation_id, dispatch_generation)); } } @@ -2199,10 +2236,10 @@ impl WorkerScheduler for ApiWorkerScheduler { let mut inner = self.inner.lock().await; let mut result = Ok(()); - for (worker_id, operation_id) in confirmed { + for (worker_id, operation_id, dispatch_generation) in confirmed { result = result.merge( inner - .worker_notify_kill_operation(&worker_id, operation_id) + .worker_notify_kill_operation(&worker_id, operation_id, dispatch_generation) .await, ); } diff --git a/nativelink-scheduler/src/worker.rs b/nativelink-scheduler/src/worker.rs index d2ef4b288b2..b1997864c3d 100644 --- a/nativelink-scheduler/src/worker.rs +++ b/nativelink-scheduler/src/worker.rs @@ -54,8 +54,8 @@ pub struct ActionInfoWithProps { #[derive(Debug)] pub enum WorkerUpdate { /// Requests that the worker begin executing this action, dispatched at - /// this scheduler time. - RunAction(Box<(OperationId, ActionInfoWithProps, WorkerTimestamp)>), + /// this scheduler time with this scheduler-wide dispatch generation. + RunAction(Box<(OperationId, ActionInfoWithProps, WorkerTimestamp, u64)>), /// Request that the worker is no longer in the pool and may discard any jobs. Disconnect, @@ -79,6 +79,15 @@ pub struct PendingActionInfoData { /// The worker's last-seen timestamp when the dispatch was sent. #[metric(help = "When this operation was dispatched to the worker.")] pub dispatched_at: WorkerTimestamp, + /// The identity of this dispatch: a scheduler-wide monotonic counter + /// allocated under the scheduler mutex at send time. `dispatched_at` + /// has second resolution and so cannot distinguish a requeue and + /// re-dispatch inside one wall-clock second; the generation is strictly + /// increasing per dispatch, so a snapshot of it names exactly one + /// dispatch. Scheduler-wide (not per-worker) so a worker that + /// disconnects and re-registers under the same id can never reissue a + /// generation an old snapshot still holds. + pub dispatch_generation: u64, /// The worker said it took the action. A worker that does not speak /// the acknowledgement never sets this, which is why the /// unacknowledged sweep is opt-in. @@ -260,9 +269,14 @@ impl Worker { pub async fn notify_update(&mut self, worker_update: WorkerUpdate) -> Result<(), Error> { match worker_update { WorkerUpdate::RunAction(action) => { - let (operation_id, action_info, dispatched_at) = *action; - self.run_action(operation_id, action_info, dispatched_at) - .await + let (operation_id, action_info, dispatched_at, dispatch_generation) = *action; + self.run_action( + operation_id, + action_info, + dispatched_at, + dispatch_generation, + ) + .await } WorkerUpdate::Disconnect => { self.metrics.notify_disconnect.inc(); @@ -366,6 +380,7 @@ impl Worker { operation_id: OperationId, action_info: ActionInfoWithProps, dispatched_at: WorkerTimestamp, + dispatch_generation: u64, ) -> Result<(), Error> { // Reject-if-exists: a same-worker retry can dispatch an op whose // earlier attempt is still resident here (e.g. an execution-deadline @@ -410,6 +425,7 @@ impl Worker { action_info, kill_requested_at: None, dispatched_at, + dispatch_generation, accepted: false, last_usage: None, }, diff --git a/nativelink-scheduler/tests/kill_revoked_redispatch_race_test.rs b/nativelink-scheduler/tests/kill_revoked_redispatch_race_test.rs new file mode 100644 index 00000000000..3b221a20b75 --- /dev/null +++ b/nativelink-scheduler/tests/kill_revoked_redispatch_race_test.rs @@ -0,0 +1,281 @@ +// 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. + +//! Deterministic reproduction of the requeue-then-redispatch TOCTOU in +//! `kill_revoked_operations` (`api_worker_scheduler.rs`): the sweep drops the +//! inner mutex between its lock-free `is_executing_on_worker` confirmation +//! and the final locked `worker_notify_kill_operation` loop. If the matcher +//! re-dispatches the same operation to the same worker inside that gap, the +//! re-guard (`contains_key` + `is_kill_requested`) passes on the fresh +//! entry (`kill_requested_at: None`) and the kill lands on live work. + +mod utils { + pub(crate) mod scheduler_utils; +} + +use core::sync::atomic::{AtomicUsize, Ordering}; +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; + +use async_trait::async_trait; +use nativelink_config::schedulers::WorkerAllocationStrategy; +use nativelink_error::Error; +use nativelink_macro::nativelink_test; +use nativelink_metric::{ + MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, +}; +use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::update_for_worker; +use nativelink_scheduler::api_worker_scheduler::ApiWorkerScheduler; +use nativelink_scheduler::platform_property_manager::PlatformPropertyManager; +use nativelink_scheduler::worker::{ActionInfoWithProps, Worker}; +use nativelink_scheduler::worker_registry::WorkerRegistry; +use nativelink_scheduler::worker_scheduler::WorkerScheduler; +use nativelink_util::action_messages::{OperationId, WorkerId}; +use nativelink_util::common::DigestInfo; +use nativelink_util::operation_state_manager::{Decline, UpdateOperationType, WorkerStateManager}; +use nativelink_util::origin_event::OriginMetadata; +use nativelink_util::platform_properties::PlatformProperties; +use nativelink_util::spawn; +use tokio::sync::{Notify, mpsc}; +use utils::scheduler_utils::make_base_action_info; + +/// A state manager scripted to answer `is_executing_on_worker` truthfully +/// for the interleaving under test, with rendezvous points so the test can +/// act inside the sweep's lock-free windows. +struct ScriptedStateManager { + op_x: OperationId, + x_calls: AtomicUsize, + y_calls: AtomicUsize, + /// Sweep reached X's first (snapshot-verification) check; the snapshot + /// of running operations has been taken by now. + x1_reached: Notify, + /// Test releases X's first check (after the decline requeued X). + x1_gate: Notify, + /// Test releases Y's first check. + y1_gate: Notify, + /// X's confirmation check has returned false; `confirmed` now contains + /// X and the sweep still holds no lock. + x2_done: Notify, + /// Test releases Y's confirmation check (the await point that models + /// the gap before the sweep re-acquires the inner mutex). + y2_gate: Notify, +} + +impl MetricsComponent for ScriptedStateManager { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } +} + +#[async_trait] +impl WorkerStateManager for ScriptedStateManager { + async fn update_operation( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + _update: UpdateOperationType, + ) -> Result<(), Error> { + // The decline requeues the operation in the real state manager; + // nothing to record here. + Ok(()) + } + + async fn is_executing_on_worker( + &self, + operation_id: &OperationId, + _worker_id: &WorkerId, + ) -> Result { + if *operation_id == self.op_x { + if self.x_calls.fetch_add(1, Ordering::SeqCst) == 0 { + self.x1_reached.notify_one(); + self.x1_gate.notified().await; + // Truthful: X was declined and requeued. + } else { + // Confirmation check. Truthful: X is still queued (the + // re-dispatch has not happened yet -- it happens right + // after this returns, in the gap before the sweep + // takes the inner lock). + self.x2_done.notify_one(); + } + Ok(false) + } else if self.y_calls.fetch_add(1, Ordering::SeqCst) == 0 { + self.y1_gate.notified().await; + // Truthful: Y was declined and requeued. + Ok(false) + } else { + // Confirmation check for Y: the sweep awaits here after + // X's confirmation returned. This is the window in + // which the matcher re-dispatches X. + self.y2_gate.notified().await; + // Y got reassigned elsewhere: skip it this pass. + Ok(true) + } + } +} + +fn make_action(digest_byte: u8) -> ActionInfoWithProps { + ActionInfoWithProps { + inner: make_base_action_info( + std::time::UNIX_EPOCH, + DigestInfo::new([digest_byte; 32], 512), + ), + platform_properties: PlatformProperties::default(), + origin_metadata: OriginMetadata::default(), + scheduler_start_execute_event_id: None, + } +} + +#[nativelink_test] +async fn kill_revoked_operations_must_not_kill_a_redispatched_operation() -> Result<(), Error> { + let worker_id = WorkerId("worker1".to_string()); + let op_x = OperationId::default(); + let op_y = OperationId::default(); + + let scripted = Arc::new(ScriptedStateManager { + op_x: op_x.clone(), + x_calls: AtomicUsize::new(0), + y_calls: AtomicUsize::new(0), + x1_reached: Notify::new(), + x1_gate: Notify::new(), + y1_gate: Notify::new(), + x2_done: Notify::new(), + y2_gate: Notify::new(), + }); + + let scheduler = ApiWorkerScheduler::new( + scripted.clone(), + Arc::new(PlatformPropertyManager::new(HashMap::new())), + WorkerAllocationStrategy::default(), + None, + None, + Arc::new(Notify::new()), + 100, // worker_timeout_s + 100, // unacknowledged_kill_timeout_s + 100, // dispatch_ack_timeout_s + Arc::new(WorkerRegistry::new()), + None, + false, + Duration::from_mins(1), + ); + + let (tx, mut rx) = mpsc::channel(16); + scheduler + .add_worker(Worker::new( + worker_id.clone(), + PlatformProperties::default(), + tx, + 0, + 4, + )) + .await?; + // Drain the initial connection message. + drop(rx.recv().await); + + // Dispatch X and Y to the worker. + scheduler + .worker_notify_run_action(worker_id.clone(), op_x.clone(), make_action(1), 0) + .await?; + scheduler + .worker_notify_run_action(worker_id.clone(), op_y.clone(), make_action(2), 0) + .await?; + // Drain the two StartAction messages. + drop(rx.recv().await); + drop(rx.recv().await); + + // The sweep runs concurrently, exactly as the worker_api_server + // interval task does. + let sweep = spawn!("kill_revoked_sweep", { + let scheduler = scheduler.clone(); + async move { scheduler.kill_revoked_operations().await } + }); + + // The sweep has snapshotted (W, X) and (W, Y) with + // kill_requested_at: None and reached its lock-free checks. + scripted.x1_reached.notified().await; + + // The worker declines both operations (e.g. the unacknowledged sweep or + // load backpressure): complete_action removes them from + // running_action_infos and the state manager requeues them. + scheduler + .update_action( + &worker_id, + &op_x, + UpdateOperationType::UpdateWithDecline(Decline { + reason: "test".to_string(), + }), + ) + .await?; + scheduler + .update_action( + &worker_id, + &op_y, + UpdateOperationType::UpdateWithDecline(Decline { + reason: "test".to_string(), + }), + ) + .await?; + + // Release the first-round checks, X strictly before Y so the sweep's + // `revoked` list is ordered [X, Y]. + scripted.x1_gate.notify_one(); + for _ in 0..16 { + tokio::task::yield_now().await; + } + scripted.y1_gate.notify_one(); + + // Wait until X's confirmation check has returned false (X is in + // `confirmed`) while the sweep is parked on Y's confirmation check, + // holding no lock. This is the TOCTOU window. + scripted.x2_done.notified().await; + + // The matching engine re-dispatches X to the same worker: a fresh, + // legitimate dispatch with kill_requested_at: None. The wall-clock + // timestamp is the SAME second as the original dispatch: `dispatched_at` + // has second resolution, so a requeue-and-redispatch inside one second + // is indistinguishable by timestamp. Only a per-dispatch generation + // tells the two dispatches apart. + scheduler + .worker_notify_run_action(worker_id.clone(), op_x.clone(), make_action(1), 0) + .await?; + match rx.recv().await.unwrap().update { + Some(update_for_worker::Update::StartAction(_)) => {} + v => panic!("Expected StartAction for the re-dispatch, got: {v:?}"), + } + + // Let the sweep finish: it re-acquires the inner lock and runs + // worker_notify_kill_operation for X. + scripted.y2_gate.notify_one(); + sweep.await.unwrap()?; + + // The fresh dispatch must not be killed: the state manager considers X + // legitimately executing on this worker again. On buggy code the + // re-guard passes on the NEW entry and a KillOperationRequest lands on + // live work (and the worker's subsequent report is dropped by + // update_action's is_kill_requested check). + match rx.try_recv() { + Err(mpsc::error::TryRecvError::Empty) => {} + Ok(update) => panic!( + "The freshly re-dispatched operation was killed by the stale sweep: {:?}", + update.update + ), + Err(err) => panic!("Worker channel unexpectedly closed: {err:?}"), + } + + Ok(()) +} From 0ac94135c1185508a78bc761623b44b8bd414605 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 15:03:07 -0400 Subject: [PATCH 23/23] migrate: worker_retry_reservation_test uses 4-tuple RunAction 1.7.3 WorkerUpdate::RunAction carries a dispatch_generation (the 4th tuple element the kill_revoked generation fix relies on); the N2 same-worker-redispatch test predated it. Co-Authored-By: Claude Opus 4.8 --- nativelink-scheduler/tests/worker_retry_reservation_test.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/nativelink-scheduler/tests/worker_retry_reservation_test.rs b/nativelink-scheduler/tests/worker_retry_reservation_test.rs index 8c734fd567a..51f33aa8f5f 100644 --- a/nativelink-scheduler/tests/worker_retry_reservation_test.rs +++ b/nativelink-scheduler/tests/worker_retry_reservation_test.rs @@ -91,6 +91,7 @@ async fn same_worker_retry_conserves_resource_budget() -> Result<(), Error> { op.clone(), action_needing(4), NOW_TIME, + 0, )))) .await?; assert_eq!( @@ -107,6 +108,7 @@ async fn same_worker_retry_conserves_resource_budget() -> Result<(), Error> { op.clone(), action_needing(4), NOW_TIME, + 1, )))) .await; // The worker may accept or reject the redundant dispatch; either way the