[READABLE DIFF] migrate-1.7.3: 18/22 verified fixes on upstream v1.7.3 (4 timing-relevant held)

nativelink-config/src/cas_server.rs

@@ -673,7 +673,7 @@
     ///
     /// Interaction with `serve_compression = "zstd"`: a preserved original
     /// takes precedence over transcoding, so a compressed push is served
     /// back in the CLIENT's original codec (no re-encode, saving CPU) and
     /// only `Compression: none` pushes are transcoded to zstd. If a
     /// preserved compressed blob is later evicted, serving degrades
     /// gracefully to the uncompressed NAR (`Compression: none`); the
@@ -724,7 +724,7 @@
     /// Maximum number of zstd NAR transcodes running concurrently when
     /// `serve_compression` is `"zstd"`. A transcode streams a whole NAR
     /// through a zstd encoder and spools the output, so unbounded fan-out
     /// of narinfo PUTs would balloon CPU and temp disk. Concurrent PUTs for
     /// the SAME NAR are additionally coalesced into a single transcode.
     ///
     /// Default: 8
@@ -762,7 +762,7 @@
     /// satisfies a `completeness_checking` `path_info_store` immediately;
     /// see the store-composition notes in `examples/nix_cache.json5`.
     /// Read-through populates the stores regardless of `read_only` (which
     /// only gates client PUTs).
     ///
     /// Default: [] (read-through disabled; a local miss is a plain 404)
     #[serde(default, skip_serializing_if = "Vec::is_empty")]
@@ -3310,7 +3310,6 @@
 
 #[cfg(test)]
 mod tests {
-    use tracing_test::traced_test;
 
     use super::*;
 

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",

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" }
 

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<Notify>,
 
+    /// 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<WorkerLoad>,
         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,24 +2160,33 @@ 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)
         };
 
         // 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
                         .is_executing_on_worker(&operation_id, &worker_id)
                         .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<ApiWorkerScheduler> {
+        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"
+        );
+    }
+}

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 = {

nativelink-scheduler/src/memory_awaited_action_db.rs

@@ -637,6 +637,20 @@ impl<I: InstantWrapper, NowFn: Fn() -> 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,