ClientKeepAlive update action ClientKeepAlive (#1580)

When the scheduler was updated to add the keep alive to the AwaitedAction
the MemoryAwaitedActionDb was not updated to set this when a ClientKeepAlive
was received.

Fix the test client_reconnect_keeps_action_alive which was not performing
the eviction due to optimisations in the filter_operations function which
then detected the issue.

Then update the ActionEvent::ClientKeepAlive event handler to update the
client keep alive timestamp in the AwaitedAction.

Fixes #1579.

nativelink-scheduler/src/awaited_action_db/awaited_action.rs

@@ -171,6 +171,7 @@ impl AwaitedAction {
     pub(crate) fn last_client_keepalive_timestamp(&self) -> SystemTime {
         self.last_client_keepalive_timestamp
     }
+
     pub(crate) fn update_client_keep_alive(&mut self, now: SystemTime) {
         self.last_client_keepalive_timestamp = now;
     }

nativelink-scheduler/src/awaited_action_db/mod.rs

@@ -15,6 +15,7 @@
 use std::cmp;
 use std::ops::Bound;
 use std::sync::Arc;
+use std::time::Duration;
 
 pub use awaited_action::{AwaitedAction, AwaitedActionSortKey};
 use futures::{Future, Stream};
@@ -25,6 +26,9 @@ use serde::{Deserialize, Serialize};
 
 mod awaited_action;
 
+/// Duration to wait before sending client keep alive messages.
+pub const CLIENT_KEEPALIVE_DURATION: Duration = Duration::from_secs(10);
+
 /// A simple enum to represent the state of an `AwaitedAction`.
 #[derive(Debug, Clone, Copy)]
 pub enum SortedAwaitedActionState {

nativelink-scheduler/src/memory_awaited_action_db.rs

@@ -15,7 +15,6 @@
 use std::collections::{BTreeMap, BTreeSet, HashMap};
 use std::ops::{Bound, RangeBounds};
 use std::sync::Arc;
-use std::time::Duration;
 
 use async_lock::Mutex;
 use futures::{FutureExt, Stream};
@@ -35,15 +34,12 @@ use tracing::{event, Level};
 
 use crate::awaited_action_db::{
     AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, SortedAwaitedAction,
-    SortedAwaitedActionState,
+    SortedAwaitedActionState, CLIENT_KEEPALIVE_DURATION,
 };
 
 /// Number of events to process per cycle.
 const MAX_ACTION_EVENTS_RX_PER_CYCLE: usize = 1024;
 
-/// Duration to wait before sending client keep alive messages.
-const CLIENT_KEEPALIVE_DURATION: Duration = Duration::from_secs(10);
-
 /// Represents a client that is currently listening to an action.
 /// When the client is dropped, it will send the `AwaitedAction` to the
 /// `event_tx` if there are other cleanups needed.
@@ -452,11 +448,21 @@ impl<I: InstantWrapper, NowFn: Fn() -> I + Clone + Send + Sync> AwaitedActionDbI
                     }
                 }
                 ActionEvent::ClientKeepAlive(client_id) => {
-                    let maybe_size = self
+                    if let Some(client_awaited_action) = self
                         .client_operation_to_awaited_action
-                        .size_for_key(&client_id)
-                        .await;
-                    if maybe_size.is_none() {
+                        .get(&client_id)
+                        .await
+                    {
+                        if let Some(awaited_action_sender) = self
+                            .operation_id_to_awaited_action
+                            .get(&client_awaited_action.operation_id)
+                        {
+                            awaited_action_sender.send_if_modified(|awaited_action| {
+                                awaited_action.update_client_keep_alive((self.now_fn)().now());
+                                false
+                            });
+                        }
+                    } else {
                         event!(
                             Level::ERROR,
                             ?client_id,

nativelink-scheduler/src/simple_scheduler.rs

@@ -38,7 +38,7 @@ use tokio_stream::StreamExt;
 use tracing::{event, info_span, Level};
 
 use crate::api_worker_scheduler::ApiWorkerScheduler;
-use crate::awaited_action_db::AwaitedActionDb;
+use crate::awaited_action_db::{AwaitedActionDb, CLIENT_KEEPALIVE_DURATION};
 use crate::platform_property_manager::PlatformPropertyManager;
 use crate::simple_scheduler_state_manager::SimpleSchedulerStateManager;
 use crate::worker::{ActionInfoWithProps, Worker, WorkerTimestamp};
@@ -368,6 +368,12 @@ impl SimpleScheduler {
         if client_action_timeout_s == 0 {
             client_action_timeout_s = DEFAULT_CLIENT_ACTION_TIMEOUT_S;
         }
+        // This matches the value of CLIENT_KEEPALIVE_DURATION which means that
+        // tasks are going to be dropped all over the place, this isn't a good
+        // setting.
+        if client_action_timeout_s <= CLIENT_KEEPALIVE_DURATION.as_secs() {
+            event!(Level::ERROR, client_action_timeout_s, "Setting client_action_timeout_s to less than the client keep alive interval is going to cause issues, please set above {}.", CLIENT_KEEPALIVE_DURATION.as_secs());
+        }
 
         let mut max_job_retries = spec.max_job_retries;
         if max_job_retries == 0 {

nativelink-scheduler/src/store_awaited_action_db.rs

@@ -38,14 +38,11 @@ use tracing::{event, Level};
 
 use crate::awaited_action_db::{
     AwaitedAction, AwaitedActionDb, AwaitedActionSubscriber, SortedAwaitedAction,
-    SortedAwaitedActionState,
+    SortedAwaitedActionState, CLIENT_KEEPALIVE_DURATION,
 };
 
 type ClientOperationId = OperationId;
 
-/// Duration to wait before sending client keep alive messages.
-const CLIENT_KEEPALIVE_DURATION: Duration = Duration::from_secs(10);
-
 /// Maximum number of retries to update client keep alive.
 const MAX_RETRIES_FOR_CLIENT_KEEPALIVE: u32 = 8;
 

nativelink-scheduler/tests/simple_scheduler_test.rs

@@ -49,7 +49,8 @@ use nativelink_util::action_messages::{
 use nativelink_util::common::DigestInfo;
 use nativelink_util::instant_wrapper::MockInstantWrapped;
 use nativelink_util::operation_state_manager::{
-    ActionStateResult, ClientStateManager, OperationFilter, UpdateOperationType,
+    ActionStateResult, ClientStateManager, OperationFilter, OperationStageFlags,
+    UpdateOperationType,
 };
 use nativelink_util::platform_properties::{PlatformProperties, PlatformPropertyValue};
 use pretty_assertions::assert_eq;
@@ -2292,18 +2293,15 @@ async fn client_reconnect_keeps_action_alive() -> Result<(), Error> {
         assert_eq!(poll!(&mut changed_fut), Poll::Pending);
         tokio::task::yield_now().await;
         // Eviction happens when someone touches the internal
-        // evicting map. So we constantly ask for some other client
-        // to trigger eviction logic.
-        assert!(scheduler
+        // evicting map.  So we constantly ask for all queued actions.
+        // Regression: https://github.com/TraceMachina/nativelink/issues/1579
+        let mut stream = scheduler
             .filter_operations(OperationFilter {
-                client_operation_id: Some(OperationId::from("dummy_client_id")),
+                stages: OperationStageFlags::Queued,
                 ..Default::default()
             })
-            .await
-            .unwrap()
-            .next()
-            .await
-            .is_none());
+            .await?;
+        while stream.next().await.is_some() {}
     }
 
     Ok(())