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