Implement MatchingEngineStateManager

nativelink-scheduler/BUILD.bazel

@@ -20,6 +20,7 @@ rust_library(
         "src/scheduler_state/awaited_action.rs",
         "src/scheduler_state/client_action_state_result.rs",
         "src/scheduler_state/completed_action.rs",
+        "src/scheduler_state/matching_engine_action_state_result.rs",
         "src/scheduler_state/metrics.rs",
         "src/scheduler_state/mod.rs",
         "src/scheduler_state/state_manager.rs",

nativelink-scheduler/src/operation_state_manager.rs

@@ -136,7 +136,7 @@ pub trait MatchingEngineStateManager {
 
     /// Update that state of an operation.
     async fn update_operation(
-        &self,
+        &mut self,
         operation_id: OperationId,
         worker_id: Option<WorkerId>,
         action_stage: Result<ActionStage, Error>,

nativelink-scheduler/src/scheduler_state/matching_engine_action_state_result.rs

@@ -0,0 +1,46 @@
+// Copyright 2024 The NativeLink Authors. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//    http://www.apache.org/licenses/LICENSE-2.0
+//
+// 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.
+
+use std::sync::Arc;
+
+use async_trait::async_trait;
+use nativelink_error::Error;
+use nativelink_util::action_messages::{ActionInfo, ActionState};
+use tokio::sync::watch;
+
+use crate::operation_state_manager::ActionStateResult;
+
+pub struct MatchingEngineActionStateResult {
+    pub action_info: Arc<ActionInfo>,
+}
+impl MatchingEngineActionStateResult {
+    pub(crate) fn new(action_info: Arc<ActionInfo>) -> Self {
+        Self { action_info }
+    }
+}
+
+#[async_trait]
+impl ActionStateResult for MatchingEngineActionStateResult {
+    async fn as_state(&self) -> Result<Arc<ActionState>, Error> {
+        unimplemented!()
+    }
+
+    async fn as_receiver(&self) -> Result<&'_ watch::Receiver<Arc<ActionState>>, Error> {
+        unimplemented!()
+    }
+
+    async fn as_action_info(&self) -> Result<Arc<ActionInfo>, Error> {
+        Ok(self.action_info.clone())
+    }
+}

nativelink-scheduler/src/scheduler_state/metrics.rs

@@ -19,6 +19,14 @@ pub(crate) struct Metrics {
     pub(crate) add_action_joined_running_action: CounterWithTime,
     pub(crate) add_action_joined_queued_action: CounterWithTime,
     pub(crate) add_action_new_action_created: CounterWithTime,
+
+    pub(crate) workers_evicted: CounterWithTime,
+    pub(crate) workers_evicted_with_running_action: CounterWithTime,
+
+    pub(crate) retry_action: CounterWithTime,
+    pub(crate) retry_action_max_attempts_reached: CounterWithTime,
+    pub(crate) retry_action_no_more_listeners: CounterWithTime,
+    pub(crate) retry_action_but_action_missing: CounterWithTime,
 }
 
 impl Metrics {
@@ -41,5 +49,45 @@ impl Metrics {
             "Stats about add_action().",
             vec![("result".into(), "new_action_created".into())],
         );
+
+        {
+            c.publish(
+                "workers_evicted_total",
+                &self.workers_evicted,
+                "The number of workers evicted from scheduler.",
+            );
+            c.publish(
+                "workers_evicted_with_running_action",
+                &self.workers_evicted_with_running_action,
+                "The number of jobs cancelled because worker was evicted from scheduler.",
+            );
+        }
+
+        {
+            c.publish_with_labels(
+                "retry_action",
+                &self.retry_action,
+                "Stats about retry_action().",
+                vec![("result".into(), "success".into())],
+            );
+            c.publish_with_labels(
+                "retry_action",
+                &self.retry_action_max_attempts_reached,
+                "Stats about retry_action().",
+                vec![("result".into(), "max_attempts_reached".into())],
+            );
+            c.publish_with_labels(
+                "retry_action",
+                &self.retry_action_no_more_listeners,
+                "Stats about retry_action().",
+                vec![("result".into(), "no_more_listeners".into())],
+            );
+            c.publish_with_labels(
+                "retry_action",
+                &self.retry_action_but_action_missing,
+                "Stats about retry_action().",
+                vec![("result".into(), "action_missing".into())],
+            );
+        }
     }
 }

nativelink-scheduler/src/scheduler_state/mod.rs

@@ -15,6 +15,7 @@
 pub(crate) mod awaited_action;
 pub(crate) mod client_action_state_result;
 pub(crate) mod completed_action;
+pub(crate) mod matching_engine_action_state_result;
 pub(crate) mod metrics;
 pub(crate) mod state_manager;
 pub(crate) mod workers;