Fix sleep await duration hang

nativelink-proto/com/github/trace_machina/nativelink/remote_execution/worker_api.proto

@@ -99,25 +99,8 @@ message ExecuteResult {
     /// that initially sent the job as part of the BRE protocol.
     string instance_name = 6;
 
-    /// The original execution digest request for this response. The scheduler knows what it
-    /// should be, but we do safety checks to ensure it really is the request we expected.
-    build.bazel.remote.execution.v2.Digest action_digest = 2;
-
-    /// The salt originally sent along with the StartExecute request. This salt is used
-    /// as a seed for cases where the execution digest should never be cached or merged
-    /// with other jobs. This salt is added to the hash function used to compute jobs that
-    /// are running or cached.
-    uint64 salt = 3;
-
-    // The digest function that was used to compute the action digest
-    // and all related blobs.
-    //
-    // If the digest function used is one of MD5, MURMUR3, SHA1, SHA256,
-    // SHA384, SHA512, or VSO, the client MAY leave this field unset. In
-    // that case the server SHOULD infer the digest function using the
-    // length of the action digest hash and the digest functions announced
-    // in the server's capabilities.
-    build.bazel.remote.execution.v2.DigestFunction.Value digest_function = 7;
+    /// The operation ID that was executed.
+    string operation_id = 8;
 
     /// The actual response data.
     oneof result {
@@ -131,7 +114,7 @@ message ExecuteResult {
         google.rpc.Status internal_error = 5;
     }
 
-    reserved 8; // NextId.
+    reserved 9; // NextId.
 }
 
 /// Result sent back from the server when a node connects.
@@ -179,14 +162,14 @@ message StartExecute {
     /// The action information used to execute job.
     build.bazel.remote.execution.v2.ExecuteRequest execute_request = 1;
 
-    /// See documentation in ExecuteResult::salt.
-    uint64 salt = 2;
+    /// Id of the operation.
+    string operation_id = 4;
 
     /// The time at which the command was added to the queue to allow population
     /// of the ActionResult.
     google.protobuf.Timestamp queued_timestamp = 3;
 
-    reserved 4; // NextId.
+    reserved 5; // NextId.
 }
 
 /// This is a special message used to save actions into the CAS that can be used

nativelink-proto/genproto/com.github.trace_machina.nativelink.remote_execution.pb.rs

@@ -60,31 +60,9 @@ pub struct ExecuteResult {
     /// / that initially sent the job as part of the BRE protocol.
     #[prost(string, tag = "6")]
     pub instance_name: ::prost::alloc::string::String,
-    /// / The original execution digest request for this response. The scheduler knows what it
-    /// / should be, but we do safety checks to ensure it really is the request we expected.
-    #[prost(message, optional, tag = "2")]
-    pub action_digest: ::core::option::Option<
-        super::super::super::super::super::build::bazel::remote::execution::v2::Digest,
-    >,
-    /// / The salt originally sent along with the StartExecute request. This salt is used
-    /// / as a seed for cases where the execution digest should never be cached or merged
-    /// / with other jobs. This salt is added to the hash function used to compute jobs that
-    /// / are running or cached.
-    #[prost(uint64, tag = "3")]
-    pub salt: u64,
-    /// The digest function that was used to compute the action digest
-    /// and all related blobs.
-    ///
-    /// If the digest function used is one of MD5, MURMUR3, SHA1, SHA256,
-    /// SHA384, SHA512, or VSO, the client MAY leave this field unset. In
-    /// that case the server SHOULD infer the digest function using the
-    /// length of the action digest hash and the digest functions announced
-    /// in the server's capabilities.
-    #[prost(
-        enumeration = "super::super::super::super::super::build::bazel::remote::execution::v2::digest_function::Value",
-        tag = "7"
-    )]
-    pub digest_function: i32,
+    /// / The operation ID that was executed.
+    #[prost(string, tag = "8")]
+    pub operation_id: ::prost::alloc::string::String,
     /// / The actual response data.
     #[prost(oneof = "execute_result::Result", tags = "4, 5")]
     pub result: ::core::option::Option<execute_result::Result>,
