From 46e058d15e6fe41b0eb16df98cc5e19c001b0b7a Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 7 Oct 2026 16:21:11 -0400 Subject: [PATCH 01/46] ci: defer the main.yaml integration-tests OCI lane off per-PR Completes the no-per-PR-OCI-builds intent: image-off-pr moved image.yaml off PRs and #2877 moved the nix integration job off PRs, but main.yaml's integration-tests (docker-compose image build + run_integration_tests.sh) fell between them and still ran on every PR. Defer to nightly + merge-to-main, matching #2880's pattern. Fork-local until folded into the upstream OCI/nightly set. --- .github/workflows/main.yaml | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/.github/workflows/main.yaml b/.github/workflows/main.yaml index 43dacaf1ea0..232952a3a33 100644 --- a/.github/workflows/main.yaml +++ b/.github/workflows/main.yaml @@ -1,17 +1,23 @@ name: CI # Controls when the workflow will run. +# +# The only live job here (integration-tests) builds the trace_machina/nativelink +# OCI image from the docker-compose Dockerfile and runs run_integration_tests.sh. +# That is a per-PR OCI build — the same class of heavy work #image-off-pr moved +# off PRs for image.yaml and #2877 moved off PRs for the nix integration job. +# This lane fell in the crack between those two, so defer it the same way: +# nightly cross-check + on every merge to main, not on every PR. Per-PR +# build/test signal is carried by the Nix "Bazel Dev" lane. on: push: branches: [ main ] paths-ignore: - '.github/styles/**' - 'web/**' - pull_request: - branches: [ main ] - paths-ignore: - - '.github/styles/**' - - 'web/**' + schedule: + - cron: '0 7 * * *' # 07:00 UTC daily (off-peak) + workflow_dispatch: permissions: read-all From ad812abc85e78d6902ec4fe41646703a5552f4b3 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 23:19:37 -0400 Subject: [PATCH 02/46] 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 69afdd02272..f9601addfc9 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 @@ -1459,10 +1464,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. @@ -1540,6 +1560,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, @@ -1576,6 +1609,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) } } @@ -1754,3 +1790,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 2eb11576a16b7d188f8a3b75afff444f9307e3ca Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 23:01:58 -0400 Subject: [PATCH 03/46] 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 99795eb35bc..f2971aad85b 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..bea779f915c --- /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 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: 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 b0257c07618..fa3e14bd8cd 100644 --- a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs @@ -1297,6 +1297,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 b6b2dc1d2a61083fa7edfac33f107ab8662a23de Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 23:01:48 -0400 Subject: [PATCH 04/46] fix(scheduler): evict the old incarnation when a worker re-registers under a live id ApiWorkerSchedulerImpl::add_worker used LruCache::put, which silently replaced a live entry when a worker re-registered under an existing id. The replaced Worker was dropped together with its running_action_infos: nothing was requeued, the old channel never heard Disconnect, and the gauges were only partially adjusted. Downstream, the replaced worker's operations stayed Executing assigned to the id; when the new incarnation was later evicted its empty list requeued nothing, and the eventual timeout requeued the operation to another worker while the original process might still be running it -- double execution. Now ApiWorkerScheduler::add_worker treats a live entry under the same id exactly like a disconnect before admitting the new incarnation: it goes through immediate_evict_worker, which drains running_action_infos via UpdateOperationType::UpdateWithDisconnect, sends Disconnect on the old channel, and balances the connect/disconnect gauges. A requeue failure is logged and does not block the reconnecting worker. Deterministically reproduced by the DST fuzzer schedule in PR #2831; the new regression test (worker re-registers mid-execution, asserts Disconnect on the old channel, re-dispatch of the same operation to the new channel, and completion) fails on unfixed main. Co-Authored-By: Claude Fable 5 --- .../src/api_worker_scheduler.rs | 37 ++++++- .../tests/simple_scheduler_test.rs | 103 ++++++++++++++++++ 2 files changed, 138 insertions(+), 2 deletions(-) diff --git a/nativelink-scheduler/src/api_worker_scheduler.rs b/nativelink-scheduler/src/api_worker_scheduler.rs index 64d7439e449..58a46fa96a2 100644 --- a/nativelink-scheduler/src/api_worker_scheduler.rs +++ b/nativelink-scheduler/src/api_worker_scheduler.rs @@ -546,8 +546,11 @@ impl ApiWorkerSchedulerImpl { fn add_worker(&mut self, worker: Worker) -> Result<(), Error> { let worker_id = worker.id.clone(); let platform_properties = worker.platform_properties.clone(); - // A replacement is one out and one in, so the gauge only moves for a - // genuinely new worker. + // The caller evicts a live entry under this id before admitting the + // new incarnation (see `ApiWorkerScheduler::add_worker`), so `put` + // does not normally replace anything. The branch below is kept as a + // defensive path: a replacement is one out and one in, so the gauge + // only moves for a genuinely new worker. let replaced = self.workers.put(worker_id.clone(), worker); // Add to capability index for fast matching @@ -2538,6 +2541,36 @@ impl WorkerScheduler for ApiWorkerScheduler { "Received request to add worker while shutting down" )); } + // A worker re-registering under an id the pool still holds is a new + // incarnation of the process, not the one the pool knows. Silently + // replacing the entry dropped the old `Worker` together with its + // `running_action_infos`: nothing was requeued, the old channel never + // heard a Disconnect, and its operations stayed Executing against + // this id until they timed out and were requeued to another worker + // while the original process might still be running them (double + // execution). Evict the old incarnation exactly like a disconnect + // first, then admit the new one. A failure to requeue must not block + // the reconnecting worker, so it is logged and the add proceeds. + if inner.workers.contains(&worker_id) { + let evict_res = inner + .immediate_evict_worker( + &worker_id, + make_err!( + Code::Internal, + "Worker {worker_id} re-registered; evicting the old incarnation" + ), + true, + WorkerDisconnectReason::Disconnected, + ) + .await; + if let Err(err) = evict_res { + warn!( + ?worker_id, + ?err, + "Failed to requeue everything the replaced worker incarnation was running" + ); + } + } let result = inner .add_worker(worker) .err_tip(|| "Error while adding worker, removing from pool"); diff --git a/nativelink-scheduler/tests/simple_scheduler_test.rs b/nativelink-scheduler/tests/simple_scheduler_test.rs index cd4237f029f..1f6f481cb40 100644 --- a/nativelink-scheduler/tests/simple_scheduler_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_test.rs @@ -1545,6 +1545,109 @@ async fn update_action_sends_completed_result_to_client_test() -> Result<(), Err Ok(()) } +/// A worker re-registering under an id the pool still holds must be treated +/// as a disconnect of the old incarnation: what it was running is requeued, +/// its channel hears Disconnect, and the new incarnation can pick the work +/// up and complete it. Silent replacement used to drop the old `Worker` with +/// its `running_action_infos`, leaving the operation Executing against an +/// incarnation that never held it (double execution after timeout). +#[nativelink_test] +async fn worker_reregistration_evicts_old_incarnation_and_requeues_test() -> Result<(), Error> { + let worker_id = WorkerId("worker_id".to_string()); + + let task_change_notify = Arc::new(Notify::new()); + let (scheduler, _worker_scheduler) = SimpleScheduler::new_with_callback( + &SimpleSpec::default(), + memory_awaited_action_db_factory( + 0, + &task_change_notify.clone(), + MockInstantWrapped::default, + ), + || async move {}, + task_change_notify, + MockInstantWrapped::default, + None, + ); + let action_digest = DigestInfo::new([99u8; 32], 512); + + let mut rx_from_old_worker = + setup_new_worker(&scheduler, worker_id.clone(), PlatformProperties::default()).await?; + let insert_timestamp = make_system_time(1); + let mut action_listener = + setup_action(&scheduler, action_digest, HashMap::new(), insert_timestamp).await?; + + let operation_id = { + // Other tests check full data. We only care if we got StartAction. + match rx_from_old_worker.recv().await.unwrap().update { + Some(update_for_worker::Update::StartAction(start_execute)) => { + assert_eq!( + action_listener.changed().await.unwrap().0.stage, + ActionStage::Executing + ); + OperationId::from(start_execute.operation_id) + } + v => panic!("Expected StartAction, got : {v:?}"), + } + }; + + // The same worker id re-registers on a fresh channel, as a restarted + // worker process does. + let mut rx_from_new_worker = + setup_new_worker(&scheduler, worker_id.clone(), PlatformProperties::default()).await?; + scheduler.do_try_match_for_test().await?; + + { + // The old incarnation's channel must hear a Disconnect. + let msg_for_worker = rx_from_old_worker.recv().await.unwrap(); + assert_eq!( + msg_for_worker, + UpdateForWorker { + update: Some(update_for_worker::Update::Disconnect(())) + } + ); + } + { + // The operation must be requeued and re-dispatched to the new + // incarnation, not left Executing against the dropped one. + match rx_from_new_worker.recv().await.unwrap().update { + Some(update_for_worker::Update::StartAction(start_execute)) => { + assert_eq!( + OperationId::from(start_execute.operation_id), + operation_id, + "expected the same operation re-dispatched to the new incarnation" + ); + assert_eq!(start_execute.worker_id, worker_id.to_string()); + } + v => panic!("Expected StartAction on the new channel, got : {v:?}"), + } + } + + // The new incarnation completes the action and the client sees it. + let action_result = ActionResult::default(); + scheduler + .update_action( + &worker_id, + &operation_id, + UpdateOperationType::UpdateWithActionStage(ActionStage::Completed( + action_result.clone(), + )), + ) + .await?; + + // The listener may see intermediate transitions from the requeue; wait + // for the final one. + let mut stage = action_listener.changed().await.unwrap().0.stage.clone(); + for _ in 0..8 { + if stage == ActionStage::Completed(action_result.clone()) { + break; + } + stage = action_listener.changed().await.unwrap().0.stage.clone(); + } + assert_eq!(stage, ActionStage::Completed(action_result)); + + Ok(()) +} + #[nativelink_test] async fn update_action_sends_completed_result_after_disconnect() -> Result<(), Error> { let worker_id = WorkerId("worker_id".to_string()); From 1f21271303f833bd8fc3253e687772df812f53cb Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 23:02:49 -0400 Subject: [PATCH 05/46] fix(scheduler): re-validate worker liveness under the eviction lock in remove_timedout_workers remove_timedout_workers had a TOCTOU window: it snapshots each worker's last_update_timestamp under the inner lock, releases the lock to consult the worker registry (an await boundary every sweep crosses), then re-acquires the lock and calls immediate_evict_worker using the stale snapshot. A keepalive processed in that window refreshes last_update_timestamp under the very same mutex, yet the sweep never re-consults it, so a demonstrably-alive worker is evicted and all of its running actions are requeued. Fix: after re-acquiring the inner lock, re-read the worker's last_update_timestamp and skip the eviction if it has moved past the timeout threshold. Because keepalives update the timestamp under the same mutex, the staleness decision and the eviction are now atomic with respect to keepalive processing. No new locking is added to any dispatch path. Kill-overdue evictions deliberately do not re-check: keepalives from a worker that refuses to acknowledge a kill are exactly the wedged-worker symptom that path exists to handle. Tests: two deterministic repros. One parks the sweep at the pre-existing await boundary via a cfg(test)-only Notify rendezvous and delivers a full worker_keep_alive_received; the other forces the interleaving with genuine lock contention (holding the registry write lock, then the inner lock) via a cfg(test)-only WorkerRegistry::write_lock_for_test accessor. Both fail without the fix and pass with it. The write_lock_for_test seam is kept: it exercises the window through real lock contention rather than the injected gate, so the two tests do not share a failure mode. Provenance: adversarial review fleet. Co-Authored-By: Claude Fable 5 --- .../src/api_worker_scheduler.rs | 323 ++++++++++++++++++ nativelink-scheduler/src/worker_registry.rs | 9 + 2 files changed, 332 insertions(+) diff --git a/nativelink-scheduler/src/api_worker_scheduler.rs b/nativelink-scheduler/src/api_worker_scheduler.rs index 64d7439e449..01b01d60f68 100644 --- a/nativelink-scheduler/src/api_worker_scheduler.rs +++ b/nativelink-scheduler/src/api_worker_scheduler.rs @@ -2075,6 +2075,14 @@ pub struct ApiWorkerScheduler { /// Woken when a worker joins or leaves this scheduler. local_fleet_change_notify: Arc, + + /// Test-only rendezvous: `remove_timedout_workers` parks here between its + /// liveness snapshot and re-acquiring the eviction lock, at an existing + /// await boundary, so tests can deterministically interleave a keepalive. + /// `.0` is signalled when the sweep reaches the boundary, `.1` lets it + /// proceed. + #[cfg(test)] + test_evict_gate: std::sync::OnceLock>, } impl ApiWorkerScheduler { @@ -2160,6 +2168,8 @@ impl ApiWorkerScheduler { local_fleet_change_notify, capacity_generation, scheduler_id, + #[cfg(test)] + test_evict_gate: std::sync::OnceLock::new(), }) } @@ -2748,10 +2758,41 @@ impl WorkerScheduler for ApiWorkerScheduler { return Ok(()); } + // Test-only hold at this pre-existing await boundary (the inner lock + // was dropped after the snapshot and is only re-acquired below); lets + // a test run a concurrent keepalive here deterministically. + #[cfg(test)] + if let Some(gate) = self.test_evict_gate.get() { + gate.0.notify_one(); + gate.1.notified().await; + } + let mut inner = self.inner.lock().await; let mut result = Ok(()); for (worker_id, kill_overdue) in &worker_ids_to_remove { + // The staleness decision above was made from a snapshot taken + // before this lock was (re-)acquired; the registry read is an + // await boundary every sweep crosses with the lock released. A + // keepalive processed in that window refreshes + // `last_update_timestamp` under this same mutex, so re-validate + // liveness here and skip eviction if the worker heartbeated past + // the threshold. Kill-overdue evictions are intentionally exempt: + // keepalives from a worker that will not acknowledge a kill are + // exactly the wedged-worker symptom being handled. + if !*kill_overdue + && inner + .workers + .peek(worker_id) + .is_some_and(|worker| worker.last_update_timestamp > timeout_threshold) + { + trace!( + ?worker_id, + timeout_threshold, + "Worker heartbeated between staleness snapshot and eviction lock; keeping" + ); + continue; + } let err = if *kill_overdue { make_err!( Code::Internal, @@ -3023,3 +3064,285 @@ mod peer_census_trusted_tests { )); } } + +#[cfg(test)] +mod timedout_worker_eviction_race_tests { + use nativelink_metric::{ + MetricFieldData, MetricKind, MetricPublishKnownKindData, MetricsComponent, + }; + use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::UpdateForWorker; + + use super::*; + use crate::worker_registry::WorkerRegistry; + + #[derive(Debug, Default)] + struct NoopWorkerStateManager; + + impl MetricsComponent for NoopWorkerStateManager { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } + } + + #[async_trait] + impl WorkerStateManager for NoopWorkerStateManager { + async fn update_operation( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + _update: UpdateOperationType, + ) -> Result<(), Error> { + Ok(()) + } + + async fn is_executing_on_worker( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + ) -> Result { + Ok(true) + } + } + + /// repro: `remove_timedout_workers` evicts a worker whose keepalive + /// arrived between the liveness snapshot and the eviction lock. + /// + /// The sweep snapshots `last_update_timestamp` under the inner lock, + /// drops the lock across the registry await, then re-locks and calls + /// `immediate_evict_worker` without re-checking the timestamp. A + /// keepalive that lands in that window (a real await boundary crossed by + /// every sweep) refreshes the worker under the same mutex, yet the + /// worker is still evicted. The test parks the sweep at that boundary + /// via a test-only gate, delivers the keepalive, and asserts the worker + /// survives; on current code the worker is evicted despite being + /// provably alive. + #[nativelink_macro::nativelink_test] + async fn keepalive_between_snapshot_and_evict_lock_still_evicts() { + const WORKER_TIMEOUT_S: u64 = 100; + const OLD_TS: WorkerTimestamp = 1_000; + const NOW_TS: WorkerTimestamp = 2_000; + // Fresh relative to NOW_TS: NOW_TS - FRESH_TS < WORKER_TIMEOUT_S. + const FRESH_TS: WorkerTimestamp = 1_999; + + let registry = Arc::new(WorkerRegistry::new()); + let scheduler = ApiWorkerScheduler::new( + Arc::new(NoopWorkerStateManager), + Arc::new(PlatformPropertyManager::new(HashMap::new())), + WorkerAllocationStrategy::default(), + None, + None, + None, + Arc::new(Notify::new()), + WORKER_TIMEOUT_S, + 10, + 0, + registry, + None, + false, + Duration::from_secs(15), + "test-scheduler".to_string(), + ); + let gate = Arc::new((Notify::new(), Notify::new())); + scheduler + .test_evict_gate + .set(gate.clone()) + .expect("gate set once"); + + let worker_id = WorkerId::from(String::from("worker1")); + // Keep rx alive so eviction's Disconnect send would succeed. + let (tx, _rx) = mpsc::channel::(16); + let worker = Worker::new( + worker_id.clone(), + PlatformProperties::default(), + tx, + OLD_TS, + 1, + ); + scheduler.add_worker(worker).await.expect("add worker"); + + // Start the sweep; it snapshots local_alive=false (OLD_TS is past + // the timeout), finds the registry stale too, and parks at the gate + // just before re-acquiring the inner lock to evict. + let sweep_scheduler = scheduler.clone(); + let sweep = nativelink_util::background_spawn!("timedout_sweep", async move { + sweep_scheduler.remove_timedout_workers(NOW_TS).await + }); + gate.0.notified().await; + + // The worker's keepalive arrives now: it acquires the inner lock, + // advances last_update_timestamp to FRESH_TS, and updates the + // registry heartbeat. The worker is provably alive. + scheduler + .worker_keep_alive_received(&worker_id, FRESH_TS, None) + .await + .expect("keepalive"); + assert!(scheduler.contains_worker_for_test(&worker_id).await); + + // Let the sweep proceed to its eviction lock. + gate.1.notify_one(); + sweep.await.expect("join").expect("sweep result"); + + // A correct sweep re-checks liveness under the lock and keeps the + // worker; current code evicts it unconditionally. + assert!( + scheduler.contains_worker_for_test(&worker_id).await, + "worker was evicted by remove_timedout_workers even though its \ + keepalive (timestamp {FRESH_TS}, well within the {WORKER_TIMEOUT_S}s \ + timeout at now={NOW_TS}) was fully processed before the eviction \ + lock was taken" + ); + } +} + +#[cfg(test)] +mod keepalive_eviction_race_tests { + //! Deterministic reproduction of the check-then-act race in + //! `remove_timedout_workers`: a keepalive that lands after the sweep's + //! unlocked registry check but before it re-acquires the inner lock is + //! never re-consulted, so a demonstrably live worker is evicted. + + use nativelink_config::schedulers::WorkerAllocationStrategy; + + use super::{ + ApiWorkerScheduler, Arc, Duration, Error, HashMap, MetricFieldData, MetricKind, + MetricPublishKnownKindData, MetricsComponent, Notify, OperationId, PlatformProperties, + UNIX_EPOCH, UpdateOperationType, Worker, WorkerId, WorkerScheduler, WorkerStateManager, + async_trait, mpsc, + }; + use crate::platform_property_manager::PlatformPropertyManager; + use crate::worker_registry::WorkerRegistry; + + #[derive(Debug)] + struct NoopWorkerStateManager; + + impl MetricsComponent for NoopWorkerStateManager { + fn publish( + &self, + _kind: MetricKind, + _field_metadata: MetricFieldData, + ) -> Result { + Ok(MetricPublishKnownKindData::Component) + } + } + + #[async_trait] + impl WorkerStateManager for NoopWorkerStateManager { + async fn update_operation( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + _update: UpdateOperationType, + ) -> Result<(), Error> { + Ok(()) + } + + async fn is_executing_on_worker( + &self, + _operation_id: &OperationId, + _worker_id: &WorkerId, + ) -> Result { + Ok(true) + } + } + + const WORKER_TIMEOUT_S: u64 = 100; + /// Worker registered (locally and in the registry) at t=0. + const REGISTERED_AT: u64 = 0; + /// The sweep and the keepalive both happen at t=200, i.e. one full + /// timeout past expiry, so both liveness sources initially say "dead". + const NOW: u64 = 200; + + #[nativelink_macro::nativelink_test] + async fn keepalive_between_registry_check_and_eviction_lock_is_ignored() { + let registry = Arc::new(WorkerRegistry::new()); + let scheduler = ApiWorkerScheduler::new( + Arc::new(NoopWorkerStateManager), + Arc::new(PlatformPropertyManager::new(HashMap::new())), + WorkerAllocationStrategy::default(), + None, + None, + None, + Arc::new(Notify::new()), + WORKER_TIMEOUT_S, + 60, // unacknowledged_kill_timeout_s + 0, // dispatch_ack_timeout_s + registry.clone(), + None, + false, // has_peers + Duration::from_secs(15), // record_ttl + "test-scheduler".to_string(), + ); + + let worker_id = WorkerId("stale_then_alive_worker".to_string()); + let (tx, _rx) = mpsc::channel(64); + scheduler + .add_worker(Worker::new( + worker_id.clone(), + PlatformProperties::default(), + tx, + REGISTERED_AT, + 4, + )) + .await + .expect("add_worker failed"); + + // Gate 1: hold the registry lock so the sweep, after snapshotting the + // (stale) local state and RELEASING the inner lock, parks at its + // unlocked `is_worker_alive` await. + let registry_guard = registry.write_lock_for_test().await; + + let sweep = nativelink_util::background_spawn!("timedout_sweep", { + let scheduler = scheduler.clone(); + async move { scheduler.remove_timedout_workers(NOW).await } + }); + // Let the sweep run: it takes the inner lock, snapshots the worker as + // locally stale, releases the lock, and parks on the registry read. + for _ in 0..8 { + tokio::task::yield_now().await; + } + + // Gate 2: take the inner lock (currently free - the sweep released it) + // so that after its registry read the sweep parks again before it can + // evict. This models inner-lock contention on a busy scheduler. + let mut inner_guard = scheduler.inner.lock().await; + + // Release gate 1: the sweep reads the registry (heartbeat still t=0, + // so "dead"), queues the worker for eviction, and parks on the inner + // lock we hold. + drop(registry_guard); + for _ in 0..8 { + tokio::task::yield_now().await; + } + + // The keepalive lands NOW, while the sweep is parked awaiting the + // inner lock. These are exactly the two steps of + // `worker_keep_alive_received` -> `refresh_worker`: refresh the local + // lifetime under the inner lock, then update the registry heartbeat. + inner_guard + .refresh_lifetime(&worker_id, NOW, None, true) + .expect("refresh_lifetime failed"); + registry + .update_worker_heartbeat(&worker_id, UNIX_EPOCH + Duration::from_secs(NOW)) + .await; + + // Release gate 2: the sweep re-acquires the inner lock and applies its + // stale decision without re-checking either liveness source. + drop(inner_guard); + sweep + .await + .expect("sweep panicked") + .expect("remove_timedout_workers failed"); + + // The worker heartbeated at t=NOW - both the local timestamp and the + // registry say it is alive - yet the stale sweep evicted it. + assert!( + scheduler.contains_worker_for_test(&worker_id).await, + "live worker (keepalive at t={NOW}, timeout {WORKER_TIMEOUT_S}s) \ + was evicted by a stale timeout decision" + ); + } +} diff --git a/nativelink-scheduler/src/worker_registry.rs b/nativelink-scheduler/src/worker_registry.rs index 08cbcce3da1..78d80eaf17f 100644 --- a/nativelink-scheduler/src/worker_registry.rs +++ b/nativelink-scheduler/src/worker_registry.rs @@ -126,6 +126,15 @@ impl WorkerRegistry { self.check_liveness(worker_id, timeout, now).await == WorkerLiveness::Alive } + /// Test-only: hold the internal write lock so that concurrent liveness + /// reads park at a deterministic await point. + #[cfg(test)] + pub(crate) async fn write_lock_for_test( + &self, + ) -> async_lock::RwLockWriteGuard<'_, HashMap> { + self.workers.write().await + } + pub async fn get_worker_last_seen(&self, worker_id: &WorkerId) -> Option { let workers = self.workers.read().await; workers.get(worker_id).copied() From 4c2261254cfb8da45115ae7edc0e4be49004c2fa Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 21:37:08 -0400 Subject: [PATCH 06/46] repro: failed client-keepalive store write still advances last_known_keepalive_ts, retiring a live client's action as abandoned Co-Authored-By: Claude Fable 5 --- .../client_keepalive_failed_write_test.rs | 257 ++++++++++++++++++ 1 file changed, 257 insertions(+) create mode 100644 nativelink-scheduler/tests/client_keepalive_failed_write_test.rs 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..9f712ed233a --- /dev/null +++ b/nativelink-scheduler/tests/client_keepalive_failed_write_test.rs @@ -0,0 +1,257 @@ +// 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. +#[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), + }); + + fn new_op_id() -> OperationId { + OperationId::from("unused-new-operation") + } + let now_fn: fn() -> SystemTime = SystemTime::now; + let op_id_fn: fn() -> OperationId = new_op_id; + let db: StoreAwaitedActionDb< + KeepaliveWriteFailsStore, + fn() -> OperationId, + SystemTime, + fn() -> SystemTime, + > = 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 415890b66be31412dea6d3519e53736f3f64d811 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 23:21:16 -0400 Subject: [PATCH 07/46] 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 | 25 ++++---- 2 files changed, 57 insertions(+), 27 deletions(-) diff --git a/nativelink-scheduler/src/store_awaited_action_db.rs b/nativelink-scheduler/src/store_awaited_action_db.rs index 01d486798ee..03eb524e4e6 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 index 9f712ed233a..d783faf2bc2 100644 --- a/nativelink-scheduler/tests/client_keepalive_failed_write_test.rs +++ b/nativelink-scheduler/tests/client_keepalive_failed_write_test.rs @@ -187,13 +187,24 @@ fn make_queued_action(operation_id: &OperationId) -> AwaitedAction { /// 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) +/// 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. +/// `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"); @@ -204,17 +215,9 @@ async fn failed_keepalive_write_must_be_retried() -> Result<(), Error> { keepalive_write_attempts: AtomicUsize::new(0), }); - fn new_op_id() -> OperationId { - OperationId::from("unused-new-operation") - } let now_fn: fn() -> SystemTime = SystemTime::now; let op_id_fn: fn() -> OperationId = new_op_id; - let db: StoreAwaitedActionDb< - KeepaliveWriteFailsStore, - fn() -> OperationId, - SystemTime, - fn() -> SystemTime, - > = StoreAwaitedActionDb::new( + let db: TestDb = StoreAwaitedActionDb::new( store.clone(), Arc::new(Notify::new()), now_fn, From d02cd1f4901a95384aec3fbee46654859c5dbf5a Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 1 Oct 2026 19:44:21 -0400 Subject: [PATCH 08/46] fix(scheduler): stamp worker liveness on keepalive receipt, before inner.lock (TL1) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Registry is the second liveness source remove_timedout_workers consults ("alive if EITHER source"), but refresh_worker stamped it only AFTER the inner.lock block — which holds the lock across an await (unacked-dispatch sweep). A punctual keepalive whose inner.lock is contended, or whose task is starved on a saturated runtime, refreshed NEITHER liveness source before the eviction sweep read them, falsely evicting a live worker. Transport/runtime-starvation liveness class, seeded by Aman #2856 (h2-read starvation); TL1 is the sibling where a read keepalive must still LAND through the lock. Fix: stamp update_worker_heartbeat on receipt (lightweight RwLock insert, no await-under-lock/store/sweep); make it monotone so an out-of-order keepalive cannot regress last_seen. Regression heartbeat_is_monotone_never_regresses is fails-without by construction; decoupling half exercised by the soak harness. NOT YET built/run on metal. Co-Authored-By: Claude Opus 4.8 --- .../src/api_worker_scheduler.rs | 19 +++++-- nativelink-scheduler/src/worker_registry.rs | 52 ++++++++++++++++++- 2 files changed, 65 insertions(+), 6 deletions(-) diff --git a/nativelink-scheduler/src/api_worker_scheduler.rs b/nativelink-scheduler/src/api_worker_scheduler.rs index 64d7439e449..d9d51f9bf01 100644 --- a/nativelink-scheduler/src/api_worker_scheduler.rs +++ b/nativelink-scheduler/src/api_worker_scheduler.rs @@ -2217,6 +2217,21 @@ impl ApiWorkerScheduler { load: Option, keepalive: bool, ) -> Result<(), Error> { + // TL1: stamp the independent registry liveness signal on RECEIPT, before + // taking `inner.lock()`. The registry is the second liveness source that + // `remove_timedout_workers` consults ("alive if EITHER source"), but it + // was previously stamped only AFTER the inner.lock block below — which + // holds the lock across an await (the unacked-dispatch sweep). A punctual + // keepalive whose inner.lock acquisition is contended, or whose task is + // starved on a saturated runtime, then refreshed NEITHER liveness source + // before the eviction sweep read them, falsely evicting a live worker. The + // registry write is a lightweight RwLock insert (no await-under-lock, no + // store, no sweep), so stamping it first decouples the monotone liveness + // FACT from the coordinated state write; it can no longer be starved. + 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 @@ -2258,10 +2273,6 @@ impl ApiWorkerScheduler { } } } - let now = UNIX_EPOCH + Duration::from_secs(timestamp); - self.worker_registry - .update_worker_heartbeat(worker_id, now) - .await; Ok(()) } diff --git a/nativelink-scheduler/src/worker_registry.rs b/nativelink-scheduler/src/worker_registry.rs index 08cbcce3da1..02be20ad301 100644 --- a/nativelink-scheduler/src/worker_registry.rs +++ b/nativelink-scheduler/src/worker_registry.rs @@ -67,10 +67,22 @@ impl WorkerRegistry { } } - /// Updates the heartbeat timestamp for a worker. + /// Advances a worker's heartbeat to `now`. Monotone: a late or out-of-order + /// keepalive never regresses `last_seen`. The liveness signal is the MAX of + /// receipt times — a monotone fact — and must not move backward and + /// manufacture a false `Stale` (pairs with the TL1 fix in + /// `ApiWorkerScheduler::refresh_worker`, which now stamps this on receipt, + /// possibly ahead of `refresh_lifetime`'s own out-of-order rejection). 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); + workers + .entry(worker_id.clone()) + .and_modify(|prev| { + if now > *prev { + *prev = now; + } + }) + .or_insert(now); trace!(?worker_id, now = %humantime::format_rfc3339(now), "FLOW: Worker heartbeat updated in registry"); } @@ -234,4 +246,40 @@ mod tests { .await ); } + + // TL1 regression (monotone half): an out-of-order / late keepalive must not + // regress `last_seen` and manufacture a false `Stale`. Fails WITHOUT the + // monotone fix (plain `insert` overwrites with the older timestamp, so the + // liveness check past the old deadline reads Stale). + #[nativelink_test] + async fn heartbeat_is_monotone_never_regresses() { + let registry = WorkerRegistry::new(); + let w = WorkerId::from(String::from("w")); + let timeout = Duration::from_secs(5); + let t0 = SystemTime::now(); + let t_late = t0.checked_add(Duration::from_secs(10)).unwrap(); + + registry.register_worker(&w, t0).await; + registry.update_worker_heartbeat(&w, t_late).await; // advance + registry.update_worker_heartbeat(&w, t0).await; // out-of-order (older) — must NOT regress + + assert_eq!( + registry.get_worker_last_seen(&w).await, + Some(t_late), + "an older keepalive regressed last_seen — manufactures false Stale (TL1)" + ); + // A check 1s past the OLD deadline (t0 + timeout) must still be Alive, + // because the real last_seen is t_late. + let past_old_deadline = t0 + .checked_add(timeout) + .unwrap() + .checked_add(Duration::from_secs(1)) + .unwrap(); + assert_eq!( + registry + .check_liveness(&w, timeout, past_old_deadline) + .await, + WorkerLiveness::Alive + ); + } } From cac91dfe2da310cebf3c50b100cd754487c52d3a Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 01:25:17 -0400 Subject: [PATCH 09/46] 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 cf6f65ddb38..55728366919 100644 --- a/nativelink-scheduler/src/worker.rs +++ b/nativelink-scheduler/src/worker.rs @@ -17,7 +17,7 @@ use std::collections::{HashMap, HashSet}; 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, @@ -416,6 +416,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 1c1f9e02974fe4aa49b2745576ff3d6d5948f4f5 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 23:21:36 -0400 Subject: [PATCH 10/46] 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 | 283 ++++++++++++++++++ 4 files changed, 354 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 25077de27cc..a576c0c808a 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/liveness_clock_step_test.rs", "tests/placement_test.rs", diff --git a/nativelink-scheduler/src/api_worker_scheduler.rs b/nativelink-scheduler/src/api_worker_scheduler.rs index 64d7439e449..6e32cbb0b0c 100644 --- a/nativelink-scheduler/src/api_worker_scheduler.rs +++ b/nativelink-scheduler/src/api_worker_scheduler.rs @@ -371,6 +371,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 @@ -1657,12 +1665,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; @@ -1939,13 +1952,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(()); @@ -2148,6 +2172,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, @@ -2824,7 +2849,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)| { @@ -2841,7 +2866,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) @@ -2849,8 +2883,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 @@ -2858,7 +2892,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!( @@ -2883,10 +2917,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. @@ -2895,7 +2932,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)); } } @@ -2905,10 +2942,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 cf6f65ddb38..e8591e80ca6 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. @@ -285,9 +294,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(); @@ -415,6 +429,7 @@ impl Worker { operation_id: OperationId, action_info: ActionInfoWithProps, dispatched_at: WorkerTimestamp, + dispatch_generation: u64, ) -> Result<(), Error> { let tx = &mut self.tx; let worker_platform_properties = &mut self.platform_properties; @@ -447,6 +462,7 @@ impl Worker { kill_requested_at: None, missed_reports: 0, 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..8e3460aeafe --- /dev/null +++ b/nativelink-scheduler/tests/kill_revoked_redispatch_race_test.rs @@ -0,0 +1,283 @@ +// 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, + None, // reservation_hold + 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), + "test-scheduler".to_string(), + ); + + 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 cb6341435aeac0fcb8fdc108d41c40174fd41ade Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 16:26:08 -0400 Subject: [PATCH 11/46] scheduler: validate placement requirements at assignment (N3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Matching computed typed resource requirements once and parked that snapshot; assignment reloaded the current action but only guarded against double-assignment. It never compared the reloaded action's requirements against the ones the matcher used to pick the worker. So a matcher could park a 4 GiB snapshot of X; a peer runs X, hits a memory-limit kill, requeues X escalated to 8 GiB; the stale 4 GiB snapshot then dispatched X to a 4 GiB-only worker. The DB said 8 GiB but StartExecute and the reservation carried 4 GiB — the escalation was silently bypassed. Fix: assign_operation now carries the placement-relevant requirement snapshot the worker was chosen for. The accept, inside the same CAS that flips the action to Executing, rejects (Aborted) if the action's current platform_properties no longer match that snapshot. A stale proposal becomes a lost race; the action stays queued and is rematched against its current requirements. Dispatch still sends the parked payload, so rejecting before send keeps DB and worker in agreement. Regression test assignment_rejects_a_stale_placement_after_escalation reproduces the two-matcher interleaving: it FAILS without the guard (the 4 GiB placement is silently accepted to Executing) and PASSES with it. Co-Authored-By: Claude Opus 4.8 --- nativelink-scheduler/src/simple_scheduler.rs | 21 ++- .../src/simple_scheduler_state_manager.rs | 40 ++++- .../simple_scheduler_state_manager_test.rs | 162 +++++++++++++++++- .../src/operation_state_manager.rs | 10 ++ 4 files changed, 220 insertions(+), 13 deletions(-) diff --git a/nativelink-scheduler/src/simple_scheduler.rs b/nativelink-scheduler/src/simple_scheduler.rs index 69afdd02272..8d164dd91e4 100644 --- a/nativelink-scheduler/src/simple_scheduler.rs +++ b/nativelink-scheduler/src/simple_scheduler.rs @@ -680,15 +680,28 @@ impl SimpleScheduler { action_state.client_operation_id.clone() }; - // Tell the matching engine that the operation is being assigned to a worker. + // Tell the matching engine that the operation is being assigned + // to a worker, validating that the requirements this worker was + // chosen for still hold. A peer that requeued the action with a + // raised reservation (memory escalation) between our read and + // now must not have our stale, smaller placement dispatched at + // it; the accept rejects such a placement with Aborted. let assign_result = matching_engine_state_manager - .assign_operation(&operation_id, Ok(&worker_id)) + .assign_operation( + &operation_id, + Ok(&worker_id), + Some(&action_info.inner.platform_properties), + ) .await .err_tip(|| "Failed to assign operation in do_try_match"); if let Err(err) = assign_result { if err.code == Code::Aborted { - // If the operation was aborted, it means that the operation was - // cancelled due to another operation being assigned to the worker. + // Aborted here means either another operation was + // assigned to the worker first, or our placement is + // stale because the action's requirements changed after + // we chose the worker. Either way nothing was sent; the + // action stays queued and is rematched next pass against + // its current requirements. return Ok(Dispatch::LostRace); } // Any other error is a real error. diff --git a/nativelink-scheduler/src/simple_scheduler_state_manager.rs b/nativelink-scheduler/src/simple_scheduler_state_manager.rs index 4defb3bced6..8a380872803 100644 --- a/nativelink-scheduler/src/simple_scheduler_state_manager.rs +++ b/nativelink-scheduler/src/simple_scheduler_state_manager.rs @@ -673,6 +673,7 @@ where Code::Unavailable, "The worker finished this operation but its result never reached the scheduler" )), + None, ) .await { @@ -1147,6 +1148,7 @@ where self.assign_operation( operation_id, Err(make_err!(Code::DeadlineExceeded, "Operation {cause}")), + None, ) .await } @@ -1156,6 +1158,7 @@ where operation_id: &OperationId, maybe_worker_id: Option<&WorkerId>, update: UpdateOperationType, + expected_platform_properties: Option<&HashMap>, ) -> Result<(), Error> { let update_type_str = match &update { UpdateOperationType::KeepAlive => "KeepAlive", @@ -1289,6 +1292,31 @@ where warn!(state = ?awaited_action.state(), "Action already assigned"); return Err(make_err!(Code::Aborted, "Action already assigned")); } + // Placement validation: the matcher chose this worker for a + // snapshot of the action's requirements. If the action's + // current requirements no longer match that snapshot — an + // escalation or other requeue raised the reservation after + // the worker was picked — the placement is stale. Accepting + // it would flip the action to Executing while the matcher + // dispatches the old, smaller reservation to the worker, + // silently bypassing the escalation. Reject with Aborted so + // the matcher treats it as a lost race and the action is + // rematched against its current requirements. + if stage == &ActionStage::Executing + && let Some(expected) = expected_platform_properties + && &awaited_action.action_info().platform_properties != expected + { + warn!( + %operation_id, + expected = ?expected, + current = ?awaited_action.action_info().platform_properties, + "Rejecting stale placement: requirements changed since the worker was chosen" + ); + return Err(make_err!( + Code::Aborted, + "Stale placement for {operation_id}: action requirements changed since the worker was chosen; rematch required" + )); + } stage.clone() } UpdateOperationType::UpdateWithError(err) => { @@ -1744,7 +1772,7 @@ where worker_id: &WorkerId, update: UpdateOperationType, ) -> Result<(), Error> { - self.inner_update_operation(operation_id, Some(worker_id), update) + self.inner_update_operation(operation_id, Some(worker_id), update, None) .await } @@ -1804,6 +1832,7 @@ where &self, operation_id: &OperationId, worker_id_or_reason_for_unassign: Result<&WorkerId, Error>, + expected_platform_properties: Option<&HashMap>, ) -> Result<(), Error> { let (maybe_worker_id, update) = match worker_id_or_reason_for_unassign { Ok(worker_id) => ( @@ -1812,8 +1841,13 @@ where ), Err(err) => (None, UpdateOperationType::UpdateWithError(err)), }; - self.inner_update_operation(operation_id, maybe_worker_id, update) - .await + self.inner_update_operation( + operation_id, + maybe_worker_id, + update, + expected_platform_properties, + ) + .await } async fn fail_queued_operation( diff --git a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs index b0257c07618..046aeca748b 100644 --- a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs @@ -47,7 +47,7 @@ use nativelink_util::common::DigestInfo; use nativelink_util::digest_hasher::DigestHasherFunc; use nativelink_util::instant_wrapper::MockInstantWrapped; use nativelink_util::operation_state_manager::{ - ActionStateResult, ClientStateManager, MatchingEngineStateManager, OperationFilter, + ActionStateResult, ClientStateManager, Escalation, MatchingEngineStateManager, OperationFilter, OperationStageFlags, UpdateOperationType, WorkerStateManager, }; use nativelink_util::platform_properties::PlatformProperties; @@ -98,12 +98,19 @@ fn make_system_time(add_time: u64) -> SystemTime { } fn action_info(started: SystemTime) -> ActionInfo { + action_info_with_properties(started, HashMap::default()) +} + +fn action_info_with_properties( + started: SystemTime, + platform_properties: HashMap, +) -> ActionInfo { let action_digest = DigestInfo::zero_digest(); ActionInfo { command_digest: action_digest, input_root_digest: action_digest, timeout: Duration::ZERO, - platform_properties: HashMap::default(), + platform_properties, priority: 0, load_timestamp: SystemTime::UNIX_EPOCH, insert_timestamp: started, @@ -599,7 +606,7 @@ async fn is_executing_on_worker_follows_the_assignment() -> Result<(), Error> { ); state_mgr - .assign_operation(&operation_id, Ok(&worker_id)) + .assign_operation(&operation_id, Ok(&worker_id), None) .await?; assert!( state_mgr @@ -617,6 +624,7 @@ async fn is_executing_on_worker_follows_the_assignment() -> Result<(), Error> { .assign_operation( &operation_id, Err(make_err!(Code::DeadlineExceeded, "timed out")), + None, ) .await?; assert!( @@ -625,7 +633,7 @@ async fn is_executing_on_worker_follows_the_assignment() -> Result<(), Error> { .await? ); state_mgr - .assign_operation(&operation_id, Ok(&other_worker_id)) + .assign_operation(&operation_id, Ok(&other_worker_id), None) .await?; assert!( !state_mgr @@ -1268,7 +1276,7 @@ async fn fail_queued_operation_leaves_an_executing_action_alone() -> Result<(), let (operation_id, client) = queued_operation_id(state_mgr.as_ref()).await?; let worker_id = WorkerId::from(String::from("worker")); state_mgr - .assign_operation(&operation_id, Ok(&worker_id)) + .assign_operation(&operation_id, Ok(&worker_id), None) .await?; let failed = state_mgr @@ -1363,7 +1371,7 @@ async fn queue_and_assign_with_client( let (state, _) = listed.as_state().await?; let operation_id = state.client_operation_id.clone(); state_mgr - .assign_operation(&operation_id, Ok(worker_id)) + .assign_operation(&operation_id, Ok(worker_id), None) .await?; Ok((operation_id, subscriber)) } @@ -1551,3 +1559,145 @@ async fn executing_sweep_lost_result_counts_against_max_job_retries() -> Result< ); Ok(()) } + +/// N3: matching and assignment must agree on the action's requirements. +/// +/// Two schedulers (two matchers) share one action DB and memory escalation +/// is enabled. Matcher A reads action X while it needs 4 GiB and picks a +/// 4 GiB worker for it. Before A releases that placement, a peer runs X, +/// hits a memory-limit kill, and requeues X escalated to 8 GiB. When A then +/// releases its stale 4 GiB placement, the assignment must reject it: its +/// only guard used to be double-assignment, so it accepted the queued record +/// and flipped X to Executing while A dispatched the parked 4 GiB reservation +/// to the worker — the DB said 8 GiB but `StartExecute` and the worker's +/// reservation carried 4 GiB, silently bypassing the escalation. The accept +/// now validates the placement-relevant requirements the worker was chosen +/// for and rejects the stale proposal, leaving X queued at 8 GiB to be +/// rematched. A fresh placement made against the current 8 GiB requirements +/// is accepted, so stored, matched, and assigned requirements all agree. +#[nativelink_test] +async fn assignment_rejects_a_stale_placement_after_escalation() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = state_manager(Arc::new(WorkerRegistry::new())); + + let four_gib = HashMap::from([("memory_kb".to_string(), "4000".to_string())]); + let eight_gib = HashMap::from([("memory_kb".to_string(), "8000".to_string())]); + + let client = state_mgr + .add_action( + OperationId::default(), + Arc::new(action_info_with_properties( + make_system_time(0), + four_gib.clone(), + )), + ) + .await?; + + // Matcher A reads X and picks a 4 GiB worker. This is the snapshot it + // will release against later; nothing is committed yet. + let operation_id = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Queued, + ..Default::default() + }, + ) + .await? + .next() + .await + .expect("the queued operation") + .as_state() + .await? + .0 + .client_operation_id + .clone(); + let matcher_a_snapshot = four_gib.clone(); + let worker_a = WorkerId::from(String::from("matcher-a-worker-4gib")); + + // A peer runs X on its own worker and the action is killed for memory; + // the peer requeues X escalated to 8 GiB. X is Queued again, needing more. + let peer_worker = WorkerId::from(String::from("peer-worker")); + state_mgr + .assign_operation(&operation_id, Ok(&peer_worker), Some(&four_gib)) + .await?; + state_mgr + .update_operation( + &operation_id, + &peer_worker, + UpdateOperationType::UpdateWithEscalation(Escalation { + property: "memory_kb".to_string(), + value: 8000, + reason: make_err!(Code::FailedPrecondition, "killed for memory"), + cpu: None, + }), + ) + .await?; + + // The store now records the escalated requirement. + let stored = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Queued, + ..Default::default() + }, + ) + .await? + .next() + .await + .expect("the requeued operation") + .as_action_info() + .await? + .0; + assert_eq!( + stored.platform_properties, eight_gib, + "the escalation raised the stored reservation to 8 GiB" + ); + + // Matcher A now releases its stale 4 GiB placement against its 4 GiB + // worker. Without the fix the assignment accepts this and flips X to + // Executing at 4 GiB, bypassing the escalation. It must be rejected. + let stale = state_mgr + .assign_operation(&operation_id, Ok(&worker_a), Some(&matcher_a_snapshot)) + .await + .expect_err("a stale placement must be rejected, not silently accepted"); + assert_eq!( + stale.code, + Code::Aborted, + "a rejected stale placement is a lost race, not a hard error: {stale}" + ); + assert_eq!( + client.as_state().await?.0.stage, + ActionStage::Queued, + "the rejected placement leaves X queued for a rematch, not Executing" + ); + + // A fresh placement made against the current 8 GiB requirement is + // accepted: stored, matched, and assigned requirements all agree. + state_mgr + .assign_operation(&operation_id, Ok(&worker_a), Some(&eight_gib)) + .await?; + assert_eq!( + client.as_state().await?.0.stage, + ActionStage::Executing, + "a placement matching the current 8 GiB requirement is accepted" + ); + let assigned = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Executing, + ..Default::default() + }, + ) + .await? + .next() + .await + .expect("the assigned operation") + .as_action_info() + .await? + .0; + assert_eq!( + assigned.platform_properties, eight_gib, + "the dispatched action carries the escalated 8 GiB reservation" + ); + Ok(()) +} diff --git a/nativelink-util/src/operation_state_manager.rs b/nativelink-util/src/operation_state_manager.rs index 4e6e2a30dee..f88f45cfb47 100644 --- a/nativelink-util/src/operation_state_manager.rs +++ b/nativelink-util/src/operation_state_manager.rs @@ -13,6 +13,7 @@ // limitations under the License. use core::pin::Pin; +use std::collections::HashMap; use std::sync::Arc; use std::time::SystemTime; @@ -202,10 +203,19 @@ pub trait MatchingEngineStateManager: Sync + Send + MetricsComponent { ) -> Result, Error>; /// Assign an operation to a worker or unassign it. + /// + /// `expected_platform_properties`, when assigning, is the untyped + /// requirement snapshot the matcher chose this worker for. The accept + /// rejects (with `Aborted`) if the action's current requirements no + /// longer match it — an escalation or other requeue changed the + /// reservation after the worker was picked, so the placement is stale + /// and must be rematched against the current requirements rather than + /// dispatched at the old reservation. `None` skips the check. async fn assign_operation( &self, operation_id: &OperationId, worker_id_or_reason_for_unassign: Result<&WorkerId, Error>, + expected_platform_properties: Option<&HashMap>, ) -> Result<(), Error>; /// Completes a queued operation with `err`, without retrying it. Does From 9f6fb3a7f580daedbbd65519fad0a0a93fc6b47a Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 19:24:49 -0400 Subject: [PATCH 12/46] 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 01d486798ee..4479ae18e6d 100644 --- a/nativelink-scheduler/src/store_awaited_action_db.rs +++ b/nativelink-scheduler/src/store_awaited_action_db.rs @@ -1003,15 +1003,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 dbe76f3fc1a..31dc0c4bd90 100644 --- a/nativelink-scheduler/tests/store_awaited_action_db_test.rs +++ b/nativelink-scheduler/tests/store_awaited_action_db_test.rs @@ -35,7 +35,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; @@ -443,6 +444,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 e71f6882fb6eebecc60747164dc188123927d8da Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 19:34:32 -0400 Subject: [PATCH 13/46] Scheduler: a lost queued record is garbage-collected once no client is listening, instead of being counted forever A queued index entry whose backing record cannot be decoded (for example serialization drift across a rolling upgrade, or a corrupted write) had no TTL and no deletion path. The sweep and the listings skip such an entry and count it (`record_awaited_action_orphan`), but nothing ever removes it, so it is re-listed, re-skipped and re-counted by every pass forever: the queue depth metric stays accurate while the queue itself never drains. The redis listing surfaces the failure as a stream error with no key attached, so the sweep could not delete the record even if it wanted to. The FT.AGGREGATE now also loads RediSearch's `__key`, and when a listed row fails to decode the store removes it under three guards, checked atomically in Lua: - only the queued-state listing opts in (`SchedulerIndexProvider::lost_record_liveness_guard_key`); - the record's client keepalive key (written with a TTL of the client action timeout) must not exist, so a record any client has touched within the timeout is never removed; - the record must still hold the version the listing saw, so a record rewritten concurrently is never removed. The row is still reported as lost, so the orphan counter records it one last time before it is gone. Co-Authored-By: Claude Opus 4.8 --- .../src/dynamic_fake_redis.rs | 72 +++++++--- .../src/store_awaited_action_db.rs | 13 ++ .../redis_store_awaited_action_db_test.rs | 93 +++++++++++++ nativelink-store/src/redis_store.rs | 127 ++++++++++++++++-- nativelink-store/tests/redis_store_test.rs | 32 +++-- nativelink-util/src/store_trait.rs | 17 +++ 6 files changed, 317 insertions(+), 37 deletions(-) diff --git a/nativelink-redis-tester/src/dynamic_fake_redis.rs b/nativelink-redis-tester/src/dynamic_fake_redis.rs index 3b602c597a9..10e4bc6b189 100644 --- a/nativelink-redis-tester/src/dynamic_fake_redis.rs +++ b/nativelink-redis-tester/src/dynamic_fake_redis.rs @@ -156,12 +156,13 @@ impl FakeRedisBackend { }; // Lazy implementation making assumptions. assert_eq!( - args[load_offset..load_offset + 4], + args[load_offset..load_offset + 5], vec![ OwnedFrame::BulkString(b"LOAD".to_vec()), - OwnedFrame::BulkString(b"2".to_vec()), + OwnedFrame::BulkString(b"3".to_vec()), OwnedFrame::BulkString(b"data".to_vec()), - OwnedFrame::BulkString(b"version".to_vec()) + OwnedFrame::BulkString(b"version".to_vec()), + OwnedFrame::BulkString(b"@__key".to_vec()) ] ); let mut results = vec![Value::Int(0)]; @@ -178,12 +179,13 @@ impl FakeRedisBackend { let desc = matches!(args.get(i + 3), Some(OwnedFrame::BulkString(d)) if d == b"DESC"); Some((field, desc)) }); - let sorted = |mut rows: Vec>| { + let sorted = |mut rows: Vec<(String, HashMap)>| { if let Some((field, desc)) = &sort_spec { - let key = |r: &HashMap| match r.get(field) { - Some(Value::BulkString(b)) => b.clone(), - _ => Vec::new(), - }; + let key = + |r: &(String, HashMap)| match r.1.get(field) { + Some(Value::BulkString(b)) => b.clone(), + _ => Vec::new(), + }; rows.sort_by_key(key); if *desc { rows.reverse(); @@ -199,16 +201,20 @@ impl FakeRedisBackend { .table .lock() .unwrap() - .values() - .filter(|f| f.contains_key("data") && f.contains_key("version")) - .cloned() + .iter() + .filter(|(_, f)| { + f.contains_key("data") && f.contains_key("version") + }) + .map(|(k, f)| (k.clone(), f.clone())) .collect(); - for fields in sorted(rows) { + for (key, fields) in sorted(rows) { results.push(Value::Array(vec![ Value::BulkString(b"data".to_vec()), fields["data"].clone(), Value::BulkString(b"version".to_vec()), fields["version"].clone(), + Value::BulkString(b"__key".to_vec()), + Value::BulkString(key.into_bytes()), ])); } } else { @@ -225,14 +231,14 @@ impl FakeRedisBackend { .table .lock() .unwrap() - .values() - .filter(|f| { + .iter() + .filter(|(_, f)| { f.get(field) == Some(&Value::BulkString(value.as_bytes().to_vec())) }) - .cloned() + .map(|(k, f)| (k.clone(), f.clone())) .collect(); - for fields in sorted(rows) { + for (key, fields) in sorted(rows) { { let mut record = vec![ Value::BulkString(b"data".to_vec()), @@ -245,6 +251,8 @@ impl FakeRedisBackend { record.push(Value::BulkString(b"version".to_vec())); record.push(version.clone()); } + record.push(Value::BulkString(b"__key".to_vec())); + record.push(Value::BulkString(key.into_bytes())); results.push(Value::Array(record)); } } @@ -354,6 +362,38 @@ impl FakeRedisBackend { } } + "EVAL" => { + // Only the delete-lost-record script is run via + // EVAL; emulate its semantics: delete KEYS[1] only + // if the liveness guard KEYS[2] does not exist and + // the record still holds version ARGV[1]. + let script = + str::from_utf8(args[0].as_bytes().expect("Script not bytes")).unwrap(); + assert!(script.contains("DEL"), "Unexpected EVAL script: {script:?}"); + assert_eq!(args[1], OwnedFrame::BulkString(b"2".to_vec())); + let key = str::from_utf8(args[2].as_bytes().expect("Key not bytes")) + .expect("Key cannot be parsed as string") + .to_string(); + let guard_key = + str::from_utf8(args[3].as_bytes().expect("Guard key not bytes")) + .expect("Guard key cannot be parsed as string") + .to_string(); + let expected_version = + args[4].as_bytes().expect("Version not bytes").to_vec(); + let mut table = self.table.lock().unwrap(); + if table.contains_key(&guard_key) { + Value::Int(-1) + } else if table.get(&key).and_then(|f| f.get("version")) + == Some(&Value::BulkString(expected_version)) + { + table.remove(&key); + debug!(%key, "Deleted key via delete-lost-record"); + Value::Int(1) + } else { + Value::Int(0) + } + } + "HMSET" => { let mut values = HashMap::new(); assert_eq!( diff --git a/nativelink-scheduler/src/store_awaited_action_db.rs b/nativelink-scheduler/src/store_awaited_action_db.rs index 01d486798ee..4b9ec3750de 100644 --- a/nativelink-scheduler/src/store_awaited_action_db.rs +++ b/nativelink-scheduler/src/store_awaited_action_db.rs @@ -662,6 +662,19 @@ impl SchedulerIndexProvider for SearchStateToAwaitedAction { fn index_value(&self) -> Cow<'_, str> { Cow::Borrowed(self.0) } + + /// A queued record whose data cannot be decoded may be removed, but + /// only while no client keepalive exists for it: the keepalive is + /// written with a TTL of the client action timeout, so its absence + /// means no client has been listening for at least that long. Scoped + /// to the queued listing; records in any other stage are left alone. + fn lost_record_liveness_guard_key(index_value: &str, record_key: &str) -> Option { + if index_value != get_state_prefix(SortedAwaitedActionState::Queued) { + return None; + } + let operation_id = record_key.strip_prefix(OPERATION_ID_TO_AWAITED_ACTION_KEY_PREFIX)?; + Some(format!("{CLIENT_KEEPALIVE_KEY_PREFIX}{operation_id}")) + } } impl SchedulerStoreDecodeTo for SearchStateToAwaitedAction { type DecodeOutput = AwaitedAction; diff --git a/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs b/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs index c13978bb319..bc7f733fc2a 100644 --- a/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs +++ b/nativelink-scheduler/tests/redis_store_awaited_action_db_test.rs @@ -35,6 +35,7 @@ use nativelink_scheduler::awaited_action_db::{ }; use nativelink_scheduler::default_scheduler_factory::retain_completed_for_s; use nativelink_scheduler::simple_scheduler::SimpleScheduler; +use nativelink_scheduler::simple_scheduler_state_manager::SimpleSchedulerStateManager; use nativelink_scheduler::store_awaited_action_db::StoreAwaitedActionDb; use nativelink_scheduler::worker::Worker; use nativelink_scheduler::worker_scheduler::WorkerScheduler; @@ -723,3 +724,95 @@ async fn unset_retention_is_the_shared_default_not_forever() -> Result<(), Error assert_eq!(retain_completed_for_s(30), 30); Ok(()) } + +/// A queued index entry whose backing record cannot be decoded (for example +/// serialization drift across a rolling upgrade) has no TTL and, before the +/// fix, no deletion path: every sweep listed it, skipped it and counted it +/// as an orphan, forever. The sweep's listing must remove it, and must leave +/// a healthy queued action alone. +#[nativelink_test] +async fn sweep_deletes_lost_queued_record() -> Result<(), Error> { + const CORRUPT_KEY: &str = "aa_corrupt-op"; + + let backend: FakeRedisBackend = FakeRedisBackend::new(); + let port = backend.clone().run().await; + let store = RedisStore::new_standard(RedisSpec { + addresses: vec![format!("redis://127.0.0.1:{port}")], + experimental_pub_sub_channel: Some("lost-record-sweep".to_string()), + ..Default::default() + }) + .await?; + backend.set_subscription_manager(store.subscription_manager().await.unwrap()); + let notify = Arc::new(Notify::new()); + let db = StoreAwaitedActionDb::new( + store, + notify, + MockInstantWrapped::default, + OperationId::default, + 60, + 60, + false, + ) + .await?; + + // A healthy queued action whose client is still alive. + let healthy = make_awaited_action("healthy-op"); + db.add_action( + OperationId::String("healthy-client-op".to_string()), + healthy.action_info().clone(), + Duration::from_mins(1), + ) + .await?; + + // The lost record: listed under the queued index, version intact, but + // its data does not decode. + backend.table.lock().unwrap().insert( + CORRUPT_KEY.to_string(), + HashMap::from([ + ( + "data".to_string(), + Value::BulkString(b"{this is not an AwaitedAction".to_vec()), + ), + ("version".to_string(), Value::BulkString(b"1".to_vec())), + ("state".to_string(), Value::BulkString(b"queued".to_vec())), + ( + "sort_key".to_string(), + Value::BulkString(b"0000000000000000".to_vec()), + ), + ]), + ); + + let state_manager = SimpleSchedulerStateManager::new( + 0_usize, + Duration::from_mins(1), + Duration::from_mins(1), + Duration::from_mins(10), + db, + MockInstantWrapped::default, + None, + ); + + // Several sweeps: without the fix the corrupt record survives every one + // of them (skip-and-count, never delete). + for _ in 0..3 { + state_manager.sweep_abandoned_queued_actions().await?; + } + + assert!( + !backend.table.lock().unwrap().contains_key(CORRUPT_KEY), + "an undecodable queued record must be deleted by the sweep's listing", + ); + assert_eq!( + backend + .table + .lock() + .unwrap() + .keys() + .filter(|k| k.starts_with("aa_")) + .count(), + 1, + "the healthy queued action must survive the sweep", + ); + + Ok(()) +} diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index 00dc08a29a4..d866e2492ac 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -52,8 +52,8 @@ use redis::cluster::ClusterClient; use redis::cluster_async::ClusterConnection; use redis::sentinel::{SentinelClient, SentinelNodeConnectionInfo, SentinelServerType}; use redis::{ - AsyncCommands, AsyncIter, Client, IntoConnectionInfo, PushInfo, ScanOptions, Script, Value, - pipe, + AsyncCommands, AsyncIter, Client, IntoConnectionInfo, PushInfo, RedisError, ScanOptions, + Script, Value, pipe, }; use serde::Deserialize; use serde::de::IntoDeserializer; @@ -94,7 +94,7 @@ const MAX_REDIS_RETRY_ATTEMPTS: u32 = 5; /// dropped/refused connection, an IO error, or a write that hit a freshly /// demoted replica. Anything else (e.g. a real protocol/logic error) is /// returned as-is. -fn is_retryable_redis_error(err: &redis::RedisError) -> bool { +fn is_retryable_redis_error(err: &RedisError) -> bool { err.is_connection_dropped() || err.is_connection_refusal() || err.is_io_error() @@ -1726,6 +1726,10 @@ const CURSOR_IDLE_MS: u64 = 30_000; const DATA_FIELD_NAME: &str = "data"; /// The name of the field in the Redis hash that stores the version. const VERSION_FIELD_NAME: &str = "version"; +/// The `RediSearch` built-in property holding the document's own key. Loaded in +/// searches so a row whose record cannot be decoded can still be identified +/// and removed rather than sitting in the index forever. +const DOC_KEY_FIELD_NAME: &str = "__key"; /// The time to live of indexes in seconds. After this time redis may delete the index. const INDEX_TTL_S: u64 = 60 * 60 * 24; // 24 hours. @@ -1783,6 +1787,48 @@ return {{ 1, new_version }} " ); +/// Lua script that deletes a record only if its version still matches the +/// version a search listed it at. Used to garbage-collect a record whose +/// `data` cannot be decoded: records are written atomically (delete and +/// re-insert in `LUA_VERSION_SET_SCRIPT`), so an undecodable record is never +/// a write in progress, and no writer can ever update it either, because +/// every update needs the current version from a decoded copy. The version +/// guard makes the removal safe even if that reasoning is ever violated: a +/// record rewritten after the listing is left alone. +pub const LUA_DELETE_LOST_RECORD_SCRIPT: &str = formatcp!( + r" +if redis.call('EXISTS', KEYS[2]) == 1 then + return -1 +end +if redis.call('HGET', KEYS[1], '{VERSION_FIELD_NAME}') == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +" +); + +/// Deletes the lost record at `key`, but only if the liveness guard at +/// `guard_key` does not exist (returns -1 if it does) and the record still +/// holds `version` (returns 0 if it does not). Returns 1 on deletion. +async fn delete_lost_record( + mut connection: C, + key: &str, + guard_key: &str, + version: i64, +) -> Result +where + C: ConnectionLike + Send, +{ + redis::cmd("EVAL") + .arg(LUA_DELETE_LOST_RECORD_SCRIPT) + .arg(2) + .arg(key) + .arg(guard_key) + .arg(version.to_string()) + .query_async(&mut connection) + .await +} + /// How long a freshly created index is given to finish its background /// scan before a read goes ahead regardless, and how often it is asked. const INDEX_READY_TIMEOUT: Duration = Duration::from_secs(5); @@ -2389,7 +2435,11 @@ where format!("@{}:{{ {} }}", K::INDEX_NAME, index_value) }, FtAggregateOptions { - load: vec![DATA_FIELD_NAME.into(), VERSION_FIELD_NAME.into()], + load: vec![ + DATA_FIELD_NAME.into(), + VERSION_FIELD_NAME.into(), + format!("@{DOC_KEY_FIELD_NAME}"), + ], cursor: FtAggregateCursor { count: self.max_count_per_cursor, max_idle: CURSOR_IDLE_MS, @@ -2427,6 +2477,10 @@ where }; let (connection_manager, connect_id) = self.connection_manager.get_connection().await?; + // For garbage-collecting a listed record that fails to decode. + let delete_connection = connection_manager.clone(); + let owned_index_value = index_value.to_string(); + let key_prefix = self.key_prefix.clone(); let stream = match run_ft_aggregate(connection_manager.clone()).await { // A demoted master answers READONLY and a dead/old master drops the // connection or times the command out. Both mean the master moved @@ -2510,7 +2564,11 @@ where Ok(stream) => Ok(stream), }?; - Ok(stream.filter_map(|result| async move { + Ok(stream.filter_map(move |result| { + let delete_connection = delete_connection.clone(); + let owned_index_value = owned_index_value.clone(); + let key_prefix = key_prefix.clone(); + async move { let raw_redis_map = match result { Ok(v) => v, Err(e) => { @@ -2534,6 +2592,7 @@ where let mut redis_map_iter = redis_map.iter(); let mut bytes_data: Option = None; let mut version: Option = None; + let mut doc_key: Option = None; while let Some(key) = redis_map_iter.next() { let value = redis_map_iter.next().unwrap(); let Value::BulkString(k) = key else { @@ -2573,6 +2632,15 @@ where }; version = Some(raw_version); } + DOC_KEY_FIELD_NAME => { + let Ok(str_v) = str::from_utf8(v) else { + return Some(Err(Error::new( + Code::Internal, + format!("Non-utf8 document key from ft_aggregate: {v:?}"), + ))); + }; + doc_key = Some(str_v.to_string()); + } other => { if K::MAYBE_SORT_KEY == Some(other) { // ignore sort keys @@ -2591,11 +2659,50 @@ where format!("Missing '{DATA_FIELD_NAME}' in ft_aggregate, got: {raw_redis_map:?}"), ))); }; - Some( - K::decode(version.unwrap_or(0), found_bytes_data) - .err_tip(|| "In RedisStore::search_by_index_prefix::decode"), - ) - })) + let err = match K::decode(version.unwrap_or(0), found_bytes_data) { + Ok(decoded) => return Some(Ok(decoded)), + Err(err) => err, + }; + // The record is in the index but its data does not decode (for + // example serialization drift across a rolling upgrade). It has + // no TTL and no deletion path, so without removal it is listed, + // skipped and counted by every sweep forever. Records are + // written atomically, so this is never a write in progress. + // Delete it only if the index opted in, its liveness guard says + // nothing is still waiting on it, and its version is unchanged + // since the listing; the row is still reported as lost so the + // caller counts it this one last time. + if err.code == Code::InvalidArgument + && let Some(doc_key) = doc_key + { + let record_key = doc_key.strip_prefix(&key_prefix).unwrap_or(&doc_key); + if let Some(guard_key) = + K::lost_record_liveness_guard_key(&owned_index_value, record_key) + { + let guard_key = format!("{key_prefix}{guard_key}"); + match delete_lost_record( + delete_connection, + &doc_key, + &guard_key, + version.unwrap_or(0), + ) + .await + { + Ok(removed) => warn!( + key = %doc_key, + removed, + "Undecodable record listed; removed if no longer live (1 = removed, 0 = version changed, -1 = liveness guard present)" + ), + Err(del_err) => warn!( + key = %doc_key, + ?del_err, + "Failed to remove a listed record whose data cannot be decoded" + ), + } + } + } + Some(Err(err).err_tip(|| "In RedisStore::search_by_index_prefix::decode")) + }})) } async fn count_by_index_prefix(&self, index: K) -> Result diff --git a/nativelink-store/tests/redis_store_test.rs b/nativelink-store/tests/redis_store_test.rs index 90f6e7dd2bb..368c69d1311 100644 --- a/nativelink-store/tests/redis_store_test.rs +++ b/nativelink-store/tests/redis_store_test.rs @@ -1411,9 +1411,10 @@ fn test_search_by_index() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1526,9 +1527,10 @@ fn test_search_by_index_descending() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1583,9 +1585,10 @@ fn test_search_by_index_retries_on_failover() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1686,9 +1689,10 @@ fn test_search_by_index_skips_docs_that_expired_mid_query() -> Result<(), Error> .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1785,7 +1789,7 @@ fn test_search_by_index_failure() -> Result<(), Error> { "Client: TEST - Client: unexpected command", "Error with ft_create in RedisStore::search_by_index_prefix(test:_content_prefix_sort_key_3e762c15)", "---", "Client: TEST - Client: unexpected command", "Error with second ft_aggregate in RedisStore::search_by_index_prefix(test:_content_prefix_sort_key_3e762c15)"].iter().map(ToString::to_string).collect())); assert!(logs_contain( - "Error calling ft.aggregate e=TEST - Client: unexpected command index=\"test:_content_prefix_sort_key_3e762c15\" query=\"*\" options=FtAggregateOptions { load: [\"data\", \"version\"], cursor: FtAggregateCursor { count: 1500, max_idle: 30000 }, sort_by: [\"@sort_key\"], sort_desc: false } all_args=[\"FT.AGGREGATE\", \"test:_content_prefix_sort_key_3e762c15\", \"*\", \"TIMEOUT\", \"10000\", \"LOAD\", \"2\", \"data\", \"version\", \"WITHCURSOR\", \"COUNT\", \"1500\", \"MAXIDLE\", \"30000\", \"SORTBY\", \"2\", \"@sort_key\", \"ASC\", \"MAX\", \"1000000\"]" + "Error calling ft.aggregate e=TEST - Client: unexpected command index=\"test:_content_prefix_sort_key_3e762c15\" query=\"*\" options=FtAggregateOptions { load: [\"data\", \"version\", \"@__key\"], cursor: FtAggregateCursor { count: 1500, max_idle: 30000 }, sort_by: [\"@sort_key\"], sort_desc: false } all_args=[\"FT.AGGREGATE\", \"test:_content_prefix_sort_key_3e762c15\", \"*\", \"TIMEOUT\", \"10000\", \"LOAD\", \"3\", \"data\", \"version\", \"@__key\", \"WITHCURSOR\", \"COUNT\", \"1500\", \"MAXIDLE\", \"30000\", \"SORTBY\", \"2\", \"@sort_key\", \"ASC\", \"MAX\", \"1000000\"]" )); Ok(()) @@ -1806,9 +1810,10 @@ fn test_search_by_index_swallows_already_exists_from_ft_create() -> Result<(), E .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1892,9 +1897,10 @@ fn test_search_by_index_preserves_other_ft_create_errors() -> Result<(), Error> .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -1974,9 +1980,10 @@ fn test_search_by_index_with_sort_key() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -2062,9 +2069,10 @@ fn test_search_by_index_resp3() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -2173,9 +2181,10 @@ fn test_search_by_index_skips_int_from_cursor_read() -> Result<(), Error> { .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) @@ -2264,9 +2273,10 @@ async fn search_by_index_prefix_reads_a_resp3_cursor_page() -> Result<(), Error> .arg("TIMEOUT") .arg(10000_u64) .arg("LOAD") - .arg(2) + .arg(3) .arg("data") .arg("version") + .arg("@__key") .arg("WITHCURSOR") .arg("COUNT") .arg(1500) diff --git a/nativelink-util/src/store_trait.rs b/nativelink-util/src/store_trait.rs index e68d030c670..92b9ad8b890 100644 --- a/nativelink-util/src/store_trait.rs +++ b/nativelink-util/src/store_trait.rs @@ -1025,6 +1025,23 @@ pub trait SchedulerIndexProvider { /// The value of the index. fn index_value(&self) -> Cow<'_, str>; + + /// Opts an index into garbage-collecting a listed record whose data + /// exists but cannot be decoded (for example serialization drift across + /// a rolling upgrade). Such a record has no TTL, and no reader or + /// writer can ever use it again, so without removal it is listed, + /// skipped and counted forever. + /// + /// Returning `Some(guard_key)` permits deletion of `record_key` only + /// while `guard_key` does not exist; the guard should be a liveness + /// marker (such as a client keepalive) so a record something still + /// cares about is never removed. `None`, the default, means a record + /// listed by this index is never deleted. Both keys are store keys + /// before any store-level key prefix is applied. + #[must_use] + fn lost_record_liveness_guard_key(_index_value: &str, _record_key: &str) -> Option { + None + } } /// Provides a key to lookup data in the store. From 747a2c91b353e30e7b309b58d46e4e8c9da84b50 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 1 Oct 2026 21:33:31 -0400 Subject: [PATCH 14/46] fix(scheduler): defuse cache-lookup drop-guard after claiming the map entry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The inflight-dedup drop-guard in CacheLookupScheduler::add_action was bound as `let _scope_guard = scope_guard;`, so it fired on *every* exit of the spawn — including normal completion, after the spawn had already removed its own entry and resolved the coalesced waiters. A concurrent request for the same cacheable key, arriving in the window between the leader's own `inflight_cache_checks.remove(...)` and the spawn's exit (dropping the guard), inserts a fresh entry. The stale guard then removes *that* entry, orphaning the second request: its oneshot sender is dropped, surfacing a spurious "tx hung up" error to the client and forcing a needless retry. Fix: defuse the guard via `ScopeGuard::into_inner` the moment the cache check returns. From that point every arm removes our own entry with no intervening `.await` (no drop point), so defusing is leak-free and the guard can no longer clobber a concurrent same-key re-insertion. The guard still provides true drop-safety for the only window that matters: a spawn death during the cache-store check, before we claim the entry. Regression test leader_drop_guard_does_not_orphan_a_later_same_key_request parks a leader in its forward to the inner scheduler (an await after it has removed its entry), sends a second same-key request into that window, then lets the leader finish. Without the fix the second request fails with "tx hung up"; with it, the second request is forwarded and answered. Co-Authored-By: Claude Opus 4.8 Co-Authored-By: Claude Opus 5.5 --- .../src/cache_lookup_scheduler.rs | 16 ++- .../tests/cache_lookup_scheduler_test.rs | 98 ++++++++++++++++++- 2 files changed, 110 insertions(+), 4 deletions(-) diff --git a/nativelink-scheduler/src/cache_lookup_scheduler.rs b/nativelink-scheduler/src/cache_lookup_scheduler.rs index 3d6ced82639..c8600b77d2b 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/tests/cache_lookup_scheduler_test.rs b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs index cbff20b9867..de24db6d328 100644 --- a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs +++ b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs @@ -32,12 +32,14 @@ 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::operation_state_manager::{ + ActionStateResult, ClientStateManager, OperationFilter, +}; use nativelink_util::store_trait::{ RemoveItemCallback, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, }; @@ -253,3 +255,95 @@ async fn find_by_client_operation_id_call_passed() -> Result<(), Error> { ); Ok(()) } + +fn queued_result( + client_operation_id: OperationId, + action_info: Arc, +) -> Box { + let (_tx, rx) = watch::channel(Arc::new(ActionState { + client_operation_id: client_operation_id.clone(), + stage: ActionStage::Queued, + action_digest: action_info.unique_qualifier.digest(), + last_transition_timestamp: SystemTime::now(), + })); + Box::new(TokioWatchActionStateResult::new( + client_operation_id, + action_info, + rx, + )) +} + +// A leader's cache-check task removes its in-flight entry before forwarding a +// miss to the inner scheduler, and that forward is an await point. A second +// request for the same key arriving in that window inserts a fresh entry. The +// leader's drop guard must not remove that fresh entry when the leader's task +// exits, or the second request's sender is dropped and it fails with +// "tx hung up". The default test runtime is single-threaded, so the +// interleaving below is deterministic. +#[nativelink_test] +async fn leader_drop_guard_does_not_orphan_a_later_same_key_request() -> Result<(), Error> { + let context = make_cache_scheduler()?; + let cache_scheduler = Arc::new(context.cache_scheduler); + let mock_scheduler = context.mock_scheduler; + // Nothing in the AC, so both requests miss and are forwarded. + let action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::new([7; 32], 77)); + + // Request A becomes the leader, misses, removes its entry and parks in the + // forward to the mock scheduler, which has not answered yet. + let op_a = OperationId::default(); + let request_a = tokio::spawn({ + let cache_scheduler = cache_scheduler.clone(); + let (op_a, action_info) = (op_a.clone(), action_info.clone()); + async move { + cache_scheduler + .add_action(op_a, action_info) + .await + .map(|_| ()) + } + }); + for _ in 0..50 { + tokio::task::yield_now().await; + } + + // Request B for the same key arrives while A is parked: A's entry is gone, + // so B inserts its own. + let op_b = OperationId::default(); + let request_b = tokio::spawn({ + let cache_scheduler = cache_scheduler.clone(); + let (op_b, action_info) = (op_b.clone(), action_info.clone()); + async move { + cache_scheduler + .add_action(op_b, action_info) + .await + .map(|_| ()) + } + }); + + // Answer A's forward. A's task then finishes while B's entry is in the map. + let (forwarded_a, _) = mock_scheduler + .expect_add_action(Ok(queued_result(op_a.clone(), action_info.clone()))) + .await; + assert_eq!(forwarded_a, op_a); + + // B must still be forwarded and answered, not orphaned. + let serve_b = tokio::spawn({ + let mock_scheduler = mock_scheduler.clone(); + let action_info = action_info.clone(); + async move { + mock_scheduler + .expect_add_action(Ok(queued_result(op_b, action_info))) + .await + } + }); + let result_b = tokio::time::timeout(std::time::Duration::from_secs(5), request_b) + .await + .expect("request B never completed") + .expect("request B panicked"); + assert!( + result_b.is_ok(), + "a same-key request that arrived while the leader was forwarding was orphaned: {result_b:?}" + ); + serve_b.await.expect("serving B panicked"); + request_a.await.expect("request A panicked")?; + Ok(()) +} From 80678933f8571830df7b4bee10be05a64fc7d068 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 7 Oct 2026 18:12:32 -0400 Subject: [PATCH 15/46] fix(filesystem_store): bound background delete concurrency with a semaphore Evicted entries' backing files are unlinked by a fire-and-forget background_spawn! in EncodedFilePath::drop, with no bound. Under an eviction storm (cache wipe, overwrite flood) that grows without limit: queued task frames consume memory and disk/inode reclamation lags behind what the async runtime can schedule. Open FDs are already bounded by a semaphore; deletes were not. Add max_concurrent_deletes. drop is sync so it cannot await a permit; it try_acquire_owned's and, when saturated, deletes synchronously inline -- which is the backpressure (the evicting thread pays the unlink cost directly) and never leaks a file. An inline_deletes metric exposes how often backpressure engages. Bounded by default (1024), unlike max_concurrent_writes: leaving an unbounded delete backlog off by default would move a latent memory/inode blowup from a high-context code review to a late-night config change nobody remembers to make. The default is generous enough that normal eviction never reaches it (so the inline path never fires in ordinary operation); it only engages under a genuine storm. 0 remains the explicit escape hatch for unlimited. The serde field default governs JSON5 configs (production); derive(Default) still yields 0, so existing Rust-constructed tests are unchanged. --- nativelink-config/src/stores.rs | 34 +++++++++++ nativelink-store/src/filesystem_store.rs | 55 ++++++++++++++++++ .../tests/filesystem_store_test.rs | 56 +++++++++++++++++++ 3 files changed, 145 insertions(+) diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index 81b6179f929..f1f1b36fdc9 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -844,6 +844,32 @@ pub struct FilesystemSpec { #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] pub max_concurrent_writes: usize, + /// Maximum number of concurrent background file-deletion operations. + /// When an entry leaves the store, its backing file is unlinked by a + /// background task. Without a bound, an eviction storm (cache wipe or a + /// flood of overwrites) can spawn these without limit: memory grows with the + /// queued task frames and disk/inode reclamation lags behind what the async + /// runtime can schedule. Bounding it applies backpressure the same way + /// `max_concurrent_writes` does — once this many deletes are in flight, a + /// further delete runs synchronously inline in the evicting path instead of + /// spawning, which throttles eviction to delete throughput and never leaks + /// a file. + /// + /// Bounded by default, unlike `max_concurrent_writes`: an unbounded delete + /// backlog is a latent memory/inode blowup, and leaving the protection off + /// by default just moves that risk from a high-context code review to a + /// 4 AM config change nobody remembers to make. The default is generous + /// enough that normal steady-state eviction never reaches it (so the inline + /// path never fires in ordinary operation); it only engages under a genuine + /// storm, which is exactly when the backpressure is wanted. + /// A value of 0 is the explicit escape hatch for unlimited. + /// Default: 1024. + #[serde( + default = "default_max_concurrent_deletes", + deserialize_with = "convert_numeric_with_shellexpand" + )] + pub max_concurrent_deletes: usize, + /// When true, advise the kernel to drop the page cache for each blob after /// it is written or read (`posix_fadvise` with `POSIX_FADV_DONTNEED`). On /// real filesystems this takes a globally serialized, all-CPU kernel path @@ -1797,6 +1823,14 @@ const fn default_read_batching_dispatch_slots() -> usize { 4 } +/// Generous enough that normal eviction never reaches it (so the synchronous +/// inline delete path never fires in ordinary operation), low enough that the +/// in-flight backlog stays bounded to a few MB under a storm. Tune per +/// deployment; 0 disables the bound. See `FilesystemSpec::max_concurrent_deletes`. +const fn default_max_concurrent_deletes() -> usize { + 1024 +} + const fn default_read_batching_max_queued_bytes() -> u64 { 32 * 1024 * 1024 // 32 MiB. } diff --git a/nativelink-store/src/filesystem_store.rs b/nativelink-store/src/filesystem_store.rs index a7925e18ce6..9ad1780f9fa 100644 --- a/nativelink-store/src/filesystem_store.rs +++ b/nativelink-store/src/filesystem_store.rs @@ -102,6 +102,17 @@ pub struct SharedContext { // deleted (similar to how it is done in tests). #[metric(help = "Number of active drop spawns")] pub active_drop_spawns: AtomicU64, + /// Count of deletes that ran synchronously inline because the delete + /// semaphore was saturated. A nonzero, growing value means eviction is + /// outrunning background delete throughput — the backpressure signal. + #[metric(help = "Number of deletes performed inline due to delete backpressure")] + pub inline_deletes: AtomicU64, + /// Bounds the number of in-flight background file deletions. `None` means + /// unlimited (the historical behavior). When the permits are exhausted, + /// a drop deletes its file synchronously inline instead of spawning, so + /// an eviction storm applies backpressure rather than spawning without + /// limit. See `FilesystemSpec::max_concurrent_deletes`. + delete_semaphore: Option>, #[metric(help = "Path to the configured temp path")] temp_path: String, #[metric(help = "Path to the configured content path")] @@ -187,6 +198,40 @@ impl Drop for EncodedFilePath { let file_path = self.get_file_path().to_os_string(); let shared_context = self.shared_context.clone(); + + // Bound the number of in-flight background deletes. `drop` is sync, so + // we cannot `.await` a permit the way the write path does; instead we + // `try_acquire_owned` (never blocks) and, when the permits are + // exhausted, delete synchronously inline. That inline path is the + // backpressure: the evicting thread pays the unlink cost directly + // instead of queueing an unbounded number of tasks, and the file is + // always removed (never leaked). `None` keeps the historical + // unbounded-spawn behavior. + let delete_permit = match &shared_context.delete_semaphore { + Some(sem) => { + if let Ok(permit) = Arc::clone(sem).try_acquire_owned() { + Some(permit) + } else { + // Saturated: delete synchronously inline, which is the + // backpressure, then bail out (no spawn). + shared_context + .inline_deletes + .fetch_add(1, Ordering::Relaxed); + match std::fs::remove_file(&file_path) { + Ok(()) => { + debug!(?file_path, "File deleted inline (delete backpressure)"); + } + Err(err) if err.kind() == std::io::ErrorKind::NotFound => { + debug!(?file_path, "File already gone, nothing to delete"); + } + Err(err) => error!(?file_path, ?err, "Failed to delete file inline"), + } + return; + } + } + None => None, + }; + // .fetch_add returns previous value, so we add one to get approximate current value let current_active_drop_spawns = shared_context .active_drop_spawns @@ -198,6 +243,9 @@ impl Drop for EncodedFilePath { "Spawned a filesystem_delete_file" ); background_spawn!("filesystem_delete_file", async move { + // Hold the permit until the delete actually completes so the bound + // reflects deletes in flight, not merely scheduled. + let _delete_permit = delete_permit; match fs::remove_file(&file_path).await { Ok(()) => debug!(?file_path, "File deleted"), // The file already being gone is the desired end state of a @@ -1433,8 +1481,15 @@ impl FilesystemStore { enabled }; + let delete_semaphore = if spec.max_concurrent_deletes > 0 { + Some(Arc::new(Semaphore::new(spec.max_concurrent_deletes))) + } else { + None + }; let shared_context = Arc::new(SharedContext { active_drop_spawns: AtomicU64::new(0), + inline_deletes: AtomicU64::new(0), + delete_semaphore, temp_path: spec.temp_path.clone(), content_path: spec.content_path.clone(), }); diff --git a/nativelink-store/tests/filesystem_store_test.rs b/nativelink-store/tests/filesystem_store_test.rs index 7f659a16123..d92383df92b 100644 --- a/nativelink-store/tests/filesystem_store_test.rs +++ b/nativelink-store/tests/filesystem_store_test.rs @@ -711,6 +711,62 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { check_storage_dir_empty(&temp_path).await } +// Bounding background deletes with `max_concurrent_deletes` must not change the +// correctness of eviction: an evicted entry's backing file is still removed. +// This exercises the bounded (semaphore) delete path — permit acquired, held +// for the lifetime of the delete, then released. The default unbounded path is +// covered by the rest of this suite; the inline-overflow fallback (which fires +// only under saturation) is exercised under real load by the furnace via the +// `inline_deletes` metric. +#[nativelink_test] +async fn bounded_delete_semaphore_still_removes_evicted_files() -> Result<(), Error> { + static DELETES_FINISHED: AtomicU32 = AtomicU32::new(0); + struct LocalHooks {} + impl FileEntryHooks for LocalHooks { + fn on_drop(_file_entry: &Fe) { + DELETES_FINISHED.fetch_add(1, Ordering::Relaxed); + } + } + + let digest1 = DigestInfo::try_new(HASH1, VALUE1.len())?; + let digest2 = DigestInfo::try_new(HASH2, VALUE2.len())?; + let content_path = make_temp_path("content_path"); + let temp_path = make_temp_path("temp_path"); + + let store = Arc::new( + FilesystemStore::>::new(&FilesystemSpec { + content_path: content_path.clone(), + temp_path: temp_path.clone(), + eviction_policy: Some(EvictionPolicy { + max_count: 1, + ..Default::default() + }), + // The bound under test. With it set, the delete routes through the + // semaphore path rather than the unconditional spawn. + max_concurrent_deletes: 1, + block_size: 1, + read_buffer_size: 1, + ..Default::default() + }) + .await?, + ); + + // Insert, then evict digest1 by inserting digest2 (max_count == 1). + store.update_oneshot(digest1, VALUE1.into()).await?; + store.update_oneshot(digest2, VALUE2.into()).await?; + + // Wait for the eviction's background delete to finish. + loop { + if DELETES_FINISHED.load(Ordering::Relaxed) == 1 { + break; + } + tokio::task::yield_now().await; + } + + // The evicted file must be gone from the temp staging area. + check_storage_dir_empty(&temp_path).await +} + // Test to ensure that if we are holding a reference to `FileEntry` and the contents are // replaced, the `FileEntry` continues to use the old data. // `FileEntry` file contents should be immutable for the lifetime of the object. From a202c4c292bd7b775f65dbbcb9d36a482b6984ed Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 19:02:30 -0400 Subject: [PATCH 16/46] store(dedup): reject torn content-chunk reads instead of serving them as success DedupStore::get_part trusted the per-chunk sizes recorded in the index and streamed each content chunk without checking it actually held as many bytes as its index entry declared. A short content chunk (a torn/partial read or an index/content size disagreement) would either panic the worker on the bytes_to_skip assert, or silently under-deliver and send EOF, serving a truncated read to the caller as success. Add a per-chunk size-consistency check: a fetched content chunk whose length does not equal the size its index entry declares now yields a NotFound error. The check is per-chunk (not a post-loop total), so a legitimate bounded request that overshoots the end of the blob still returns the available bytes. Regression tests: get_part_short_chunk_is_error_test (fails without the fix: "got Ok with 10 bytes") and get_part_over_length_request_still_ok_test (guards against over-correction). Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/dedup_store.rs | 17 +++++ nativelink-store/tests/dedup_store_test.rs | 83 +++++++++++++++++++++- 2 files changed, 99 insertions(+), 1 deletion(-) diff --git a/nativelink-store/src/dedup_store.rs b/nativelink-store/src/dedup_store.rs index e3b437ca00a..105be62d076 100644 --- a/nativelink-store/src/dedup_store.rs +++ b/nativelink-store/src/dedup_store.rs @@ -352,6 +352,23 @@ impl StoreDriver for DedupStore { .await .err_tip(|| "Failed to get_part in content_store in dedup_store")?; + // Content chunks are content-addressed and whole: a chunk must + // contain exactly the number of bytes its index entry declares. + // If the content store returns a short chunk (a torn/partial + // read, or an index/content size disagreement) the streaming + // loop below would either panic on the `bytes_to_skip` assert or + // silently under-deliver and send EOF as success. Reject the + // mismatch here so it surfaces as an error instead. + let declared = index_entry.size_bytes(); + let actual = data.len() as u64; + if actual != declared { + return Err(make_err!( + Code::NotFound, + "DedupStore content chunk {} is {actual} bytes but its index entry declares {declared}; refusing to serve a torn read as success", + index_entry, + )); + } + Result::<_, Error>::Ok(data) }) .buffered(self.max_concurrent_fetch_per_get); diff --git a/nativelink-store/tests/dedup_store_test.rs b/nativelink-store/tests/dedup_store_test.rs index e72f2905a62..f45f693ce0d 100644 --- a/nativelink-store/tests/dedup_store_test.rs +++ b/nativelink-store/tests/dedup_store_test.rs @@ -16,7 +16,8 @@ use nativelink_config::stores::{DedupSpec, MemorySpec, StoreSpec}; use nativelink_error::{Code, Error, ResultExt}; use nativelink_macro::nativelink_test; use nativelink_store::cas_utils::ZERO_BYTE_DIGESTS; -use nativelink_store::dedup_store::DedupStore; +use nativelink_store::compression_store::WincodeConfig; +use nativelink_store::dedup_store::{DedupIndex, DedupStore}; use nativelink_store::memory_store::MemoryStore; use nativelink_util::common::DigestInfo; use nativelink_util::store_trait::{Store, StoreLike}; @@ -445,3 +446,83 @@ async fn empty_cas_upload_does_not_publish_dedup_index() -> Result<(), Error> { assert_eq!(content.len_for_test(), 0); Ok(()) } + +// Regression: a content chunk must contain exactly the number of bytes its +// index entry declares. If the content store returns a short chunk (a torn or +// partial read, or an index/content size disagreement), `DedupStore::get_part` +// previously either panicked on the `bytes_to_skip` assert or silently +// under-delivered and sent EOF, serving a truncated read as success. This test +// hand-builds an index whose single entry declares more bytes than the stored +// content and asserts the read now surfaces an error. +#[nativelink_test] +async fn get_part_short_chunk_is_error_test() -> Result<(), Error> { + const DECLARED_SIZE: usize = 100; + const ACTUAL_SIZE: usize = 10; + + let index_store = MemoryStore::new(&MemorySpec::default()); + let content_store = MemoryStore::new(&MemorySpec::default()); + let store = DedupStore::new( + &make_default_config(), + Store::new(index_store.clone()), + Store::new(content_store.clone()), + )?; + + let chunk_digest = DigestInfo::try_new(VALID_HASH2, DECLARED_SIZE).unwrap(); + content_store + .update_oneshot(chunk_digest, vec![7u8; ACTUAL_SIZE].into()) + .await + .err_tip(|| "Failed to write short content chunk")?; + + let serialized = wincode::config::serialize( + &DedupIndex { + entries: vec![chunk_digest], + }, + WincodeConfig::new(), + ) + .unwrap(); + let top_digest = DigestInfo::try_new(VALID_HASH1, DECLARED_SIZE).unwrap(); + Store::new(index_store.clone()) + .update_oneshot(top_digest, serialized.into()) + .await + .err_tip(|| "Failed to write index entry")?; + + let result = store + .get_part_unchunked(top_digest, 0, Some(DECLARED_SIZE as u64)) + .await; + assert!( + result.is_err(), + "Expected short-chunk read to error, got Ok with {} bytes", + result.map_or(0, |d| d.len()) + ); + Ok(()) +} + +// Guard against over-correction: a bounded request for MORE bytes than the blob +// actually contains must still succeed and return the available bytes, matching +// the clamp-to-available semantics of the underlying stores. The +// size-consistency check must not turn a legitimate over-length request into an +// error. +#[nativelink_test] +async fn get_part_over_length_request_still_ok_test() -> Result<(), Error> { + const DATA_SIZE: usize = 64 * 1024; + + let store = DedupStore::new( + &make_default_config(), + Store::new(MemoryStore::new(&MemorySpec::default())), + Store::new(MemoryStore::new(&MemorySpec::default())), + )?; + + let original_data = make_random_data(DATA_SIZE); + let digest = DigestInfo::try_new(VALID_HASH1, DATA_SIZE).unwrap(); + store + .update_oneshot(digest, original_data.clone().into()) + .await + .err_tip(|| "Failed to write data to dedup store")?; + + let rt = store + .get_part_unchunked(digest, 0, Some((DATA_SIZE as u64) * 2)) + .await + .err_tip(|| "Over-length bounded read should still succeed")?; + assert_eq!(rt, original_data, "Expected available bytes to be returned"); + Ok(()) +} From 0506e038df9e92c7abca58d9960b54d6d4a5f379 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 21 Oct 2026 17:19:02 +0000 Subject: [PATCH 17/46] test(scheduler): run client_keepalive_clobber_test under Bazel The test file was only picked up by Cargo's tests/ auto-discovery; the Bazel integration suite lists its sources explicitly, so CI never ran it. Register it, and drop the dev-dependency on mock_instant, which is already a regular dependency of the crate. Co-Authored-By: Claude Opus 5.5 --- nativelink-scheduler/BUILD.bazel | 1 + nativelink-scheduler/Cargo.toml | 1 - 2 files changed, 1 insertion(+), 1 deletion(-) diff --git a/nativelink-scheduler/BUILD.bazel b/nativelink-scheduler/BUILD.bazel index 25077de27cc..0d56cfa399a 100644 --- a/nativelink-scheduler/BUILD.bazel +++ b/nativelink-scheduler/BUILD.bazel @@ -74,6 +74,7 @@ rust_test_suite( "tests/admin_test.rs", "tests/backpressure_test.rs", "tests/cache_lookup_scheduler_test.rs", + "tests/client_keepalive_clobber_test.rs", "tests/dispatch_ack_test.rs", "tests/escalation_test.rs", "tests/execution_active_count_test.rs", diff --git a/nativelink-scheduler/Cargo.toml b/nativelink-scheduler/Cargo.toml index f2971aad85b..99795eb35bc 100644 --- a/nativelink-scheduler/Cargo.toml +++ b/nativelink-scheduler/Cargo.toml @@ -53,7 +53,6 @@ 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" } From 607e546ec25a5e50e5a0e1e2da11e38e82d59850 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 21 Oct 2026 17:19:03 +0000 Subject: [PATCH 18/46] test(scheduler): run client_keepalive_failed_write_test under Bazel The test file was only picked up by Cargo's tests/ auto-discovery; the Bazel integration suite lists its sources explicitly, so CI never ran it. Co-Authored-By: Claude Opus 5.5 --- nativelink-scheduler/BUILD.bazel | 1 + 1 file changed, 1 insertion(+) diff --git a/nativelink-scheduler/BUILD.bazel b/nativelink-scheduler/BUILD.bazel index 25077de27cc..8c4e1994b79 100644 --- a/nativelink-scheduler/BUILD.bazel +++ b/nativelink-scheduler/BUILD.bazel @@ -74,6 +74,7 @@ rust_test_suite( "tests/admin_test.rs", "tests/backpressure_test.rs", "tests/cache_lookup_scheduler_test.rs", + "tests/client_keepalive_failed_write_test.rs", "tests/dispatch_ack_test.rs", "tests/escalation_test.rs", "tests/execution_active_count_test.rs", From c3c9ae163239e505ac6234e02b6c0df50061f450 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 15 Oct 2026 12:23:46 -0400 Subject: [PATCH 19/46] store(dedup): drop a content chunk whose size does not match its digest Rejecting a wrong-size chunk with NotFound stopped the truncated read, but left the bad chunk in place. `has()` sums the sizes the content store reports, so the blob stayed "present", and `update()` skips any chunk the content store already has, so the client's re-upload kept the bad copy. The blob could never be read again. On a size mismatch, `get_part` now makes a best-effort `remove` of the chunk from the content store, as `VerifyStore` does for a copy that fails verification. The blob then reads as missing and the next upload writes the chunk again. The error stays NotFound, not DataLoss, so a `FastSlowStore` above can still fall back to its slow tier. Tests now cover the error code, the removal, a read that starts inside a short chunk (the old assert panic), a chunk longer than declared, and an end-to-end truncate / fail / re-upload / read cycle. The removal, has() and repair assertions fail without this change. Co-Authored-By: Claude Opus 5.5 --- nativelink-store/src/dedup_store.rs | 42 +++++- nativelink-store/tests/dedup_store_test.rs | 144 +++++++++++++++++---- 2 files changed, 157 insertions(+), 29 deletions(-) diff --git a/nativelink-store/src/dedup_store.rs b/nativelink-store/src/dedup_store.rs index 105be62d076..c2f4a6307c9 100644 --- a/nativelink-store/src/dedup_store.rs +++ b/nativelink-store/src/dedup_store.rs @@ -354,14 +354,46 @@ impl StoreDriver for DedupStore { // Content chunks are content-addressed and whole: a chunk must // contain exactly the number of bytes its index entry declares. - // If the content store returns a short chunk (a torn/partial - // read, or an index/content size disagreement) the streaming - // loop below would either panic on the `bytes_to_skip` assert or - // silently under-deliver and send EOF as success. Reject the - // mismatch here so it surfaces as an error instead. + // If the content store returns a chunk of any other size (a + // torn or partial read, or an index/content size disagreement) + // the streaming loop below would either panic on the + // `bytes_to_skip` assert or silently under-deliver and send EOF + // as success. Reject the mismatch here so it surfaces as an + // error instead. let declared = index_entry.size_bytes(); let actual = data.len() as u64; if actual != declared { + // Every blob that references this chunk is unreadable while + // it stays, and nothing else would notice: `has()` sums the + // sizes the content store reports, and `update()` skips any + // chunk the content store already has, so a re-upload would + // keep the bad copy. Drop it, as `VerifyStore` drops a copy + // that fails verification, so `has()` reports those blobs + // missing and the next upload writes the chunk again. This + // is best effort: a content store that does not support + // `remove` keeps the chunk. + match self + .content_store + .as_store_driver_pin() + .remove(index_entry.into()) + .await + { + Ok(removed) => warn!( + %index_entry, + actual, + declared, + removed, + "Dropped a dedup content chunk whose size does not match its digest" + ), + Err(err) => warn!( + %index_entry, + ?err, + "Could not drop a dedup content chunk whose size does not match its digest" + ), + } + // NotFound rather than the DataLoss `VerifyStore` returns: a + // `FastSlowStore` above falls back to its slow tier only on + // NotFound, and the slow tier may hold a good copy. return Err(make_err!( Code::NotFound, "DedupStore content chunk {} is {actual} bytes but its index entry declares {declared}; refusing to serve a torn read as success", diff --git a/nativelink-store/tests/dedup_store_test.rs b/nativelink-store/tests/dedup_store_test.rs index f45f693ce0d..22850e323f4 100644 --- a/nativelink-store/tests/dedup_store_test.rs +++ b/nativelink-store/tests/dedup_store_test.rs @@ -12,6 +12,8 @@ // See the License for the specific language governing permissions and // limitations under the License. +use std::sync::Arc; + use nativelink_config::stores::{DedupSpec, MemorySpec, StoreSpec}; use nativelink_error::{Code, Error, ResultExt}; use nativelink_macro::nativelink_test; @@ -447,18 +449,14 @@ async fn empty_cas_upload_does_not_publish_dedup_index() -> Result<(), Error> { Ok(()) } -// Regression: a content chunk must contain exactly the number of bytes its -// index entry declares. If the content store returns a short chunk (a torn or -// partial read, or an index/content size disagreement), `DedupStore::get_part` -// previously either panicked on the `bytes_to_skip` assert or silently -// under-delivered and sent EOF, serving a truncated read as success. This test -// hand-builds an index whose single entry declares more bytes than the stored -// content and asserts the read now surfaces an error. -#[nativelink_test] -async fn get_part_short_chunk_is_error_test() -> Result<(), Error> { - const DECLARED_SIZE: usize = 100; - const ACTUAL_SIZE: usize = 10; - +/// Builds a `DedupStore` whose index for `VALID_HASH1` names one content chunk +/// of `declared_size` bytes, and stores `stored` under that chunk's digest. A +/// `stored` of any other length is a chunk whose bytes disagree with its index +/// entry, as a torn write or a misbehaving content store would leave it. +async fn store_with_one_chunk( + declared_size: usize, + stored: Vec, +) -> Result<(Arc, Arc, DigestInfo, DigestInfo), Error> { let index_store = MemoryStore::new(&MemorySpec::default()); let content_store = MemoryStore::new(&MemorySpec::default()); let store = DedupStore::new( @@ -467,11 +465,11 @@ async fn get_part_short_chunk_is_error_test() -> Result<(), Error> { Store::new(content_store.clone()), )?; - let chunk_digest = DigestInfo::try_new(VALID_HASH2, DECLARED_SIZE).unwrap(); + let chunk_digest = DigestInfo::try_new(VALID_HASH2, declared_size).unwrap(); content_store - .update_oneshot(chunk_digest, vec![7u8; ACTUAL_SIZE].into()) + .update_oneshot(chunk_digest, stored.into()) .await - .err_tip(|| "Failed to write short content chunk")?; + .err_tip(|| "Failed to write content chunk")?; let serialized = wincode::config::serialize( &DedupIndex { @@ -480,20 +478,118 @@ async fn get_part_short_chunk_is_error_test() -> Result<(), Error> { WincodeConfig::new(), ) .unwrap(); - let top_digest = DigestInfo::try_new(VALID_HASH1, DECLARED_SIZE).unwrap(); - Store::new(index_store.clone()) + let top_digest = DigestInfo::try_new(VALID_HASH1, declared_size).unwrap(); + index_store .update_oneshot(top_digest, serialized.into()) .await .err_tip(|| "Failed to write index entry")?; + Ok((store, content_store, top_digest, chunk_digest)) +} - let result = store - .get_part_unchunked(top_digest, 0, Some(DECLARED_SIZE as u64)) - .await; - assert!( - result.is_err(), - "Expected short-chunk read to error, got Ok with {} bytes", - result.map_or(0, |d| d.len()) +// A content chunk must contain exactly the number of bytes its index entry +// declares. A short chunk used to be served as a truncated read that ended in a +// normal EOF. The read must fail with NotFound, and the bad chunk must be gone +// afterwards so `has()` stops reporting the blob as present. +#[nativelink_test] +async fn get_part_short_chunk_is_error_test() -> Result<(), Error> { + let (store, content_store, top_digest, chunk_digest) = + store_with_one_chunk(100, vec![7u8; 10]).await?; + + let result = store.get_part_unchunked(top_digest, 0, Some(100)).await; + let err = result.expect_err("a short chunk must not be served as a successful read"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + + assert_eq!( + content_store.has(chunk_digest).await?, + None, + "the short chunk must be dropped from the content store" ); + assert_eq!( + store.has(top_digest).await?, + None, + "a blob whose chunk was dropped must no longer be reported as present" + ); + Ok(()) +} + +// A read that starts inside a short chunk used to trip the `bytes_to_skip` +// assert and panic. It must fail like any other read of that chunk. +#[nativelink_test] +async fn get_part_offset_inside_short_chunk_is_error_test() -> Result<(), Error> { + let (store, _content_store, top_digest, _chunk_digest) = + store_with_one_chunk(100, vec![7u8; 10]).await?; + + let result = store.get_part_unchunked(top_digest, 50, Some(10)).await; + let err = result.expect_err("a read inside a short chunk must fail, not panic"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + Ok(()) +} + +// A chunk longer than its index entry declares is just as wrong as a short one. +#[nativelink_test] +async fn get_part_long_chunk_is_error_test() -> Result<(), Error> { + let (store, content_store, top_digest, chunk_digest) = + store_with_one_chunk(100, vec![7u8; 150]).await?; + + let result = store.get_part_unchunked(top_digest, 0, None).await; + let err = result.expect_err("a long chunk must not be served as a successful read"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + assert_eq!(content_store.has(chunk_digest).await?, None); + Ok(()) +} + +// End to end: a blob with one corrupted chunk fails to read, is then reported +// missing, and the client's re-upload makes it readable again. Before the bad +// chunk was dropped, `update()` skipped it because the content store still had +// a value under its digest, so the blob stayed unreadable forever. +#[nativelink_test] +async fn short_chunk_is_rewritten_by_the_next_upload_test() -> Result<(), Error> { + let index_store = MemoryStore::new(&MemorySpec::default()); + let content_store = MemoryStore::new(&MemorySpec::default()); + let store = DedupStore::new( + &make_default_config(), + Store::new(index_store.clone()), + Store::new(content_store.clone()), + )?; + + let original_data = make_random_data(MEGABYTE_SZ); + let digest = DigestInfo::try_new(VALID_HASH1, MEGABYTE_SZ).unwrap(); + store + .update_oneshot(digest, original_data.clone().into()) + .await + .err_tip(|| "Failed to write data to dedup store")?; + + // Truncate one chunk from the middle of the blob in place. + let index_bytes = index_store.get_part_unchunked(digest, 0, None).await?; + let index: DedupIndex = + wincode::config::deserialize(&index_bytes, WincodeConfig::new()).unwrap(); + assert!(index.entries.len() > 2, "expected several chunks"); + let victim = index.entries[index.entries.len() / 2]; + let victim_bytes = content_store.get_part_unchunked(victim, 0, None).await?; + content_store + .update_oneshot(victim, victim_bytes.slice(..victim_bytes.len() / 2)) + .await?; + + let err = store + .get_part_unchunked(digest, 0, None) + .await + .expect_err("a blob with a truncated chunk must not read back"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + assert_eq!( + store.has(digest).await?, + None, + "the blob must now be missing" + ); + + store + .update_oneshot(digest, original_data.clone().into()) + .await + .err_tip(|| "Failed to re-upload data to dedup store")?; + let rt_data = store + .get_part_unchunked(digest, 0, None) + .await + .err_tip(|| "The re-uploaded blob must read back")?; + assert_eq!(rt_data, original_data, "Expected round trip data to match"); Ok(()) } From 3584d241b02b5c3cfbda33771b6130b949be6fa7 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 10:48:21 -0400 Subject: [PATCH 20/46] 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/oci_store_test.rs | 2 +- nativelink-store/tests/ontap_s3_store_test.rs | 63 ++++++++++++++++++- nativelink-store/tests/s3_store_test.rs | 61 +++++++++++++++++- 8 files changed, 215 insertions(+), 11 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 9fe7d161d56..b59ca03e852 100644 --- a/nativelink-store/src/gcs_store.rs +++ b/nativelink-store/src/gcs_store.rs @@ -436,8 +436,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 224fc2df52e..72d795ee85c 100644 --- a/nativelink-store/src/ontap_s3_store.rs +++ b/nativelink-store/src/ontap_s3_store.rs @@ -622,8 +622,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 34e532a542a..6bd31287b8a 100644 --- a/nativelink-store/src/s3_store.rs +++ b/nativelink-store/src/s3_store.rs @@ -637,8 +637,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 fce8d76b6a6..8a7ccfaa45b 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(()) @@ -1043,3 +1043,52 @@ async fn incomplete_upload_never_publishes_gcs_object() -> 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/oci_store_test.rs b/nativelink-store/tests/oci_store_test.rs index d2eb5cadbd5..1d092465fd1 100644 --- a/nativelink-store/tests/oci_store_test.rs +++ b/nativelink-store/tests/oci_store_test.rs @@ -250,7 +250,7 @@ async fn ranged_get_sends_range_header_path_style() -> Result<(), Error> { .uri(format!( "https://{HOST}/{BUCKET}/{VALID_HASH1}-1000?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() diff --git a/nativelink-store/tests/ontap_s3_store_test.rs b/nativelink-store/tests/ontap_s3_store_test.rs index bbcd72ef25f..83a48caa2ab 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 @@ -773,3 +773,62 @@ async fn incomplete_small_upload_never_sends_ontap_put() -> 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.{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 a9f4b7ee185..f4f193c2d3a 100644 --- a/nativelink-store/tests/s3_store_test.rs +++ b/nativelink-store/tests/s3_store_test.rs @@ -326,7 +326,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() @@ -564,7 +564,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() @@ -1097,3 +1097,60 @@ async fn incomplete_multipart_upload_aborts_before_s3_commit() -> Result<(), Err assert_eq!(methods, vec!["POST", "DELETE"]); 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 5ac552526cda8053f84658b7d6171f73d1a651a7 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Sat, 18 Jul 2026 18:49:43 -0400 Subject: [PATCH 21/46] s3_store: size multipart parts at 64 MiB by default, not the 5 MiB floor MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The part-size math targeted MIN_MULTIPART_SIZE (5 MiB), so a large blob was split into thousands of tiny parts — a 12 GiB blob became ~2300 concurrent 5 MiB PUTs. That churns HTTP connection pools (aws-smithy "connection never set" storms); on a good link the retries win slowly, on a lossy one they exhaust and the slow-store write fails, which write-through then cascades into dropping the successful fast-store write too. Size parts at a configurable multipart_part_size (default 64 MiB), growing only as needed to stay under the 10,000-part ceiling. 12 GiB -> ~192 parts. The same knob is wired into the ONTAP S3 store, whose previous formula divided the object size by the 5 MiB floor — a part count — and clamped that as a byte size, so every blob used 5 MiB parts. Larger parts mean more buffered bytes, so this also makes the buffering bounds real, in both stores. The existing select loop drained the bounded part channel into an unbounded FuturesUnordered, so the documented per-upload bound was never enforced; the recv arm is now gated on in-flight uploads, bounding one upload's buffer at roughly (2 * multipart_max_concurrent_uploads + 1) * part_size. Across uploads, a store-wide admission semaphore (max_concurrent_multipart_uploads, default 4, whole-upload permits so it cannot deadlock the part loop) bounds the aggregate near 5 GiB at the defaults. The single-request threshold intentionally stays at 5 MiB: the single-PUT path replays from a max_retry_buffer_per_request-sized buffer (default 5 MiB), so raising the threshold would break retries for mid-size uploads. Mid-size blobs pay two extra small requests. The existing multipart mock tests pin multipart_part_size to the old 5 MiB so their part boundaries stay small (which also exercises the knob); a new test locks the default to a single part for a blob between 5 MiB and 64 MiB. Co-Authored-By: Claude Fable 5 --- nativelink-config/src/stores.rs | 38 +++++++ nativelink-store/src/ontap_s3_store.rs | 62 ++++++++++- nativelink-store/src/s3_store.rs | 97 +++++++++++++--- .../ontap_s3_existence_cache_store_test.rs | 2 + nativelink-store/tests/ontap_s3_store_test.rs | 6 + nativelink-store/tests/s3_store_test.rs | 104 ++++++++++++++++++ 6 files changed, 292 insertions(+), 17 deletions(-) diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index 81b6179f929..7c073892a2f 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1398,6 +1398,42 @@ pub struct CommonObjectSpec { )] pub multipart_max_concurrent_uploads: Option, + /// Target size of each part in a multipart upload, growing only as + /// needed to stay under the provider's 10,000-part ceiling. Parts sized + /// at the 5 MiB provider floor turn a multi-GiB blob into thousands of + /// tiny concurrent requests, which churns HTTP connection pools; a + /// larger part keeps the request count low. Each queued or in-flight + /// part is buffered in memory, so one multipart upload buffers up to + /// roughly `(2 * multipart_max_concurrent_uploads + 1) * + /// multipart_part_size`; `max_concurrent_multipart_uploads` bounds the + /// aggregate across uploads. Values are clamped to the provider's + /// 5 MiB minimum and 5 GiB maximum part sizes. + /// + /// Currently honored by the S3 and ONTAP S3 providers. + /// + /// Default: 64MB. + #[serde( + default, + deserialize_with = "convert_optional_data_size_with_shellexpand" + )] + pub multipart_part_size: Option, + + /// Store-wide cap on the number of blobs concurrently uploaded via the + /// multipart path. Each multipart upload can buffer roughly + /// `(2 * multipart_max_concurrent_uploads + 1) * multipart_part_size`, + /// so this bounds the store's aggregate upload buffer (about 5 GiB at + /// the defaults). Uploads over the 5 MiB multipart threshold wait for a + /// slot; smaller blobs are unaffected. + /// + /// Currently honored by the S3 and ONTAP S3 providers. + /// + /// Default: 4. + #[serde( + default, + deserialize_with = "convert_optional_numeric_with_shellexpand" + )] + pub max_concurrent_multipart_uploads: Option, + /// Allow unencrypted HTTP connections. Only use this for local testing. /// /// Default: false @@ -1502,6 +1538,8 @@ impl Default for CommonObjectSpec { consider_expired_after_s: 0, max_retry_buffer_per_request: None, multipart_max_concurrent_uploads: None, + multipart_part_size: None, + max_concurrent_multipart_uploads: None, insecure_allow_http: false, disable_http2: false, throttle_retry_delay_s: default_throttle_retry_delay_s(), diff --git a/nativelink-store/src/ontap_s3_store.rs b/nativelink-store/src/ontap_s3_store.rs index 224fc2df52e..39c0fc673f4 100644 --- a/nativelink-store/src/ontap_s3_store.rs +++ b/nativelink-store/src/ontap_s3_store.rs @@ -50,6 +50,7 @@ use rustls::{ClientConfig, RootCertStore}; use rustls_pki_types::CertificateDer; use rustls_pki_types::pem::PemObject; use sha2::{Digest, Sha256}; +use tokio::sync::Semaphore; use tokio::time::sleep; use tracing::{Level, event}; @@ -71,6 +72,16 @@ const DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST: usize = 20 * 1024 * 1024; // 20MB // Default limit for concurrent part uploads per multipart upload const DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS: usize = 10; +// Default target size for each multipart part; see the `multipart_part_size` +// config docs. Sizing parts at the 5 MiB floor maximises the part count and +// drowns the endpoint in tiny concurrent PUTs on multi-GiB blobs. +const DEFAULT_MULTIPART_PART_SIZE: u64 = 64 * 1024 * 1024; // 64MB. + +// Default store-wide admission cap on concurrent multipart uploads; see the +// `max_concurrent_multipart_uploads` config docs. A permit is held for the +// whole upload (never per part), so this cannot deadlock the part loop. +const DEFAULT_MAX_CONCURRENT_MULTIPART_UPLOADS: usize = 4; + #[derive(Debug, MetricsComponent)] pub struct OntapS3Store { s3_client: Arc, @@ -86,6 +97,11 @@ pub struct OntapS3Store { max_retry_buffer_per_request: usize, #[metric(help = "The number of concurrent uploads allowed for multipart uploads")] multipart_max_concurrent_uploads: usize, + #[metric(help = "The target size of each part in a multipart upload")] + multipart_part_size: u64, + // Store-wide admission control for concurrent multipart uploads. + // Bounds aggregate buffered memory. + multipart_upload_slots: Arc, remove_callbacks: Mutex>, } @@ -212,6 +228,19 @@ where .common .multipart_max_concurrent_uploads .unwrap_or(DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS), + multipart_part_size: spec + .common + .multipart_part_size + .map_or(DEFAULT_MULTIPART_PART_SIZE, |size| { + u64::try_from(size).unwrap_or(u64::MAX) + }) + .clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE), + multipart_upload_slots: Arc::new(Semaphore::new( + spec.common + .max_concurrent_multipart_uploads + .unwrap_or(DEFAULT_MAX_CONCURRENT_MULTIPART_UPLOADS) + .max(1), + )), remove_callbacks: Mutex::new(vec![]), })) } @@ -425,6 +454,16 @@ where } // Handle multipart upload for large files + // Admission control: cap concurrent multipart uploads store-wide so a + // burst of large writes cannot multiply the per-upload buffer into an + // OOM. Held (RAII) for the whole upload; released on drop. + let _multipart_slot = self + .multipart_upload_slots + .clone() + .acquire_owned() + .await + .map_err(|e| make_err!(Code::Internal, "multipart upload semaphore closed: {e}"))?; + let upload_id = &self .retrier .retry(unfold((), move |()| async move { @@ -458,10 +497,23 @@ where })) .await?; - let bytes_per_upload_part = - (max_size / (MIN_MULTIPART_SIZE - 1)).clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE); + // Size each part at the configured target, growing it only as needed + // to keep the part count within the 10,000-part ceiling for very + // large objects. (The previous formula divided the object size by the + // 5 MiB floor — a part COUNT — and clamped that as a byte size, so + // every blob used 5 MiB parts.) + let bytes_per_upload_part = cmp::max( + self.multipart_part_size, + max_size.div_ceil(MAX_UPLOAD_PARTS as u64), + ) + .clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE); let upload_parts = move || async move { + // Together with the gated `recv` in the select loop below, this + // bounds buffered part data for one upload at roughly + // `(2 * multipart_max_concurrent_uploads + 1) * bytes_per_upload_part`: + // once the channel is full and N uploads are in flight, the + // `tx.send` above blocks and the reader stops consuming. let (tx, mut rx) = tokio::sync::mpsc::channel(self.multipart_max_concurrent_uploads); let read_stream_fut = ( @@ -548,7 +600,11 @@ where total_uploaded = result?; }, Some(upload_result) = upload_futures.next() => completed_parts.push(upload_result?), - Some(fut) = rx.recv() => upload_futures.push(fut), + // Gated so the bounded channel exerts real backpressure: + // draining it eagerly into the unbounded `FuturesUnordered` + // would free channel slots and let the reader buffer far + // ahead of the uploads. + Some(fut) = rx.recv(), if upload_futures.len() < self.multipart_max_concurrent_uploads => upload_futures.push(fut), } } diff --git a/nativelink-store/src/s3_store.rs b/nativelink-store/src/s3_store.rs index 34e532a542a..3e0b7a982d6 100644 --- a/nativelink-store/src/s3_store.rs +++ b/nativelink-store/src/s3_store.rs @@ -51,7 +51,7 @@ use nativelink_util::store_trait::{ RemoveCallback, StoreDriver, StoreKey, StoreOptimizations, UploadSizeInfo, }; use parking_lot::Mutex; -use tokio::sync::mpsc; +use tokio::sync::{Semaphore, mpsc}; use tokio::time::sleep; use tracing::{error, info}; @@ -75,6 +75,17 @@ const MAX_MULTIPART_SIZE: u64 = 5 * 1024 * 1024 * 1024; // 5GB. // Note: Type 'u64' chosen to simplify calculations const MAX_UPLOAD_PARTS: u64 = 10_000; +// Default target size for each multipart part. Deliberately far above the +// 5 MiB floor: sizing parts at the floor turns a multi-GiB blob into +// thousands of tiny PUTs, which churns and poisons the HTTP connection pool +// (aws-smithy "connection never set" storms) and makes large uploads slow and +// flaky — on a lossy link the retries exhaust and the write fails outright. +// 64 MiB keeps a 12 GiB blob at ~192 parts; combined with the concurrency +// limit it bounds in-flight memory to +// `multipart_max_concurrent_uploads * multipart_part_size`. +// Note: If you change this, adjust the docs in the config. +const DEFAULT_MULTIPART_PART_SIZE: u64 = 64 * 1024 * 1024; // 64MB. + // Default max buffer size for retrying upload requests. // Note: If you change this, adjust the docs in the config. const DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST: usize = 5 * 1024 * 1024; // 5MB. @@ -83,6 +94,16 @@ const DEFAULT_MAX_RETRY_BUFFER_PER_REQUEST: usize = 5 * 1024 * 1024; // 5MB. // Note: If you change this, adjust the docs in the config. const DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS: usize = 10; +// Default store-wide admission cap on concurrent multipart uploads. The +// per-upload concurrency and part size bound the buffer of a single upload +// at roughly `(2 * multipart_max_concurrent_uploads + 1) * multipart_part_size` +// (~1.3 GiB at the defaults); without a store-wide limit, N simultaneous +// large writes multiply that footprint unbounded. Four concurrent uploads +// bounds the default worst case near 5 GiB. A permit is held for the whole +// upload (never per part), so this cannot deadlock the inner part loop. +// Note: If you change this, adjust the docs in the config. +const DEFAULT_MAX_CONCURRENT_MULTIPART_UPLOADS: usize = 4; + #[derive(Debug, MetricsComponent)] pub struct S3Store { s3_client: Arc, @@ -98,6 +119,12 @@ pub struct S3Store { max_retry_buffer_per_request: usize, #[metric(help = "The number of concurrent uploads allowed for multipart uploads")] multipart_max_concurrent_uploads: usize, + #[metric(help = "The target size of each part in a multipart upload")] + multipart_part_size: u64, + // Store-wide admission control for concurrent multipart uploads (see + // DEFAULT_MAX_CONCURRENT_MULTIPART_UPLOADS). Bounds aggregate buffered + // memory. + multipart_upload_slots: Arc, remove_callbacks: Mutex>, } @@ -212,6 +239,19 @@ where .common .multipart_max_concurrent_uploads .unwrap_or(DEFAULT_MULTIPART_MAX_CONCURRENT_UPLOADS), + multipart_part_size: spec + .common + .multipart_part_size + .map_or(DEFAULT_MULTIPART_PART_SIZE, |size| { + u64::try_from(size).unwrap_or(u64::MAX) + }) + .clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE), + multipart_upload_slots: Arc::new(Semaphore::new( + spec.common + .max_concurrent_multipart_uploads + .unwrap_or(DEFAULT_MAX_CONCURRENT_MULTIPART_UPLOADS) + .max(1), + )), remove_callbacks: Mutex::new(Vec::new()), })) } @@ -422,6 +462,20 @@ where .await; } + // Admission control: cap concurrent multipart uploads store-wide so a + // burst of large writes cannot multiply the per-upload buffer + // (`multipart_max_concurrent_uploads * multipart_part_size`) into an + // OOM. The permit is held (RAII) for the whole upload and acquired + // after the single-request early-return above, so small blobs never + // wait on a slot; it is whole-upload admission, never per-part, so it + // cannot deadlock the part loop below. + let _multipart_slot = self + .multipart_upload_slots + .clone() + .acquire_owned() + .await + .map_err(|e| make_err!(Code::Internal, "multipart upload semaphore closed: {e}"))?; + let upload_id = &self .retrier .retry(unfold((), move |()| async move { @@ -458,26 +512,37 @@ where // S3 requires us to upload in parts if the size is greater than 5GB. The part size must be at least // 5MB (except last part) and can have up to 10,000 parts. - // Calculate of number of chunks if we upload in 5MB chucks (min chunk size), clamping to - // 10,000 parts and correcting for lossy integer division. This provides the - let chunk_count = (max_size / MIN_MULTIPART_SIZE).clamp(0, MAX_UPLOAD_PARTS - 1) + 1; + // Size each part at the configured target, growing it only as needed + // to keep the part count within the 10,000-part ceiling for very + // large objects. Sizing at the 5 MiB floor (the old behaviour) + // maximised the part count and drowned providers in tiny concurrent + // PUTs. + let bytes_per_upload_part = cmp::max( + self.multipart_part_size, + max_size.div_ceil(MAX_UPLOAD_PARTS), + ) + .clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE); - // Using clamped first approximation of number of chunks, calculate byte count of each - // chunk, excluding last chunk, clamping to min/max upload size 5MB, 5GB. - let bytes_per_upload_part = - (max_size / chunk_count).clamp(MIN_MULTIPART_SIZE, MAX_MULTIPART_SIZE); + // Number of parts this yields (last part may be short). + let chunk_count = max_size.div_ceil(bytes_per_upload_part); - // Sanity check before continuing. - if !(MIN_MULTIPART_SIZE..MAX_MULTIPART_SIZE).contains(&bytes_per_upload_part) { + // Sanity check before continuing. Inclusive of MAX_MULTIPART_SIZE: + // the construction-time clamp allows a configured part size of + // exactly the 5 GiB provider maximum. + if !(MIN_MULTIPART_SIZE..=MAX_MULTIPART_SIZE).contains(&bytes_per_upload_part) { return Err(make_err!( Code::FailedPrecondition, - "Failed to calculate file chuck size (min, max, calc): {MIN_MULTIPART_SIZE}, {MAX_MULTIPART_SIZE}, {bytes_per_upload_part}", + "Failed to calculate file chunk size (min, max, calc): {MIN_MULTIPART_SIZE}, {MAX_MULTIPART_SIZE}, {bytes_per_upload_part}", )); } let upload_parts = move || async move { - // This will ensure we only have `multipart_max_concurrent_uploads` * `bytes_per_upload_part` - // bytes in memory at any given time waiting to be uploaded. + // Together with the gated `recv` in the select loop below, this + // bounds buffered part data for one upload at roughly + // `(2 * multipart_max_concurrent_uploads + 1) * bytes_per_upload_part`: + // up to N parts uploading, N queued in the channel, and one being + // read. Once the channel is full and N uploads are in flight, the + // `tx.send` above blocks and the reader stops consuming. let (tx, mut rx) = mpsc::channel(self.multipart_max_concurrent_uploads); let read_stream_fut = async move { @@ -559,7 +624,11 @@ where total_uploaded = result?; }, // Return error or wait for other futures. Some(upload_result) = upload_futures.next() => completed_parts.push(upload_result?), - Some(fut) = rx.recv() => upload_futures.push(fut), + // Gated so the bounded channel exerts real backpressure: + // draining it eagerly into the unbounded `FuturesUnordered` + // would free channel slots and let the reader buffer far + // ahead of the uploads. + Some(fut) = rx.recv(), if upload_futures.len() < self.multipart_max_concurrent_uploads => upload_futures.push(fut), } } diff --git a/nativelink-store/tests/ontap_s3_existence_cache_store_test.rs b/nativelink-store/tests/ontap_s3_existence_cache_store_test.rs index e195c7df687..299453eb209 100644 --- a/nativelink-store/tests/ontap_s3_existence_cache_store_test.rs +++ b/nativelink-store/tests/ontap_s3_existence_cache_store_test.rs @@ -69,6 +69,8 @@ async fn create_test_store(mock_client: StaticReplayClient) -> Result Result<(), Error> { bucket: BUCKET_NAME.to_string(), vserver_name: VSERVER_NAME.to_string(), endpoint: "https://example.com".to_string(), + common: CommonObjectSpec { + // Keep the historical 5 MiB parts so the mocked part + // boundaries below stay small. + multipart_part_size: Some(5 * 1024 * 1024), + ..Default::default() + }, ..Default::default() }), s3_client, diff --git a/nativelink-store/tests/s3_store_test.rs b/nativelink-store/tests/s3_store_test.rs index a9f4b7ee185..89592e35f16 100644 --- a/nativelink-store/tests/s3_store_test.rs +++ b/nativelink-store/tests/s3_store_test.rs @@ -426,6 +426,98 @@ async fn get_part_simple_retries() -> Result<(), Error> { Ok(()) } +// A blob over the 5 MiB multipart threshold but under the default 64 MiB +// part target must upload as ONE part. Sizing parts at the 5 MiB floor (the +// old behaviour) maximised the part count, drowning providers in tiny +// concurrent PUTs on multi-GiB blobs. +#[nativelink_test] +async fn multipart_default_part_size_uploads_one_part() -> Result<(), Error> { + const MIN_MULTIPART_SIZE: usize = 5 * 1024 * 1024; // 5mb. + const AC_ENTRY_SIZE: usize = MIN_MULTIPART_SIZE * 2 + 50; + + let mut send_data: Vec = Vec::with_capacity(AC_ENTRY_SIZE); + for i in 0..send_data.capacity() { + send_data.push(u8::try_from((i * 3) % 256).expect("modulo 256 always fits in u8")); + } + let digest = DigestInfo::try_new(VALID_HASH1, send_data.len())?; + + 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}?uploads", + )) + .method("POST") + .body(SdkBody::empty()) + .unwrap(), + http::Response::builder() + .status(StatusCode::OK) + .body(SdkBody::from(concat!( + "", + "Dummy-uploadid", + "" + ))) + .unwrap(), + ), + ReplayEvent::new( + http::Request::builder() + .uri(format!( + "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?x-id=UploadPart&partNumber=1&uploadId=Dummy-uploadid", + )) + .method("PUT") + .header("content-type", "application/octet-stream") + .header("content-length", AC_ENTRY_SIZE.to_string()) + .body(SdkBody::from(&send_data[..])) + .unwrap(), + http::Response::builder() + .status(StatusCode::OK) + .body(SdkBody::empty()) + .unwrap(), + ), + ReplayEvent::new( + http::Request::builder() + .uri(format!( + "https://{BUCKET_NAME}.s3.{REGION}.amazonaws.com/{VALID_HASH1}-{AC_ENTRY_SIZE}?uploadId=Dummy-uploadid", + )) + .method("POST") + .header("content-length", "138") + .body(SdkBody::from(concat!( + "", + "1", + "", + ))) + .unwrap(), + http::Response::builder() + .status(StatusCode::OK) + .body(SdkBody::from(concat!( + "", + "", + ))) + .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, + )?; + store + .update_oneshot(digest, send_data.clone().into()) + .await?; + mock_client.assert_requests_match(&[]); + Ok(()) +} + #[nativelink_test] async fn multipart_update_large_cas() -> Result<(), Error> { // Same as in s3_store. @@ -538,6 +630,12 @@ async fn multipart_update_large_cas() -> Result<(), Error> { let store = S3Store::new_with_client_and_jitter( &ExperimentalAwsSpec { bucket: BUCKET_NAME.to_string(), + common: CommonObjectSpec { + // Keep the historical 5 MiB parts so the mocked part + // boundaries below stay small. + multipart_part_size: Some(5 * 1024 * 1024), + ..Default::default() + }, ..Default::default() }, s3_client, @@ -890,6 +988,12 @@ async fn multipart_chunk_size_clamp_min() -> Result<(), Error> { let store = S3Store::new_with_client_and_jitter( &ExperimentalAwsSpec { bucket: BUCKET_NAME.to_string(), + common: CommonObjectSpec { + // Keep the historical 5 MiB parts so the mocked part + // boundaries below stay small. + multipart_part_size: Some(5 * 1024 * 1024), + ..Default::default() + }, ..Default::default() }, s3_client, From 11f83898111d8d6320c3368aa3ae60babac63014 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 11:28:56 -0400 Subject: [PATCH 22/46] fix(store): mongo update_data version-CAS must return retryable Ok(None), not Err ExperimentalMongoStore::update_data used find_one_and_update(...).upsert(true). On a version-CAS conflict (a concurrent writer already bumped the version) the {_id, version} filter matches nothing, so upsert tries to INSERT a doc with the existing _id and MongoDB returns a DuplicateKeyError (E11000), surfaced as Err(Internal). The version-CAS contract (and the Redis backend, which returns Ok(None)) require a conflict to be retryable; instead concurrent updates on a Mongo-backed record spuriously failed. Fix: .upsert(current_version == 0) so only a genuine first-create upserts; a non-matching filter on an existing record now returns Ok(None) (retryable). Adds a deterministic fails-without regression and corrects an existing (ignored) scheduler test that had pinned the buggy Err behavior. Co-Authored-By: Claude Opus 4.8 --- 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 8a7bbb0096f..185b9f855b4 100644 --- a/nativelink-store/src/mongo_store.rs +++ b/nativelink-store/src/mongo_store.rs @@ -975,7 +975,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 d6133631254..3933ec5b554 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"); } @@ -1056,3 +1059,50 @@ async fn incomplete_upload_preserves_mongo_value() -> Result<(), Error> { } 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 af50b05732380c9f165487f6a6780827a6c6fd3e Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 15 Oct 2026 12:29:24 -0400 Subject: [PATCH 23/46] completeness: keep CAS error codes, and test every miss path Three follow-ups to the change that stopped backend errors being served as misses. - A failure in the CAS existence check (the output files, stdout, stderr) was still rewrapped as `Internal`, which retry policies treat as permanent. It now keeps its own code, so a transient CAS fault is retryable like an AC one. `incomplete_entries_counter` now counts only genuine misses, not propagated errors. - Tests cover the two misses that must stay `NotFound` and were untested: an absent action-result record and an evicted Tree blob. Both fail if every lookup error is propagated. The existing get-path miss now checks for `NotFound` exactly, and a new test checks a CAS failure keeps its code (it fails against the old `Internal` wrapping). - The scheduler regression test is bounded, so a regression fails in 5 s instead of hanging, and its fake store uses its configured code. Comments no longer claim undecodable records propagate: `ac_utils` maps decode failures to `NotFound`, so they stay misses. Co-Authored-By: Claude Opus 5.5 --- .../src/cache_lookup_scheduler.rs | 8 +- .../tests/cache_lookup_scheduler_test.rs | 37 ++--- .../src/completeness_checking_store.rs | 34 +++-- .../tests/completeness_checking_store_test.rs | 133 +++++++++++++++--- 4 files changed, 161 insertions(+), 51 deletions(-) diff --git a/nativelink-scheduler/src/cache_lookup_scheduler.rs b/nativelink-scheduler/src/cache_lookup_scheduler.rs index 3d6ced82639..0f9abd399ad 100644 --- a/nativelink-scheduler/src/cache_lookup_scheduler.rs +++ b/nativelink-scheduler/src/cache_lookup_scheduler.rs @@ -303,10 +303,10 @@ impl CacheLookupScheduler { } Err(err) => { // A NotFound just means we need to execute the action. Any - // other lookup failure (an unavailable AC backend, a decode - // fault) degrades to execution too: a cache fault must never - // fail the request, only cost a redundant execution. The - // error itself still reaches operators via the log. + // other lookup failure (an unavailable AC backend, a timeout) + // degrades to execution too: a cache fault must never fail + // the request, only cost a redundant execution. The error + // itself still reaches operators via the log. if err.code != Code::NotFound { warn!(?err, "Cache lookup failed; falling through to execution"); } diff --git a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs index cbff20b9867..71272994e9a 100644 --- a/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs +++ b/nativelink-scheduler/tests/cache_lookup_scheduler_test.rs @@ -48,11 +48,10 @@ use tokio::{self}; use tokio_stream::StreamExt; use utils::scheduler_utils::{TokioWatchActionStateResult, make_base_action_info}; -/// An AC store whose reads always fail with a non-`NotFound` backend error. +/// An AC store whose operations always fail with a non-`NotFound` backend error. #[derive(Debug, MetricsComponent)] struct AlwaysErrStore { - #[metric(help = "Error code every operation fails with")] - code: u32, + code: Code, } #[async_trait] @@ -66,7 +65,7 @@ impl StoreDriver for AlwaysErrStore { _keys: &[StoreKey<'_>], _results: &mut [Option], ) -> Result<(), Error> { - Err(make_err!(Code::Unavailable, "injected AC backend failure")) + Err(make_err!(self.code, "injected AC backend failure")) } async fn update( @@ -75,7 +74,7 @@ impl StoreDriver for AlwaysErrStore { _reader: DropCloserReadHalf, _size_info: UploadSizeInfo, ) -> Result { - Err(make_err!(Code::Unavailable, "injected AC backend failure")) + Err(make_err!(self.code, "injected AC backend failure")) } async fn get_part( @@ -85,7 +84,7 @@ impl StoreDriver for AlwaysErrStore { _offset: u64, _length: Option, ) -> Result<(), Error> { - Err(make_err!(Code::Unavailable, "injected AC backend failure")) + Err(make_err!(self.code, "injected AC backend failure")) } fn inner_store(&self, _key: Option) -> &dyn StoreDriver { @@ -199,14 +198,14 @@ async fn add_action_handles_skip_cache() -> Result<(), Error> { Ok(()) } -/// A faulting AC backend (unavailable Redis, decode error) must degrade the +/// A faulting AC backend (an unavailable Redis, a timeout) must degrade the /// cache lookup to execution, not fail the Execute request: a cache fault /// costs a redundant execution, never a failed build. #[nativelink_test] async fn ac_fault_falls_through_to_execution() -> Result<(), Error> { let mock_scheduler = Arc::new(MockActionScheduler::new()); let ac_store = Store::new(Arc::new(AlwaysErrStore { - code: Code::Unavailable as u32, + code: Code::Unavailable, })); let cache_scheduler = CacheLookupScheduler::new(ac_store, mock_scheduler.clone())?; let action_info = make_base_action_info(UNIX_EPOCH, DigestInfo::new([7; 32], 123)); @@ -218,14 +217,20 @@ async fn ac_fault_falls_through_to_execution() -> Result<(), Error> { last_transition_timestamp: SystemTime::now(), })); let client_operation_id = OperationId::default(); - let (add_result, _) = join!( - cache_scheduler.add_action(client_operation_id.clone(), action_info.clone()), - mock_scheduler.expect_add_action(Ok(Box::new(TokioWatchActionStateResult::new( - client_operation_id, - action_info, - forward_watch_channel_rx - )))) - ); + // If the fault failed the request instead, the mock would never see the + // action and `expect_add_action` would wait forever; bound it. + let (add_result, _) = tokio::time::timeout(std::time::Duration::from_secs(5), async { + join!( + cache_scheduler.add_action(client_operation_id.clone(), action_info.clone()), + mock_scheduler.expect_add_action(Ok(Box::new(TokioWatchActionStateResult::new( + client_operation_id, + action_info, + forward_watch_channel_rx + )))) + ) + }) + .await + .expect("the AC fault was not forwarded to execution"); add_result.expect("AC backend fault must fall through to execution, not fail the client"); Ok(()) } diff --git a/nativelink-store/src/completeness_checking_store.rs b/nativelink-store/src/completeness_checking_store.rs index 9cc282d5d09..5a3e9dfd829 100644 --- a/nativelink-store/src/completeness_checking_store.rs +++ b/nativelink-store/src/completeness_checking_store.rs @@ -297,28 +297,34 @@ impl CompletenessCheckingStore { // Poll both futures at the same time. select! { r = check_existence_fut => { - return Err(make_err!( - Code::Internal, - "CompletenessCheckingStore's check_existence_fut ended unexpectedly {r:?}" - )); + // Before `done` is set, this future only finishes when a CAS + // existence check fails. Keep that error's code, so a + // transient CAS fault stays retryable instead of reaching + // the client as Internal. + return Err(match r { + Err(err) => err, + Ok(()) => make_err!( + Code::Internal, + "CompletenessCheckingStore's check_existence_fut ended unexpectedly" + ), + }); } maybe_result = futures.next() => { match maybe_result { Some(Ok(())) => self.complete_entries_counter.inc(), Some(Err((err, i))) => { - self.incomplete_entries_counter.inc(); - // Only a genuine `NotFound` may be laundered into a "missing" - // slot (which callers surface as `Code::NotFound`). That covers - // a truly absent action-result record and a clean completeness - // miss — the intended eviction-coherence behavior (an evicted - // referent must still 404). ANY other code is a real backend - // read/decode failure; collapsing it into `None` would make an - // I/O fault indistinguishable from absence and get served as a - // 404. Propagate it so the caller can tell "genuinely absent" - // from "could not determine". + // Only a `NotFound` becomes a miss: an absent action + // result, or an absent Tree blob. Records that fail to + // decode also arrive as `NotFound` (`ac_utils` maps + // them), so they stay misses too. Any other code is a + // backend failure, such as an unavailable store or a + // timeout, and turning it into a miss would serve an + // outage as a 404 and throw away a valid cached + // result. Propagate it instead. if err.code != Code::NotFound { return Err(err); } + self.incomplete_entries_counter.inc(); state_mux.lock().results[i] = None; } None => { diff --git a/nativelink-store/tests/completeness_checking_store_test.rs b/nativelink-store/tests/completeness_checking_store_test.rs index ed3dcd2c38e..01ba27d89b1 100644 --- a/nativelink-store/tests/completeness_checking_store_test.rs +++ b/nativelink-store/tests/completeness_checking_store_test.rs @@ -43,6 +43,21 @@ const STDOUT: DigestInfo = DigestInfo::new([5u8; 32], 0); const STDERR: DigestInfo = DigestInfo::new([6u8; 32], 0); async fn setup() -> Result<(Arc, Arc, DigestInfo), Error> { + let (ac_store, cas_store, action_result_digest, _tree_digest) = setup_with_tree().await?; + Ok((ac_store, cas_store, action_result_digest)) +} + +/// Like `setup`, but also returns the digest of the `Tree` the action result's +/// output directory points at. +async fn setup_with_tree() -> Result< + ( + Arc, + Arc, + DigestInfo, + DigestInfo, + ), + Error, +> { let backend_store = Store::new(MemoryStore::new(&MemorySpec::default())); let cas_store = MemoryStore::new(&MemorySpec::default()); let ac_store = @@ -114,7 +129,7 @@ async fn setup() -> Result<(Arc, Arc, Di ) .await?; - Ok((ac_store, cas_store, action_result_digest)) + Ok((ac_store, cas_store, action_result_digest, tree_digest)) } #[nativelink_test] @@ -241,12 +256,14 @@ async fn verify_completeness_get() -> Result<(), Error> { cas_store.remove_entry(OUTPUT_FILE.into()).await; - assert!( - ac_store - .get_part_unchunked(action_result_digest, 0, None) - .await - .is_err(), - ".get() should fail with item missing in CAS", + let err = ac_store + .get_part_unchunked(action_result_digest, 0, None) + .await + .expect_err(".get() should fail with item missing in CAS"); + assert_eq!( + err.code, + Code::NotFound, + "a missing referent must stay a NotFound miss: {err:?}" ); } @@ -320,14 +337,14 @@ impl StoreDriver for AlwaysErrStore { default_health_status_indicator!(AlwaysErrStore); -/// Regression: a real backend read/decode failure on the action-result record -/// must propagate with its underlying, non-`NotFound` code — it must NOT be -/// collapsed into `Code::NotFound`. A `NotFound` miss tells the client the -/// result does not exist and to re-execute; a backend fault (unavailable -/// Redis, decode failure) is transient and retryable, and serving it as a -/// clean miss silently discards a valid cached result. The genuine -/// eviction-miss path (referent absent) is covered by -/// `verify_has_function_call_checks_cas` and stays `NotFound`. +/// Regression: a backend failure reading the action-result record must +/// propagate with its own code, not be turned into `Code::NotFound`. A +/// `NotFound` miss tells the client the result does not exist and to +/// re-execute; a backend fault (an unavailable Redis, a timeout) is transient +/// and retryable, and serving it as a miss throws away a valid cached result. +/// The misses that must stay `NotFound` are covered by +/// `verify_has_function_call_checks_cas`, `absent_record_is_a_miss` and +/// `absent_tree_is_a_miss`. #[nativelink_test] async fn backend_error_on_record_propagates_as_non_not_found() -> Result<(), Error> { let make_store = || { @@ -339,7 +356,7 @@ async fn backend_error_on_record_propagates_as_non_not_found() -> Result<(), Err }; let some_digest = DigestInfo::new([9u8; 32], 123); - // get_part path (nix narinfo GET / Bazel AC get). + // get_part path (an AC GetActionResult). let get_err = make_store() .get_part_unchunked(some_digest, 0, None) .await @@ -355,7 +372,7 @@ async fn backend_error_on_record_propagates_as_non_not_found() -> Result<(), Err "backend fault must propagate with its real code" ); - // has path (nix narinfo HEAD / Bazel AC has). + // has path. let has_err = make_store() .has(some_digest) .await @@ -373,3 +390,85 @@ async fn backend_error_on_record_propagates_as_non_not_found() -> Result<(), Err Ok(()) } + +/// An action result that is not in the AC at all is a plain miss. +#[nativelink_test] +async fn absent_record_is_a_miss() -> Result<(), Error> { + let (ac_store, _cas_store, _action_result_digest) = setup().await?; + let absent = DigestInfo::new([8u8; 32], 64); + + assert_eq!(ac_store.has(absent).await?, None); + let err = ac_store + .get_part_unchunked(absent, 0, None) + .await + .expect_err("an absent action result must not read back"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + Ok(()) +} + +/// An action result whose output directory's `Tree` blob has been evicted from +/// the CAS is incomplete, so it is a miss, not an error. +#[nativelink_test] +async fn absent_tree_is_a_miss() -> Result<(), Error> { + let (ac_store, cas_store, action_result_digest, tree_digest) = setup_with_tree().await?; + cas_store.remove_entry(tree_digest.into()).await; + + assert_eq!(ac_store.has(action_result_digest).await?, None); + let err = ac_store + .get_part_unchunked(action_result_digest, 0, None) + .await + .expect_err("an action result with an evicted Tree must not read back"); + assert_eq!(err.code, Code::NotFound, "unexpected error: {err:?}"); + Ok(()) +} + +/// A failure checking the referenced outputs in the CAS keeps its own code. It +/// used to reach the client as `Internal`, which retry policies treat as +/// permanent. +#[nativelink_test] +async fn cas_backend_error_propagates_with_its_code() -> Result<(), Error> { + let backend_store = Store::new(MemoryStore::new(&MemorySpec::default())); + let cas_store = Store::new(Arc::new(AlwaysErrStore { + code: Code::Unavailable, + })); + let ac_store = CompletenessCheckingStore::new(backend_store, cas_store); + + // No output directories, so the only CAS call is the existence check of + // the output file, stdout and stderr. + let action_result = ProtoActionResult { + output_files: vec![OutputFile { + digest: Some(OUTPUT_FILE.into()), + ..Default::default() + }], + stdout_digest: Some(STDOUT.into()), + stderr_digest: Some(STDERR.into()), + ..Default::default() + }; + let action_result_digest = serialize_and_upload_message( + &action_result, + ac_store.as_pin(), + &mut DigestHasherFunc::Blake3.hasher(), + ) + .await?; + + let has_err = ac_store + .has(action_result_digest) + .await + .expect_err("a CAS failure must not be swallowed"); + assert_eq!( + has_err.code, + Code::Unavailable, + "unexpected error: {has_err:?}" + ); + + let get_err = ac_store + .get_part_unchunked(action_result_digest, 0, None) + .await + .expect_err("a CAS failure must not be swallowed"); + assert_eq!( + get_err.code, + Code::Unavailable, + "unexpected error: {get_err:?}" + ); + Ok(()) +} From 81beca251d32887b5949c9f1f5a437f9433a21a6 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 15 Oct 2026 12:43:47 -0400 Subject: [PATCH 24/46] redis: retry the eviction subscription properly, and keep reconnect off the lock Follow-ups from review of the subscription-lifecycle fix. reconnect: subscribing the new connection to patterns recorded mid-reconnect happened while holding the connection write lock, so every get_connection (every Redis operation) waited on those round trips. Subscribing now happens without the lock; the lock is taken only to check nothing is still missing and to publish, repeating if a pattern arrived in between. A pattern recorded after the final check still waits for the lock in get_connection and then subscribes on the new connection, so the original race stays closed. Eviction subscription retry: - The task holds a weak reference, so a dropped store ends it instead of being kept alive and retried forever. - It logs the first failure as an error, then only at powers of two, so a permanent failure (an ACL forbidding PSUBSCRIBE) doesn't log every 5 s. - On a connection-type failure it re-resolves the connection, as the store's other retry paths do, instead of retrying a dead handle. - Reading notify-keyspace-events stays advisory, now logged at warn with its own message so it can't be mistaken for the subscription failing. Tests: - The existing init test only covered the advisory CONFIG GET and passed without any retry. A new test fails the first PSUBSCRIBE and checks the subscription is still established; it fails if the retry loop gives up. - A new test checks a dropped store ends the retry task; it fails if the task holds a strong reference. - A new unit test checks get_connection is not blocked while reconnect catches up; it fails against the previous locking. - The reconnect race test uses #[nativelink_test] instead of a hand-built runtime, so the clippy and lint exemptions are gone, and the test names now describe the fixed behaviour. - The fake Redis gains make_fake_redis_with_first_responses, which answers a command differently the first time only. Co-Authored-By: Claude Opus 5.5 --- nativelink-redis-tester/src/fake_redis.rs | 50 ++ nativelink-redis-tester/src/lib.rs | 3 +- nativelink-store/BUILD.bazel | 1 + nativelink-store/src/redis_store.rs | 566 ++++++++++++--------- nativelink-store/tests/redis_store_test.rs | 187 +++++-- 5 files changed, 528 insertions(+), 279 deletions(-) diff --git a/nativelink-redis-tester/src/fake_redis.rs b/nativelink-redis-tester/src/fake_redis.rs index 95a49685b87..4f38592314f 100644 --- a/nativelink-redis-tester/src/fake_redis.rs +++ b/nativelink-redis-tester/src/fake_redis.rs @@ -15,6 +15,8 @@ use core::fmt::Write; use core::hash::BuildHasher; use std::collections::HashMap; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; use nativelink_util::background_spawn; use redis::Value; @@ -292,3 +294,51 @@ pub async fn make_fake_redis_with_responses u16 { make_fake_redis_with_multiple_responses(vec![responses]).await } + +/// Like [`make_fake_redis_with_responses`], except that the first request +/// matching a key of `first_responses` gets that key's response instead, once. +/// Every later request gets its normal response. Use this to test that a +/// caller retries after a transient error. +pub async fn make_fake_redis_with_first_responses< + B: BuildHasher + Clone + Send + 'static + Sync, +>( + responses: HashMap, + first_responses: HashMap, +) -> u16 { + let first_responses: Vec<(String, String, Arc)> = first_responses + .into_iter() + .map(|(key, value)| (key, value, Arc::new(AtomicBool::new(false)))) + .collect(); + let handler = move |buf: &[u8]| -> String { + let str_buf = String::from_utf8_lossy(buf).into_owned(); + for (key, value, used) in &first_responses { + if str_buf.starts_with(key) && !used.swap(true, Ordering::SeqCst) { + info!("Responding once to {}", str_buf.replace("\r\n", "\\r\\n")); + return value.clone(); + } + } + for (key, value) in &responses { + if str_buf.starts_with(key) { + info!("Responding to {}", str_buf.replace("\r\n", "\\r\\n")); + return value.clone(); + } + } + warn!( + "Unknown command: {}", + str_buf.chars().take(1000).collect::() + ); + String::new() + }; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let (listener_ready_tx, listener_ready_rx) = oneshot::channel::<()>(); + background_spawn!("listener", async move { + fake_redis_internal(listener, listener_ready_tx, vec![handler]).await; + }); + listener_ready_rx + .await + .expect("Expected successful listener boot"); + info!(port, "Fake redis booted"); + port +} diff --git a/nativelink-redis-tester/src/lib.rs b/nativelink-redis-tester/src/lib.rs index 896d384df96..65ec4874674 100644 --- a/nativelink-redis-tester/src/lib.rs +++ b/nativelink-redis-tester/src/lib.rs @@ -20,7 +20,8 @@ mod read_only_redis; pub use dynamic_fake_redis::{FakeRedisBackend, SubscriptionManagerNotify}; pub use fake_redis::{ 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, + fake_redis_sentinel_stream, fake_redis_stream, make_fake_redis_with_first_responses, + make_fake_redis_with_responses, }; pub use pubsub::MockPubSub; pub use read_only_redis::ReadOnlyRedis; diff --git a/nativelink-store/BUILD.bazel b/nativelink-store/BUILD.bazel index a7715695201..618a926a3ff 100644 --- a/nativelink-store/BUILD.bazel +++ b/nativelink-store/BUILD.bazel @@ -266,6 +266,7 @@ rust_test( "@crates//:redis", "@crates//:serde_json", "@crates//:sha2", + "@crates//:tracing-test", ], ) diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index 00dc08a29a4..eed0526e4a4 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -90,6 +90,14 @@ const DEFAULT_RETRY_DELAY: f32 = 0.1; /// paths so a transient error doesn't fail an otherwise healthy request. const MAX_REDIS_RETRY_ATTEMPTS: u32 = 5; +/// First delay between attempts to subscribe to eviction events after a +/// failure. It doubles after each failure, up to +/// [`EVICTION_SUBSCRIBE_MAX_DELAY`]. +const EVICTION_SUBSCRIBE_INITIAL_DELAY: Duration = Duration::from_millis(50); + +/// Longest delay between attempts to subscribe to eviction events. +const EVICTION_SUBSCRIBE_MAX_DELAY: Duration = Duration::from_secs(5); + /// Whether a Redis error is worth re-resolving the master and retrying: a /// dropped/refused connection, an IO error, or a write that hit a freshly /// demoted replica. Anything else (e.g. a real protocol/logic error) is @@ -270,6 +278,16 @@ where Ok(()) } + /// The recorded subscriptions that are not in `subscribed`. + fn subscriptions_not_in(&self, subscribed: &HashSet) -> Vec { + self.subscriptions + .lock() + .iter() + .filter(|pattern| !subscribed.contains(*pattern)) + .cloned() + .collect() + } + async fn new(connect_func: Box>) -> Result { let connection_manager = connect_func().await?; let update_if_version_matches_script = Script::new(LUA_VERSION_SET_SCRIPT); @@ -320,36 +338,27 @@ impl RedisManager for StandardRedisManager let mut connection_manager = (self.connect_func)().await?; let new_uuid = Uuid::new_v4(); self.configure(&mut connection_manager).await?; - let subscriptions = { - let guard = self.subscriptions.lock(); - guard.iter().map(Clone::clone).collect::>() - }; - for subscription in &subscriptions { - connection_manager.psubscribe(subscription).await?; - } - 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 { + // Subscribe the new connection to every recorded pattern before it is + // published, so a pattern is never recorded as active but missing from + // the live connection. `psubscribe` records its pattern before it calls + // `get_connection`, so a pattern recorded while this runs is either + // picked up by a later pass of this loop, or is recorded after the + // final check, in which case its `get_connection` waits for the write + // lock below and then subscribes on the new connection itself. The + // subscribing happens without the write lock, so `get_connection` + // readers keep using the old handle meanwhile; the lock is held only to + // check that nothing is missing and to publish. + let mut subscribed = HashSet::new(); + loop { + for subscription in self.subscriptions_not_in(&subscribed) { connection_manager.psubscribe(&subscription).await?; - replayed.insert(subscription); + subscribed.insert(subscription); + } + let mut guard = self.connection_manager.write().await; + if self.subscriptions_not_in(&subscribed).is_empty() { + *guard = (connection_manager.clone(), new_uuid); + break; } - *guard = (connection_manager.clone(), new_uuid); } info!(old = %uuid, new = %new_uuid, "StandardRedisManager re-resolved the Redis master"); Ok((connection_manager, new_uuid)) @@ -609,6 +618,86 @@ where } } + /// Subscribes to Redis keyspace events, so evictions and expiries reach the + /// remove callbacks. Succeeds at most once; later calls return at once. + async fn subscribe_to_eviction_events(&self) -> Result<(), Error> { + self.has_remove_callback_subscribe + .get_or_try_init(|| async { + self.warn_if_eviction_events_are_disabled().await; + let (_, uuid) = self.connection_manager.get_connection().await?; + if let Err(err) = self.connection_manager.psubscribe("__key*__:*").await { + // A dropped connection or a failover in progress: re-resolve + // the connection, as the other retry paths in this store do, + // so the next attempt does not reuse a dead one. + if matches!(err.code, Code::Unavailable | Code::DeadlineExceeded) + && let Err(reconnect_err) = self.connection_manager.reconnect(uuid).await + { + debug!(?reconnect_err, "Reconnect after a failed PSUBSCRIBE failed"); + } + return Err(err); + } + Ok(()) + }) + .await?; + Ok(()) + } + + /// Logs if the server's `notify-keyspace-events` setting will not deliver + /// the eviction events the remove callbacks need. Advisory only: the + /// subscription is attempted either way, including when the setting + /// cannot be read (some managed Redis services block `CONFIG`). + async fn warn_if_eviction_events_are_disabled(&self) { + let cfg = async { + let mut client = self.get_client().await?; + redis::cmd("CONFIG") + .arg("GET") + .arg("notify-keyspace-events") + .query_async::>(&mut client.connection_manager) + .await + .map_err(|e| Error::from(e).append("Parsing notify-keyspace-events")) + } + .await; + let cfg = match cfg { + Ok(cfg) => cfg, + Err(err) => { + warn!( + ?err, + "Could not read notify-keyspace-events; subscribing to eviction events anyway" + ); + return; + } + }; + if cfg.len() != 1 { + warn!(?cfg, "Got multiple items for CONFIG GET, expected one"); + return; + } + let Some((_, events_cfg)) = cfg.first() else { + return; + }; + 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" + ); + } + async fn get_client(&self) -> Result, Error> { let local_client_permits = self.client_permits.clone(); let remaining = local_client_permits.available_permits(); @@ -1550,67 +1639,55 @@ where fn register_remove_callback(self: Arc, callback: RemoveCallback) -> Result<(), Error> { debug!(?callback, "New callback"); + // The task holds only a weak reference, so a store that is dropped + // while the subscription is still being retried ends the task instead + // of being kept alive by it. + let weak_self = Arc::downgrade(&self); background_spawn!("remove_callback_subscribe", async move { - self.remove_callbacks.lock().await.push(callback); - // 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); + let Some(store) = weak_self.upgrade() else { + return; + }; + store.remove_callbacks.lock().await.push(callback); + drop(store); + // This function has already returned Ok(()) to its caller (an + // `ExistenceCacheStore` relies on these callbacks to invalidate its + // entries), so the subscription is retried until it succeeds. A + // single transient error, such as -LOADING during a restart or + // failover, used to leave remove callbacks off for good. + let mut retry_delay = EVICTION_SUBSCRIBE_INITIAL_DELAY; + let mut failures: u32 = 0; 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" + let Some(store) = weak_self.upgrade() else { + debug!("RedisStore dropped before it subscribed to eviction events"); + return; + }; + let Err(err) = store.subscribe_to_eviction_events().await else { + if failures > 0 { + info!( + failures, + "Subscribed to Redis eviction events after retrying" ); - sleep(retry_delay).await; - retry_delay = (retry_delay * 2).min(Duration::from_secs(5)); } + return; + }; + drop(store); + failures += 1; + // Log the first failure as an error, then only at powers of + // two, so a permanent failure (for example an ACL that forbids + // PSUBSCRIBE) does not log every few seconds forever. + if failures == 1 { + error!( + ?err, + "Could not subscribe to Redis eviction events; retrying until it succeeds" + ); + } else if failures.is_power_of_two() { + warn!( + ?err, + failures, "Still unable to subscribe to Redis eviction events" + ); } + sleep(retry_delay).await; + retry_delay = (retry_delay * 2).min(EVICTION_SUBSCRIBE_MAX_DELAY); } }); Ok(()) @@ -2711,33 +2788,21 @@ 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 { +mod psubscribe_reconnect_tests { use core::time::Duration; - use std::collections::HashSet; + use std::collections::{HashMap, HashSet}; use std::sync::Arc; + use nativelink_macro::nativelink_test; + use nativelink_util::{background_spawn, spawn}; use redis::aio::{ConnectionManager, ConnectionManagerConfig}; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::sync::Notify; + use tokio::sync::mpsc::{UnboundedReceiver, unbounded_channel}; + use tokio::time::timeout; - use super::{RedisManager, StandardRedisManager}; + use super::{RedisConnectFn, RedisManager, StandardRedisManager}; /// Reads one CRLF-terminated line from the stream. async fn read_line(stream: &mut TcpStream) -> Option { @@ -2757,7 +2822,7 @@ mod psubscribe_reconnect_race_tests { Some(String::from_utf8_lossy(&line).into_owned()) } - /// Parses one RESP command (array of bulk strings) from the client. + /// Parses one RESP command (an 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('*') { @@ -2775,21 +2840,51 @@ mod psubscribe_reconnect_race_tests { Some(parts) } + /// A `PSUBSCRIBE` reply the fake server holds back on the reconnect + /// connection until the test releases it. + #[derive(Default)] + struct HeldReply { + /// Notified when the `PSUBSCRIBE` arrives. + arrived: Notify, + /// The reply is sent once this is notified. + release: Notify, + } + + #[derive(Default)] struct FakeServerState { - /// Sha1 to return for SCRIPT LOAD (must match the client's own hash). + /// 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 each connection, by accept order, was subscribed 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, + /// Replies to hold back on connection 1, the one `reconnect` opens, + /// keyed by pattern. + held_on_reconnect: HashMap>, + } + + impl FakeServerState { + fn new(held_patterns: &[&str]) -> Arc { + Arc::new(Self { + script_hash: redis::Script::new(super::LUA_VERSION_SET_SCRIPT) + .get_hash() + .to_string(), + held_on_reconnect: held_patterns + .iter() + .map(|pattern| ((*pattern).to_string(), Arc::default())) + .collect(), + ..Self::default() + }) + } + + fn held(&self, pattern: &str) -> &HeldReply { + &self.held_on_reconnect[pattern] + } + + fn patterns_on(&self, conn_idx: usize) -> Vec { + self.patterns_per_conn.lock()[conn_idx].clone() + } } 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() @@ -2812,13 +2907,11 @@ mod psubscribe_reconnect_race_tests { "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; + if conn_idx == 1 + && let Some(held) = state.held_on_reconnect.get(&pattern) + { + held.arrived.notify_one(); + held.release.notified().await; } let reply = format!( ">3\r\n$10\r\npsubscribe\r\n${}\r\n{}\r\n:1\r\n", @@ -2834,121 +2927,134 @@ mod psubscribe_reconnect_race_tests { } } - /// 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; - } - }); + /// Starts the fake server and returns a manager connected to it, the + /// generation id of its first connection, and the receiver for server + /// pushes, which must stay alive for the length of the test. + async fn connect( + state: &Arc, + ) -> ( + Arc>, + uuid::Uuid, + UnboundedReceiver, + ) { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let port = listener.local_addr().unwrap().port(); + let server_state = state.clone(); + background_spawn!("fake_redis_accept", async move { + let mut conn_idx = 0usize; + loop { + let (stream, _) = listener.accept().await.unwrap(); + server_state.patterns_per_conn.lock().push(Vec::new()); + background_spawn!( + "fake_redis_conn", + handle_conn(stream, conn_idx, server_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(), - ) + let (push_tx, push_rx) = unbounded_channel(); + let url = format!("redis://127.0.0.1:{port}/?protocol=resp3"); + let connect_func: Box> = 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) + }) + }); + let manager = Arc::new(StandardRedisManager::new(connect_func).await.unwrap()); + let (_, uuid) = manager.get_connection().await.unwrap(); + (manager, uuid, push_rx) + } + + /// A `psubscribe` that runs while a `reconnect` is subscribing its new + /// connection must end up subscribed on that new connection. + /// + /// The interleaving, with every step at a real await point: + /// 1. The manager connects (connection 0) and subscribes to `pre`. + /// 2. `reconnect` opens connection 1 and parks while subscribing it to + /// `pre`, because the fake server holds that reply. + /// 3. The test subscribes to `pat`. `get_connection` still returns + /// connection 0, so `pat` lands there. + /// 4. The server releases the reply and `reconnect` publishes connection 1. + /// + /// Before the fix, `reconnect` only replayed a snapshot taken before step 3, + /// so the live connection never subscribed to `pat` and its events were lost + /// until some later reconnect. + #[nativelink_test] + async fn psubscribe_during_reconnect_lands_on_new_connection() { + let state = FakeServerState::new(&["pre"]); + let (manager, uuid0, _push_rx) = connect(&state).await; + manager.psubscribe("pre").await.unwrap(); + + let reconnect = spawn!("reconnect", { + let manager = manager.clone(); + async move { manager.reconnect(uuid0).await } + }); + timeout(Duration::from_secs(5), state.held("pre").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" - ); + .expect("reconnect never started subscribing its 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" - ); + manager.psubscribe("pat").await.unwrap(); - // ...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", - ); + state.held("pre").release.notify_one(); + let (_new_conn, new_uuid) = reconnect.await.unwrap().unwrap(); + assert_ne!(uuid0, new_uuid, "reconnect must publish a new connection"); + + let recorded: HashSet = manager.subscriptions.lock().clone(); + assert!( + recorded.contains("pat"), + "pattern must be recorded as active" + ); + let live = state.patterns_on(1); + assert!( + live.contains(&"pat".to_string()), + "pattern \"pat\" is recorded as active but the live connection is only \ + subscribed to {live:?}" + ); + } + + /// While `reconnect` subscribes patterns recorded after its first pass, it + /// must not hold the connection write lock, or every `get_connection` + /// (every Redis operation) waits on those network round trips. + #[nativelink_test] + async fn reconnect_catch_up_does_not_block_get_connection() { + let state = FakeServerState::new(&["pre", "pat"]); + let (manager, uuid0, _push_rx) = connect(&state).await; + manager.psubscribe("pre").await.unwrap(); + + let reconnect = spawn!("reconnect", { + let manager = manager.clone(); + async move { manager.reconnect(uuid0).await } }); + timeout(Duration::from_secs(5), state.held("pre").arrived.notified()) + .await + .expect("reconnect never started subscribing its new connection"); + manager.psubscribe("pat").await.unwrap(); + state.held("pre").release.notify_one(); + + // `reconnect` now subscribes connection 1 to `pat`, and the server holds + // that reply. Readers must still get the old connection meanwhile. + timeout(Duration::from_secs(5), state.held("pat").arrived.notified()) + .await + .expect("reconnect never caught up on the pattern recorded mid-reconnect"); + let (_conn, uuid) = timeout(Duration::from_secs(2), manager.get_connection()) + .await + .expect("get_connection blocked while reconnect was subscribing") + .unwrap(); + assert_eq!(uuid, uuid0, "the new connection must not be published yet"); + + state.held("pat").release.notify_one(); + let (_new_conn, new_uuid) = reconnect.await.unwrap().unwrap(); + assert_ne!(uuid0, new_uuid, "reconnect must publish a new connection"); + assert!(state.patterns_on(1).contains(&"pat".to_string())); } } diff --git a/nativelink-store/tests/redis_store_test.rs b/nativelink-store/tests/redis_store_test.rs index 90f6e7dd2bb..8e8977fadf3 100644 --- a/nativelink-store/tests/redis_store_test.rs +++ b/nativelink-store/tests/redis_store_test.rs @@ -26,13 +26,13 @@ use nativelink_macro::nativelink_test; use nativelink_redis_tester::{ 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, + make_fake_redis_with_first_responses, make_fake_redis_with_responses, }; use nativelink_store::cas_utils::ZERO_BYTE_DIGESTS; use nativelink_store::existence_cache_store::ExistenceCacheStore; use nativelink_store::redis_store::{ ClusterRedisManager, DEFAULT_MAX_CHUNK_UPLOADS_PER_UPDATE, DEFAULT_MAX_COUNT_PER_CURSOR, - LUA_VERSION_SET_SCRIPT, RedisStore, RedisSubscriptionManager, decode_key, + LUA_VERSION_SET_SCRIPT, RedisStore, RedisSubscriptionManager, StandardRedisManager, decode_key, }; use nativelink_util::buf_channel::make_buf_channel_pair; use nativelink_util::common::DigestInfo; @@ -44,6 +44,7 @@ use nativelink_util::store_trait::{ TrueValue, UploadSizeInfo, }; use pretty_assertions::assert_eq; +use redis::aio::ConnectionManager; use redis::{PushInfo, RedisError, Value, make_extension_error}; use redis_test::{MockCmd, MockRedisConnection}; use tokio::sync::mpsc::UnboundedReceiver; @@ -2706,7 +2707,10 @@ where 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); - store.register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + // Keep `store` alive: the background subscription holds only a weak reference. + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; timeout(Duration::from_secs(3), async move { loop { @@ -2755,21 +2759,16 @@ 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. +// `register_remove_callback` returns Ok(()) at once and sets up the +// `__key*__:*` subscription in the background. Reading `notify-keyspace-events` +// first is only advisory, so a server that refuses `CONFIG GET` (here with a +// transient `-LOADING`) must still end up subscribed. This used to leave remove +// callbacks off for the life of the process. #[nativelink_test] -async fn transient_init_error_permanently_disables_remove_callbacks() -> Result<(), Error> { +async fn remove_callback_subscribes_even_if_config_cannot_be_read() -> 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") @@ -2785,6 +2784,125 @@ async fn transient_init_error_permanently_disables_remove_callbacks() -> Result< let redis_port = make_fake_redis_with_responses(responses) .instrument(redis_span) .await; + let store = remove_callback_test_store(redis_port).await; + // Keep `store` alive: the background subscription holds only a weak reference. + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + + wait_for_log( + logs_contain, + "Could not read notify-keyspace-events; subscribing to eviction events anyway", + ) + .await; + wait_for_log( + logs_contain, + "new psubscribe complete pattern=\"__key*__:*\" new_subscription=true", + ) + .await; + Ok(()) +} + +// A `PSUBSCRIBE` that fails once with a transient error must be retried until +// the subscription is established. +#[nativelink_test] +async fn remove_callback_subscription_is_retried_after_a_failed_psubscribe() -> Result<(), Error> { + let redis_span = info_span!("redis"); + + let mut responses = add_lua_version_script(fake_redis_stream()); + add_to_response( + &mut responses, + redis::cmd("CONFIG") + .arg("GET") + .arg("notify-keyspace-events"), + vec![Value::Map(vec![( + Value::BulkString(b"notify-keyspace-events".into()), + Value::BulkString(b"KA".into()), + )])], + ); + add_to_response( + &mut responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + vec![Value::Nil], + ); + let mut first_responses = HashMap::new(); + add_to_response_raw( + &mut first_responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + "-LOADING Redis is loading the dataset in memory\r\n".to_string(), + ); + let redis_port = make_fake_redis_with_first_responses(responses, first_responses) + .instrument(redis_span) + .await; + let store = remove_callback_test_store(redis_port).await; + // Keep `store` alive: the background subscription holds only a weak reference. + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + + wait_for_log( + logs_contain, + "Could not subscribe to Redis eviction events; retrying until it succeeds", + ) + .await; + wait_for_log( + logs_contain, + "new psubscribe complete pattern=\"__key*__:*\" new_subscription=true", + ) + .await; + wait_for_log( + logs_contain, + "Subscribed to Redis eviction events after retrying failures=1", + ) + .await; + Ok(()) +} + +// The subscription retry task holds only a weak reference to the store, so a +// store dropped while its subscription keeps failing (here every PSUBSCRIBE is +// refused) ends the task instead of being kept alive and retried forever. +#[nativelink_test] +async fn remove_callback_retry_stops_when_the_store_is_dropped() -> Result<(), Error> { + let redis_span = info_span!("redis"); + + let mut responses = add_lua_version_script(fake_redis_stream()); + add_to_response_raw( + &mut responses, + redis::cmd("CONFIG") + .arg("GET") + .arg("notify-keyspace-events"), + "-NOPERM this user has no permissions to run the 'config|get' command\r\n".to_string(), + ); + add_to_response_raw( + &mut responses, + redis::cmd("PSUBSCRIBE").arg("__key*__:*"), + "-NOPERM this user has no permissions to access one of the channels\r\n".to_string(), + ); + let redis_port = make_fake_redis_with_responses(responses) + .instrument(redis_span) + .await; + let store = remove_callback_test_store(redis_port).await; + store + .clone() + .register_remove_callback(Arc::new(LoggingRemoveCallback {}))?; + + wait_for_log( + logs_contain, + "Could not subscribe to Redis eviction events; retrying until it succeeds", + ) + .await; + drop(store); + wait_for_log( + logs_contain, + "RedisStore dropped before it subscribed to eviction events", + ) + .await; + Ok(()) +} + +async fn remove_callback_test_store( + redis_port: u16, +) -> Arc>> { let spec = RedisSpec { addresses: vec![format!("redis://127.0.0.1:{redis_port}/")], mode: RedisMode::Standard, @@ -2793,46 +2911,19 @@ async fn transient_init_error_permanently_disables_remove_callbacks() -> Result< 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 {}))?; + Arc::new(raw_store) +} - // Deterministic sync point: wait until the background init task has hit - // the transient error and logged it. +/// Waits up to 5 s for `needle` to appear in the calling test's logs. +/// `logs_contain` is the function `#[nativelink_test]` defines in each test. +async fn wait_for_log(logs_contain: impl Fn(&str) -> bool + Send + Sync, needle: &str) { timeout(Duration::from_secs(5), async { - loop { - if logs_contain("Error while trying to initialise remove_callback_subscribe") { - break; - } + while !logs_contain(needle) { 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(()) + .unwrap_or_else(|_| panic!("timed out waiting for log line: {needle}")); } #[nativelink_test] From 5c5d8044312148643cf25719d07fb8030a182a85 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 20:57:45 -0400 Subject: [PATCH 25/46] redis: pin the C10 torn-read hole with a reproduction test (RFC-track residual) PR #2841 added length + existence guards to the chunked Redis read path. Astra's C10 audit flagged that length checking cannot establish stable value identity: if a key's value is replaced by a *different* value of the same total length between chunk reads, the reader stitches bytes from two values and serves a torn blob as success. STRLEN is identical the whole way through, so no length-based guard can catch it. This confirms the hole is reachable and ships a reproduction, rather than a shallow length patch that would not close it. `get_part` reads a value larger than `read_chunk_size` with one `GETRANGE` per chunk, each a separate round-trip; values are mutable (`update` writes a temp key then `RENAME`s it onto the final key, discarding whatever was there). Content addressing makes this a non-issue for CAS, but the AC / scheduler / arbitrary-key uses of this same driver mutate values legitimately, so an equal-length replacement interleaved between two chunk reads tears the result. The new test drives exactly that interleaving through the mock connection and asserts the current code returns head-of-A spliced onto tail-of-B as Ok. Why no point fix here: closing equal-length replacement requires a snapshot-consistent read across more than one round-trip. On plain Redis strings that is only achievable by (a) reading the whole range in one atomic command (EVAL / MULTI), which materializes the entire value in a single reply and is unbounded for large blobs, or (b) attaching a version/identity to the value that `update` writes and `get_part` checks across chunks -- a change to the write path and the key/value schema. Both change the store read/write contract broadly, so the complete fix is RFC-track rather than a local patch. The single-chunk read path is already atomic and unaffected. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/tests/redis_store_test.rs | 88 ++++++++++++++++++++++ 1 file changed, 88 insertions(+) diff --git a/nativelink-store/tests/redis_store_test.rs b/nativelink-store/tests/redis_store_test.rs index 90f6e7dd2bb..e2702b4245a 100644 --- a/nativelink-store/tests/redis_store_test.rs +++ b/nativelink-store/tests/redis_store_test.rs @@ -566,6 +566,94 @@ async fn test_large_downloads_are_chunked() -> Result<(), Error> { Ok(()) } +// C10 (read integrity): equal-length in-place replacement tears a chunked read. +// +// `get_part` reads a value larger than `read_chunk_size` with one `GETRANGE` +// per chunk, each a separate round-trip. The value under a key is mutable: +// `update` lands data on a temp key and `RENAME`s it onto the final key, +// discarding whatever was there (CAS content-addressing makes this a non-issue +// for CAS, but the AC and arbitrary-key uses of this same driver mutate values +// legitimately). If a writer replaces the value between two chunk reads with a +// *different value of the same total length*, the reader stitches bytes from +// two different values and returns them as success. +// +// A length guard cannot catch this: STRLEN is identical before, during, and +// after the replacement. This test pins the current (buggy) behavior so the +// hole is captured in-tree and so any future fix has a fail->pass target. +// When a fix lands that establishes stable value identity across chunks, this +// test's final assertion flips from "returns torn bytes as Ok" to "returns an +// error", and the comment/name should change with it. +#[nativelink_test] +async fn c10_equal_length_replacement_tears_chunked_read() -> Result<(), Error> { + const READ_CHUNK_SIZE: usize = 1024; + // Two values of identical total length. Value A is all 'A's, value B all + // 'B's. The reader asks for A (via the first chunk) and then, mid-read, + // the key is swapped to B, so the second chunk comes from B. + let total_len = READ_CHUNK_SIZE + 128; + let value_a = Bytes::from(vec![b'A'; total_len]); + let value_b = Bytes::from(vec![b'B'; total_len]); + assert_eq!( + value_a.len(), + value_b.len(), + "same length is the whole point" + ); + + let digest = DigestInfo::try_new(VALID_HASH1, 1)?; + let real_key = format!("{digest}"); + + let commands = vec![ + // get_part goes straight to GETRANGE per chunk (no up-front STRLEN): + // the existence/length pipeline lives in has_with_results, a separate + // call. STRLEN would be identical for A and B anyway -- the swap is + // indistinguishable by length. + // First chunk: bytes [0, READ_CHUNK_SIZE) of value A. + MockCmd::new( + redis::cmd("GETRANGE") + .arg(real_key.clone()) + .arg(0) + .arg(READ_CHUNK_SIZE.try_into().unwrap_or(i64::MAX) - 1), + Ok(Value::BulkString(value_a.slice(..READ_CHUNK_SIZE).into())), + ), + // *** The value is replaced by B here (equal length). *** + // Second chunk: bytes [READ_CHUNK_SIZE, total_len) now come from B. + MockCmd::new( + redis::cmd("GETRANGE") + .arg(real_key) + .arg(READ_CHUNK_SIZE.try_into().unwrap_or(i64::MAX)) + .arg(total_len.try_into().unwrap_or(i64::MAX) - 1), + Ok(Value::BulkString(value_b.slice(READ_CHUNK_SIZE..).into())), + ), + ]; + + let store = make_mock_store(commands).await; + + let get_result = store + .get_part_unchunked(digest, 0, Some(total_len as u64)) + .await; + + // The stitched result is neither A nor B: it is the head of A spliced onto + // the tail of B -- a torn blob served as success. This is the C10 hole. + let torn = get_result.expect("current code returns the torn blob as Ok"); + let mut expected_torn = Vec::with_capacity(total_len); + expected_torn.extend_from_slice(&value_a[..READ_CHUNK_SIZE]); + expected_torn.extend_from_slice(&value_b[READ_CHUNK_SIZE..]); + assert_eq!( + torn, + Bytes::from(expected_torn), + "C10: read is torn across an equal-length replacement -- head from A, tail from B", + ); + assert_ne!( + torn, value_a, + "torn blob is not the value the reader asked for" + ); + assert_ne!( + torn, value_b, + "torn blob is not the replacement value either" + ); + + Ok(()) +} + #[nativelink_test] async fn yield_between_sending_packets_in_update() -> Result<(), Error> { let data_p1 = Bytes::from(vec![b'A'; DEFAULT_READ_CHUNK_SIZE + 512]); From b0f3b0067c28d8aea88605abf0f11caab88c9813 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Mon, 3 Aug 2026 15:13:00 -0400 Subject: [PATCH 26/46] fix(store): dead shard-ring peer fails fast instead of wedging every RPC MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The production wedge: a ShardStore ring over grpc_store backends hangs EVERY CAS RPC when one peer is dead or restarting, with no error and no timeout — the operational workaround was restarting the frontend. Where the hang lived: ConnectionManager retries a dead endpoint's connect forever in the background (by design — never give up on an endpoint), but connection() requests queue unbounded in waiting_connections while it does, so no error ever reaches the retry/error machinery. rpc_timeout_s was only applied on the ByteStream.Write path; reads and unary RPCs had no deadline at all. The fix, two layers: - connection_manager: track per-endpoint connect health. When the last connect attempt to every endpoint has failed and no channel is available, fail queued and incoming connection() requests fast with Unavailable instead of queuing them forever. A background reconnect success clears the flag and service resumes. - grpc_store: apply the existing rpc_timeout_s knob to every RPC — unary calls, and read-stream establishment through the first response chunk — not just writes. A timeout maps to DeadlineExceeded, which is retryable. Tests (repro first; both failed against the old code by hanging past a generous outer timeout): - connection_manager_test::dead_endpoint_fails_fast_instead_of_queuing_forever - shard_ring_dead_peer_test: an in-process two-shard lab ring (one live fake CAS peer, one dead port) proving reads and multi-key existence checks settle with Unavailable/DeadlineExceeded while the live shard keeps serving. Co-Authored-By: Claude Fable 5 --- .../vocabularies/TraceMachina/accept.txt | 1 + nativelink-config/src/stores.rs | 12 +- nativelink-store/BUILD.bazel | 1 + nativelink-store/src/grpc_store.rs | 43 +- .../tests/shard_ring_dead_peer_test.rs | 450 ++++++++++++++++++ nativelink-util/BUILD.bazel | 1 + nativelink-util/src/connection_manager.rs | 101 +++- .../tests/connection_manager_test.rs | 257 ++++++++++ 8 files changed, 847 insertions(+), 19 deletions(-) create mode 100644 nativelink-store/tests/shard_ring_dead_peer_test.rs diff --git a/.github/styles/config/vocabularies/TraceMachina/accept.txt b/.github/styles/config/vocabularies/TraceMachina/accept.txt index c0832d0ef15..c68d4ee5fd3 100644 --- a/.github/styles/config/vocabularies/TraceMachina/accept.txt +++ b/.github/styles/config/vocabularies/TraceMachina/accept.txt @@ -493,3 +493,4 @@ Flybody MuJoCo connectome simulatable +unary diff --git a/nativelink-config/src/stores.rs b/nativelink-config/src/stores.rs index 81b6179f929..8abaa14c390 100644 --- a/nativelink-config/src/stores.rs +++ b/nativelink-config/src/stores.rs @@ -1628,12 +1628,18 @@ pub struct GrpcSpec { #[serde(default, deserialize_with = "convert_numeric_with_shellexpand")] pub connections_per_endpoint: usize, - /// Maximum time (seconds) allowed for a single RPC request (e.g. a - /// `ByteStream.Write` call) before it is cancelled. + /// Maximum time (seconds) allowed for a single RPC request before it + /// is cancelled with `DeadlineExceeded` (retryable). This applies to + /// every RPC this store makes: unary calls (e.g. `FindMissingBlobs`, + /// `BatchReadBlobs`), `ByteStream.Write` calls, and the establishment + /// of read streams through their first response chunk. The body of an + /// established read stream is not bounded by this deadline. /// /// A value of 0 (the default) disables the per-RPC timeout. Dead /// connections are still detected by the HTTP/2 and TCP keepalive - /// mechanisms configured on each endpoint. + /// mechanisms configured on each endpoint, and endpoints that cannot + /// be connected to at all fail fast with `Unavailable` rather than + /// queuing requests indefinitely. /// /// For large uploads (multi-GB), either leave this at 0 or set it /// large enough to accommodate the full transfer time. diff --git a/nativelink-store/BUILD.bazel b/nativelink-store/BUILD.bazel index a7715695201..75e7e2c8430 100644 --- a/nativelink-store/BUILD.bazel +++ b/nativelink-store/BUILD.bazel @@ -166,6 +166,7 @@ rust_test_suite( "tests/redis_store_test.rs", "tests/ref_store_test.rs", "tests/s3_store_test.rs", + "tests/shard_ring_dead_peer_test.rs", "tests/shard_store_test.rs", "tests/size_partitioning_store_test.rs", "tests/store_manager_test.rs", diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index 451ad956cd2..3fa7efd31cd 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -334,6 +334,32 @@ impl GrpcStore { })) } + /// Bounds `fut` by the configured per-RPC deadline (`rpc_timeout_s`). + /// A timeout maps to `DeadlineExceeded`, which is retryable, so the + /// existing retry machinery can route around a peer that accepted the + /// request but will never answer (e.g. a restarting shard). + async fn with_rpc_timeout( + &self, + context: &'static str, + fut: impl Future> + Send, + ) -> Result { + // Boxed so this wrapper adds a pointer to the caller's future, not a + // second inline copy of `fut` — the RPC futures it wraps are large, + // and embedding them twice measurably deepens poll stacks. + let mut fut = Box::pin(fut); + if self.rpc_timeout.is_zero() { + return fut.as_mut().await; + } + match tokio::time::timeout(self.rpc_timeout, fut.as_mut()).await { + Ok(result) => result, + Err(_elapsed) => Err(make_err!( + Code::DeadlineExceeded, + "{context} timed out after {}s", + self.rpc_timeout.as_secs() + )), + } + } + async fn perform_request(&self, input: I, mut request: F) -> Result where F: FnMut(I) -> Fut + Send + Copy, @@ -345,7 +371,7 @@ impl GrpcStore { .retry(unfold(input, move |input| async move { let input_clone = input.clone(); Some(( - request(input_clone) + self.with_rpc_timeout("GrpcStore RPC", request(input_clone)) .await .map_or_else(RetryResult::Retry, RetryResult::Ok), input, @@ -1454,11 +1480,14 @@ impl GrpcStore { }; let mut stream = match self - .read_internal(ReadRequest { - resource_name, - read_offset: 0, - read_limit: 0, - }) + .with_rpc_timeout( + "GrpcStore::get_part_compressed read", + self.read_internal(ReadRequest { + resource_name, + read_offset: 0, + read_limit: 0, + }), + ) .await { Ok(stream) => stream, @@ -2043,7 +2072,7 @@ impl StoreDriver for GrpcStore { read_limit: local_state.read_limit, }; let mut stream = match self - .read_internal(request) + .with_rpc_timeout("GrpcStore::get_part read", self.read_internal(request)) .await .err_tip(|| "in GrpcStore::get_part()") { diff --git a/nativelink-store/tests/shard_ring_dead_peer_test.rs b/nativelink-store/tests/shard_ring_dead_peer_test.rs new file mode 100644 index 00000000000..238b8f0d6e0 --- /dev/null +++ b/nativelink-store/tests/shard_ring_dead_peer_test.rs @@ -0,0 +1,450 @@ +// 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. + +//! Lab reproduction of the production "shard-ring wedge": a `ShardStore` +//! ring over `GrpcStore` backends (peer CAS shards) where one peer is dead +//! (killed or never started). Before the fix, every CAS RPC that touched +//! the dead shard hung forever: the peer's `ConnectionManager` retried the +//! connect in the background while `connection()` requests queued +//! unbounded, so no error ever reached the retry machinery. The +//! operational workaround was `systemctl restart nativelink` on the +//! frontend. +//! +//! These tests must fail (by generous outer timeout) against the +//! pre-fix code and pass after dead peers yield fast `Unavailable`. + +use core::pin::Pin; +use core::time::Duration; +use std::collections::HashMap; +use std::sync::Arc; + +use async_lock::Mutex; +use bytes::Bytes; +use futures::{Stream, StreamExt}; +use nativelink_config::stores::{ + GrpcEndpoint, GrpcSpec, MemorySpec, Retry, ShardConfig, ShardSpec, StoreSpec, StoreType, +}; +use nativelink_error::{Code, Error}; +use nativelink_macro::nativelink_test; +use nativelink_proto::build::bazel::remote::execution::v2::content_addressable_storage_server::{ + ContentAddressableStorage, ContentAddressableStorageServer, +}; +use nativelink_proto::build::bazel::remote::execution::v2::{ + BatchReadBlobsRequest, BatchReadBlobsResponse, BatchUpdateBlobsRequest, + BatchUpdateBlobsResponse, FindMissingBlobsRequest, FindMissingBlobsResponse, GetTreeRequest, + GetTreeResponse, SpliceBlobRequest, SpliceBlobResponse, SplitBlobRequest, SplitBlobResponse, +}; +use nativelink_proto::google::bytestream::byte_stream_server::{ByteStream, ByteStreamServer}; +use nativelink_proto::google::bytestream::{ + QueryWriteStatusRequest, QueryWriteStatusResponse, ReadRequest, ReadResponse, WriteRequest, + WriteResponse, +}; +use nativelink_store::grpc_store::GrpcStore; +use nativelink_store::memory_store::MemoryStore; +use nativelink_store::shard_store::ShardStore; +use nativelink_util::background_spawn; +use nativelink_util::common::DigestInfo; +use nativelink_util::store_trait::{Store, StoreLike}; +use tokio::time::timeout; +use tonic::transport::Server; +use tonic::transport::server::TcpIncoming; +use tonic::{Request, Response, Status, Streaming}; + +/// Generous bound: far longer than any configured timeout/retry budget in +/// this test, far shorter than "forever". If an operation does not settle +/// within this, it is the production hang. +const GENEROUS_TIMEOUT: Duration = Duration::from_secs(20); + +/// A minimal in-process CAS shard: blobs live in a shared map, +/// `FindMissingBlobs`, `ByteStream.Read` and `ByteStream.Write` behave like +/// a real peer. +#[derive(Debug, Clone)] +struct FakeShardServer { + blobs: Arc>>, +} + +type GetTreeStream = Pin> + Send + 'static>>; + +#[tonic::async_trait] +impl ContentAddressableStorage for FakeShardServer { + type GetTreeStream = GetTreeStream; + + async fn find_missing_blobs( + &self, + grpc_request: Request, + ) -> Result, Status> { + let request = grpc_request.into_inner(); + let blobs = self.blobs.lock().await; + let missing_blob_digests = request + .blob_digests + .into_iter() + .filter(|digest| !blobs.contains_key(&digest.hash)) + .collect(); + Ok(Response::new(FindMissingBlobsResponse { + missing_blob_digests, + })) + } + + #[allow(clippy::unimplemented)] + async fn batch_update_blobs( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn batch_read_blobs( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn get_tree( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn split_blob( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } + + #[allow(clippy::unimplemented)] + async fn splice_blob( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } +} + +type ReadStream = Pin> + Send + 'static>>; + +/// Extracts the blob hash from a `ByteStream` resource name of the form +/// `[{instance}/]uploads/{uuid}/blobs/{fn}/{hash}/{size}` or +/// `[{instance}/]blobs/{fn}/{hash}/{size}`. +fn hash_from_resource_name(resource_name: &str) -> String { + let mut components = resource_name.rsplit('/'); + let _size = components.next(); + components.next().unwrap_or_default().to_string() +} + +#[tonic::async_trait] +impl ByteStream for FakeShardServer { + type ReadStream = ReadStream; + + async fn read( + &self, + grpc_request: Request, + ) -> Result, Status> { + let request = grpc_request.into_inner(); + let hash = hash_from_resource_name(&request.resource_name); + let data = self + .blobs + .lock() + .await + .get(&hash) + .cloned() + .ok_or_else(|| Status::not_found(format!("Blob {hash} not seeded")))?; + let offset = usize::try_from(request.read_offset) + .map_err(|_| Status::invalid_argument("Bad read_offset"))?; + if offset > data.len() { + return Err(Status::out_of_range("read_offset past end of blob")); + } + let stream: Self::ReadStream = Box::pin(futures::stream::iter(vec![Ok(ReadResponse { + data: data.slice(offset..), + })])); + Ok(Response::new(stream)) + } + + async fn write( + &self, + grpc_request: Request>, + ) -> Result, Status> { + let mut stream = grpc_request.into_inner(); + let mut hash = String::new(); + let mut data = Vec::new(); + while let Some(request) = stream.next().await { + let request = request?; + if !request.resource_name.is_empty() { + hash = hash_from_resource_name(&request.resource_name); + } + data.extend_from_slice(&request.data); + if request.finish_write { + break; + } + } + let committed_size = + i64::try_from(data.len()).map_err(|_| Status::invalid_argument("Blob too large"))?; + self.blobs.lock().await.insert(hash, Bytes::from(data)); + Ok(Response::new(WriteResponse { committed_size })) + } + + #[allow(clippy::unimplemented)] + async fn query_write_status( + &self, + _grpc_request: Request, + ) -> Result, Status> { + unimplemented!(); + } +} + +/// Spawns a live in-process shard peer and returns its port and blob map. +fn spawn_live_peer() -> (u16, Arc>>) { + let server = FakeShardServer { + blobs: Arc::new(Mutex::new(HashMap::new())), + }; + let blobs = server.blobs.clone(); + let listener = TcpIncoming::bind("127.0.0.1:0".parse().unwrap()).unwrap(); + let port = listener.local_addr().unwrap().port(); + let cas_service = ContentAddressableStorageServer::new(server.clone()); + let bytestream_service = ByteStreamServer::new(server); + background_spawn!("shard_ring_dead_peer_test_server", async move { + Server::builder() + .add_service(cas_service) + .add_service(bytestream_service) + .serve_with_incoming(listener) + .await + .unwrap(); + }); + (port, blobs) +} + +/// Reserves a port with nothing listening on it: the dead peer (killed or +/// never started). Connect attempts fail with `ECONNREFUSED`. +fn dead_peer_port() -> u16 { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap().port() +} + +fn grpc_spec_for_port(port: u16) -> GrpcSpec { + GrpcSpec { + instance_name: String::new(), + endpoints: vec![GrpcEndpoint { + address: format!("http://127.0.0.1:{port}"), + tls_config: None, + concurrency_limit: None, + connect_timeout_s: 1, + tcp_keepalive_s: 0, + http2_keepalive_interval_s: 0, + http2_keepalive_timeout_s: 0, + }], + store_type: StoreType::Cas, + retry: Retry { + max_retries: 1, + delay: 0.05, + jitter: 0.0, + retry_on_errors: None, + }, + max_concurrent_requests: 0, + connections_per_endpoint: 0, + rpc_timeout_s: 5, + use_legacy_resource_names: false, + headers: HashMap::new(), + forward_headers: vec![], + experimental_read_batching: None, + experimental_remote_cache_compression: None, + } +} + +/// Mirrors `ShardStore`'s routing for a two-shard equal-weight ring: the +/// digest's xor-folded key lands in shard 0 iff it is `<= u32::MAX / 2`. +fn shard_index_for(digest: &DigestInfo) -> usize { + let packed_hash = digest.packed_hash(); + let size_bytes = digest.size_bytes().to_le_bytes(); + let mut key: u32 = 0; + for chunk in 0..8 { + key ^= u32::from_le_bytes(packed_hash[chunk * 4..chunk * 4 + 4].try_into().unwrap()); + } + key ^= u32::from_le_bytes(size_bytes[0..4].try_into().unwrap()); + key ^= u32::from_le_bytes(size_bytes[4..8].try_into().unwrap()); + usize::from(key > u32::MAX / 2) +} + +/// Finds a digest (with 64-byte content) that routes to `shard`. +fn digest_for_shard(shard: usize) -> (DigestInfo, Bytes) { + for i in 0..1024usize { + let hash = format!("{i:064x}"); + let content = Bytes::from(format!("{i:0>64}")); + let digest = DigestInfo::try_new(&hash, content.len()).unwrap(); + if shard_index_for(&digest) == shard { + return (digest, content); + } + } + unreachable!("no digest found for shard {shard} in 1024 candidates"); +} + +struct Ring { + shard_store: Arc, + live_shard: usize, + live_blobs: Arc>>, +} + +/// Builds a two-shard equal-weight ring: shard 0 is a live in-process +/// peer, shard 1 is dead (nothing listening). +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))?; + 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. + let shard_store = ShardStore::new( + &ShardSpec { + stores: (0..2) + .map(|_| ShardConfig { + store: StoreSpec::Memory(MemorySpec::default()), + weight: Some(1), + }) + .collect(), + }, + vec![Store::new(live_store), Store::new(dead_store)], + )?; + Ok(Ring { + shard_store, + live_shard: 0, + live_blobs, + }) +} + +// Guard for `shard_index_for`: it re-implements `ShardStore`'s xor-fold +// routing, which would silently diverge if upstream ever changes the ring +// hash — the dead-peer tests above would then route probes to the wrong +// shard and test nothing. This builds a REAL `ShardStore` over two +// instrumented `MemoryStore` backends, writes each probe digest through +// it, and asserts the blob landed on the shard `shard_index_for` +// predicted. If the ring hash changes upstream, this fails loudly. +#[nativelink_test] +async fn shard_index_for_matches_real_shard_store_routing() -> Result<(), Error> { + let backends: Vec> = (0..2) + .map(|_| MemoryStore::new(&MemorySpec::default())) + .collect(); + let shard_store = ShardStore::new( + &ShardSpec { + stores: (0..2) + .map(|_| ShardConfig { + store: StoreSpec::Memory(MemorySpec::default()), + weight: Some(1), + }) + .collect(), + }, + backends + .iter() + .map(|backend| Store::new(backend.clone())) + .collect(), + )?; + + for shard in 0..2usize { + let (digest, content) = digest_for_shard(shard); + shard_store.update_oneshot(digest, content).await?; + for (backend_index, backend) in backends.iter().enumerate() { + let present = backend.has(digest).await?.is_some(); + assert_eq!( + present, + backend_index == shard, + "probe digest {digest} predicted for shard {shard} but presence on \ + backend {backend_index} was {present}: `shard_index_for` no longer \ + matches ShardStore's ring hash", + ); + } + } + Ok(()) +} + +// A read routed to the dead shard must settle with an error (which the +// retry machinery upstream can handle) instead of hanging forever. +#[nativelink_test] +async fn read_routed_to_dead_shard_errors_instead_of_hanging() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, _) = digest_for_shard(1 - ring.live_shard); + + let result = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read routed to the dead shard hung past the generous timeout: the shard-ring wedge"); + let err = result.expect_err("read from a dead shard cannot succeed"); + assert!( + matches!(err.code, Code::Unavailable | Code::DeadlineExceeded), + "expected Unavailable/DeadlineExceeded from the dead shard, got: {err:?}", + ); + Ok(()) +} + +// The production symptom: a multi-key existence check (FindMissingBlobs) +// fans out to every shard, so ONE dead peer wedged EVERY CAS RPC through +// the frontend. It must settle with an error instead. +#[nativelink_test] +async fn has_with_dead_shard_errors_instead_of_hanging() -> Result<(), Error> { + let ring = make_ring().await?; + let (live_digest, _) = digest_for_shard(ring.live_shard); + let (dead_digest, _) = digest_for_shard(1 - ring.live_shard); + + let result = timeout( + GENEROUS_TIMEOUT, + ring.shard_store + .has_many(&[live_digest.into(), dead_digest.into()]), + ) + .await + .expect("multi-key has() with a dead shard hung past the generous timeout"); + let err = result.expect_err("has() touching a dead shard cannot succeed"); + assert!( + matches!(err.code, Code::Unavailable | Code::DeadlineExceeded), + "expected Unavailable/DeadlineExceeded from the dead shard, got: {err:?}", + ); + Ok(()) +} + +// The live shard must keep serving while its ring-mate is dead: routing +// around the dead peer stays possible for keys that don't touch it. +#[nativelink_test] +async fn live_shard_still_serves_while_peer_dead() -> Result<(), Error> { + let ring = make_ring().await?; + let (live_digest, content) = digest_for_shard(ring.live_shard); + + timeout( + GENEROUS_TIMEOUT, + ring.shard_store + .update_oneshot(live_digest, content.clone()), + ) + .await + .expect("write to the live shard hung")?; + assert!( + ring.live_blobs + .lock() + .await + .contains_key(&live_digest.packed_hash().to_string()), + "blob did not land on the live peer", + ); + + let data = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(live_digest, 0, None), + ) + .await + .expect("read from the live shard hung")?; + assert_eq!(data, content, "live shard returned wrong content"); + Ok(()) +} diff --git a/nativelink-util/BUILD.bazel b/nativelink-util/BUILD.bazel index bc008ef5ea0..fa453f15375 100644 --- a/nativelink-util/BUILD.bazel +++ b/nativelink-util/BUILD.bazel @@ -106,6 +106,7 @@ rust_test_suite( "tests/cache_size_metrics_test.rs", "tests/channel_body_for_tests_test.rs", "tests/common_test.rs", + "tests/connection_manager_test.rs", "tests/evicting_map_test.rs", "tests/fastcdc_test.rs", "tests/fs_test.rs", diff --git a/nativelink-util/src/connection_manager.rs b/nativelink-util/src/connection_manager.rs index 753e7cb956a..145f84a790a 100644 --- a/nativelink-util/src/connection_manager.rs +++ b/nativelink-util/src/connection_manager.rs @@ -35,7 +35,7 @@ use crate::retry::{self, Retrier, RetryResult}; #[derive(Debug)] pub struct ConnectionManager { // The channel to request connections from the worker. - worker_tx: mpsc::Sender<(String, oneshot::Sender)>, + worker_tx: mpsc::Sender<(String, oneshot::Sender>)>, } /// The index into `ConnectionManagerWorker::endpoints`. @@ -94,6 +94,21 @@ struct ConnectionManagerWorker { /// The endpoints to establish Channels and the identifier of the last /// connection attempt to that endpoint. endpoints: Vec<(ConnectionIndex, Endpoint)>, + /// Per-endpoint connect health: `true` once a connect retry cycle to + /// that endpoint has failed, cleared when a connect succeeds. When + /// every endpoint is marked failed and no channel is available, the + /// endpoint set is unreachable — requests fail fast with + /// `Unavailable` instead of queuing forever while the background + /// reconnect loop runs (the "dead peer wedges every RPC" hang). + endpoint_connect_failed: Vec, + /// Number of currently established channels, whether sitting in + /// `available_channels` or checked out to in-flight `Connection`s. + /// `available_channels` alone cannot distinguish "everything is down" + /// from "everything is busy": borrowed channels are not in that queue, + /// so a busy-but-healthy pool plus one transient connect failure must + /// not be declared dead. Incremented on connect success, decremented + /// when a channel is torn down on transport error. + live_channels: usize, /// The channel used to communicate between a Connection and the worker. connection_tx: mpsc::UnboundedSender, /// Gates the maximum number of in-flight `Connection` objects. @@ -112,7 +127,7 @@ struct ConnectionManagerWorker { /// Connected channels that are available for use. available_channels: VecDeque, /// Requests for a Channel when available - (reason, request) - waiting_connections: VecDeque<(String, oneshot::Sender)>, + waiting_connections: VecDeque<(String, oneshot::Sender>)>, /// The retry configuration for connecting to an Endpoint, on failure will /// restart the retrier after a 1 second delay. retrier: Retrier, @@ -140,10 +155,11 @@ impl ConnectionManager { // which defeats the object since there would be no backpressure // applied. Therefore it makes sense for this to be unbounded. let (connection_tx, connection_rx) = mpsc::unbounded_channel(); - let endpoints = endpoints + let endpoints: Vec<(ConnectionIndex, Endpoint)> = endpoints .into_iter() .map(|endpoint| (0, endpoint)) .collect(); + let endpoint_connect_failed = vec![false; endpoints.len()]; // Zero means unlimited, which becomes a semaphore holding // Semaphore::MAX_PERMITS. That is a sentinel, not a free-slot count, @@ -159,6 +175,8 @@ impl ConnectionManager { } let worker = ConnectionManagerWorker { endpoints, + endpoint_connect_failed, + live_channels: 0, available_connections: Arc::new(Semaphore::new(max_concurrent_requests)), bounded_connections, connection_tx, @@ -190,7 +208,7 @@ impl ConnectionManager { .await .map_err(|err| make_err!(Code::Unavailable, "Requesting a new connection: {err:?}"))?; rx.await - .map_err(|err| make_err!(Code::Unavailable, "Waiting for a new connection: {err:?}")) + .map_err(|err| make_err!(Code::Unavailable, "Waiting for a new connection: {err:?}"))? } } @@ -198,7 +216,7 @@ impl ConnectionManagerWorker { async fn service_requests( mut self, connections_per_endpoint: usize, - mut worker_rx: mpsc::Receiver<(String, oneshot::Sender)>, + mut worker_rx: mpsc::Receiver<(String, oneshot::Sender>)>, mut connection_rx: mpsc::UnboundedReceiver, ) { // Make the initial set of connections, connection failures will be @@ -252,6 +270,13 @@ impl ConnectionManagerWorker { fn handle_connected(&mut self, connection_result: IndexedChannel) { match connection_result { Ok(established_channel) => { + if let Some(failed) = self + .endpoint_connect_failed + .get_mut(established_channel.identifier.endpoint_index) + { + *failed = false; + } + self.live_channels += 1; self.available_channels.push_back(established_channel); self.maybe_available_connection(); } @@ -259,11 +284,52 @@ impl ConnectionManagerWorker { // beginning of the retry period. Never want to be in a // situation where we give up on an Endpoint forever. Err((identifier, _)) => { + if let Some(failed) = self + .endpoint_connect_failed + .get_mut(identifier.endpoint_index) + { + *failed = true; + } + // If this failure means the whole endpoint set is now + // unreachable, waiting requests would hang until a peer + // comes back — fail them fast so the caller's retry/error + // machinery can handle it. + self.fail_requests_if_all_endpoints_down(); self.connect_endpoint(identifier.endpoint_index, Some(identifier.connection_index)); } } } + /// True when no channel is live (available or checked out) and the + /// last connect attempt to every endpoint failed: nothing can service + /// a request until a background reconnect succeeds. + fn all_endpoints_down(&self) -> bool { + self.live_channels == 0 + && !self.endpoint_connect_failed.is_empty() + && self.endpoint_connect_failed.iter().all(|failed| *failed) + } + + fn unavailable_error(&self) -> Error { + make_err!( + Code::Unavailable, + "All {} gRPC endpoint(s) are unreachable (connect attempts failed); failing fast instead of queuing", + self.endpoints.len() + ) + } + + fn fail_requests_if_all_endpoints_down(&mut self) { + if !self.all_endpoints_down() { + return; + } + while let Some((reason, tx)) = self.waiting_connections.pop_front() { + warn!( + reason, + "ConnectionManager: all endpoints down, failing queued connection request" + ); + drop(tx.send(Err(self.unavailable_error()))); + } + } + fn connect_endpoint(&mut self, endpoint_index: usize, connection_index: Option) { let Some((current_connection_index, endpoint)) = self.endpoints.get_mut(endpoint_index) else { @@ -333,7 +399,7 @@ impl ConnectionManagerWorker { } // This must never be made async otherwise the select may cancel it. - fn handle_worker(&mut self, reason: String, tx: oneshot::Sender) { + fn handle_worker(&mut self, reason: String, tx: oneshot::Sender>) { let maybe_permit = self.available_connections.clone().try_acquire_owned().ok(); if let Some(permit) = maybe_permit && let Some(channel) = self.available_channels.pop_front() @@ -341,6 +407,12 @@ impl ConnectionManagerWorker { debug!(reason, "ConnectionManager: request running"); record_connection_acquired("grpc", self.free_connection_slots(), false); self.provide_channel(channel, tx, permit); + } else if self.all_endpoints_down() { + warn!( + reason, + "ConnectionManager: all endpoints down, failing connection request" + ); + drop(tx.send(Err(self.unavailable_error()))); } else { debug!( available_permits = self.available_connections.available_permits(), @@ -357,15 +429,15 @@ impl ConnectionManagerWorker { fn provide_channel( &self, channel: EstablishedChannel, - tx: oneshot::Sender, + tx: oneshot::Sender>, permit: OwnedSemaphorePermit, ) { - drop(tx.send(Connection { + drop(tx.send(Ok(Connection { tx: self.connection_tx.clone(), pending_channel: Some(channel.channel.clone()), channel, _permit: permit, - })); + }))); } fn maybe_available_connection(&mut self) { @@ -413,6 +485,17 @@ impl ConnectionManagerWorker { original_length != self.available_channels.len() }; if should_reconnect { + // The channel this identifier referred to is gone; + // saturate in case of duplicate error reports for the + // same pending channel. + self.live_channels = self.live_channels.saturating_sub(1); + // If this was the last live channel and every endpoint's + // last connect cycle had already failed, the pool just + // became unreachable — flush queued waiters now instead + // of leaving them parked until the reconnect scheduled + // below fails a full connect cycle later (new requests + // already fail fast in `handle_worker` in this state). + self.fail_requests_if_all_endpoints_down(); self.connect_endpoint(identifier.endpoint_index, None); } } diff --git a/nativelink-util/tests/connection_manager_test.rs b/nativelink-util/tests/connection_manager_test.rs index 286ea578311..f293b0201dd 100644 --- a/nativelink-util/tests/connection_manager_test.rs +++ b/nativelink-util/tests/connection_manager_test.rs @@ -161,6 +161,263 @@ async fn aborted_caller_future_does_not_leak_permits() -> Result<(), Error> { Ok(()) } +/// The production "shard-ring wedge": an endpoint that is dead (nothing +/// listening — killed or never started) must yield a fast `Unavailable` +/// error from `connection()` instead of queuing the request forever while +/// the worker retries the connect in the background. Before the fix this +/// test failed: `cm.connection()` never resolved and the outer timeout +/// elapsed. +#[nativelink_test] +async fn dead_endpoint_fails_fast_instead_of_queuing_forever() -> Result<(), Error> { + // Bind a port and immediately drop the listener so nothing is + // listening there — connects get ECONNREFUSED, the "peer killed or + // never started" case observed in production. + let dead_port = { + let listener = std::net::TcpListener::bind("127.0.0.1:0").unwrap(); + listener.local_addr().unwrap().port() + }; + let endpoint = Endpoint::from_shared(format!("http://127.0.0.1:{dead_port}")) + .unwrap() + .connect_timeout(Duration::from_secs(1)); + + let cm = ConnectionManager::new(vec![endpoint], 1, 2, Retry::default(), no_jitter()); + + let result = timeout(Duration::from_secs(10), cm.connection("dead-peer".into())) + .await + .expect("connection() to a dead endpoint hung past 10s instead of erroring"); + assert_eq!( + result.err().map(|err| err.code), + Some(nativelink_error::Code::Unavailable), + "expected fast Unavailable from a dead endpoint", + ); + Ok(()) +} + +/// The false-negative guard on the fast-fail gating: a busy-but-healthy +/// pool must NOT be declared dead. One endpoint, two connection slots: a +/// single-accept proxy in front of a real gRPC server lets exactly one +/// connect succeed, then drops its listener so the second slot's connect +/// is refused (`ECONNREFUSED`), marking the endpoint's last connect cycle +/// as failed. We hold the one live connection (so `available_channels` is +/// empty) and issue a new `connection()` request: with gating keyed only +/// on `available_channels` + connect health this errored `Unavailable` +/// immediately; with live-channel tracking it queues, and resolves as +/// soon as the held connection is released. +#[nativelink_test] +async fn busy_pool_with_transient_connect_failure_does_not_fail_fast() -> Result<(), Error> { + // Real gRPC server the proxied connection terminates at. + let backend = fake_grpc_server_endpoint().await; + let backend_uri = backend.uri().clone(); + let backend_addr = format!( + "{}:{}", + backend_uri.host().unwrap(), + backend_uri.port_u16().unwrap() + ); + + // Proxy that accepts exactly ONE connection, pipes it to the backend, + // then drops the listener — every later connect is refused. + let proxy_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_port = proxy_listener.local_addr().unwrap().port(); + background_spawn!("single_accept_proxy", async move { + let (mut inbound, _) = proxy_listener.accept().await.unwrap(); + // Drop the listener so the second connect slot gets ECONNREFUSED. + drop(proxy_listener); + let mut outbound = tokio::net::TcpStream::connect(backend_addr).await.unwrap(); + drop(tokio::io::copy_bidirectional(&mut inbound, &mut outbound).await); + }); + + let endpoint = Endpoint::from_shared(format!("http://127.0.0.1:{proxy_port}")) + .unwrap() + .connect_timeout(Duration::from_secs(1)); + let cm = Arc::new(ConnectionManager::new( + vec![endpoint], + /* connections_per_endpoint = */ 2, + /* max_concurrent_requests = */ 2, + Retry::default(), + no_jitter(), + )); + + // Take (and hold) the one connection that could be established. This + // may need a moment while the connect race settles; the winner is the + // proxied slot, the loser marks the endpoint's connect cycle failed. + let held = timeout(Duration::from_secs(10), cm.connection("held-live".into())) + .await + .expect("acquiring the single live connection hung")?; + + // Give the losing connect slot ample time to fail its cycle (default + // retry is a single ~1s-connect-timeout attempt) and set the + // per-endpoint failed flag; pre-fix this is what armed the trap. + tokio::time::sleep(Duration::from_secs(3)).await; + + // A new request must queue, not fail fast: the pool is busy, not dead. + let cm_for_queued = Arc::clone(&cm); + let queued = + tokio::spawn(async move { cm_for_queued.connection("queued-while-busy".into()).await }); + tokio::time::sleep(Duration::from_millis(500)).await; + assert_eq!( + queued.is_finished(), + false, + "connection() failed fast while a live connection was checked out — \ + busy-but-healthy pool wrongly declared dead", + ); + + // Releasing the held connection must let the queued request through. + drop(held); + let conn = timeout(Duration::from_secs(5), queued) + .await + .expect("queued request did not resolve after the live connection was released") + .unwrap()?; + drop(conn); + Ok(()) +} + +/// The teardown-to-zero window: when the LAST live channel is torn down +/// by an in-use transport error while every endpoint's last connect cycle +/// had already failed, queued waiters must be flushed with `Unavailable` +/// at that instant — not parked until the freshly scheduled reconnect +/// fails a whole connect cycle later (`handle_worker` already fails new +/// requests fast in that state; queued ones must not be second-class). +/// +/// Setup: one endpoint, two connection slots, both piped through a proxy +/// to a real gRPC server; `max_concurrent_requests` of 2 with both +/// permits held so returned channel clones can never be handed to the +/// queued request. The retrier is configured so a connect cycle to the +/// (killed) endpoint takes ~3s: after severing both connections, an RPC +/// on the second held connection tears its channel down and its +/// reconnect cycle fails, arming the endpoint's connect-failed flag +/// while the first channel is still live. A request is then queued, and +/// an RPC on the first held connection tears the last live channel down +/// to zero. Post-fix the queued request errors `Unavailable` right +/// there; pre-fix nothing flushed it until the next connect cycle failed +/// ~3s later, so the 1.5s deadline below distinguishes the two. +#[nativelink_test] +async fn last_live_channel_teardown_flushes_queued_requests() -> Result<(), Error> { + let backend = fake_grpc_server_endpoint().await; + let backend_uri = backend.uri().clone(); + let backend_addr = format!( + "{}:{}", + backend_uri.host().unwrap(), + backend_uri.port_u16().unwrap() + ); + + let proxy_listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let proxy_port = proxy_listener.local_addr().unwrap().port(); + let (kill_tx, kill_rx) = tokio::sync::watch::channel(false); + background_spawn!("teardown_test_proxy", async move { + // Pipe exactly the two initial connections (one per slot). On the + // kill signal, drop the listener first (so reconnects are refused, + // failing their connect cycles) and then sever both pipes. + for _ in 0..2 { + let (mut inbound, _) = proxy_listener.accept().await.unwrap(); + let mut rx = kill_rx.clone(); + let addr = backend_addr.clone(); + background_spawn!("teardown_test_pipe", async move { + let mut outbound = tokio::net::TcpStream::connect(addr).await.unwrap(); + tokio::select! { + _ = tokio::io::copy_bidirectional(&mut inbound, &mut outbound) => {} + _ = rx.changed() => {} + } + }); + } + let mut rx = kill_rx.clone(); + drop(rx.changed().await); + drop(proxy_listener); + // Pipes observe the same signal and sever their connections. + }); + + let endpoint = Endpoint::from_shared(format!("http://127.0.0.1:{proxy_port}")) + .unwrap() + .connect_timeout(Duration::from_secs(1)); + // delay 1.5 with one retry makes a failing connect cycle take one + // ~3s backoff sleep between two instant ECONNREFUSED attempts: long + // enough that "flushed at teardown" and "flushed by the next connect + // failure" are cleanly distinguishable below. + let retry = Retry { + max_retries: 1, + delay: 1.5, + ..Default::default() + }; + let cm = Arc::new(ConnectionManager::new( + vec![endpoint], + /* connections_per_endpoint = */ 2, + /* max_concurrent_requests = */ 2, + retry, + no_jitter(), + )); + + // Hold both permits (and both live channels). + let held1 = timeout(Duration::from_secs(10), cm.connection("held-1".into())) + .await + .expect("acquiring the first piped connection hung")?; + let held2 = timeout(Duration::from_secs(10), cm.connection("held-2".into())) + .await + .expect("acquiring the second piped connection hung")?; + + // Kill the endpoint: refuse new connects and sever both connections. + tokio::time::sleep(Duration::from_millis(500)).await; + kill_tx.send(true).unwrap(); + tokio::time::sleep(Duration::from_millis(100)).await; + + // Tear down the second channel with an in-use error; keep the client + // (and its permit) alive. Its reconnect cycle fails ~3s later, + // arming the endpoint's connect-failed flag while the first channel + // keeps live_channels at 1. + let mut client2 = + nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient::new(held2); + let rpc2 = timeout( + Duration::from_secs(5), + client2.query_write_status(QueryWriteStatusRequest { + resource_name: "poke-2".into(), + }), + ) + .await + .expect("RPC on the severed second connection hung"); + assert_eq!(rpc2.is_err(), true, "RPC on a severed connection succeeded"); + tokio::time::sleep(Duration::from_secs(4)).await; + + // Queue a request; it must wait (one channel still live, pool busy). + let cm_for_queued = Arc::clone(&cm); + let queued = + tokio::spawn(async move { cm_for_queued.connection("queued-behind-held".into()).await }); + tokio::time::sleep(Duration::from_millis(300)).await; + assert_eq!( + queued.is_finished(), + false, + "queued request resolved while a channel was live and permits were held", + ); + + // Tear down the LAST live channel with an in-use error. + let mut client1 = + nativelink_proto::google::bytestream::byte_stream_client::ByteStreamClient::new(held1); + let rpc1 = timeout( + Duration::from_secs(5), + client1.query_write_status(QueryWriteStatusRequest { + resource_name: "poke-1".into(), + }), + ) + .await + .expect("RPC on the severed first connection hung"); + assert_eq!(rpc1.is_err(), true, "RPC on a severed connection succeeded"); + + // The queued request must be flushed with Unavailable at the + // teardown itself — the next connect-cycle failure is ~3s away. + let queued_result = timeout(Duration::from_millis(1500), queued) + .await + .expect( + "queued request still parked after the last live channel was torn down \ + with all endpoints' connect cycles failed", + ) + .unwrap(); + assert_eq!( + queued_result.err().map(|err| err.code), + Some(nativelink_error::Code::Unavailable), + "expected Unavailable for the queued request after teardown-to-zero", + ); + drop(client1); + drop(client2); + Ok(()) +} + #[nativelink_test] async fn extra_request_above_max_blocks_until_a_release() -> Result<(), Error> { const MAX_CONCURRENT: usize = 2; From 59937482409134cbcfa6b5b9f917a79d35af9c08 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Sat, 4 Jul 2026 10:10:34 -0400 Subject: [PATCH 27/46] nativelink: add --check for offline config validation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `nativelink ` previously validated a config only by booting it (binding listeners, connecting to backends), so there was no way to check a config in CI without a live environment. Add `nativelink --check `: it parses the config (already running check_store_conflict) and resolves every store/scheduler name reference — the typo'd-reference class that otherwise only surfaces at boot — then exits without constructing the store manager, binding a socket, connecting to any backend, or creating store directories. Success prints a one-line summary to stdout (exit 0); unresolved references and duplicate store/scheduler names are collected and reported to stderr, each with its config path (exit 1). CasConfig::validate_references walks workers and server services, and recurses ref_store references for the transitive closure of stores that are actually wired to a consumer. It is conservative to avoid false positives: optional/defaulted refs are checked only when present, and store-internal refs of unwired stores are skipped — so all 17 shipped example configs pass --check (proven offline via strace: zero network connects), while a dangling reference is named and fails. Co-Authored-By: Claude Fable 5 --- .../vocabularies/TraceMachina/accept.txt | 3 + nativelink-config/src/cas_server.rs | 695 +++++++++++++++++- nativelink-config/tests/json5_test.rs | 5 +- src/bin/nativelink.rs | 53 +- 4 files changed, 749 insertions(+), 7 deletions(-) diff --git a/.github/styles/config/vocabularies/TraceMachina/accept.txt b/.github/styles/config/vocabularies/TraceMachina/accept.txt index c0832d0ef15..71a32a34411 100644 --- a/.github/styles/config/vocabularies/TraceMachina/accept.txt +++ b/.github/styles/config/vocabularies/TraceMachina/accept.txt @@ -493,3 +493,6 @@ Flybody MuJoCo connectome simulatable +[Rr]eachability +recursed +recursing diff --git a/nativelink-config/src/cas_server.rs b/nativelink-config/src/cas_server.rs index 1b4aac223fe..c0cc48c0afc 100644 --- a/nativelink-config/src/cas_server.rs +++ b/nativelink-config/src/cas_server.rs @@ -12,7 +12,7 @@ // See the License for the specific language governing permissions and // limitations under the License. -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err}; #[cfg(feature = "dev-schema")] @@ -20,7 +20,7 @@ use schemars::JsonSchema; use serde::{Deserialize, Serialize}; use tracing::warn; -use crate::schedulers::SchedulerSpec; +use crate::schedulers::{ExperimentalSimpleSchedulerBackend, SchedulerSpec}; use crate::serde_utils::{ convert_boolean_with_shellexpand, convert_data_size_with_shellexpand, convert_duration_with_shellexpand, convert_numeric_with_shellexpand, @@ -2007,6 +2007,347 @@ impl CasConfig { } Ok(()) } + + /// Offline structural validation of every store/scheduler name reference + /// in this config. This is a pure, read-only check intended for + /// `nativelink --check`: it constructs no stores, binds no sockets, opens + /// no connections, and creates no directories. + /// + /// Every store name referenced by a worker or by a server service, and + /// every scheduler name referenced by a service, must resolve to a store + /// declared in `stores` / a scheduler declared in `schedulers`. Store + /// references made *inside* a declared store's spec (the `ref_store` + /// wrappers such as `fast_slow`, `dedup`, `completeness_checking`, ...) + /// are also resolved, but only for the transitive closure of stores that + /// are actually wired to a worker or service. Store definitions that are + /// never referenced by any consumer (for example a pure store-catalog + /// sample config) are intentionally not walked, so their illustrative + /// placeholder references are never flagged. + /// + /// All problems are collected and reported together, each with a path + /// such as `servers[0].services.cas[main].cas_store references + /// undefined store 'CAS_STORE'`, rather than failing on the first one. + /// + /// Duplicate store names and duplicate scheduler names are also reported, + /// because [`crate::stores`] / the store manager resolve names last-wins + /// and would otherwise silently drop the earlier definition. + /// + /// # Errors + /// + /// Returns `Err(Code::InvalidArgument)` if any referenced store/scheduler + /// name is undeclared, or if a store/scheduler name is declared more than + /// once. + pub fn validate_references(&self) -> Result<(), Error> { + let mut problems: Vec = Vec::new(); + + // Declared registries. Collision on insert is a duplicate name. + let mut store_names: HashSet<&str> = HashSet::new(); + for store in &self.stores { + if !store_names.insert(store.name.as_str()) { + problems.push(format!("duplicate store name '{}'", store.name)); + } + } + let mut scheduler_names: HashSet<&str> = HashSet::new(); + for scheduler in self.schedulers.iter().flatten() { + if !scheduler_names.insert(scheduler.name.as_str()) { + problems.push(format!("duplicate scheduler name '{}'", scheduler.name)); + } + } + + // Only validate a reference category when its registry is non-empty. + // A config that declares zero stores (or zero schedulers) is a + // fragment or legacy-format sample whose backing definitions live + // elsewhere; every reference would then trivially be "unresolved" and + // flagging them would be a false positive on an otherwise-valid + // sample. Real deployments always declare their stores, so genuine + // typos are still caught. + let mut checker = ReferenceChecker { + check_stores: !store_names.is_empty(), + check_schedulers: !scheduler_names.is_empty(), + store_names, + scheduler_names, + problems, + used_stores: Vec::new(), + }; + + for (worker_idx, worker) in self.workers.iter().flatten().enumerate() { + let WorkerConfig::Local(local) = worker; + checker.store( + &format!("workers[{worker_idx}].local.cas_fast_slow_store"), + &local.cas_fast_slow_store, + ); + if let Some(ac_store) = &local.upload_action_result.ac_store { + checker.store( + &format!("workers[{worker_idx}].local.upload_action_result.ac_store"), + ac_store, + ); + } + if let Some(historical) = &local.upload_action_result.historical_results_store { + checker.store( + &format!( + "workers[{worker_idx}].local.upload_action_result.historical_results_store" + ), + historical, + ); + } + } + + if let Some(origin_events) = &self.experimental_origin_events { + checker.store( + "experimental_origin_events.publisher.store", + &origin_events.publisher.store, + ); + } + + for (server_idx, server) in self.servers.iter().enumerate() { + let Some(services) = &server.services else { + continue; + }; + let prefix = format!("servers[{server_idx}].services"); + for entry in services.cas.iter().flatten() { + let name = &entry.instance_name; + checker.store(&format!("{prefix}.cas[{name}].cas_store"), &entry.cas_store); + if let Some(index_store) = entry + .experimental_chunking + .as_ref() + .and_then(|chunking| chunking.index_store.as_ref()) + { + checker.store( + &format!("{prefix}.cas[{name}].experimental_chunking.index_store"), + index_store, + ); + } + } + for entry in services.ac.iter().flatten() { + let name = &entry.instance_name; + checker.store(&format!("{prefix}.ac[{name}].ac_store"), &entry.ac_store); + } + for entry in services.execution.iter().flatten() { + let name = &entry.instance_name; + checker.store( + &format!("{prefix}.execution[{name}].cas_store"), + &entry.cas_store, + ); + checker.scheduler( + &format!("{prefix}.execution[{name}].scheduler"), + &entry.scheduler, + ); + } + for entry in services.capabilities.iter().flatten() { + if let Some(remote_execution) = &entry.remote_execution { + let name = &entry.instance_name; + checker.scheduler( + &format!("{prefix}.capabilities[{name}].remote_execution.scheduler"), + &remote_execution.scheduler, + ); + } + } + for entry in services.bytestream.iter().flatten() { + let name = &entry.instance_name; + checker.store( + &format!("{prefix}.bytestream[{name}].cas_store"), + &entry.cas_store, + ); + } + for entry in services.fetch.iter().flatten() { + let name = &entry.instance_name; + checker.store( + &format!("{prefix}.fetch[{name}].fetch_store"), + &entry.fetch_store, + ); + } + for entry in services.push.iter().flatten() { + let name = &entry.instance_name; + checker.store( + &format!("{prefix}.push[{name}].push_store"), + &entry.push_store, + ); + } + if let Some(worker_api) = &services.worker_api { + checker.scheduler( + &format!("{prefix}.worker_api.scheduler"), + &worker_api.scheduler, + ); + } + if let Some(bep) = &services.experimental_bep { + checker.store(&format!("{prefix}.experimental_bep.store"), &bep.store); + } + } + + // Scheduler specs themselves reference stores (and nest further + // schedulers); walk every declared scheduler so those references are + // resolved and their stores seed the `ref_store` reachability walk. + for scheduler in self.schedulers.iter().flatten() { + checker.scheduler_spec(&format!("schedulers[{}]", scheduler.name), &scheduler.spec); + } + + // Resolve `ref_store` references reachable from the wired-up stores. + if checker.check_stores { + let store_specs: HashMap<&str, &StoreSpec> = self + .stores + .iter() + .map(|store| (store.name.as_str(), &store.spec)) + .collect(); + checker.walk_used_stores(&store_specs); + } + + if checker.problems.is_empty() { + Ok(()) + } else { + Err(make_err!( + Code::InvalidArgument, + "configuration reference validation failed:\n - {}", + checker.problems.join("\n - ") + )) + } + } +} + +/// Recursively collects every `ref_store` name that appears anywhere inside a +/// single [`StoreSpec`] tree (the wrapper stores nest other `StoreSpec`s). The +/// match is exhaustive on purpose so a newly added wrapping store forces this +/// to be revisited. +fn collect_ref_store_names<'a>(spec: &'a StoreSpec, out: &mut Vec<&'a StoreRefName>) { + match spec { + StoreSpec::RefStore(ref_spec) => out.push(&ref_spec.name), + StoreSpec::CacheMetrics(inner) => collect_ref_store_names(&inner.backend, out), + StoreSpec::Verify(inner) => collect_ref_store_names(&inner.backend, out), + StoreSpec::CompletenessChecking(inner) => { + collect_ref_store_names(&inner.backend, out); + collect_ref_store_names(&inner.cas_store, out); + } + StoreSpec::Compression(inner) => collect_ref_store_names(&inner.backend, out), + StoreSpec::Dedup(inner) => { + collect_ref_store_names(&inner.index_store, out); + collect_ref_store_names(&inner.content_store, out); + } + StoreSpec::ExistenceCache(inner) => collect_ref_store_names(&inner.backend, out), + StoreSpec::FastSlow(inner) => { + collect_ref_store_names(&inner.fast, out); + collect_ref_store_names(&inner.slow, out); + } + StoreSpec::Shard(inner) => { + for shard in &inner.stores { + collect_ref_store_names(&shard.store, out); + } + } + StoreSpec::SizePartitioning(inner) => { + collect_ref_store_names(&inner.lower_store, out); + collect_ref_store_names(&inner.upper_store, out); + } + // Leaf stores and stores whose backends are concrete (non-`StoreSpec`) + // specs carry no nested `ref_store` names. + StoreSpec::Memory(_) + | StoreSpec::ExperimentalCloudObjectStore(_) + | StoreSpec::OntapS3ExistenceCache(_) + | StoreSpec::Filesystem(_) + | StoreSpec::Grpc(_) + | StoreSpec::RedisStore(_) + | StoreSpec::Noop(_) + | StoreSpec::ExperimentalMongo(_) => {} + } +} + +/// Accumulates unresolved-reference problems while walking a [`CasConfig`]. +struct ReferenceChecker<'a> { + store_names: HashSet<&'a str>, + scheduler_names: HashSet<&'a str>, + check_stores: bool, + check_schedulers: bool, + problems: Vec, + /// Declared store names reached from a worker/service, used to seed the + /// `ref_store` reachability walk. + used_stores: Vec<&'a str>, +} + +impl<'a> ReferenceChecker<'a> { + /// Records that `name` (found at `path`) must resolve to a declared store. + fn store(&mut self, path: &str, name: &'a str) { + if !self.check_stores { + return; + } + if self.store_names.contains(name) { + self.used_stores.push(name); + } else { + self.problems + .push(format!("{path} references undefined store '{name}'")); + } + } + + /// Records that `name` (found at `path`) must resolve to a declared + /// scheduler. + fn scheduler(&mut self, path: &str, name: &str) { + if !self.check_schedulers { + return; + } + if !self.scheduler_names.contains(name) { + self.problems + .push(format!("{path} references undefined scheduler '{name}'")); + } + } + + /// Walks a [`SchedulerSpec`] tree, resolving every store reference it + /// carries and recursing into nested schedulers. The match is exhaustive + /// on purpose so a newly added scheduler variant forces this to be + /// revisited (same design as [`collect_ref_store_names`]). + fn scheduler_spec(&mut self, path: &str, spec: &'a SchedulerSpec) { + match spec { + SchedulerSpec::Simple(simple) => match &simple.experimental_backend { + Some(ExperimentalSimpleSchedulerBackend::Redis(redis)) => { + self.store( + &format!("{path}.simple.experimental_backend.redis.redis_store"), + &redis.redis_store, + ); + } + Some(ExperimentalSimpleSchedulerBackend::Memory) | None => {} + }, + SchedulerSpec::CacheLookup(inner) => { + self.store(&format!("{path}.cache_lookup.ac_store"), &inner.ac_store); + self.scheduler_spec(&format!("{path}.cache_lookup.scheduler"), &inner.scheduler); + } + SchedulerSpec::PropertyModifier(inner) => { + self.scheduler_spec( + &format!("{path}.property_modifier.scheduler"), + &inner.scheduler, + ); + } + SchedulerSpec::HistoricalResource(inner) => { + self.scheduler_spec( + &format!("{path}.historical_resource.scheduler"), + &inner.scheduler, + ); + } + // Carries no store references and no nested scheduler. + SchedulerSpec::Grpc(_) => {} + } + } + + /// Walks the transitive closure of stores reachable from + /// [`Self::used_stores`], resolving each nested `ref_store` against the + /// declared store set. A visited set bounds the walk against `ref_store` + /// cycles. + fn walk_used_stores(&mut self, store_specs: &HashMap<&'a str, &'a StoreSpec>) { + let mut visited: HashSet<&'a str> = HashSet::new(); + while let Some(name) = self.used_stores.pop() { + if !visited.insert(name) { + continue; + } + let Some(&spec) = store_specs.get(name) else { + continue; + }; + let mut refs: Vec<&'a StoreRefName> = Vec::new(); + collect_ref_store_names(spec, &mut refs); + for ref_name in refs { + if self.store_names.contains(ref_name.as_str()) { + self.used_stores.push(ref_name.as_str()); + } else { + self.problems.push(format!( + "store '{name}' references undefined store '{ref_name}'" + )); + } + } + } + } } #[cfg(test)] @@ -2290,4 +2631,354 @@ mod tests { assert!(logs_contain("instance_name=\"opt-out\"")); assert!(logs_contain(CasConfig::ZSTD_COMPRESSION_DOCS_URL)); } + + // ---------------------------------------------------------------------- + // `validate_references` (offline `nativelink --check`) tests. + // ---------------------------------------------------------------------- + + fn parse_cas_config(json5: &str) -> CasConfig { + serde_json5::from_str(json5).expect("valid CasConfig json5") + } + + /// A fully-resolved config that exercises worker, service and nested + /// `ref_store` reference sites. `WRAPPED` wraps a `ref_store` to `CAS`, + /// so the reachability walk must resolve it. + const VALID_CONFIG: &str = r#"{ + stores: [ + { name: "CAS", memory: {} }, + { name: "AC", memory: {} }, + { name: "WRAPPED", fast_slow: { + fast: { ref_store: { name: "CAS" } }, + slow: { noop: {} }, + } }, + ], + schedulers: [{ name: "SCHED", simple: {} }], + workers: [{ + local: { + worker_api_endpoint: { uri: "grpc://127.0.0.1:50061" }, + cas_fast_slow_store: "WRAPPED", + upload_action_result: { ac_store: "AC" }, + work_directory: "/tmp/work", + platform_properties: {}, + }, + }], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { + cas: [{ instance_name: "main", cas_store: "CAS" }], + ac: [{ instance_name: "main", ac_store: "AC" }], + execution: [{ instance_name: "main", cas_store: "CAS", scheduler: "SCHED" }], + capabilities: [{ instance_name: "main", remote_execution: { scheduler: "SCHED" } }], + bytestream: [{ instance_name: "main", cas_store: "CAS" }], + worker_api: { scheduler: "SCHED" }, + }, + }], + }"#; + + #[test] + fn validate_references_accepts_resolved_config() { + parse_cas_config(VALID_CONFIG) + .validate_references() + .expect("every reference resolves"); + } + + #[test] + fn validate_references_flags_undefined_service_store() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { cas: [{ instance_name: "main", cas_store: "CAS_TYPO" }] }, + }], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + assert_eq!(err.code, Code::InvalidArgument); + let message = err.messages.join("\n"); + assert!( + message.contains("CAS_TYPO"), + "error must name the bad store ref: {message}" + ); + assert!( + message.contains("servers[0].services.cas[main].cas_store"), + "error must name the reference site: {message}" + ); + } + + #[test] + fn validate_references_flags_undefined_scheduler() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + schedulers: [{ name: "SCHED", simple: {} }], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { + execution: [{ + instance_name: "main", + cas_store: "CAS", + scheduler: "SCHED_TYPO", + }], + }, + }], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + assert_eq!(err.code, Code::InvalidArgument); + let message = err.messages.join("\n"); + assert!( + message.contains("SCHED_TYPO"), + "error must name the bad scheduler ref: {message}" + ); + assert!( + message.contains("scheduler"), + "error must name the reference site: {message}" + ); + } + + #[test] + fn validate_references_flags_undefined_worker_store() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + workers: [{ + local: { + worker_api_endpoint: { uri: "grpc://127.0.0.1:50061" }, + cas_fast_slow_store: "WORKER_TYPO", + work_directory: "/tmp/work", + platform_properties: {}, + }, + }], + servers: [], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("WORKER_TYPO") + && message.contains("workers[0].local.cas_fast_slow_store"), + "error must name the bad worker store ref and site: {message}" + ); + } + + #[test] + fn validate_references_flags_nested_ref_store_typo() { + // `WRAPPED` is wired to a service, so its nested `ref_store` typo must + // be resolved and flagged via the reachability walk. + let cfg = parse_cas_config( + r#"{ + stores: [ + { name: "CAS", memory: {} }, + { name: "WRAPPED", fast_slow: { + fast: { ref_store: { name: "CAS_TYPO" } }, + slow: { noop: {} }, + } }, + ], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { cas: [{ instance_name: "main", cas_store: "WRAPPED" }] }, + }], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("CAS_TYPO") && message.contains("store 'WRAPPED'"), + "error must name the nested bad ref and its owning store: {message}" + ); + } + + #[test] + fn validate_references_flags_undefined_scheduler_redis_store() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + schedulers: [{ + name: "SCHED", + simple: { + experimental_backend: { + redis: { redis_store: "REDIS_TYPO" }, + }, + }, + }], + servers: [], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("REDIS_TYPO") + && message + .contains("schedulers[SCHED].simple.experimental_backend.redis.redis_store"), + "error must name the bad redis store ref and its site: {message}" + ); + } + + #[test] + fn validate_references_flags_undefined_cache_lookup_ac_store() { + // The nested scheduler inside `cache_lookup` must be recursed into, + // and its `ac_store` reference resolved. + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + schedulers: [{ + name: "SCHED", + property_modifier: { + modifications: [], + scheduler: { + cache_lookup: { + ac_store: "AC_TYPO", + scheduler: { simple: {} }, + }, + }, + }, + }], + servers: [], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("AC_TYPO") + && message.contains( + "schedulers[SCHED].property_modifier.scheduler.cache_lookup.ac_store" + ), + "error must name the nested bad ac_store ref and its site: {message}" + ); + } + + #[test] + fn validate_references_flags_undefined_chunking_index_store() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { + cas: [{ + instance_name: "main", + cas_store: "CAS", + experimental_chunking: { index_store: "INDEX_TYPO" }, + }], + }, + }], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("INDEX_TYPO") + && message + .contains("servers[0].services.cas[main].experimental_chunking.index_store"), + "error must name the bad index store ref and its site: {message}" + ); + } + + #[test] + fn validate_references_flags_duplicate_store_name() { + let cfg = parse_cas_config( + r#"{ + stores: [ + { name: "CAS", memory: {} }, + { name: "CAS", memory: {} }, + ], + servers: [], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("duplicate store name 'CAS'"), + "error must report the duplicate store: {message}" + ); + } + + #[test] + fn validate_references_flags_duplicate_scheduler_name() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + schedulers: [ + { name: "SCHED", simple: {} }, + { name: "SCHED", simple: {} }, + ], + servers: [], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + assert!( + message.contains("duplicate scheduler name 'SCHED'"), + "error must report the duplicate scheduler: {message}" + ); + } + + #[test] + fn validate_references_collects_all_problems() { + let cfg = parse_cas_config( + r#"{ + stores: [{ name: "CAS", memory: {} }], + schedulers: [{ name: "SCHED", simple: {} }], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { + cas: [{ instance_name: "main", cas_store: "BAD_CAS" }], + ac: [{ instance_name: "main", ac_store: "BAD_AC" }], + }, + }], + }"#, + ); + let err = cfg.validate_references().unwrap_err(); + let message = err.messages.join("\n"); + // Reporting is not fail-fast: both bad references appear together. + assert!( + message.contains("BAD_CAS") && message.contains("BAD_AC"), + "error must collect every unresolved reference: {message}" + ); + } + + #[test] + fn validate_references_skips_empty_store_registry() { + // A fragment/legacy sample with no declared stores must not be + // flagged: every reference would trivially be unresolved, which would + // be a false positive on an otherwise-valid sample. + let cfg = parse_cas_config( + r#"{ + stores: [], + servers: [{ + listener: { http: { socket_address: "0.0.0.0:50051" } }, + services: { + cas: [{ instance_name: "", cas_store: "CAS_MAIN_STORE" }], + execution: [{ + instance_name: "", + cas_store: "WORKER_STORE", + scheduler: "MAIN_SCHEDULER", + }], + }, + }], + }"#, + ); + cfg.validate_references() + .expect("empty registries defer to boot, not a false positive"); + } + + #[test] + fn validate_references_ignores_unreferenced_store_internal_refs() { + // A pure store-catalog config (no servers, no workers) may carry + // illustrative placeholder `ref_store` names; because nothing wires + // those stores to a consumer, their internal refs are not walked and + // must not be flagged. + let cfg = parse_cas_config( + r#"{ + stores: [ + { name: "REAL", memory: {} }, + { name: "SHOWCASE", ref_store: { name: "PLACEHOLDER_NOT_DECLARED" } }, + ], + servers: [], + }"#, + ); + cfg.validate_references() + .expect("unreferenced catalog stores are not walked"); + } } diff --git a/nativelink-config/tests/json5_test.rs b/nativelink-config/tests/json5_test.rs index f1dadde10ab..3f2b4ccdd24 100644 --- a/nativelink-config/tests/json5_test.rs +++ b/nativelink-config/tests/json5_test.rs @@ -38,8 +38,11 @@ fn test_example_parsing() { if !config_file.contains(".json5") { continue; } - CasConfig::try_from_json5_file(&config_file) + let config = CasConfig::try_from_json5_file(&config_file) .unwrap_or_else(|e| panic!("Error while reading {config_file}: {e}")); + config + .validate_references() + .unwrap_or_else(|e| panic!("Unresolved reference in {config_file}: {e}")); found_at_least_one_entry = true; } diff --git a/src/bin/nativelink.rs b/src/bin/nativelink.rs index 4a796002009..7ee55120aeb 100644 --- a/src/bin/nativelink.rs +++ b/src/bin/nativelink.rs @@ -128,6 +128,10 @@ struct Args { /// Config file to use. #[clap(value_parser)] config_file: String, + + /// Validate the configuration and exit without starting any services. + #[clap(long)] + check: bool, } trait RoutesExt { @@ -881,14 +885,57 @@ async fn inner_main( Ok(()) } -fn get_config() -> Result { +fn get_config() -> Result<(Args, CasConfig), Error> { let args = Args::parse(); - CasConfig::try_from_json5_file(&args.config_file) + let cfg = CasConfig::try_from_json5_file(&args.config_file) + .err_tip(|| format!("While loading config file {}", args.config_file))?; + Ok((args, cfg)) } fn main() -> Result<(), Box> { install_default_rustls_crypto_provider(); + // Parse args and load the config *before* building any runtime so that + // `--check` can run a fully offline validation and exit without starting + // a serving runtime, binding a socket, connecting to a backend, or + // creating any store directories. + let (args, mut cfg) = get_config()?; + + if args.check { + return match cfg.validate_references() { + Ok(()) => { + let num_schedulers = cfg.schedulers.as_ref().map_or(0, Vec::len); + #[allow( + clippy::print_stdout, + reason = "`--check` reports its success summary on stdout for CI" + )] + { + println!( + "OK: {} — {} stores, {} schedulers, {} servers, all references resolve", + args.config_file, + cfg.stores.len(), + num_schedulers, + cfg.servers.len(), + ); + } + Ok(()) + } + Err(err) => { + #[allow( + clippy::print_stderr, + reason = "`--check` reports validation failures on stderr for CI" + )] + { + eprintln!("FAIL: {} — configuration is invalid", args.config_file); + for message in &err.messages { + eprintln!("{message}"); + } + } + std::process::exit(1); + } + }; + } + // Set QoS to USER_INITIATED on the main thread *before* the tokio // runtime is built so the spawned worker threads inherit P-core // scheduling preference via pthread QoS inheritance on Apple @@ -910,8 +957,6 @@ fn main() -> Result<(), Box> { #[expect(clippy::disallowed_methods, reason = "tracing init on main runtime")] runtime.block_on(async { tokio::spawn(async { init_tracing().await }).await? })?; - let mut cfg = get_config()?; - let global_cfg = if let Some(global_cfg) = &mut cfg.global { if global_cfg.max_open_files == 0 { global_cfg.max_open_files = fs::DEFAULT_OPEN_FILE_LIMIT; From 25dd77fe84d3facb33d370b853c21a85fac837ca Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 2 Oct 2026 16:08:10 -0400 Subject: [PATCH 28/46] fix(store): wake in-flight slow-write waiters when any upload completes the digest An existence check that found a slow-store write in flight parked on that specific writer's completion mutex. A ByteStream upload that is interrupted (cancel, half-close, client deadline) is parked for resumption, keeping its FastSlowStore::update() future -- and thus the locked mutex -- alive. A later registration for the same key overwrote the map entry, so completing the digest under a fresh upload id never woke the original waiters: they starved until their client deadline even though the blob was present and new lookups succeeded. Replace the per-writer mutex with a per-key slot shared by all concurrent writers of that key, broadcasting Pending/Done over a watch channel. Any writer completing the digest wakes every waiter; the entry is removed (and waiters released with "not written") only when the last writer deregisters, which also fixes the stale-remove where a second writer's guard evicted a slot the first writer still occupied. Regression: stalled_upload_waiter_released_by_fresh_upload_of_same_digest -- fails on v1.7.3 (waiter hangs past a 2s deadline), passes with the fix; also asserts a genuinely-absent digest still resolves promptly. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/fast_slow_store.rs | 103 ++++++++++++++---- .../tests/fast_slow_store_test.rs | 67 ++++++++++++ 2 files changed, 149 insertions(+), 21 deletions(-) diff --git a/nativelink-store/src/fast_slow_store.rs b/nativelink-store/src/fast_slow_store.rs index 08d91800559..d1b97bd6a9f 100644 --- a/nativelink-store/src/fast_slow_store.rs +++ b/nativelink-store/src/fast_slow_store.rs @@ -77,7 +77,33 @@ pub struct FastSlowStore { // concurrent writer that has not yet finished pushing to the slow store // is still visible to a concurrent existence check, preventing redundant // duplicate uploads of the same blob. - in_flight_slow_writes: Mutex, Arc>>>, + // + // Each key maps to a single shared slot that *all* concurrent writers of + // that key register on. This is load-bearing for liveness: waiters key on + // the digest, not on any individual upload, so a second upload (e.g. a + // retry under a fresh ByteStream upload id after the first was + // interrupted and parked for resumption) completing the blob wakes + // waiters that subscribed while the first upload was in flight. + in_flight_slow_writes: Mutex, InFlightSlowWrite>>, +} + +/// State of an in-flight slow-store write, broadcast to existence-check +/// waiters through a `watch` channel. +#[derive(Clone, Copy, Debug)] +enum InFlightWriteState { + /// At least one writer is still streaming this key to the slow store. + Pending, + /// The write finished: `Some(size)` if a writer completed successfully, + /// `None` if every writer gave up without completing. + Done(MaybeSize), +} + +/// Per-key slot shared by every concurrent writer of that key. +#[derive(Debug)] +struct InFlightSlowWrite { + tx: tokio::sync::watch::Sender, + /// Number of live [`InFlightSlowWriteGuard`]s registered on this slot. + writers: usize, } // This guard ensures that the populating_digests is cleared even if the future @@ -104,19 +130,22 @@ impl LoaderGuard<'_> { } } -/// Cancel-safe RAII guard that removes an entry from -/// `FastSlowStore::in_flight_slow_writes` when dropped. This ensures the -/// map does not leak entries if the surrounding `update()` future is -/// cancelled before the slow-store write completes. +/// Cancel-safe RAII guard that deregisters a writer from the key's shared +/// slot in `FastSlowStore::in_flight_slow_writes` when dropped. The map entry +/// is removed only when the *last* writer for the key deregisters, so a +/// parked/cancelled upload cannot orphan waiters nor evict a slot a newer +/// upload of the same key is still using. struct InFlightSlowWriteGuard { weak_store: Weak, key: Option>, - write_complete_guard: tokio::sync::OwnedMutexGuard, + tx: tokio::sync::watch::Sender, } impl InFlightSlowWriteGuard { - fn complete(mut self, size: MaybeSize) { - *self.write_complete_guard = size; + /// Marks the write done and wakes all waiters; deregistration happens in + /// `Drop`. + fn complete(self, size: MaybeSize) { + self.tx.send_replace(InFlightWriteState::Done(size)); } } @@ -128,7 +157,27 @@ impl Drop for InFlightSlowWriteGuard { let Some(key) = self.key.take() else { return; }; - store.in_flight_slow_writes.lock().remove(&key); + let mut map = store.in_flight_slow_writes.lock(); + let std::collections::hash_map::Entry::Occupied(mut occupied_entry) = map.entry(key) else { + return; + }; + // Only touch the slot this guard registered on (defense against a + // same-key slot that was replaced out from under us). + if !occupied_entry.get().tx.same_channel(&self.tx) { + return; + } + let slot = occupied_entry.get_mut(); + slot.writers -= 1; + if slot.writers == 0 { + // Last writer for this key. If nobody completed the write, wake + // waiters with "not written" so they do not hang forever. + slot.tx.send_modify(|state| { + if matches!(state, InFlightWriteState::Pending) { + *state = InFlightWriteState::Done(None); + } + }); + occupied_entry.remove(); + } } } @@ -175,18 +224,18 @@ impl FastSlowStore { fn register_in_flight_slow_write(&self, key: StoreKey<'_>) -> InFlightSlowWriteGuard { let owned = key.into_owned(); - let write_complete = Arc::new(tokio::sync::Mutex::new(None)); - let write_complete_guard = write_complete - .clone() - .try_lock_owned() - .expect("Newly created mutex is locked"); - self.in_flight_slow_writes - .lock() - .insert(owned.borrow().into_owned(), write_complete); + let mut map = self.in_flight_slow_writes.lock(); + let slot = map + .entry(owned.borrow().into_owned()) + .or_insert_with(|| InFlightSlowWrite { + tx: tokio::sync::watch::channel(InFlightWriteState::Pending).0, + writers: 0, + }); + slot.writers += 1; InFlightSlowWriteGuard { weak_store: self.weak_self.clone(), key: Some(owned), - write_complete_guard, + tx: slot.tx.clone(), } } @@ -560,9 +609,21 @@ impl StoreDriver for FastSlowStore { for (i, (k, result)) in key.iter().zip(results.iter_mut()).enumerate() { if result.is_none() { let owned = k.borrow().into_owned(); - if let Some(maybe_size) = in_flight.get(&owned) { - let maybe_size = maybe_size.clone(); - in_flight_futs.push(async move { (i, *maybe_size.lock().await) }); + if let Some(slot) = in_flight.get(&owned) { + let mut rx = slot.tx.subscribe(); + in_flight_futs.push(async move { + loop { + let state = *rx.borrow_and_update(); + match state { + InFlightWriteState::Done(size) => return (i, size), + InFlightWriteState::Pending => {} + } + if rx.changed().await.is_err() { + // All writers dropped without completing. + return (i, None); + } + } + }); } } } diff --git a/nativelink-store/tests/fast_slow_store_test.rs b/nativelink-store/tests/fast_slow_store_test.rs index 35a04d2511c..1514b815f67 100644 --- a/nativelink-store/tests/fast_slow_store_test.rs +++ b/nativelink-store/tests/fast_slow_store_test.rs @@ -2347,3 +2347,70 @@ async fn incomplete_slow_read_never_populates_memory_or_filesystem() -> Result<( } Ok(()) } + +// Regression test for N14: an interrupted ByteStream upload is parked for +// resumption, keeping its FastSlowStore::update() future (and in-flight +// slow-write registration) alive. Existence-check waiters that subscribed to +// that in-flight write must be woken when the SAME digest is completed by a +// DIFFERENT upload (fresh upload id), not starve until their client deadline. +#[nativelink_test] +async fn stalled_upload_waiter_released_by_fresh_upload_of_same_digest() -> Result<(), Error> { + let (fast_slow_store, _fast_store, _slow_store) = make_stores(); + let data = make_random_data(100); + let digest = DigestInfo::try_new(VALID_HASH, data.len())?; + + // Upload A: starts streaming, then stalls forever mid-stream (simulates a + // cancelled/half-closed ByteStream upload parked for resumption). We keep + // `stalled_tx` alive so the update() future stays pending. + let (mut stalled_tx, stalled_rx) = make_buf_channel_pair(); + let store_a = fast_slow_store.clone(); + let upload_a = tokio::spawn(async move { + store_a + .update(digest, stalled_rx, UploadSizeInfo::ExactSize(100)) + .await + }); + stalled_tx.send(Bytes::from(make_random_data(10))).await?; + // Let upload A register its in-flight slow write. + tokio::task::yield_now().await; + + // Waiter: FindMissingBlobs-style existence check that parks on the + // in-flight write of upload A. + let store_w = fast_slow_store.clone(); + let waiter = tokio::spawn(async move { store_w.has(digest).await }); + // Let the waiter subscribe to the in-flight write. + tokio::time::sleep(Duration::from_millis(50)).await; + assert!(!waiter.is_finished(), "Waiter should be parked on upload A"); + + // Upload B: completes the same digest under a "fresh upload id". + fast_slow_store + .update_oneshot(digest, data.clone().into()) + .await?; + + // The waiter must be released promptly by upload B's completion. + let result = tokio::time::timeout(Duration::from_secs(2), waiter) + .await + .map_err(|_| { + make_err!( + Code::DeadlineExceeded, + "Waiter hung after same digest was completed by a fresh upload" + ) + })? + .expect("waiter task panicked")?; + assert_eq!( + result, + Some(100), + "Waiter should see the completed blob size" + ); + + // A genuinely-absent digest must still resolve promptly once its only + // (stalled) writer goes away. + drop(stalled_tx); + drop(upload_a.await); + let absent = DigestInfo::try_new(VALID_HASH, 12345)?; + let absent_result = tokio::time::timeout(Duration::from_secs(2), fast_slow_store.has(absent)) + .await + .map_err(|_| make_err!(Code::DeadlineExceeded, "Absent digest check hung"))??; + assert_eq!(absent_result, None); + + Ok(()) +} From 2b5e9de65f2235467e700c6b38accc1eb1475a19 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 21 Oct 2026 17:25:22 +0000 Subject: [PATCH 29/46] scheduler: register worker_retry_reservation_test in the Bazel suite The N2 regression test was added under nativelink-scheduler/tests/ but never listed in the rust_test_suite, so Cargo ran it and Bazel (and CI) did not. List it alongside the other integration tests. Co-Authored-By: Claude Opus 5.5 --- nativelink-scheduler/BUILD.bazel | 1 + 1 file changed, 1 insertion(+) diff --git a/nativelink-scheduler/BUILD.bazel b/nativelink-scheduler/BUILD.bazel index 25077de27cc..114db264c1b 100644 --- a/nativelink-scheduler/BUILD.bazel +++ b/nativelink-scheduler/BUILD.bazel @@ -94,6 +94,7 @@ rust_test_suite( "tests/unsatisfiable_action_test.rs", "tests/worker_capability_index_test.rs", "tests/worker_registry_test.rs", + "tests/worker_retry_reservation_test.rs", "tests/worker_running_operations_report_test.rs", ], compile_data = [ From 7d610f7e1c455b904e2d97fe634819ef843d28f2 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 21 Oct 2026 17:26:45 +0000 Subject: [PATCH 30/46] docs(config-reference): regenerate main for `max_concurrent_deletes` Regenerated web/apps/docs/content/docs/reference/nativelink-config/main.mdx from the nativelink-config crate (build-schema + schema-to-mdx), as AGENTS.md requires when a config field or its doc comment changes. --- .../docs/content/docs/reference/nativelink-config/main.mdx | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx index 5d97a6a58b5..cb58419344b 100644 --- a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx +++ b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx @@ -5,7 +5,7 @@ full: true --- {/* AUTOGENERATED — do not edit by hand. - Source: nativelink-config @ main (6d6266cf) + Source: nativelink-config @ main (80678933f8) Regenerate from web/: bun --filter @nativelink/docs gen:config-reference */} @@ -830,6 +830,7 @@ See [`experimental_cloud_object_store`](#experimental_cloud_object_store) for de | `eviction_policy` | [EvictionPolicy](#evictionpolicy) | — | — | Policy used to evict items out of the store. Failure to set this value will cause items to never be removed from the store causing infinite memory usage. | | `block_size` | integer (uint64) | — | 4kb | The block size of the filesystem for the running machine value is used to determine an entry's actual size on disk consumed For a 4KB block size filesystem, a 1B file actually consumes 4KB | | `max_concurrent_writes` | integer (uint) | — | unlimited | Maximum number of concurrent write operations allowed. Each write involves streaming data to a temp file and calling `sync_all()`, which can saturate disk I/O when many writes happen simultaneously. Limiting concurrency prevents disk saturation from blocking the async runtime. A value of 0 means unlimited (no concurrency limit). | +| `max_concurrent_deletes` | integer (uint) | — | 1024 | Maximum number of concurrent background file-deletion operations. When an entry leaves the store, its backing file is unlinked by a background task. Without a bound, an eviction storm (cache wipe or a flood of overwrites) can spawn these without limit: memory grows with the queued task frames and disk/inode reclamation lags behind what the async runtime can schedule. Bounding it applies backpressure the same way `max_concurrent_writes` does — once this many deletes are in flight, a further delete runs synchronously inline in the evicting path instead of spawning, which throttles eviction to delete throughput and never leaks a file. | | `evict_page_cache` | boolean | — | false | When true, advise the kernel to drop the page cache for each blob after it is written or read (`posix_fadvise` with `POSIX_FADV_DONTNEED`). On real filesystems this takes a globally serialized, all-CPU kernel path that stalls on many-core hosts, and it evicts the page-cache tier that fast-disk deployments rely on. Leave off unless you specifically want to keep this store's I/O out of the page cache. | ## RefSpec From 279253c4ed76a384809f8715b6dc4a0735d4d9fe Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 21 Oct 2026 17:26:45 +0000 Subject: [PATCH 31/46] docs(config-reference): regenerate main for `multipart_part_size` and `max_concurrent_multipart_uploads` Regenerated web/apps/docs/content/docs/reference/nativelink-config/main.mdx from the nativelink-config crate (build-schema + schema-to-mdx), as AGENTS.md requires when a config field or its doc comment changes. --- .../docs/content/docs/reference/nativelink-config/main.mdx | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx index 5d97a6a58b5..32fcd5642e7 100644 --- a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx +++ b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx @@ -5,7 +5,7 @@ full: true --- {/* AUTOGENERATED — do not edit by hand. - Source: nativelink-config @ main (6d6266cf) + Source: nativelink-config @ main (5ac552526c) Regenerate from web/: bun --filter @nativelink/docs gen:config-reference */} @@ -1126,6 +1126,8 @@ The total delay is additive, so this example produces 9.525 to 15.875 s of total | `consider_expired_after_s` | integer (uint32) | — | 0 | If the number of seconds since the `last_modified` time of the object is greater than this value, the object will not be considered "existing". This allows for external tools to delete objects that have not been uploaded in a long time. If a client receives a `NotFound` the client should re-upload the object. | | `max_retry_buffer_per_request` | integer (uint) | — | 5MB | The maximum buffer size to retain in case of a retryable error during upload. Setting this to zero will disable upload buffering; this means that in the event of a failure during upload, the entire upload will be aborted and the client will likely receive an error. | | `multipart_max_concurrent_uploads` | integer (uint) | — | 10 | Maximum number of concurrent `UploadPart` requests per `MultipartUpload`. | +| `multipart_part_size` | integer (uint) | — | 64MB | Target size of each part in a multipart upload, growing only as needed to stay under the provider's 10,000-part ceiling. Parts sized at the 5 MiB provider floor turn a multi-GiB blob into thousands of tiny concurrent requests, which churns HTTP connection pools; a larger part keeps the request count low. Each queued or in-flight part is buffered in memory, so one multipart upload buffers up to roughly `(2 * multipart_max_concurrent_uploads + 1) * multipart_part_size`; `max_concurrent_multipart_uploads` bounds the aggregate across uploads. Values are clamped to the provider's 5 MiB minimum and 5 GiB maximum part sizes. | +| `max_concurrent_multipart_uploads` | integer (uint) | — | 4 | Store-wide cap on the number of blobs concurrently uploaded via the multipart path. Each multipart upload can buffer roughly `(2 * multipart_max_concurrent_uploads + 1) * multipart_part_size`, so this bounds the store's aggregate upload buffer (about 5 GiB at the defaults). Uploads over the 5 MiB multipart threshold wait for a slot; smaller blobs are unaffected. | | `insecure_allow_http` | boolean | — | false | Allow unencrypted HTTP connections. Only use this for local testing. | | `disable_http2` | boolean | — | false | Disable HTTP/2 connections and only use HTTP/1.1. Default client configuration will have HTTP/1.1 and HTTP/2 enabled for connection schemes. HTTP/2 should be disabled if environments have poor support or performance related to HTTP/2. Safe to keep default unless underlying network environment, S3, or GCS API servers specify otherwise. | | `throttle_retry_delay_s` | number (float) | — | 2.0 | Minimum number of seconds to wait after a response that asked the client to slow down (HTTP 503 `ServerBusy` or 429) when it carries no `Retry-After` header.0 | From f786d769ef1647565c2e1bcf922e72713e961f08 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 21 Oct 2026 17:27:03 +0000 Subject: [PATCH 32/46] docs(config-reference): regenerate main for `rpc_timeout_s` Regenerated web/apps/docs/content/docs/reference/nativelink-config/main.mdx from the nativelink-config crate (build-schema + schema-to-mdx), as AGENTS.md requires when a config field or its doc comment changes. --- .../docs/content/docs/reference/nativelink-config/main.mdx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx index 5d97a6a58b5..cdd2255366b 100644 --- a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx +++ b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx @@ -5,7 +5,7 @@ full: true --- {/* AUTOGENERATED — do not edit by hand. - Source: nativelink-config @ main (6d6266cf) + Source: nativelink-config @ main (b0f3b0067c) Regenerate from web/: bun --filter @nativelink/docs gen:config-reference */} @@ -856,7 +856,7 @@ See [`experimental_cloud_object_store`](#experimental_cloud_object_store) for de | `retry` | [Retry](#retry) | — | — | Retry configuration to use when a network request fails. | | `max_concurrent_requests` | integer (uint) | — | `0` | Limit the number of simultaneous upstream requests to this many. A value of zero is treated as unlimited. If the limit is reached the request is queued. | | `connections_per_endpoint` | integer (uint) | — | 1 | The number of connections to make to each specified endpoint to balance the load over multiple TCP connections. | -| `rpc_timeout_s` | integer (uint64) | — | 0 (disabled) | Maximum time (seconds) allowed for a single RPC request (e.g. a `ByteStream.Write` call) before it is cancelled. | +| `rpc_timeout_s` | integer (uint64) | — | 0 (disabled) | Maximum time (seconds) allowed for a single RPC request before it is cancelled with `DeadlineExceeded` (retryable). This applies to every RPC this store makes: unary calls (e.g. `FindMissingBlobs`, `BatchReadBlobs`), `ByteStream.Write` calls, and the establishment of read streams through their first response chunk. The body of an established read stream is not bounded by this deadline. | | `use_legacy_resource_names` | boolean | — | false | Use legacy `ByteStream` resource name format, omitting the digest function component from the path. | | `headers` | map of string to string | — | — | Static headers to attach to every outgoing gRPC request sent to this store's upstream endpoints. Useful for fixed authentication tokens (e.g. `{"authorization": "Bearer "}`) and other static metadata. | | `forward_headers` | array of string | — | — | Header names to forward from the incoming client request to every outgoing upstream request. The header value is taken from the client request that triggered this store operation. Use this to pass through dynamic credentials such as JWT tokens sent by build clients. | From 690fff9ecdf18f9ebdf215b6d9eec899bf895edc Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 1 Oct 2026 21:58:14 -0400 Subject: [PATCH 33/46] 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 | 11 ++ nativelink-store/tests/grpc_store_test.rs | 119 +++++++++++++++++++--- 2 files changed, 114 insertions(+), 16 deletions(-) diff --git a/nativelink-store/src/grpc_store.rs b/nativelink-store/src/grpc_store.rs index 451ad956cd2..79ae9e09cfa 100644 --- a/nativelink-store/src/grpc_store.rs +++ b/nativelink-store/src/grpc_store.rs @@ -2087,6 +2087,17 @@ 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 a11c6fb518e..156dde6e6fa 100644 --- a/nativelink-store/tests/grpc_store_test.rs +++ b/nativelink-store/tests/grpc_store_test.rs @@ -123,6 +123,11 @@ struct FakeStreamServer { fail_mid_upload: Arc>, failure_code: tonic::Code, require_zero_start: bool, + /// 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 { @@ -136,6 +141,8 @@ impl FakeStreamServer { fail_mid_upload: Arc::new(Mutex::new(0)), failure_code: tonic::Code::Unavailable, require_zero_start: false, + read_payload: Arc::new(RAW_INPUT.as_bytes().to_vec()), + first_read_partial: None, } } @@ -149,10 +156,6 @@ impl FakeStreamServer { type ReadStream = Pin> + Send + 'static>>; -struct ReaderState { - responded: bool, -} - #[tonic::async_trait] impl ByteStream for FakeStreamServer { type ReadStream = ReadStream; @@ -177,19 +180,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))) } @@ -1108,3 +1131,67 @@ async fn incomplete_oneshot_cas_upload_never_sends_remote_write() -> Result<(), assert!(server.write_requests.lock().await.is_empty()); 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 cfeb5a271b555531bbca8928a75ba7a1110995b2 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Sat, 3 Oct 2026 00:32:32 -0400 Subject: [PATCH 34/46] fix(store): bounded read failover around dead shards in ShardStore ShardStore routes blobs over a weighted consistent-hash ring built once at construction (cumulative weights, binary search). A dead shard keeps its ring weight forever, so its share of the keyspace (~10% in our fleet) keeps routing reads to a node that answers nothing but Unavailable/DeadlineExceeded; every CAS lookup or fetch for those keys fails (or, before the connection-manager fix, wedged) until an operator restarts the frontend. Fix: read-path failover with per-shard health tracking. - Each shard carries a ShardHealth (AtomicU32 consecutive-failure count + AtomicU64 last-failure instant, nanos since store creation). A shard becomes "suspect" after 3 consecutive connectivity failures (Unavailable/DeadlineExceeded) and is skipped as a first read target for a 5s cooldown, after which it is probed again. Any success resets the counter. - get_part and has_with_results fail over in ring order on a connectivity error from the primary (or skip a suspect primary), trying each shard at most once. get_part only fails over when no bytes have reached the client (writer reborrowed across attempts), and never lets a failover shard's NotFound mask the owning shard's connectivity error. If every shard is suspect, the primary is probed directly rather than failing without touching the network. - The write path (update) is unchanged: writes always target the key's home shard so a flaky peer cannot cause duplicate or misplaced blobs; a failed write errors exactly as before. Tests (nativelink-store/tests/shard_ring_dead_peer_test.rs): reads and existence checks for keys owned by a dead shard now succeed via a live replica, multi-key has() no longer fails wholesale when one bucket's shard is down, and a revived shard serves its own keys again after the suspect cooldown. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/src/shard_store.rs | 212 ++++++++++++++++-- .../tests/shard_ring_dead_peer_test.rs | 135 +++++++++-- 2 files changed, 316 insertions(+), 31 deletions(-) diff --git a/nativelink-store/src/shard_store.rs b/nativelink-store/src/shard_store.rs index 2f249448199..56f644e69e5 100644 --- a/nativelink-store/src/shard_store.rs +++ b/nativelink-store/src/shard_store.rs @@ -15,6 +15,8 @@ use core::hash::Hasher; use core::ops::BitXor; use core::pin::Pin; +use core::sync::atomic::{AtomicU32, AtomicU64, Ordering}; +use core::time::Duration; use std::hash::DefaultHasher; use std::sync::Arc; @@ -22,13 +24,39 @@ use async_trait::async_trait; use futures::future::try_join_all; use futures::stream::{FuturesUnordered, TryStreamExt}; use nativelink_config::stores::ShardSpec; -use nativelink_error::{Error, ResultExt, error_if}; +use nativelink_error::{Code, Error, ResultExt, error_if}; 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, Store, StoreDriver, StoreKey, StoreLike, UploadSizeInfo, }; +use tokio::time::Instant; + +/// Number of consecutive connectivity failures after which a shard is +/// considered suspect and skipped as the first read target. +const SUSPECT_FAILURE_THRESHOLD: u32 = 3; + +/// How long a suspect shard is skipped on the read path before it is +/// probed again. +const SUSPECT_COOLDOWN: Duration = Duration::from_secs(5); + +/// Errors that indicate the shard itself is unreachable (as opposed to a +/// definitive answer like `NotFound`), making read failover safe. +const fn is_connectivity_error(err: &Error) -> bool { + matches!(err.code, Code::Unavailable | Code::DeadlineExceeded) +} + +/// Read-path health state for one shard. Writes never consult this: a +/// blob's home shard is fixed by the ring, so writes always target it. +#[derive(Debug, Default)] +struct ShardHealth { + /// Consecutive connectivity failures observed on this shard. + consecutive_failures: AtomicU32, + /// Nanoseconds (since `ShardStore` creation) of the most recent + /// connectivity failure. Zero means "never failed". + last_failure_nanos: AtomicU64, +} #[derive(Debug, MetricsComponent)] struct StoreAndWeight { @@ -36,6 +64,7 @@ struct StoreAndWeight { weight: u32, #[metric(help = "The underlying store")] store: Store, + health: ShardHealth, } #[derive(Debug, MetricsComponent)] @@ -47,6 +76,8 @@ pub struct ShardStore { help = "The weights and stores that are used to determine which store to use" )] weights_and_stores: Vec, + /// Anchor for `ShardHealth::last_failure_nanos` timestamps. + created_at: Instant, } impl ShardStore { @@ -85,8 +116,107 @@ impl ShardStore { weights_and_stores: weights .into_iter() .zip(stores) - .map(|(weight, store)| StoreAndWeight { weight, store }) + .map(|(weight, store)| StoreAndWeight { + weight, + store, + health: ShardHealth::default(), + }) .collect(), + created_at: Instant::now(), + })) + } + + fn now_nanos(&self) -> u64 { + u64::try_from(self.created_at.elapsed().as_nanos()).unwrap_or(u64::MAX) + } + + /// A shard is suspect once it has accumulated enough consecutive + /// connectivity failures and its cooldown has not yet elapsed. After + /// the cooldown it is probed again (treated as healthy). + fn is_suspect(&self, store_idx: usize) -> bool { + let health = &self.weights_and_stores[store_idx].health; + if health.consecutive_failures.load(Ordering::Relaxed) < SUSPECT_FAILURE_THRESHOLD { + return false; + } + let last_failure_nanos = health.last_failure_nanos.load(Ordering::Relaxed); + let cooldown_nanos = u64::try_from(SUSPECT_COOLDOWN.as_nanos()).unwrap_or(u64::MAX); + self.now_nanos().saturating_sub(last_failure_nanos) < cooldown_nanos + } + + fn note_success(&self, store_idx: usize) { + self.weights_and_stores[store_idx] + .health + .consecutive_failures + .store(0, Ordering::Relaxed); + } + + fn note_connectivity_failure(&self, store_idx: usize) { + let health = &self.weights_and_stores[store_idx].health; + health.consecutive_failures.fetch_add(1, Ordering::Relaxed); + health + .last_failure_nanos + .store(self.now_nanos(), Ordering::Relaxed); + } + + /// Runs `has_with_results` against the ring starting at `primary_idx`, + /// failing over in ring order on connectivity errors (each shard is + /// tried at most once). Suspect shards are skipped unless every shard + /// is suspect, in which case the primary is probed directly. + async fn has_on_ring( + &self, + primary_idx: usize, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + let num_stores = self.weights_and_stores.len(); + let mut last_err = None; + let mut attempted_any = false; + for attempt in 0..num_stores { + let store_idx = (primary_idx + attempt) % num_stores; + if self.is_suspect(store_idx) { + continue; + } + attempted_any = true; + let store = &self.weights_and_stores[store_idx].store; + match store.has_with_results(keys, results).await { + Ok(()) => { + self.note_success(store_idx); + return Ok(()); + } + Err(err) if is_connectivity_error(&err) => { + self.note_connectivity_failure(store_idx); + results.fill(None); + last_err = Some(err); + } + Err(err) => { + return Err(err) + .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}"); + } + } + } + if !attempted_any { + // Every shard is suspect: probe the primary rather than fail + // without touching the network. + let store = &self.weights_and_stores[primary_idx].store; + match store.has_with_results(keys, results).await { + Ok(()) => { + self.note_success(primary_idx); + return Ok(()); + } + Err(err) => { + if is_connectivity_error(&err) { + self.note_connectivity_failure(primary_idx); + } + return Err(err) + .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}"); + } + } + } + Err(last_err.unwrap_or_else(|| { + Error::new( + Code::Unavailable, + "All shards unavailable in ShardStore::has_with_results()".to_string(), + ) })) } @@ -167,11 +297,7 @@ impl StoreDriver for ShardStore { if keys.len() == 1 { // Hot path: It is very common to lookup only one key. let store_idx = self.get_store_index(&keys[0]); - let store = &self.weights_and_stores[store_idx].store; - return store - .has_with_results(keys, results) - .await - .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}}"); + return self.has_on_ring(store_idx, keys, results).await; } let mut keys_for_store: Vec<(KeyIdxVec, KeyVec)> = self .weights_and_stores @@ -192,12 +318,9 @@ impl StoreDriver for ShardStore { .into_iter() .enumerate() .map(|(store_idx, (key_idxs, keys))| async move { - let store = &self.weights_and_stores[store_idx].store; let mut inner_results = vec![None; keys.len()]; - store - .has_with_results(&keys, &mut inner_results) - .await - .err_tip(|| "In ShardStore::has_with_results() for store {store_idx}")?; + self.has_on_ring(store_idx, &keys, &mut inner_results) + .await?; Result::<_, Error>::Ok((key_idxs, inner_results)) }) .collect(); @@ -231,11 +354,66 @@ impl StoreDriver for ShardStore { offset: u64, length: Option, ) -> Result<(), Error> { - let store = self.get_store(&key); - store - .get_part(key, writer, offset, length) - .await - .err_tip(|| "In ShardStore::get_part()") + let primary_idx = self.get_store_index(&key); + let num_stores = self.weights_and_stores.len(); + let bytes_written_start = writer.get_bytes_written(); + let mut last_connectivity_err = None; + let mut last_other_err = None; + let mut attempted_any = false; + for attempt in 0..num_stores { + let store_idx = (primary_idx + attempt) % num_stores; + if self.is_suspect(store_idx) { + continue; + } + attempted_any = true; + let store = &self.weights_and_stores[store_idx].store; + match store + .get_part(key.borrow(), &mut *writer, offset, length) + .await + { + Ok(()) => { + self.note_success(store_idx); + return Ok(()); + } + Err(err) if is_connectivity_error(&err) => { + self.note_connectivity_failure(store_idx); + // Failover is only safe if nothing reached the client. + if writer.get_bytes_written() != bytes_written_start { + return Err(err).err_tip(|| "In ShardStore::get_part()"); + } + last_connectivity_err = Some(err); + } + Err(err) if err.code == Code::NotFound && store_idx != primary_idx => { + // A failover shard legitimately may not have the blob; + // keep looking, but never let it mask a connectivity + // error from the shard that owns the key. + self.note_success(store_idx); + last_other_err = Some(err); + } + Err(err) => return Err(err).err_tip(|| "In ShardStore::get_part()"), + } + } + if !attempted_any { + // Every shard is suspect: probe the primary rather than fail + // without touching the network. + let store = &self.weights_and_stores[primary_idx].store; + let result = store.get_part(key, writer, offset, length).await; + match &result { + Ok(()) => self.note_success(primary_idx), + Err(err) if is_connectivity_error(err) => { + self.note_connectivity_failure(primary_idx); + } + Err(_) => {} + } + return result.err_tip(|| "In ShardStore::get_part()"); + } + Err(last_connectivity_err.or(last_other_err).unwrap_or_else(|| { + Error::new( + Code::Unavailable, + "All shards unavailable in ShardStore::get_part()".to_string(), + ) + })) + .err_tip(|| "In ShardStore::get_part()") } fn inner_store(&self, key: Option) -> &'_ dyn StoreDriver { diff --git a/nativelink-store/tests/shard_ring_dead_peer_test.rs b/nativelink-store/tests/shard_ring_dead_peer_test.rs index 238b8f0d6e0..db790cd873a 100644 --- a/nativelink-store/tests/shard_ring_dead_peer_test.rs +++ b/nativelink-store/tests/shard_ring_dead_peer_test.rs @@ -208,13 +208,14 @@ impl ByteStream for FakeShardServer { } } -/// Spawns a live in-process shard peer and returns its port and blob map. -fn spawn_live_peer() -> (u16, Arc>>) { +/// Spawns a live in-process shard peer on `port` (0 for ephemeral) and +/// returns its bound port and blob map. +fn spawn_live_peer_on(port: u16) -> (u16, Arc>>) { let server = FakeShardServer { blobs: Arc::new(Mutex::new(HashMap::new())), }; let blobs = server.blobs.clone(); - let listener = TcpIncoming::bind("127.0.0.1:0".parse().unwrap()).unwrap(); + let listener = TcpIncoming::bind(format!("127.0.0.1:{port}").parse().unwrap()).unwrap(); let port = listener.local_addr().unwrap().port(); let cas_service = ContentAddressableStorageServer::new(server.clone()); let bytestream_service = ByteStreamServer::new(server); @@ -297,12 +298,13 @@ struct Ring { shard_store: Arc, live_shard: usize, live_blobs: Arc>>, + dead_port: u16, } /// Builds a two-shard equal-weight ring: shard 0 is a live in-process /// peer, shard 1 is dead (nothing listening). async fn make_ring() -> Result { - let (live_port, live_blobs) = spawn_live_peer(); + let (live_port, live_blobs) = spawn_live_peer_on(0); let dead_port = dead_peer_port(); let live_store = GrpcStore::new(&grpc_spec_for_port(live_port))?; @@ -325,6 +327,7 @@ async fn make_ring() -> Result { shard_store, live_shard: 0, live_blobs, + dead_port, }) } @@ -395,28 +398,132 @@ async fn read_routed_to_dead_shard_errors_instead_of_hanging() -> Result<(), Err // The production symptom: a multi-key existence check (FindMissingBlobs) // fans out to every shard, so ONE dead peer wedged EVERY CAS RPC through -// the frontend. It must settle with an error instead. +// the frontend. With read failover the dead shard's bucket is retried on +// the live shard, so the whole call succeeds instead of erroring. #[nativelink_test] -async fn has_with_dead_shard_errors_instead_of_hanging() -> Result<(), Error> { +async fn has_with_dead_shard_fails_over_instead_of_erroring() -> Result<(), Error> { let ring = make_ring().await?; - let (live_digest, _) = digest_for_shard(ring.live_shard); - let (dead_digest, _) = digest_for_shard(1 - ring.live_shard); + let (live_digest, live_content) = digest_for_shard(ring.live_shard); + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + // Both blobs are present on the live shard (as a replica would be). + { + let mut blobs = ring.live_blobs.lock().await; + blobs.insert(live_digest.packed_hash().to_string(), live_content.clone()); + blobs.insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + } - let result = timeout( + let results = timeout( GENEROUS_TIMEOUT, ring.shard_store .has_many(&[live_digest.into(), dead_digest.into()]), ) .await - .expect("multi-key has() with a dead shard hung past the generous timeout"); - let err = result.expect_err("has() touching a dead shard cannot succeed"); - assert!( - matches!(err.code, Code::Unavailable | Code::DeadlineExceeded), - "expected Unavailable/DeadlineExceeded from the dead shard, got: {err:?}", + .expect("multi-key has() with a dead shard hung past the generous timeout")?; + assert_eq!( + results, + vec![ + Some(live_content.len() as u64), + Some(dead_content.len() as u64) + ], + "both keys must resolve: the dead shard's bucket via failover", + ); + Ok(()) +} + +// A single-key has() routed to the dead shard must fail over to the live +// shard that holds the blob instead of surfacing Unavailable. +#[nativelink_test] +async fn has_routed_to_dead_shard_fails_over_to_live_replica() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + ring.live_blobs + .lock() + .await + .insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + + let result = timeout(GENEROUS_TIMEOUT, ring.shard_store.has(dead_digest)) + .await + .expect("single-key has() routed to the dead shard hung")?; + assert_eq!( + result, + Some(dead_content.len() as u64), + "has() must find the blob on the live shard via failover", + ); + Ok(()) +} + +// A read routed to the dead shard must fail over to the live shard that +// holds a replica of the blob and return its content. +#[nativelink_test] +async fn get_part_routed_to_dead_shard_fails_over_to_live_replica() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + ring.live_blobs + .lock() + .await + .insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + + let data = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read routed to the dead shard hung")?; + assert_eq!( + data, dead_content, + "failover read must return the live shard's replica", ); Ok(()) } +// Health recovery: after enough consecutive failures the dead shard is +// marked suspect and skipped; once the peer comes back and the cooldown +// elapses it is probed again and serves its own keys directly. +#[nativelink_test] +async fn dead_shard_recovers_after_cooldown() -> Result<(), Error> { + let ring = make_ring().await?; + let (dead_digest, dead_content) = digest_for_shard(1 - ring.live_shard); + + // Trip the suspect threshold: each read fails over to the live shard + // (which lacks the blob), so the call errors, but each attempt counts + // a connectivity failure against the dead shard. + for _ in 0..3 { + let result = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read with dead shard hung"); + assert!(result.is_err(), "blob exists nowhere yet; read must fail"); + } + + // Revive the dead peer on its original port, now holding the blob. + let (_, revived_blobs) = spawn_live_peer_on(ring.dead_port); + revived_blobs + .lock() + .await + .insert(dead_digest.packed_hash().to_string(), dead_content.clone()); + + // Wait out the suspect cooldown so the shard is probed again. + tokio::time::sleep(Duration::from_secs(6)).await; + + let data = timeout( + GENEROUS_TIMEOUT, + ring.shard_store.get_part_unchunked(dead_digest, 0, None), + ) + .await + .expect("read after recovery hung")?; + assert_eq!(data, dead_content, "revived shard must serve its own key"); + + // The success must have reset the failure counter: an immediate + // follow-up (well inside any cooldown) is also served directly. + let result = timeout(GENEROUS_TIMEOUT, ring.shard_store.has(dead_digest)) + .await + .expect("has() after recovery hung")?; + assert_eq!(result, Some(dead_content.len() as u64)); + Ok(()) +} + // The live shard must keep serving while its ring-mate is dead: routing // around the dead peer stays possible for keys that don't touch it. #[nativelink_test] From 9009570165653cf16f6b7aa58f1f5a5733ebc692 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Sat, 3 Oct 2026 00:26:26 -0400 Subject: [PATCH 35/46] feat(cas): make the BatchReadBlobs per-blob timeout configurable BatchReadBlobs applies a hardcoded 30 s deadline to each blob read so a single stalled read cannot pin the whole batch until the client's overall deadline. Under cold large-tree reads against a slow backing tier -- e.g. an R2/S3 object store behind a fast store doing writeback backfill on read -- large blobs legitimately take longer than 30 s to materialize, and the hardcoded deadline fails the read with DeadlineExceeded even though the transfer is making progress, failing the entire batch. Add an optional per-instance `batch_read_per_blob_timeout_s` field to the CAS service config. Unset preserves today's 30 s default exactly; deployments with slow cold tiers can raise it. The simplest defensible design -- a configurable base rather than a size-scaled formula -- was chosen deliberately: the operator knows their tier's throughput floor, and a scaling heuristic baked into the server is hard to tune and harder to remove later. BatchUpdateBlobs keeps the fixed default. The new regression test stalls the store for 2 minutes and shows the read succeeds with a configured 3 minute timeout; with the fix reverted the configured value is ignored and the read fails with DeadlineExceeded at 30 s (the existing default-timeout test continues to cover the unset case). Co-Authored-By: Claude Opus 4.8 --- nativelink-config/src/cas_server.rs | 11 +++ nativelink-service/src/cas_server.rs | 27 +++++++- nativelink-service/tests/cas_server_test.rs | 67 ++++++++++++++++++- .../tests/grpc_wire_compression_test.rs | 1 + 4 files changed, 100 insertions(+), 6 deletions(-) diff --git a/nativelink-config/src/cas_server.rs b/nativelink-config/src/cas_server.rs index 1b4aac223fe..bc1561a294b 100644 --- a/nativelink-config/src/cas_server.rs +++ b/nativelink-config/src/cas_server.rs @@ -148,6 +148,17 @@ pub struct CasStoreConfig { /// and behavior is identical to when this option did not exist. #[serde(default)] pub experimental_chunking: Option, + + /// Per-blob timeout in seconds applied inside `BatchReadBlobs`. Each + /// blob read in a batch races this deadline so a single stalled read + /// cannot pin the whole batch until the client's overall deadline. + /// Cold reads through a slow backing tier (e.g. an object store that + /// populates a local cache on read) can legitimately take longer than + /// the default for large blobs; raise this value for such deployments. + /// + /// Default: 30 seconds. + #[serde(default)] + pub batch_read_per_blob_timeout_s: Option, } #[derive(Deserialize, Serialize, Debug, Clone)] diff --git a/nativelink-service/src/cas_server.rs b/nativelink-service/src/cas_server.rs index 953c2bc6e5a..fd3d4832896 100644 --- a/nativelink-service/src/cas_server.rs +++ b/nativelink-service/src/cas_server.rs @@ -173,6 +173,10 @@ impl ChunkingInstance { #[derive(Debug)] pub struct CasServer { stores: HashMap, + /// Per-instance per-blob deadline applied inside `BatchReadBlobs`. + /// Defaults to `BATCH_PER_BLOB_TIMEOUT`; configurable via + /// `batch_read_per_blob_timeout_s`. + batch_read_per_blob_timeouts: HashMap, remote_cache_compression_instances: RemoteCacheCompressionInstances, chunking_instances: HashMap, chunking_metrics: ChunkingMetrics, @@ -180,7 +184,9 @@ pub struct CasServer { type GetTreeStream = Pin> + Send + 'static>>; -/// Per-blob deadline applied inside `BatchReadBlobs` / `BatchUpdateBlobs`. +/// Default per-blob deadline applied inside `BatchReadBlobs` / +/// `BatchUpdateBlobs`. For `BatchReadBlobs` this can be overridden per +/// instance via `batch_read_per_blob_timeout_s`. const BATCH_PER_BLOB_TIMEOUT: Duration = Duration::from_secs(30); /// Maximum size of a single chunk accepted in a `SpliceBlob` request. @@ -209,11 +215,20 @@ impl CasServer { remote_cache_compression_instances: &RemoteCacheCompressionInstances, ) -> Result { let mut stores = HashMap::with_capacity(configs.len()); + let mut batch_read_per_blob_timeouts = HashMap::with_capacity(configs.len()); let mut chunking_instances = HashMap::new(); for config in configs { let store = store_manager.get_store(&config.cas_store).ok_or_else(|| { make_input_err!("'cas_store': '{}' does not exist", config.cas_store) })?; + batch_read_per_blob_timeouts.insert( + config.instance_name.clone(), + config + .batch_read_per_blob_timeout_s + .map_or(BATCH_PER_BLOB_TIMEOUT, |secs| { + Duration::from_secs(secs.into()) + }), + ); if let Some(chunking_config) = &config.experimental_chunking { let avg_chunk_size_bytes = chunking_config .validated_avg_chunk_size_bytes() @@ -273,6 +288,7 @@ impl CasServer { } Ok(Self { stores, + batch_read_per_blob_timeouts, remote_cache_compression_instances: remote_cache_compression_instances.clone(), chunking_instances, chunking_metrics: ChunkingMetrics::default(), @@ -411,6 +427,11 @@ impl CasServer { } let store_ref = &store; + let per_blob_timeout = self + .batch_read_per_blob_timeouts + .get(instance_name) + .copied() + .unwrap_or(BATCH_PER_BLOB_TIMEOUT); let remote_cache_compression_enabled = self .remote_cache_compression_instances .enabled_for(instance_name); @@ -429,7 +450,7 @@ impl CasServer { // Apply a per-blob deadline so one slow read does not // make the whole batch hit the client's overall deadline. let result = match tokio::time::timeout( - BATCH_PER_BLOB_TIMEOUT, + per_blob_timeout, store_ref.get_part_unchunked(digest_copy, 0, None), ) .await @@ -438,7 +459,7 @@ impl CasServer { Err(_elapsed) => Err(make_err!( Code::DeadlineExceeded, "BatchReadBlobs per-blob timeout ({} s) elapsed for digest {}", - BATCH_PER_BLOB_TIMEOUT.as_secs(), + per_blob_timeout.as_secs(), digest_copy, )), }; diff --git a/nativelink-service/tests/cas_server_test.rs b/nativelink-service/tests/cas_server_test.rs index 055bfe9a884..3ee5477c25f 100644 --- a/nativelink-service/tests/cas_server_test.rs +++ b/nativelink-service/tests/cas_server_test.rs @@ -78,6 +78,7 @@ fn make_cas_server(store_manager: &StoreManager) -> Result { config: nativelink_config::cas_server::CasStoreConfig { cas_store: "main_cas".to_string(), experimental_chunking: None, + batch_read_per_blob_timeout_s: None, }, }], store_manager, @@ -92,6 +93,7 @@ fn make_cas_server_with_zstd(store_manager: &StoreManager) -> Result, _key: StoreKey<'_>, - _writer: &mut DropCloserWriteHalf, + writer: &mut DropCloserWriteHalf, _offset: u64, _length: Option, ) -> Result<(), Error> { tokio::time::sleep(self.delay).await; - Ok(()) + if !self.data.is_empty() { + writer.send(self.data.clone()).await?; + } + writer.send_eof() } fn inner_store(&self, _digest: Option) -> &'_ dyn StoreDriver { @@ -760,14 +767,23 @@ impl StoreDriver for StallStore { default_health_status_indicator!(StallStore); fn make_cas_server_with_stall_store(delay: Duration) -> Result { + make_cas_server_with_stall_store_and_timeout(delay, Bytes::new(), None) +} + +fn make_cas_server_with_stall_store_and_timeout( + delay: Duration, + data: Bytes, + batch_read_per_blob_timeout_s: Option, +) -> Result { let store_manager = Arc::new(StoreManager::new()); - store_manager.add_store("main_cas", Store::new(Arc::new(StallStore { delay })))?; + store_manager.add_store("main_cas", Store::new(Arc::new(StallStore { delay, data })))?; CasServer::new( &[WithInstanceName { instance_name: INSTANCE_NAME.to_string(), config: nativelink_config::cas_server::CasStoreConfig { cas_store: "main_cas".to_string(), experimental_chunking: None, + batch_read_per_blob_timeout_s, }, }], &store_manager, @@ -858,6 +874,47 @@ async fn batch_read_blobs_per_blob_timeout_returns_deadline_exceeded() Ok(()) } +#[nativelink_test(start_paused = true)] +async fn batch_read_blobs_respects_configured_per_blob_timeout() +-> Result<(), Box> { + const VALUE: &str = "1"; + + // The store stalls for 2 minutes — well beyond the 30 s default that + // would fail the read — but the instance is configured with a 3 minute + // per-blob timeout, so the read must succeed. + let cas_server = make_cas_server_with_stall_store_and_timeout( + Duration::from_mins(2), + Bytes::from_static(VALUE.as_bytes()), + Some(180), + )?; + + let digest = Digest { + hash: HASH1.to_string(), + size_bytes: VALUE.len().try_into().unwrap_or(i64::MAX), + }; + let raw_response = cas_server + .batch_read_blobs(Request::new(BatchReadBlobsRequest { + instance_name: INSTANCE_NAME.to_string(), + digests: vec![digest.clone()], + acceptable_compressors: vec![compressor::Value::Identity.into()], + digest_function: digest_function::Value::Sha256.into(), + })) + .await; + + let response = raw_response.unwrap().into_inner(); + assert_eq!(response.responses.len(), 1); + let entry = &response.responses[0]; + assert_eq!(entry.digest.as_ref(), Some(&digest)); + let status = entry.status.as_ref().expect("status set"); + assert_eq!( + status.code, + Code::Ok as i32, + "expected Ok with a raised per-blob timeout, got status: {status:?}", + ); + assert_eq!(entry.data, VALUE.as_bytes()); + Ok(()) +} + #[nativelink_test] async fn batch_update_blobs_zstd_compressed() -> Result<(), Box> { let store_manager = make_store_manager().await?; @@ -1061,6 +1118,7 @@ fn make_chunking_cas_server_with_avg( instance_name: INSTANCE_NAME.to_string(), config: nativelink_config::cas_server::CasStoreConfig { cas_store: "main_cas".to_string(), + batch_read_per_blob_timeout_s: None, experimental_chunking: Some(nativelink_config::cas_server::CasChunkingConfig { index_store: Some("chunk_index".to_string()), avg_chunk_size_bytes, @@ -1516,6 +1574,7 @@ async fn chunking_rejects_index_store_same_as_cas_store() -> Result<(), Box Result<(), Box Result<(), Box Date: Wed, 21 Oct 2026 17:30:19 +0000 Subject: [PATCH 36/46] docs(config-reference): regenerate main for `batch_read_per_blob_timeout_s` Regenerated web/apps/docs/content/docs/reference/nativelink-config/main.mdx from the nativelink-config crate (build-schema + schema-to-mdx), as AGENTS.md requires when a config field or its doc comment changes. --- .../docs/content/docs/reference/nativelink-config/main.mdx | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx index 5d97a6a58b5..3ac09ba144d 100644 --- a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx +++ b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx @@ -5,7 +5,7 @@ full: true --- {/* AUTOGENERATED — do not edit by hand. - Source: nativelink-config @ main (6d6266cf) + Source: nativelink-config @ main (9009570165) Regenerate from web/: bun --filter @nativelink/docs gen:config-reference */} @@ -1491,6 +1491,7 @@ eternal 3600 s), so this maps them to classes without a CAS fetch. | `instance_name` | string | — | `""` | Used when the config references `instance_name` in the protocol. | | `cas_store` | string | Yes | — | The store name referenced in the `stores` map in the main config. This store name referenced here may be reused multiple times. | | `experimental_chunking` | [CasChunkingConfig](#caschunkingconfig) | — | not set — chunking RPCs are rejected, nothing is advertised, | Optional and experimental: enables the REAPI `SplitBlob`/`SpliceBlob` RPCs used by content-defined chunking clients (e.g. Bazel's `--experimental_remote_cache_chunking`). When set, the capabilities service advertises blob split/splice support and `FastCDC` 2020 parameters for this instance. When `cas_store` is a grpc store the RPCs are forwarded to the backend (which must support chunking with matching parameters); otherwise they are served locally. | +| `batch_read_per_blob_timeout_s` | integer (uint32) | — | 30 seconds | Per-blob timeout in seconds applied inside `BatchReadBlobs`. Each blob read in a batch races this deadline so a single stalled read cannot pin the whole batch until the client's overall deadline. Cold reads through a slow backing tier (e.g. an object store that populates a local cache on read) can legitimately take longer than the default for large blobs; raise this value for such deployments. | ## ActionCacheServiceConfig From e97e796ae20948b86d7a0b38f67237eb6ad6f4b8 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 01:34:10 -0400 Subject: [PATCH 37/46] 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 | 146 ++++++++++++++++++ 2 files changed, 192 insertions(+), 4 deletions(-) diff --git a/nativelink-worker/src/running_actions_manager.rs b/nativelink-worker/src/running_actions_manager.rs index 155f28fe3e3..d631248e61f 100644 --- a/nativelink-worker/src/running_actions_manager.rs +++ b/nativelink-worker/src/running_actions_manager.rs @@ -4977,6 +4977,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<()>, @@ -5096,6 +5105,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 { @@ -5539,7 +5549,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) @@ -5554,6 +5570,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()) { @@ -5598,12 +5626,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 07a0abe9941..2020e3bc5cd 100644 --- a/nativelink-worker/tests/running_actions_manager_test.rs +++ b/nativelink-worker/tests/running_actions_manager_test.rs @@ -1740,6 +1740,152 @@ 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, + #[cfg(target_os = "linux")] + readonly_input_mounts: Vec::new(), + 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 cfadc1562a7bf2862790bddab33d43fe8defecc1 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Fri, 17 Jul 2026 22:36:28 -0400 Subject: [PATCH 38/46] buf_channel: widen in-flight buffer 2 -> 32 to avoid write-stall resets MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The buf channel between "read from client" and "write to store" held only 2 chunks (an explicitly-arbitrary magic number). During a large ByteStream.Write forwarded to a slow shard CAS, that near-zero buffer stalls the producer the instant the consumer lags, which stops the server draining its inbound stream, empties the HTTP/2 flow-control window, and resets the connection mid-stream — dropping multi-GiB toolchain uploads. Name it a const and widen to 32; tokio's bounded channel doesn't preallocate, so the capacity costs memory only while data is actually queued (exactly during a downstream stall). Complements the transport-error -> Unavailable coding fix (5de7c286) and the HTTP/2 adaptive-window server config. Two filesystem_store tests (file_continues_to_stream_on_content_replace_test, file_gets_cleans_up_on_cache_eviction) relied on the old 2-slot capacity to keep the store's get() stalled with the file open after the reader took one byte; with a 10-byte value and read_buffer_size 1 the whole blob now fits in the channel, get() finishes, and the temp file is gone before the assertion. They now stream a 100-byte value, which is longer than the channel holds. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../tests/filesystem_store_test.rs | 56 +++++++++++++++---- nativelink-util/src/buf_channel.rs | 16 ++++-- 2 files changed, 55 insertions(+), 17 deletions(-) diff --git a/nativelink-store/tests/filesystem_store_test.rs b/nativelink-store/tests/filesystem_store_test.rs index 7f659a16123..2c6582c858b 100644 --- a/nativelink-store/tests/filesystem_store_test.rs +++ b/nativelink-store/tests/filesystem_store_test.rs @@ -359,6 +359,35 @@ const HASH1: &str = "0123456789abcdef000000000000000000010000000000000123456789a const HASH2: &str = "0123456789abcdef000000000000000000020000000000000123456789abcdef"; const VALUE1: &str = "0123456789"; const VALUE2: &str = "9876543210"; + +// Longer, in `read_buffer_size: 1` chunks, than `make_buf_channel_pair`'s +// in-flight capacity, so a reader that has taken only the first byte keeps the +// store's `get` stalled with the file still open. Tests that need an open +// stream on a file use these instead of `VALUE1`/`VALUE2`. +const LONG_VALUE1: &str = concat!( + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", + "0123456789", +); +const LONG_VALUE2: &str = concat!( + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", + "9876543210", +); const STRING_NAME: &str = "String_Filename"; #[nativelink_test] @@ -519,7 +548,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error> } } - let digest1 = DigestInfo::try_new(HASH1, VALUE1.len())?; + let digest1 = DigestInfo::try_new(HASH1, LONG_VALUE1.len())?; let content_path = make_temp_path("content_path"); let temp_path = make_temp_path("temp_path"); @@ -539,7 +568,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error> ); // Insert data into store. - store.update_oneshot(digest1, VALUE1.into()).await?; + store.update_oneshot(digest1, LONG_VALUE1.into()).await?; let (writer, mut reader) = make_buf_channel_pair(); let store_clone = store.clone(); @@ -557,13 +586,13 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error> .err_tip(|| "Error reading first byte")?; assert_eq!( first_byte[0], - VALUE1.as_bytes()[0], + LONG_VALUE1.as_bytes()[0], "Expected first byte to match" ); } // Replace content. - store.update_oneshot(digest1, VALUE2.into()).await?; + store.update_oneshot(digest1, LONG_VALUE2.into()).await?; // Ensure we let any background tasks finish. tokio::task::yield_now().await; @@ -582,7 +611,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error> let data = read_file_contents(path.as_os_str()).await?; assert_eq!( &data[..], - VALUE1.as_bytes(), + LONG_VALUE1.as_bytes(), "Expected file content to match" ); } @@ -599,7 +628,7 @@ async fn file_continues_to_stream_on_content_replace_test() -> Result<(), Error> assert_eq!( &remaining_file_data, - &VALUE1.as_bytes()[1..], + &LONG_VALUE1.as_bytes()[1..], "Expected file content to match" ); @@ -627,8 +656,8 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { } } - let digest1 = DigestInfo::try_new(HASH1, VALUE1.len())?; - let digest2 = DigestInfo::try_new(HASH2, VALUE2.len())?; + let digest1 = DigestInfo::try_new(HASH1, LONG_VALUE1.len())?; + let digest2 = DigestInfo::try_new(HASH2, LONG_VALUE2.len())?; let content_path = make_temp_path("content_path"); let temp_path = make_temp_path("temp_path"); @@ -648,7 +677,10 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { ); // Insert data into store. - store.update_oneshot(digest1, VALUE1.into()).await.unwrap(); + store + .update_oneshot(digest1, LONG_VALUE1.into()) + .await + .unwrap(); let mut reader = { let (writer, reader) = make_buf_channel_pair(); @@ -664,7 +696,7 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { assert!(reader.peek().await.is_ok(), "Could not peek into reader"); // Insert new content. This will evict the old item. - store.update_oneshot(digest2, VALUE2.into()).await?; + store.update_oneshot(digest2, LONG_VALUE2.into()).await?; // Ensure we let any background tasks finish. tokio::task::yield_now().await; @@ -683,7 +715,7 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { let data = read_file_contents(path.as_os_str()).await?; assert_eq!( &data[..], - VALUE1.as_bytes(), + LONG_VALUE1.as_bytes(), "Expected file content to match" ); } @@ -698,7 +730,7 @@ async fn file_gets_cleans_up_on_cache_eviction() -> Result<(), Error> { .await .err_tip(|| "Error reading remaining bytes")?; - assert_eq!(&reader_data, VALUE1, "Expected file content to match"); + assert_eq!(&reader_data, LONG_VALUE1, "Expected file content to match"); loop { if DELETES_FINISHED.load(Ordering::Relaxed) == 1 { diff --git a/nativelink-util/src/buf_channel.rs b/nativelink-util/src/buf_channel.rs index d60f8fa362c..741d1245611 100644 --- a/nativelink-util/src/buf_channel.rs +++ b/nativelink-util/src/buf_channel.rs @@ -27,6 +27,16 @@ use tracing::warn; const ZERO_DATA: Bytes = Bytes::new(); +/// Chunks buffered in flight per channel. Small enough that a fast producer +/// still applies backpressure promptly, but not so small (it was 2) that a +/// momentarily-slow consumer — e.g. a shard CAS mid-forward of a multi-GiB blob +/// — stalls the producer immediately. That instant stall stops the server +/// draining its inbound stream, the HTTP/2 flow-control window empties, and the +/// connection resets mid-stream (dropping large cacheable uploads). tokio's +/// bounded channel does not preallocate, so this capacity costs memory only +/// while data is actually queued (i.e. exactly during a downstream stall). +const CHANNEL_BUFFER_CHUNKS: usize = 32; + /// Create a channel pair that can be used to transport buffer objects around to /// different components. This wrapper is used because the streams give some /// utility like managing EOF in a more friendly way, ensure if no EOF is received @@ -34,11 +44,7 @@ const ZERO_DATA: Bytes = Bytes::new(); /// the number of bytes sent. #[must_use] pub fn make_buf_channel_pair() -> (DropCloserWriteHalf, DropCloserReadHalf) { - // We allow up to 2 items in the buffer at any given time. There is no major - // reason behind this magic number other than thinking it will be nice to give - // a little time for another thread to wake up and consume data if another - // thread is pumping large amounts of data into the channel. - let (tx, rx) = mpsc::channel(2); + let (tx, rx) = mpsc::channel(CHANNEL_BUFFER_CHUNKS); let eof_sent = Arc::new(AtomicBool::new(false)); ( DropCloserWriteHalf { From e77973981c08e372a070dcc5e02114ffec662d13 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Tue, 29 Sep 2026 06:14:16 -0400 Subject: [PATCH 39/46] =?UTF-8?q?sha2:=20enable=20`asm`=20=E2=80=94=20hard?= =?UTF-8?q?ware=20SHA-256=20with=20runtime=20detection?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RustCrypto gates the aarch64 hardware SHA-2 backend behind the `asm` feature; enabling it pulls sha2-asm (cc-compiled) and, via cpufeatures, runtime-detects the HWCAP sha2 bit. One binary then uses on-core hardware SHA-256 where present and the software path elsewhere. Digest hashing is a hot path for the CAS. Co-Authored-By: Claude Fable 5 --- Cargo.lock | 10 ++++++++++ Cargo.toml | 5 ++++- MODULE.bazel.lock | 1 + 3 files changed, 15 insertions(+), 1 deletion(-) diff --git a/Cargo.lock b/Cargo.lock index 1f9474622e9..8a4ef6748b5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4604,6 +4604,7 @@ dependencies = [ "cfg-if", "cpufeatures 0.2.17", "digest 0.10.7", + "sha2-asm", ] [[package]] @@ -4617,6 +4618,15 @@ dependencies = [ "digest 0.11.3", ] +[[package]] +name = "sha2-asm" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b845214d6175804686b2bd482bcffe96651bb2d1200742b712003504a2dac1ab" +dependencies = [ + "cc", +] + [[package]] name = "sharded-slab" version = "0.1.7" diff --git a/Cargo.toml b/Cargo.toml index 0f20bfb61f3..97074514c7b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -65,7 +65,10 @@ rand = { version = "0.9.0", default-features = false, features = [ rustls-pki-types = { version = "1.13.1", features = [ "std", ], default-features = false } -sha2 = { version = "0.10.8", default-features = false } +# `asm` unlocks the RustCrypto hardware SHA-2 backends (aarch64 is gated behind +# it) with runtime HWCAP detection — hardware SHA on-core, software fallback +# elsewhere. +sha2 = { version = "0.10.8", default-features = false, features = ["asm"] } tokio = { version = "1.52.2", features = [ "fs", "io-util", diff --git a/MODULE.bazel.lock b/MODULE.bazel.lock index fceb1b0e9e9..cfa8dfef3b9 100644 --- a/MODULE.bazel.lock +++ b/MODULE.bazel.lock @@ -1152,6 +1152,7 @@ "sha1_0.10.6": "{\"dependencies\":[{\"name\":\"cfg-if\",\"req\":\"^1.0\"},{\"name\":\"cpufeatures\",\"req\":\"^0.2\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86\\\", target_arch = \\\"x86_64\\\"))\"},{\"name\":\"digest\",\"req\":\"^0.10.7\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.10.7\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.2.2\"},{\"name\":\"sha1-asm\",\"optional\":true,\"req\":\"^0.5\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86\\\", target_arch = \\\"x86_64\\\"))\"}],\"features\":{\"asm\":[\"sha1-asm\"],\"compress\":[],\"default\":[\"std\"],\"force-soft\":[],\"loongarch64_asm\":[],\"oid\":[\"digest/oid\"],\"std\":[\"digest/std\"]}}", "sha1_0.11.0": "{\"dependencies\":[{\"name\":\"cfg-if\",\"req\":\"^1.0\"},{\"name\":\"cpufeatures\",\"req\":\"^0.3\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86\\\", target_arch = \\\"x86_64\\\"))\"},{\"name\":\"digest\",\"req\":\"^0.11\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.11\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^1\"}],\"features\":{\"alloc\":[\"digest/alloc\"],\"default\":[\"alloc\",\"oid\"],\"oid\":[\"digest/oid\"],\"zeroize\":[\"digest/zeroize\"]}}", "sha1_smol_1.0.1": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"openssl\",\"req\":\"^0.10\"},{\"kind\":\"dev\",\"name\":\"rand\",\"req\":\"^0.4\"},{\"name\":\"serde\",\"optional\":true,\"req\":\"^1.0\"},{\"kind\":\"dev\",\"name\":\"serde_json\",\"req\":\"^1.0\"}],\"features\":{\"alloc\":[],\"std\":[\"alloc\"]}}", + "sha2-asm_0.6.4": "{\"dependencies\":[{\"kind\":\"build\",\"name\":\"cc\",\"req\":\"^1.0\"}],\"features\":{}}", "sha2_0.10.9": "{\"dependencies\":[{\"name\":\"cfg-if\",\"req\":\"^1.0\"},{\"name\":\"cpufeatures\",\"req\":\"^0.2\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86_64\\\", target_arch = \\\"x86\\\"))\"},{\"name\":\"digest\",\"req\":\"^0.10.7\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.10.7\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^0.2.2\"},{\"name\":\"sha2-asm\",\"optional\":true,\"req\":\"^0.6.1\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86_64\\\", target_arch = \\\"x86\\\"))\"}],\"features\":{\"asm\":[\"sha2-asm\"],\"asm-aarch64\":[\"asm\"],\"compress\":[],\"default\":[\"std\"],\"force-soft\":[],\"force-soft-compact\":[],\"loongarch64_asm\":[],\"oid\":[\"digest/oid\"],\"std\":[\"digest/std\"]}}", "sha2_0.11.0": "{\"dependencies\":[{\"name\":\"cfg-if\",\"req\":\"^1\"},{\"name\":\"cpufeatures\",\"req\":\"^0.3\",\"target\":\"cfg(any(target_arch = \\\"aarch64\\\", target_arch = \\\"x86_64\\\", target_arch = \\\"x86\\\"))\"},{\"name\":\"digest\",\"req\":\"^0.11\"},{\"features\":[\"dev\"],\"kind\":\"dev\",\"name\":\"digest\",\"req\":\"^0.11\"},{\"kind\":\"dev\",\"name\":\"hex-literal\",\"req\":\"^1\"}],\"features\":{\"alloc\":[\"digest/alloc\"],\"default\":[\"alloc\",\"oid\"],\"oid\":[\"digest/oid\"],\"zeroize\":[\"digest/zeroize\"]}}", "sharded-slab_0.1.7": "{\"dependencies\":[{\"kind\":\"dev\",\"name\":\"criterion\",\"req\":\"^0.3\"},{\"kind\":\"dev\",\"name\":\"indexmap\",\"req\":\"^1\"},{\"name\":\"lazy_static\",\"req\":\"^1\"},{\"features\":[\"checkpoint\"],\"name\":\"loom\",\"optional\":true,\"req\":\"^0.5\",\"target\":\"cfg(loom)\"},{\"features\":[\"checkpoint\"],\"kind\":\"dev\",\"name\":\"loom\",\"req\":\"^0.5\",\"target\":\"cfg(loom)\"},{\"kind\":\"dev\",\"name\":\"memory-stats\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"proptest\",\"req\":\"^1\"},{\"kind\":\"dev\",\"name\":\"slab\",\"req\":\"^0.4.2\"}],\"features\":{}}", From a8cd409af1c163984f69d1564e59febe4464ac64 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Wed, 30 Sep 2026 01:27:20 -0400 Subject: [PATCH 40/46] 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. Co-Authored-By: Claude Opus 4.8 --- nativelink-store/tests/filesystem_store_test.rs | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/nativelink-store/tests/filesystem_store_test.rs b/nativelink-store/tests/filesystem_store_test.rs index 7f659a16123..44e7300d109 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" From 0953a633f3d9e31652496d87d904a214d98a8a7a Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 22 Oct 2026 10:14:31 +0000 Subject: [PATCH 41/46] test(fast_slow_store): show the v1.8.0 populate failure from a filesystem slow tier Since 086289a063 (#2974) FastSlowStore's populate checks the bytes it received from the slow tier against the size it declared to the fast tier, and it takes that size from the slow tier's has(). A FilesystemStore reports its block-rounded footprint there (4096 for a 17-byte blob at the default block size), so every read that has to fill the fast tier from a filesystem slow tier fails with DataLoss: Expected 4096 upload bytes, received 17 Slow store ended before the cache fill was complete unless the blob is a block multiple. That is the dedup index store in nativelink-config/examples/filesystem_cas.json5 and the action cache in kubernetes/nativelink/nativelink-config.json5, which shipped broken in v1.8.0. This commit adds only tests, so CI goes red on it; the fix follows. - populate_from_filesystem_slow_tier_at_default_block_size: the shipped shape, with a real FilesystemStore slow tier. - populate_does_not_trust_slow_store_has_size_for_digest: a slow tier whose has() reports less (a compressed stream) or more (a block-rounded footprint) than the blob. - populate_keeps_entry_whose_size_is_not_its_digests: passes today and guards the fix against the obvious simplification of declaring the digest's size, which an AC or dedup index entry does not have. Co-Authored-By: Claude Fable 5.1 --- .../tests/fast_slow_store_test.rs | 194 ++++++++++++++++++ 1 file changed, 194 insertions(+) diff --git a/nativelink-store/tests/fast_slow_store_test.rs b/nativelink-store/tests/fast_slow_store_test.rs index 35a04d2511c..1f0d1f3d0e1 100644 --- a/nativelink-store/tests/fast_slow_store_test.rs +++ b/nativelink-store/tests/fast_slow_store_test.rs @@ -2347,3 +2347,197 @@ async fn incomplete_slow_read_never_populates_memory_or_filesystem() -> Result<( } Ok(()) } + +/// A filesystem store, which refuses a write whose length is not its declared +/// `ExactSize`, with one-byte blocks so `has()` reports the stored length. +async fn make_filesystem_fast_store() -> Result<(Arc, FilesystemSpec), Error> { + let spec = FilesystemSpec { + content_path: make_temp_path("content_path"), + temp_path: make_temp_path("temp_path"), + block_size: 1, + ..Default::default() + }; + Ok((FilesystemStore::::new(&spec).await?, spec)) +} + +/// Reads `key` through a `FastSlowStore` whose fast tier is `fast` and whose +/// slow tier alone holds `data`, then checks the fast tier was filled with it. +async fn check_populates_filesystem_fast_store( + fast: Arc, + fast_spec: FilesystemSpec, + slow: Store, + key: StoreKey<'_>, + data: &Bytes, +) -> Result<(), Error> { + let fast = Store::new(fast); + let store = Store::new(FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Filesystem(fast_spec), + slow: StoreSpec::Memory(MemorySpec::default()), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + fast.clone(), + slow.clone(), + )); + slow.update_oneshot(key.borrow(), data.clone()).await?; + assert_eq!(fast.has(key.borrow()).await?, None); + + assert_eq!( + store.get_part_unchunked(key.borrow(), 0, None).await?, + data, + "{key}" + ); + assert_eq!( + fast.has(key.borrow()).await?, + Some(data.len() as u64), + "{key}" + ); + assert_eq!(fast.get_part_unchunked(key.borrow(), 0, None).await?, data); + Ok(()) +} + +/// A slow store whose `has()` reports `reported_size` for whatever it holds, +/// the way a tier that reports a physical size does. +#[derive(MetricsComponent)] +struct MisreportingSlowStore { + inner: Arc, + reported_size: u64, +} + +#[async_trait] +impl StoreDriver for MisreportingSlowStore { + async fn post_init(self: Arc) -> Result<(), Error> { + Ok(()) + } + + async fn has_with_results( + self: Pin<&Self>, + keys: &[StoreKey<'_>], + results: &mut [Option], + ) -> Result<(), Error> { + Pin::new(self.inner.as_ref()) + .has_with_results(keys, results) + .await?; + for result in results.iter_mut().filter(|result| result.is_some()) { + *result = Some(self.reported_size); + } + Ok(()) + } + + async fn update( + self: Pin<&Self>, + key: StoreKey<'_>, + reader: DropCloserReadHalf, + size_info: UploadSizeInfo, + ) -> Result { + Pin::new(self.inner.as_ref()) + .update(key, reader, size_info) + .await + } + + async fn get_part( + self: Pin<&Self>, + key: StoreKey<'_>, + writer: &mut DropCloserWriteHalf, + offset: u64, + length: Option, + ) -> Result<(), Error> { + Pin::new(self.inner.as_ref()) + .get_part(key, writer, offset, length) + .await + } + + fn inner_store(&self, _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> { + Ok(()) + } +} + +default_health_status_indicator!(MisreportingSlowStore); + +// The populate cannot take a digest's size from the slow tier's `has()` on +// trust: a tier that reports less (a compressed stream) or more (a +// block-rounded footprint) than the blob must still fill the fast tier with +// the blob as it is. +#[nativelink_test] +async fn populate_does_not_trust_slow_store_has_size_for_digest() -> Result<(), Error> { + let data = Bytes::from(make_random_data(3000)); + let digest = DigestInfo::try_new(VALID_HASH, data.len())?; + for reported_size in [1, 4096] { + let slow = Store::new(Arc::new(MisreportingSlowStore { + inner: MemoryStore::new(&MemorySpec::default()), + reported_size, + })); + let (fast, fast_spec) = make_filesystem_fast_store().await?; + check_populates_filesystem_fast_store(fast, fast_spec, slow, digest.into(), &data).await?; + } + Ok(()) +} + +// An action cache entry is keyed by the action's digest and a dedup index entry +// by its blob's, so a digest's size is not always the size of what it names. +// The slow tier's `has()` is right here and the digest is not, so the populate +// must not declare the digest's size as exact either. +#[nativelink_test] +async fn populate_keeps_entry_whose_size_is_not_its_digests() -> Result<(), Error> { + let data = Bytes::from(make_random_data(3000)); + let digest = DigestInfo::try_new(VALID_HASH, 100)?; + let slow = Store::new(MemoryStore::new(&MemorySpec::default())); + let (fast, fast_spec) = make_filesystem_fast_store().await?; + check_populates_filesystem_fast_store(fast, fast_spec, slow, digest.into(), &data).await +} + +// The shipped shape of the bug: a filesystem slow tier at the default 4 KiB +// block size reports its block-rounded footprint from `has()` (4096 for a +// 17-byte blob), and the populate declares that to the fast tier as an +// `ExactSize`. Since the populate checks the bytes it received against that +// size, every read that has to fill the fast tier from such a slow tier fails +// with DataLoss unless the blob happens to be a block multiple. This is the +// dedup index store in `nativelink-config/examples/filesystem_cas.json5` and +// the action cache in `kubernetes/nativelink/nativelink-config.json5`. +#[nativelink_test] +async fn populate_from_filesystem_slow_tier_at_default_block_size() -> Result<(), Error> { + let slow_spec = FilesystemSpec { + content_path: make_temp_path("slow_content_path"), + temp_path: make_temp_path("slow_temp_path"), + ..Default::default() + }; + let slow = Store::new(FilesystemStore::::new(&slow_spec).await?); + let fast = Store::new(MemoryStore::new(&MemorySpec::default())); + let data = Bytes::from_static(b"seventeen bytes!!"); + let digest = DigestInfo::try_new(VALID_HASH, data.len())?; + slow.update_oneshot(digest, data.clone()).await?; + assert_eq!( + slow.has(digest).await?, + Some(4096), + "the filesystem store reports its block-rounded size" + ); + let store = Store::new(FastSlowStore::new( + &FastSlowSpec { + fast: StoreSpec::Memory(MemorySpec::default()), + slow: StoreSpec::Filesystem(slow_spec), + fast_direction: StoreDirection::default(), + slow_direction: StoreDirection::default(), + bypass_dedup_threshold_bytes: 0, + }, + fast.clone(), + slow.clone(), + )); + let got = store.get_part_unchunked(digest, 0, None).await?; + assert_eq!(got, data); + assert_eq!(fast.has(digest).await?, Some(data.len() as u64)); + Ok(()) +} From c6a62b71e1b1a0410c8ddc8350b2f7e9d3a3935d Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 22 Oct 2026 10:15:06 +0000 Subject: [PATCH 42/46] fix(fast_slow_store): don't declare a disputed size exact when filling the fast tier FastSlowStore's populate took the size it declares to the fast tier, as ExactSize, from the slow tier's has(). Since 086289a063 (#2974) the populate also checks the bytes it received against that size, so any slow tier whose has() is not the blob's byte count fails every fill: FilesystemStore reports its block-rounded footprint (4096 for a 17-byte blob at the default block size), dedup sums those footprints, and GrpcStore in AC mode reports u64::MAX. Memory over filesystem, as in the filesystem_cas.json5 dedup index store and the Kubernetes example's action cache, fails every read that misses the fast tier. Declaring ExactSize(digest.size_bytes()) for every digest key would fix a CAS but break the stores where a digest's size is not the size of what it names: an action cache keys an ActionResult by the action's digest, and a dedup index store keys a blob's index by the blob's digest. A filesystem fast tier would then refuse every such entry. So for a digest key the populate now checks one answer against the other. When has() agrees with the digest's size, that size is declared exact. When they disagree, neither is trusted and the populate declares MaxSize of the larger, so the fast tier accepts the bytes it actually gets and the received-vs-declared check (which only applies to ExactSize) cannot misfire. A string key, which has no digest, keeps has()'s answer but as a MaxSize for the same reason. This is the follow-up #2953 reverted #2881 to wait for; #2974 relanded the check without it. Co-Authored-By: Claude Fable 5.1 --- nativelink-store/src/fast_slow_store.rs | 53 +++++++++++++++++-------- 1 file changed, 36 insertions(+), 17 deletions(-) diff --git a/nativelink-store/src/fast_slow_store.rs b/nativelink-store/src/fast_slow_store.rs index 08d91800559..bba1b3471a2 100644 --- a/nativelink-store/src/fast_slow_store.rs +++ b/nativelink-store/src/fast_slow_store.rs @@ -198,6 +198,25 @@ impl FastSlowStore { } } + /// The size to declare to the fast store when filling it with `key`, given + /// the size the slow store's `has()` reported. + /// + /// For a digest there are two answers, and when they agree the size is + /// exact. When they disagree this store cannot tell which is wrong: a slow + /// tier can report a physical size (a filesystem store reports its + /// block-rounded footprint), and an action cache or a dedup index keys an + /// entry by a digest whose size is not the entry's. Declaring either as + /// exact makes a fast store that checks `ExactSize`, as the filesystem + /// store does, refuse a correct blob, so the larger is declared as a bound. + fn fast_store_upload_size(key: &StoreKey<'_>, slow_size: u64) -> UploadSizeInfo { + match Self::digest_size_bytes(key) { + Some(digest_size) if digest_size != slow_size => { + UploadSizeInfo::MaxSize(max(digest_size, slow_size)) + } + _ => UploadSizeInfo::ExactSize(slow_size), + } + } + /// Whether a read should skip dedup and hit the slow store directly. A /// threshold of 0 disables the bypass so every read goes through dedup. fn should_bypass_dedup(&self, key: &StoreKey<'_>) -> bool { @@ -301,23 +320,23 @@ impl FastSlowStore { ); UploadSizeInfo::MaxSize(u64::MAX) } else { - UploadSizeInfo::ExactSize(self - .slow_store - .has(key.borrow()) - .await - .err_tip(|| "Failed to run has() on slow store")? - .ok_or_else(|| { - name_missing_digest( - make_err!( - Code::NotFound, - "Object {} not found in either fast or slow store. \ - If using multiple workers, ensure all workers share the same CAS storage path.", - key.as_str() - ), - &key, - ) - })? - ) + let slow_size = self + .slow_store + .has(key.borrow()) + .await + .err_tip(|| "Failed to run has() on slow store")? + .ok_or_else(|| { + name_missing_digest( + make_err!( + Code::NotFound, + "Object {} not found in either fast or slow store. \ + If using multiple workers, ensure all workers share the same CAS storage path.", + key.as_str() + ), + &key, + ) + })?; + Self::fast_store_upload_size(&key, slow_size) }; let send_range = offset..length.map_or(u64::MAX, |length| length + offset); From 0cbe03960765b4c7db19856addeaaccfa209b48e Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 22 Oct 2026 11:46:53 +0000 Subject: [PATCH 43/46] docs(config-reference): regenerate main after merging the soak set Co-Authored-By: Claude Fable 5.1 --- .../docs/content/docs/reference/nativelink-config/main.mdx | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx index 3ac09ba144d..97c19799f9e 100644 --- a/web/apps/docs/content/docs/reference/nativelink-config/main.mdx +++ b/web/apps/docs/content/docs/reference/nativelink-config/main.mdx @@ -5,7 +5,7 @@ full: true --- {/* AUTOGENERATED — do not edit by hand. - Source: nativelink-config @ main (9009570165) + Source: nativelink-config @ main (acc4210f3d) Regenerate from web/: bun --filter @nativelink/docs gen:config-reference */} @@ -830,6 +830,7 @@ See [`experimental_cloud_object_store`](#experimental_cloud_object_store) for de | `eviction_policy` | [EvictionPolicy](#evictionpolicy) | — | — | Policy used to evict items out of the store. Failure to set this value will cause items to never be removed from the store causing infinite memory usage. | | `block_size` | integer (uint64) | — | 4kb | The block size of the filesystem for the running machine value is used to determine an entry's actual size on disk consumed For a 4KB block size filesystem, a 1B file actually consumes 4KB | | `max_concurrent_writes` | integer (uint) | — | unlimited | Maximum number of concurrent write operations allowed. Each write involves streaming data to a temp file and calling `sync_all()`, which can saturate disk I/O when many writes happen simultaneously. Limiting concurrency prevents disk saturation from blocking the async runtime. A value of 0 means unlimited (no concurrency limit). | +| `max_concurrent_deletes` | integer (uint) | — | 1024 | Maximum number of concurrent background file-deletion operations. When an entry leaves the store, its backing file is unlinked by a background task. Without a bound, an eviction storm (cache wipe or a flood of overwrites) can spawn these without limit: memory grows with the queued task frames and disk/inode reclamation lags behind what the async runtime can schedule. Bounding it applies backpressure the same way `max_concurrent_writes` does — once this many deletes are in flight, a further delete runs synchronously inline in the evicting path instead of spawning, which throttles eviction to delete throughput and never leaks a file. | | `evict_page_cache` | boolean | — | false | When true, advise the kernel to drop the page cache for each blob after it is written or read (`posix_fadvise` with `POSIX_FADV_DONTNEED`). On real filesystems this takes a globally serialized, all-CPU kernel path that stalls on many-core hosts, and it evicts the page-cache tier that fast-disk deployments rely on. Leave off unless you specifically want to keep this store's I/O out of the page cache. | ## RefSpec @@ -856,7 +857,7 @@ See [`experimental_cloud_object_store`](#experimental_cloud_object_store) for de | `retry` | [Retry](#retry) | — | — | Retry configuration to use when a network request fails. | | `max_concurrent_requests` | integer (uint) | — | `0` | Limit the number of simultaneous upstream requests to this many. A value of zero is treated as unlimited. If the limit is reached the request is queued. | | `connections_per_endpoint` | integer (uint) | — | 1 | The number of connections to make to each specified endpoint to balance the load over multiple TCP connections. | -| `rpc_timeout_s` | integer (uint64) | — | 0 (disabled) | Maximum time (seconds) allowed for a single RPC request (e.g. a `ByteStream.Write` call) before it is cancelled. | +| `rpc_timeout_s` | integer (uint64) | — | 0 (disabled) | Maximum time (seconds) allowed for a single RPC request before it is cancelled with `DeadlineExceeded` (retryable). This applies to every RPC this store makes: unary calls (e.g. `FindMissingBlobs`, `BatchReadBlobs`), `ByteStream.Write` calls, and the establishment of read streams through their first response chunk. The body of an established read stream is not bounded by this deadline. | | `use_legacy_resource_names` | boolean | — | false | Use legacy `ByteStream` resource name format, omitting the digest function component from the path. | | `headers` | map of string to string | — | — | Static headers to attach to every outgoing gRPC request sent to this store's upstream endpoints. Useful for fixed authentication tokens (e.g. `{"authorization": "Bearer "}`) and other static metadata. | | `forward_headers` | array of string | — | — | Header names to forward from the incoming client request to every outgoing upstream request. The header value is taken from the client request that triggered this store operation. Use this to pass through dynamic credentials such as JWT tokens sent by build clients. | @@ -1126,6 +1127,8 @@ The total delay is additive, so this example produces 9.525 to 15.875 s of total | `consider_expired_after_s` | integer (uint32) | — | 0 | If the number of seconds since the `last_modified` time of the object is greater than this value, the object will not be considered "existing". This allows for external tools to delete objects that have not been uploaded in a long time. If a client receives a `NotFound` the client should re-upload the object. | | `max_retry_buffer_per_request` | integer (uint) | — | 5MB | The maximum buffer size to retain in case of a retryable error during upload. Setting this to zero will disable upload buffering; this means that in the event of a failure during upload, the entire upload will be aborted and the client will likely receive an error. | | `multipart_max_concurrent_uploads` | integer (uint) | — | 10 | Maximum number of concurrent `UploadPart` requests per `MultipartUpload`. | +| `multipart_part_size` | integer (uint) | — | 64MB | Target size of each part in a multipart upload, growing only as needed to stay under the provider's 10,000-part ceiling. Parts sized at the 5 MiB provider floor turn a multi-GiB blob into thousands of tiny concurrent requests, which churns HTTP connection pools; a larger part keeps the request count low. Each queued or in-flight part is buffered in memory, so one multipart upload buffers up to roughly `(2 * multipart_max_concurrent_uploads + 1) * multipart_part_size`; `max_concurrent_multipart_uploads` bounds the aggregate across uploads. Values are clamped to the provider's 5 MiB minimum and 5 GiB maximum part sizes. | +| `max_concurrent_multipart_uploads` | integer (uint) | — | 4 | Store-wide cap on the number of blobs concurrently uploaded via the multipart path. Each multipart upload can buffer roughly `(2 * multipart_max_concurrent_uploads + 1) * multipart_part_size`, so this bounds the store's aggregate upload buffer (about 5 GiB at the defaults). Uploads over the 5 MiB multipart threshold wait for a slot; smaller blobs are unaffected. | | `insecure_allow_http` | boolean | — | false | Allow unencrypted HTTP connections. Only use this for local testing. | | `disable_http2` | boolean | — | false | Disable HTTP/2 connections and only use HTTP/1.1. Default client configuration will have HTTP/1.1 and HTTP/2 enabled for connection schemes. HTTP/2 should be disabled if environments have poor support or performance related to HTTP/2. Safe to keep default unless underlying network environment, S3, or GCS API servers specify otherwise. | | `throttle_retry_delay_s` | number (float) | — | 2.0 | Minimum number of seconds to wait after a response that asked the client to slow down (HTTP 503 `ServerBusy` or 429) when it carries no `Retry-After` header.0 | From eb93a28e4604060cc223959d492b24b3b092a729 Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 22 Oct 2026 11:55:11 +0000 Subject: [PATCH 44/46] test: repair three test files after merging the soak set - fast_slow_store_test: the merge of the N14 inflight-waiter test had landed inside the filesystem-slow-tier test's body; close that test. - simple_scheduler_state_manager_test: rebuilt from the first parent plus the N3 placement-version hunks, so both the concurrent-listing test and the stale-placement test are whole. - worker_retry_reservation_test: RunAction carries a dispatch generation since the N3 merge; pass 0. Co-Authored-By: Claude Fable 5.1 --- .../simple_scheduler_state_manager_test.rs | 32 ++++++++++++------- .../tests/worker_retry_reservation_test.rs | 2 ++ .../tests/fast_slow_store_test.rs | 12 ++++--- 3 files changed, 29 insertions(+), 17 deletions(-) diff --git a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs index f95faa7fc8e..9e1831095a9 100644 --- a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs @@ -1921,6 +1921,26 @@ async fn the_listing_reads_records_concurrently_in_the_stores_order() -> Result< OperationFilter { stages: OperationStageFlags::Queued, order_by_priority_direction: Some(OrderDirection::Desc), + ..Default::default() + }, + ) + .await? + .collect::>() + .await; + let mut priorities = Vec::with_capacity(listed.len()); + for operation in &listed { + priorities.push(operation.as_action_info().await?.0.priority); + } + let expected: Vec = (0..QUEUE).rev().collect(); + assert_eq!(priorities, expected, "the listing keeps the store's order"); + assert_eq!( + gauge.peak.load(Ordering::SeqCst), + LISTING_READ_CONCURRENCY, + "the reads run this many at a time" + ); + Ok(()) +} + /// N3: matching and assignment must agree on the action's requirements. /// /// Two schedulers (two matchers) share one action DB and memory escalation @@ -1964,18 +1984,6 @@ async fn assignment_rejects_a_stale_placement_after_escalation() -> Result<(), E }, ) .await? - .collect::>() - .await; - let mut priorities = Vec::with_capacity(listed.len()); - for operation in &listed { - priorities.push(operation.as_action_info().await?.0.priority); - } - let expected: Vec = (0..QUEUE).rev().collect(); - assert_eq!(priorities, expected, "the listing keeps the store's order"); - assert_eq!( - gauge.peak.load(Ordering::SeqCst), - LISTING_READ_CONCURRENCY, - "the reads run this many at a time" .next() .await .expect("the queued operation") diff --git a/nativelink-scheduler/tests/worker_retry_reservation_test.rs b/nativelink-scheduler/tests/worker_retry_reservation_test.rs index 8c734fd567a..c7cd4a8f41c 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, // dispatch generation )))) .await?; assert_eq!( @@ -107,6 +108,7 @@ async fn same_worker_retry_conserves_resource_budget() -> Result<(), Error> { op.clone(), action_needing(4), NOW_TIME, + 0, // dispatch generation )))) .await; // The worker may accept or reject the redundant dispatch; either way the diff --git a/nativelink-store/tests/fast_slow_store_test.rs b/nativelink-store/tests/fast_slow_store_test.rs index 5b3ba03a5bd..6ffdc9d8f75 100644 --- a/nativelink-store/tests/fast_slow_store_test.rs +++ b/nativelink-store/tests/fast_slow_store_test.rs @@ -2584,11 +2584,10 @@ async fn populate_from_filesystem_slow_tier_at_default_block_size() -> Result<() let data = Bytes::from_static(b"seventeen bytes!!"); let digest = DigestInfo::try_new(VALID_HASH, data.len())?; slow.update_oneshot(digest, data.clone()).await?; - assert_eq!( - slow.has(digest).await?, - Some(4096), - "the filesystem store reports its block-rounded size" - ); + // Before #3006 a filesystem store's `has()` reported its block-rounded + // size on disk (4096 here); since then it reports the content length. + // Either way the populate must not refuse the 17 bytes it gets. + assert!(slow.has(digest).await?.is_some()); let store = Store::new(FastSlowStore::new( &FastSlowSpec { fast: StoreSpec::Memory(MemorySpec::default()), @@ -2603,6 +2602,9 @@ async fn populate_from_filesystem_slow_tier_at_default_block_size() -> Result<() let got = store.get_part_unchunked(digest, 0, None).await?; assert_eq!(got, data); assert_eq!(fast.has(digest).await?, Some(data.len() as u64)); + Ok(()) +} + // Regression test for N14: an interrupted ByteStream upload is parked for // resumption, keeping its FastSlowStore::update() future (and in-flight // slow-write registration) alive. Existence-check waiters that subscribed to From 2a23827baf1bb62633b3cdd39279a2851f2a147e Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 22 Oct 2026 13:07:11 +0000 Subject: [PATCH 45/46] fix(scheduler): read the memory store's listing whole, never under the lock inside the stream Since 4edf8224ea (#3001) the filter pass reads listed rows 64 at a time with `buffered`, which keeps asking the listing for the next row while it hands rows downstream. The memory store's listing took the store's lock inside the stream (a `ChunkedStream` whose chunk fetch awaits `inner.lock()`), so under load the next chunk fetch was left parked inside the stream, waiting for the lock, while the matcher wrote for a row it had been handed, on the same task, and waited for the same lock. Once the parked wait had starved past async-lock's fairness threshold the lock was reserved for it, nothing polled it, and the scheduler stopped for good: every state update blocked, actions the workers had finished and uploaded "timed out with no update from their worker", and nothing dispatched again. On railgun a `bazel build //...` of this repository through a three-worker fleet wedged within 15 s on #3001 and on main, every time; v1.7.7, #2954 and v1.8.0 run it clean. The listing is now read whole under one hold of the lock, taken in the `async fn` before the stream exists, and the stream only hands out what was read. The chunk fetch already read the whole range in one go, so the snapshot is what the stream yielded before; the lock just moves out of it. Rows are still handed on one at a time as they are read, which the matcher's mid-pass placement relies on, and `buffered` keeps its concurrency for the Redis store, whose reads are round trips. The test pins what the fix relies on: a row added after the listing was requested is not in it. The lazy listing read it at the first poll. Co-Authored-By: Claude Fable 5.1 --- .../src/memory_awaited_action_db.rs | 104 +++++++----------- .../simple_scheduler_state_manager_test.rs | 45 ++++++++ 2 files changed, 85 insertions(+), 64 deletions(-) diff --git a/nativelink-scheduler/src/memory_awaited_action_db.rs b/nativelink-scheduler/src/memory_awaited_action_db.rs index 384b0237277..26fae954a8d 100644 --- a/nativelink-scheduler/src/memory_awaited_action_db.rs +++ b/nativelink-scheduler/src/memory_awaited_action_db.rs @@ -19,14 +19,13 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::Arc; use async_lock::Mutex; -use futures::{FutureExt, Stream}; +use futures::{FutureExt, Stream, stream}; use nativelink_config::stores::EvictionPolicy; use nativelink_error::{Code, Error, ResultExt, error_if, make_err}; use nativelink_metric::MetricsComponent; use nativelink_util::action_messages::{ ActionInfo, ActionStage, ActionUniqueKey, ActionUniqueQualifier, OperationId, }; -use nativelink_util::chunked_stream::ChunkedStream; use nativelink_util::evicting_map::{EvictingMap, LenEntry}; use nativelink_util::instant_wrapper::InstantWrapper; use nativelink_util::metrics::{ @@ -979,28 +978,25 @@ impl I + Clone + Send + Sync + 'static> Awaite .await } - fn get_all_awaited_actions( + /// The listing is read whole, under one hold of the lock taken here, + /// before the stream exists, and the stream then only hands out what + /// was read. The lock must not be taken inside the stream: a reader + /// that prefetches rows (the filter pass reads 64 at a time) would be + /// left with a lock wait parked inside the stream while it hands rows + /// on, and if the reader then writes, on the same task, it waits for + /// the same lock; the release goes to the parked wait, which nothing + /// polls, and the scheduler stops for good. + async fn get_all_awaited_actions( &self, - ) -> impl Future>, Error>> - { - std::future::ready(Ok(ChunkedStream::new( - Bound::Unbounded, - Bound::Unbounded, - move |start, end, mut output| async move { - let inner = self.inner.lock().await; - let mut maybe_new_start = None; - - for (operation_id, item) in - inner.get_awaited_actions_range(start.as_ref(), end.as_ref()) - { - output.push_back(item); - maybe_new_start = Some(operation_id); - } - - Ok(maybe_new_start - .map(|new_start| ((Bound::Excluded(new_start.clone()), end), output))) - }, - ))) + ) -> Result> + Send, Error> { + let listed: Vec> = { + let inner = self.inner.lock().await; + inner + .get_awaited_actions_range(Bound::Unbounded, Bound::Unbounded) + .map(|(_, item)| Ok(item)) + .collect() + }; + Ok(stream::iter(listed)) } async fn get_by_operation_id( @@ -1010,52 +1006,32 @@ impl I + Clone + Send + Sync + 'static> Awaite Ok(self.inner.lock().await.get_by_operation_id(operation_id)) } - fn get_range_of_actions( + /// Read whole under one hold of the lock, like `get_all_awaited_actions`, + /// and for the same reason. + async fn get_range_of_actions( &self, state: SortedAwaitedActionState, start: Bound, end: Bound, desc: bool, - ) -> impl Future> + Send, Error>> - { - std::future::ready(Ok(ChunkedStream::new( - start, - end, - move |start, end, mut output| async move { - let inner = self.inner.lock().await; - let mut done = true; - let mut new_start = start.as_ref(); - let mut new_end = end.as_ref(); - - let iterator = inner - .get_range_of_actions(state, (start.as_ref(), end.as_ref())) - .map(|res| res.err_tip(|| "In AwaitedActionDb::get_range_of_actions")); - - // TODO(palfrey) This should probably use the `.left()/right()` pattern, - // but that doesn't exist in the std or any libraries we use. - if desc { - for result in iterator.rev() { - let (sorted_awaited_action, item) = - result.err_tip(|| "In AwaitedActionDb::get_range_of_actions")?; - output.push_back(item); - new_end = Bound::Excluded(sorted_awaited_action); - done = false; - } - } else { - for result in iterator { - let (sorted_awaited_action, item) = - result.err_tip(|| "In AwaitedActionDb::get_range_of_actions")?; - output.push_back(item); - new_start = Bound::Excluded(sorted_awaited_action); - done = false; - } - } - if done { - return Ok(None); - } - Ok(Some(((new_start.cloned(), new_end.cloned()), output))) - }, - ))) + ) -> Result> + Send, Error> { + let listed: Vec> = { + let inner = self.inner.lock().await; + let iterator = inner + .get_range_of_actions(state, (start.as_ref(), end.as_ref())) + .map(|res| { + res.map(|(_, item)| item) + .err_tip(|| "In AwaitedActionDb::get_range_of_actions") + }); + // TODO(palfrey) This should probably use the `.left()/right()` pattern, + // but that doesn't exist in the std or any libraries we use. + if desc { + iterator.rev().collect() + } else { + iterator.collect() + } + }; + Ok(stream::iter(listed)) } async fn update_awaited_action(&self, new_awaited_action: AwaitedAction) -> Result<(), Error> { diff --git a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs index 1686aa6f203..1ede82e5d86 100644 --- a/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs +++ b/nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs @@ -1778,3 +1778,48 @@ async fn the_listing_reads_records_concurrently_in_the_stores_order() -> Result< ); Ok(()) } + +/// The memory store's listing is read whole, under one hold of the store's +/// lock taken when the listing is requested, and the stream then only hands +/// out what was read. The lock must not be taken inside the stream: the +/// filter pass reads rows 64 at a time, so while it hands rows on it keeps +/// asking the listing for more, and a lock wait that asking starts is +/// parked inside the stream, polled again only when the next row is asked +/// for. Once that wait has been starved past async-lock's fairness +/// threshold the lock is reserved for it, and when the matcher then writes +/// for a row it was handed, on the same task, it waits for the same lock: +/// nothing polls the parked wait, and the scheduler stops for good. Reading +/// lazily was how that wait got into the stream; a row added after the +/// listing was requested is therefore not in it. +#[nativelink_test] +async fn the_listing_is_read_whole_when_it_is_requested() -> Result<(), Error> { + MockClock::set_time(Duration::from_secs(NOW_TIME)); + let state_mgr = state_manager(Arc::new(WorkerRegistry::new())); + // Kept so the clients count as waiting. + let _clients = add_prioritised_actions(state_mgr.as_ref(), 3).await?; + + let listed = MatchingEngineStateManager::filter_operations( + state_mgr.as_ref(), + OperationFilter { + stages: OperationStageFlags::Queued, + order_by_priority_direction: Some(OrderDirection::Desc), + ..Default::default() + }, + ) + .await?; + // Added after the listing was requested, before any of it was read. + let _late = state_mgr + .add_action( + OperationId::default(), + Arc::new(action_info(make_system_time(0))), + ) + .await?; + + let listed: Vec<_> = listed.collect().await; + assert_eq!( + listed.len(), + 3, + "the listing holds what was there when it was requested" + ); + Ok(()) +} From fcb785f4df86f61ac910f10cd8c208aaeb70182d Mon Sep 17 00:00:00 2001 From: b7r6 Date: Thu, 22 Oct 2026 14:10:25 +0000 Subject: [PATCH 46/46] redis: require Send of the connection in RedisStore's inherent impl `warn_if_eviction_events_are_disabled` holds a `ClientWithPermit` across an await, which the nightly clippy lane (`future_not_send`) rejects while `C` is only `Sync` there. Every connection type used is `Send`, and the `StoreDriver` impl already requires it. Co-Authored-By: Claude Fable 5.1 --- nativelink-store/src/redis_store.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nativelink-store/src/redis_store.rs b/nativelink-store/src/redis_store.rs index d0a3766e231..52c0e96e508 100644 --- a/nativelink-store/src/redis_store.rs +++ b/nativelink-store/src/redis_store.rs @@ -541,8 +541,8 @@ pub fn decode_key<'a>( impl RedisStore where - C: ConnectionLike + Clone + Sync, - M: RedisManager + Sync, + C: ConnectionLike + Clone + Send + Sync, + M: RedisManager + Send + Sync, { /// Used for testing when determinism is required. #[expect(clippy::too_many_arguments)]