StateManager can now be notified of noone listeneing (#1093)

ActionStateResult is now wired up to ActionListener allowing
it to be notified of Drop calls. This will be used to do client
operation id cleanups.

nativelink-scheduler/src/action_scheduler.rs

@@ -29,13 +29,10 @@ pub trait ActionListener: Sync + Send + Unpin {
     /// Returns the client operation id.
     fn client_operation_id(&self) -> &ClientOperationId;
 
-    /// Returns the current action state.
-    fn action_state(&self) -> Arc<ActionState>;
-
     /// Waits for the action state to change.
     fn changed(
         &mut self,
-    ) -> Pin<Box<dyn Future<Output = Result<Arc<ActionState>, Error>> + Send + Sync + '_>>;
+    ) -> Pin<Box<dyn Future<Output = Result<Arc<ActionState>, Error>> + Send + '_>>;
 }
 
 /// ActionScheduler interface is responsible for interactions between the scheduler

nativelink-scheduler/src/cache_lookup_scheduler.rs

@@ -116,13 +116,9 @@ impl ActionListener for CachedActionListener {
         &self.client_operation_id
     }
 
-    fn action_state(&self) -> Arc<ActionState> {
-        self.action_state.clone()
-    }
-
     fn changed(
         &mut self,
-    ) -> Pin<Box<dyn Future<Output = Result<Arc<ActionState>, Error>> + Send + Sync + '_>> {
+    ) -> Pin<Box<dyn Future<Output = Result<Arc<ActionState>, Error>> + Send + '_>> {
         Box::pin(async { Ok(self.action_state.clone()) })
     }
 }

nativelink-scheduler/src/default_action_listener.rs

@@ -39,31 +39,29 @@ impl DefaultActionListener {
             action_state,
         }
     }
+
+    pub async fn changed(&mut self) -> Result<Arc<ActionState>, Error> {
+        self.action_state.changed().await.map_or_else(
+            |e| {
+                Err(make_err!(
+                    Code::Internal,
+                    "Sender of ActionState went away unexpectedly - {e:?}"
+                ))
+            },
+            |()| Ok(self.action_state.borrow_and_update().clone()),
+        )
+    }
 }
 
 impl ActionListener for DefaultActionListener {
     fn client_operation_id(&self) -> &ClientOperationId {
         &self.client_operation_id
     }
 
-    fn action_state(&self) -> Arc<ActionState> {
-        self.action_state.borrow().clone()
-    }
-
     fn changed(
         &mut self,
-    ) -> Pin<Box<dyn Future<Output = Result<Arc<ActionState>, Error>> + Send + Sync + '_>> {
-        Box::pin(async move {
-            self.action_state.changed().await.map_or_else(
-                |e| {
-                    Err(make_err!(
-                        Code::Internal,
-                        "Sender of ActionState went away unexpectedly - {e:?}"
-                    ))
-                },
-                |()| Ok(self.action_state.borrow_and_update().clone()),
-            )
-        })
+    ) -> Pin<Box<dyn Future<Output = Result<Arc<ActionState>, Error>> + Send + '_>> {
+        Box::pin(self.changed())
     }
 }
 

nativelink-scheduler/src/scheduler_state/client_action_state_result.rs

@@ -47,6 +47,6 @@ impl ActionStateResult for ClientActionStateResult {
     }
 
     async fn as_action_info(&self) -> Result<Arc<ActionInfo>, Error> {
-        unimplemented!()
+        todo!()
     }
 }

nativelink-scheduler/src/simple_scheduler.rs

@@ -17,20 +17,19 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 use futures::{Future, Stream};
-use nativelink_error::{Error, ResultExt};
+use nativelink_error::{make_err, Code, Error, ResultExt};
 use nativelink_util::action_messages::{
-    ActionInfo, ActionStage, ClientOperationId, OperationId, WorkerId,
+    ActionInfo, ActionStage, ActionState, ClientOperationId, OperationId, WorkerId,
 };
 use nativelink_util::metrics_utils::Registry;
 use nativelink_util::spawn;
 use nativelink_util::task::JoinHandleDropGuard;
-use tokio::sync::Notify;
+use tokio::sync::{watch, Notify};
 use tokio::time::Duration;
 use tokio_stream::StreamExt;
 use tracing::{event, Level};
 
 use crate::action_scheduler::{ActionListener, ActionScheduler};
