Repository: TraceMachina/nativelink (language: rust)

Issue:
Action times out if client is listening
The addition of client_action_timeout_s means that an action times out if a client isn't actively polling for it.  However, the RBE WaitExecution allows passive polling of an action.  This is not considered and therefore long actions time out an signal down this channel that they have timed out.  If there is an active WaitExecution for an action then it should not be subject to the client action timeout.


Discussion / hints:
Wow! You found the smoking gun (I think). Thank you so much Chris!

We will look into this right away. 
I think this is handled by the code in `OperationSubscriber::changed` for `store_awaited_action_db`, however this same logic doesn't apply in `memory_awaited_action_db`.
This might fix it:
```
diff --git a/nativelink-scheduler/src/memory_awaited_action_db.rs b/nativelink-scheduler/src/memory_awaited_action_db.rs
index ef703329..55b797a0 100644
--- a/nativelink-scheduler/src/memory_awaited_action_db.rs
+++ b/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,
```

Fix patch:
diff --git a/nativelink-scheduler/src/awaited_action_db/awaited_action.rs b/nativelink-scheduler/src/awaited_action_db/awaited_action.rs
index dc7be982a..fd8d3bc8d 100644
--- a/nativelink-scheduler/src/awaited_action_db/awaited_action.rs
+++ b/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;
     }
diff --git a/nativelink-scheduler/src/awaited_action_db/mod.rs b/nativelink-scheduler/src/awaited_action_db/mod.rs
index 6126c582f..fba678220 100644
--- a/nativelink-scheduler/src/awaited_action_db/mod.rs
+++ b/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 {
diff --git a/nativelink-scheduler/src/memory_awaited_action_db.rs b/nativelink-scheduler/src/memory_awaited_action_db.rs
index ef7033296..628c69975 100644
--- a/nativelink-scheduler/src/memory_awaited_action_db.rs
+++ b/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,
diff --git a/nativelink-scheduler/src/simple_scheduler.rs b/nativelink-scheduler/src/simple_scheduler.rs
index c91db82fc..2826ae99d 100644
--- a/nativelink-scheduler/src/simple_scheduler.rs
+++ b/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 {
diff --git a/nativelink-scheduler/src/store_awaited_action_db.rs b/nativelink-scheduler/src/store_awaited_action_db.rs
index f326456c4..d40a3199c 100644
--- a/nativelink-scheduler/src/store_awaited_action_db.rs
+++ b/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;
 


Test patch:
diff --git a/nativelink-scheduler/tests/simple_scheduler_test.rs b/nativelink-scheduler/tests/simple_scheduler_test.rs
index 9644d9131..bca6b58cb 100644
--- a/nativelink-scheduler/tests/simple_scheduler_test.rs
+++ b/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(())