fix(scheduler): client keep-alive stamps survive concurrent versioned writes

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

nativelink-scheduler/tests/client_keepalive_clobber_test.rs

@@ -0,0 +1,156 @@
+// Copyright 2026 The NativeLink Authors. All rights reserved.
+//
+// Licensed under the Functional Source License, Version 1.1, Apache 2.0 Future License (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    See LICENSE file for details
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Repro: client keep-alive writes bypass the version CAS in the memory
+//! backend, so a concurrent read-modify-write (a worker update) that
+//! snapshotted the action before the keep-alive silently reverts
+//! `last_client_keepalive_timestamp` — and passes the version check, so the
+//! `Code::Aborted` retry defense in `SimpleSchedulerStateManager` never
+//! fires. Repeated indefinitely, the client-timeout sweep kills an action
+//! whose client is alive and sending keep-alives.
+
+use core::time::Duration;
+use std::collections::HashMap;
+use std::sync::Arc;
+use std::time::SystemTime;
+
+use mock_instant::thread_local::MockClock;
+use nativelink_config::stores::EvictionPolicy;
+use nativelink_error::Error;
+use nativelink_macro::nativelink_test;
+use nativelink_scheduler::awaited_action_db::{
+    AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, CLIENT_KEEPALIVE_DURATION,
+};
+use nativelink_scheduler::memory_awaited_action_db::MemoryAwaitedActionDb;
+use nativelink_util::action_messages::{
+    ActionInfo, ActionUniqueKey, ActionUniqueQualifier, OperationId,
+};
+use nativelink_util::common::DigestInfo;
+use nativelink_util::digest_hasher::DigestHasherFunc;
+use nativelink_util::instant_wrapper::MockInstantWrapped;
+use pretty_assertions::assert_eq;
+use tokio::sync::Notify;
+
+const INSTANCE_NAME: &str = "keepalive_clobber_instance";
+
+fn make_action_info() -> Arc<ActionInfo> {
+    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(&current);
+            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(())
+}

nativelink-scheduler/tests/simple_scheduler_state_manager_test.rs

@@ -1290,6 +1290,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<T: AwaitedActionDb>(Arc<T>);
+
+impl<T: AwaitedActionDb> MetricsComponent for SharedDb<T> {
+    fn publish(
+        &self,
+        kind: MetricKind,
+        field_metadata: MetricFieldData,
+    ) -> Result<MetricPublishKnownKindData, nativelink_metric::Error> {
+        self.0.publish(kind, field_metadata)
+    }
+}
+
+impl<T: AwaitedActionDb> AwaitedActionDb for SharedDb<T> {
+    type Subscriber = T::Subscriber;
+
+    async fn get_awaited_action_by_id(
+        &self,
+        client_operation_id: &OperationId,
+    ) -> Result<Option<Self::Subscriber>, Error> {
+        self.0.get_awaited_action_by_id(client_operation_id).await
+    }
+
+    async fn get_all_awaited_actions(
+        &self,
+    ) -> Result<impl Stream<Item = Result<Self::Subscriber, Error>> + Send, Error> {
+        self.0.get_all_awaited_actions().await
+    }
+
+    async fn get_by_operation_id(
+        &self,
+        operation_id: &OperationId,
+    ) -> Result<Option<Self::Subscriber>, Error> {
+        self.0.get_by_operation_id(operation_id).await
+    }
+
+    async fn get_range_of_actions(
+        &self,
+        state: SortedAwaitedActionState,
+        start: Bound<SortedAwaitedAction>,
+        end: Bound<SortedAwaitedAction>,
+        desc: bool,
+    ) -> Result<impl Stream<Item = Result<Self::Subscriber, Error>> + 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<ActionInfo>,
+        no_event_action_timeout: Duration,
+    ) -> Result<Self::Subscriber, Error> {
+        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));