@@ -168,9 +146,9 @@ pub struct StartExecute {
     pub execute_request: ::core::option::Option<
         super::super::super::super::super::build::bazel::remote::execution::v2::ExecuteRequest,
     >,
-    /// / See documentation in ExecuteResult::salt.
-    #[prost(uint64, tag = "2")]
-    pub salt: u64,
+    /// / Id of the operation.
+    #[prost(string, tag = "4")]
+    pub operation_id: ::prost::alloc::string::String,
     /// / The time at which the command was added to the queue to allow population
     /// / of the ActionResult.
     #[prost(message, optional, tag = "3")]

nativelink-scheduler/BUILD.bazel

@@ -11,6 +11,7 @@ rust_library(
     srcs = [
         "src/action_scheduler.rs",
         "src/cache_lookup_scheduler.rs",
+        "src/default_action_listener.rs",
         "src/default_scheduler_factory.rs",
         "src/grpc_scheduler.rs",
         "src/lib.rs",
@@ -55,6 +56,7 @@ rust_library(
         "@crates//:scopeguard",
         "@crates//:serde",
         "@crates//:serde_json",
+        "@crates//:static_assertions",
         "@crates//:tokio",
         "@crates//:tokio-stream",
         "@crates//:tonic",

nativelink-scheduler/Cargo.toml

@@ -33,6 +33,7 @@ redis = { version = "0.25.2", features = ["aio", "tokio", "json"] }
 serde = "1.0.203"
 redis-macros = "0.3.0"
 serde_json = "1.0.117"
+static_assertions = "1.1.0"
 
 [dev-dependencies]
 nativelink-macro = { path = "../nativelink-macro" }

nativelink-scheduler/src/action_scheduler.rs

@@ -12,16 +12,29 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::pin::Pin;
 use std::sync::Arc;
 
 use async_trait::async_trait;
+use futures::Future;
 use nativelink_error::Error;
-use nativelink_util::action_messages::{ActionInfo, ActionInfoHashKey, ActionState};
+use nativelink_util::action_messages::{ActionInfo, ActionState, ClientOperationId};
 use nativelink_util::metrics_utils::Registry;
-use tokio::sync::watch;
 
 use crate::platform_property_manager::PlatformPropertyManager;
 
+/// ActionListener interface is responsible for interfacing with clients
+/// that are interested in the state of an action.
+pub trait ActionListener: Sync + Send + Unpin {
+    /// Returns the client operation id.
+    fn client_operation_id(&self) -> &ClientOperationId;
+
+    /// Waits for the action state to change.
+    fn changed(
+        &mut self,
+    ) -> Pin<Box<dyn Future<Output = Result<Arc<ActionState>, Error>> + Send + '_>>;
+}
+
 /// ActionScheduler interface is responsible for interactions between the scheduler
 /// and action related operations.
 #[async_trait]
@@ -35,17 +48,15 @@ pub trait ActionScheduler: Sync + Send + Unpin {
     /// Adds an action to the scheduler for remote execution.
     async fn add_action(
         &self,
+        client_operation_id: ClientOperationId,
         action_info: ActionInfo,
-    ) -> Result<watch::Receiver<Arc<ActionState>>, Error>;
+    ) -> Result<Pin<Box<dyn ActionListener>>, Error>;
 
     /// Find an existing action by its name.
-    async fn find_existing_action(
+    async fn find_by_client_operation_id(
         &self,
-        unique_qualifier: &ActionInfoHashKey,
-    ) -> Option<watch::Receiver<Arc<ActionState>>>;
-
-    /// Cleans up the cache of recently completed actions.
-    async fn clean_recently_completed_actions(&self);
+        client_operation_id: &ClientOperationId,
+    ) -> Result<Option<Pin<Box<dyn ActionListener>>>, Error>;
 
     /// Register the metrics for the action scheduler.
     fn register_metrics(self: Arc<Self>, _registry: &mut Registry) {}