-use crate::default_action_listener::DefaultActionListener;
 use crate::operation_state_manager::{
     ActionStateResult, ClientStateManager, MatchingEngineStateManager, OperationFilter,
     OperationStageFlags,
@@ -53,6 +52,60 @@ const DEFAULT_RETAIN_COMPLETED_FOR_S: u64 = 60;
 /// If this changes, remember to change the documentation in the config.
 const DEFAULT_MAX_JOB_RETRIES: usize = 3;
 
+struct SimpleSchedulerActionListener {
+    client_operation_id: ClientOperationId,
+    action_state_result: Arc<dyn ActionStateResult>,
+    maybe_receiver: Option<watch::Receiver<Arc<ActionState>>>,
+}
+
+impl SimpleSchedulerActionListener {
+    fn new(
+        client_operation_id: ClientOperationId,
+        action_state_result: Arc<dyn ActionStateResult>,
+    ) -> Self {
+        Self {
+            client_operation_id,
+            action_state_result,
+            maybe_receiver: None,
+        }
+    }
+}
+
+impl ActionListener for SimpleSchedulerActionListener {
+    fn client_operation_id(&self) -> &ClientOperationId {
+        &self.client_operation_id
+    }
+
+    fn changed(
+        &mut self,
+    ) -> Pin<Box<dyn Future<Output = Result<Arc<ActionState>, Error>> + Send + '_>> {
+        let action_state_result = self.action_state_result.clone();
+        Box::pin(async move {
+            let receiver = match &mut self.maybe_receiver {
+                Some(receiver) => receiver,
+                None => {
+                    let mut receiver = action_state_result
+                        .as_receiver()
+                        .await
+                        .err_tip(|| "In SimpleSchedulerActionListener::changed getting receiver")?
+                        .into_owned();
+                    receiver.mark_changed();
+                    self.maybe_receiver = Some(receiver.clone());
+                    self.maybe_receiver.as_mut().unwrap()
+                }
+            };
+            receiver.changed().await.map_err(|_| {
+                make_err!(
+                    Code::Internal,
+                    "Sender hungup in SimpleSchedulerActionListener::changed()"
+                )
+            })?;
+            let result = receiver.borrow().clone();
+            Ok(result)
+        })
+    }
+}
+
 /// Engine used to manage the queued/running tasks and relationship with
 /// the worker nodes. All state on how the workers and actions are interacting
 /// should be held in this struct.
@@ -89,15 +142,11 @@ impl SimpleScheduler {
             .client_state_manager
             .add_action(client_operation_id.clone(), action_info)
             .await?;
-        add_action_result
-            .as_receiver()
-            .await
-            .map(move |receiver| -> Pin<Box<dyn ActionListener>> {
-                Box::pin(DefaultActionListener::new(
-                    client_operation_id,
-                    receiver.into_owned(),
-                ))
-            })
+
+        Ok(Box::pin(SimpleSchedulerActionListener::new(
+            client_operation_id,
+            add_action_result,
+        )))
     }
 
     async fn clean_recently_completed_actions(&self) {
@@ -127,16 +176,12 @@ impl SimpleScheduler {
 
         let mut stream = filter_result
             .err_tip(|| "In SimpleScheduler::find_by_client_operation_id getting filter result")?;
-        let Some(result) = stream.next().await else {
+        let Some(action_state_result) = stream.next().await else {
             return Ok(None);
         };
-        Ok(Some(Box::pin(DefaultActionListener::new(
+        Ok(Some(Box::pin(SimpleSchedulerActionListener::new(
             client_operation_id.clone(),
-            result
-                .as_receiver()
-                .await
-                .err_tip(|| "In SimpleScheduler::find_by_client_operation_id getting receiver")?
-                .into_owned(),
+            action_state_result,
         ))))
     }
 

nativelink-scheduler/tests/simple_scheduler_test.rs

@@ -393,14 +393,14 @@ async fn remove_worker_reschedules_multiple_running_job_test() -> Result<(), Err
     {
         let expected_action_stage = ActionStage::Executing;
         // Client should get notification saying it's being executed.
-        let action_state = client2_action_listener.action_state();
+        let action_state = client1_action_listener.changed().await.unwrap();
         // We now know the name of the action so populate it.
         assert_eq!(&action_state.stage, &expected_action_stage);
     }
     {
         let expected_action_stage = ActionStage::Executing;
         // Client should get notification saying it's being executed.
-        let action_state = client2_action_listener.action_state();
+        let action_state = client2_action_listener.changed().await.unwrap();
         // We now know the name of the action so populate it.
         assert_eq!(&action_state.stage, &expected_action_stage);
     }

nativelink-service/src/execution_server.rs

@@ -165,7 +165,7 @@ pub struct ExecutionServer {
     instance_infos: HashMap<InstanceName, InstanceInfo>,
 }
 
-type ExecuteStream = Pin<Box<dyn Stream<Item = Result<Operation, Status>> + Send + Sync + 'static>>;
+type ExecuteStream = Pin<Box<dyn Stream<Item = Result<Operation, Status>> + Send + 'static>>;
 
 impl ExecutionServer {
     pub fn new(