ClientKeepAlive update action ClientKeepAlive

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/memory_awaited_action_db.rs

@@ -452,11 +452,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

@@ -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 <= 10 {
+            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 10.");
+        }
 
         let mut max_job_retries = spec.max_job_retries;
         if max_job_retries == 0 {

nativelink-scheduler/tests/simple_scheduler_test.rs

@@ -49,7 +49,7 @@ 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 +2292,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(())