Introduce execution_completion_behaviour: one_shot_always for workers.

Author: dkostyrev

nativelink-config/src/cas_server.rs

@@ -803,6 +803,17 @@ pub struct LocalWorkerConfig {
     /// them from CAS for every action.
     /// Default: None (directory cache disabled)
     pub directory_cache: Option<DirectoryCacheConfig>,
+
+    #[serde(default)]
+    pub execution_completion_behaviour: ExecutionCompletionBehaviour,
+}
+
+#[derive(Deserialize, Serialize, Debug, Default, Copy, Clone)]
+#[serde(rename_all = "snake_case")]
+pub enum ExecutionCompletionBehaviour {
+    #[default]
+    Default,
+    OneShotAlways,
 }
 
 #[derive(Deserialize, Serialize, Debug, Clone)]

nativelink-worker/src/local_worker.rs

@@ -18,18 +18,17 @@ use core::sync::atomic::{AtomicU64, Ordering};
 use core::time::Duration;
 use std::process::Stdio;
 use std::sync::{Arc, Weak};
-
 use futures::future::BoxFuture;
 use futures::stream::FuturesUnordered;
-use futures::{Future, FutureExt, StreamExt, TryFutureExt, select};
-use nativelink_config::cas_server::LocalWorkerConfig;
-use nativelink_error::{Code, Error, ResultExt, make_err, make_input_err};
+use futures::{select, Future, FutureExt, StreamExt, TryFutureExt};
+use nativelink_config::cas_server::{ExecutionCompletionBehaviour, LocalWorkerConfig};
+use nativelink_error::{make_err, make_input_err, Code, Error, ResultExt};
 use nativelink_metric::{MetricsComponent, RootMetricsComponent};
 use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::update_for_worker::Update;
 use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::worker_api_client::WorkerApiClient;
 use nativelink_proto::com::github::trace_machina::nativelink::remote_execution::{
-    ExecuteComplete, ExecuteResult, GoingAwayRequest, KeepAliveRequest, UpdateForWorker,
-    execute_result,
+    execute_result, ExecuteComplete, ExecuteResult, GoingAwayRequest, KeepAliveRequest,
+    UpdateForWorker,
 };
 use nativelink_store::fast_slow_store::FastSlowStore;
 use nativelink_util::action_messages::{ActionResult, ActionStage, OperationId};
@@ -42,10 +41,11 @@ use nativelink_util::{spawn, tls_utils};
 use opentelemetry::context::Context;
 use tokio::process;
 use tokio::sync::{broadcast, mpsc};
+use tokio::sync::broadcast::{Receiver, Sender};
 use tokio::time::sleep;
 use tokio_stream::wrappers::UnboundedReceiverStream;
 use tonic::Streaming;
-use tracing::{Level, debug, error, event, info, info_span, instrument, warn};
+use tracing::{debug, error, event, info, info_span, instrument, warn, Level};
 
 use crate::running_actions_manager::{
     ExecutionConfiguration, Metrics as RunningActionManagerMetrics, RunningAction,
@@ -82,6 +82,7 @@ struct LocalWorkerImpl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsM
     // on by the scheduler.
     actions_in_transit: Arc<AtomicU64>,
     metrics: Arc<Metrics>,
+    shutdown_tx: Sender<ShutdownGuard>,
 }
 
 async fn preconditions_met(precondition_script: Option<String>) -> Result<(), Error> {
@@ -123,6 +124,7 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
         worker_id: String,
         running_actions_manager: Arc<U>,
         metrics: Arc<Metrics>,
+        shutdown_tx: Sender<ShutdownGuard>,
     ) -> Self {
         Self {
             config,
@@ -135,6 +137,7 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
             // on by the scheduler.
             actions_in_transit: Arc::new(AtomicU64::new(0)),
             metrics,
+            shutdown_tx,
         }
     }
 
@@ -184,6 +187,8 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
 
         let (add_future_channel, add_future_rx) = mpsc::unbounded_channel();
         let mut add_future_rx = UnboundedReceiverStream::new(add_future_rx).fuse();
+        let (inner_shutdown_channel, inner_shutdown_rx) = mpsc::unbounded_channel();
+        let mut inner_shutdown_rx = UnboundedReceiverStream::new(inner_shutdown_rx).fuse();
 
         let mut update_for_worker_stream = update_for_worker_stream.fuse();
         // A notify which is triggered every time actions_in_flight is subtracted.
@@ -193,6 +198,9 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
         let actions_in_flight = Arc::new(AtomicU64::new(0));
         // Set to true when shutting down, this stops any new StartAction.
         let mut shutting_down = false;
+        // Channel to signal when shutdown is complete (GoingAway sent, ready to exit).
+        let (shutdown_complete_tx, shutdown_complete_rx) = mpsc::unbounded_channel::<()>();
+        let mut shutdown_complete_rx = UnboundedReceiverStream::new(shutdown_complete_rx).fuse();
 
         loop {
             select! {
@@ -342,6 +350,7 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
                             self.actions_in_transit.fetch_add(1, Ordering::Release);
 
                             let add_future_channel = add_future_channel.clone();
+                            let inner_shutdown_channel = inner_shutdown_channel.clone();
 
                             info_span!(
                                 "worker_start_action_ctx",
@@ -364,7 +373,16 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
                                             error!(?err, "Error executing action");
                                         }
                                         add_future_channel
-                                            .send(make_publish_future(res).then(move |res| {
+                                            .send(make_publish_future(res)
+                                        .then(move |res| {
+                                                match self.config.execution_completion_behaviour {
+                                                    ExecutionCompletionBehaviour::OneShotAlways => {
+                                                        inner_shutdown_channel.send(()).ok();
+                                                    }
+                                                    ExecutionCompletionBehaviour::Default => {
+                                                        // Do nothing
+                                                    }
+                                                }
                                                 actions_in_flight.fetch_sub(1, Ordering::Release);
                                                 actions_notify.notify_one();
                                                 core::future::ready(res)
@@ -388,13 +406,23 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
                     let fut = res.err_tip(|| "New future stream receives should never be closed")?;
                     futures.push(fut);
                 },
+                _ = inner_shutdown_rx.next() => {
+                    warn!("Shutting down worker because of inner shutdown signal",);
+                    let guard = ShutdownGuard::default();
+                    drop(self.shutdown_tx.send(guard.clone()));
+                }
                 res = futures.next() => res.err_tip(|| "Keep-alive should always pending. Likely unable to send data to scheduler")??,
+                _ = shutdown_complete_rx.next() => {
+                    info!("Shutdown complete, exiting worker loop");
+                    return Ok(());
+                },
                 complete_msg = shutdown_rx.recv().fuse() => {
                     warn!("Worker loop received shutdown signal. Shutting down worker...",);
                     let mut grpc_client = self.grpc_client.clone();
                     let shutdown_guard = complete_msg.map_err(|e| make_err!(Code::Internal, "Failed to receive shutdown message: {e:?}"))?;
                     let actions_in_flight = actions_in_flight.clone();
                     let actions_notify = actions_notify.clone();
+                    let shutdown_complete_tx = shutdown_complete_tx.clone();
                     let shutdown_future = async move {
                         // Wait for in-flight operations to be fully completed.
                         while actions_in_flight.load(Ordering::Acquire) > 0 {
@@ -408,6 +436,8 @@ impl<'a, T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorke
                         }
                         // Allow shutdown to occur now.
                         drop(shutdown_guard);
+                        // Signal that shutdown is complete.
+                        let _ = shutdown_complete_tx.send(());
                         Ok::<(), Error>(())
                     };
                     futures.push(shutdown_future.boxed());
@@ -638,7 +668,8 @@ impl<T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorker<T,
     #[instrument(skip(self), level = Level::INFO)]
     pub async fn run(
         mut self,
-        mut shutdown_rx: broadcast::Receiver<ShutdownGuard>,
+        shutdown_tx: Sender<ShutdownGuard>,
+        mut shutdown_rx: Receiver<ShutdownGuard>,
     ) -> Result<(), Error> {
         let sleep_fn = self
             .sleep_fn
@@ -673,6 +704,7 @@ impl<T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorker<T,
                         worker_id,
                         self.running_actions_manager.clone(),
                         self.metrics.clone(),
+                        shutdown_tx.clone(),
                     ),
                     update_for_worker_stream,
                 ),
@@ -683,30 +715,37 @@ impl<T: WorkerApiClientTrait + 'static, U: RunningActionsManager> LocalWorker<T,
             );
 
             // Now listen for connections and run all other services.
-            if let Err(err) = inner.run(update_for_worker_stream, &mut shutdown_rx).await {
-                'no_more_actions: {
-                    // Ensure there are no actions in transit before we try to kill
-                    // all our actions.
-                    const ITERATIONS: usize = 1_000;
-
-                    const ERROR_MSG: &str = "Actions in transit did not reach zero before we disconnected from the scheduler";
-
-                    let sleep_duration = ACTIONS_IN_TRANSIT_TIMEOUT_S / ITERATIONS as f32;
-                    for _ in 0..ITERATIONS {
-                        if inner.actions_in_transit.load(Ordering::Acquire) == 0 {
-                            break 'no_more_actions;
+            match inner.run(update_for_worker_stream, &mut shutdown_rx).await {
+                Ok(()) => {
+                    // Graceful shutdown completed, return without retrying.
+                    info!("Worker completed graceful shutdown");
+                    return Ok(());
+                }
+                Err(err) => {
+                    'no_more_actions: {
+                        // Ensure there are no actions in transit before we try to kill
+                        // all our actions.
+                        const ITERATIONS: usize = 1_000;
+
+                        const ERROR_MSG: &str = "Actions in transit did not reach zero before we disconnected from the scheduler";
+
+                        let sleep_duration = ACTIONS_IN_TRANSIT_TIMEOUT_S / ITERATIONS as f32;
+                        for _ in 0..ITERATIONS {
+                            if inner.actions_in_transit.load(Ordering::Acquire) == 0 {
+                                break 'no_more_actions;
+                            }
+                            (sleep_fn_pin)(Duration::from_secs_f32(sleep_duration)).await;
                         }
-                        (sleep_fn_pin)(Duration::from_secs_f32(sleep_duration)).await;
+                        error!(ERROR_MSG);
+                        return Err(err.append(ERROR_MSG));
                     }
-                    error!(ERROR_MSG);
-                    return Err(err.append(ERROR_MSG));
-                }
-                error!(?err, "Worker disconnected from scheduler");
-                // Kill off any existing actions because if we re-connect, we'll
-                // get some more and it might resource lock us.
-                self.running_actions_manager.kill_all().await;
+                    error!(?err, "Worker disconnected from scheduler");
+                    // Kill off any existing actions because if we re-connect, we'll
+                    // get some more and it might resource lock us.
+                    self.running_actions_manager.kill_all().await;
 
-                (error_handler)(err).await; // Try to connect again.
+                    (error_handler)(err).await; // Try to connect again.
+                }
             }
         }
         // Unreachable.

nativelink-worker/tests/local_worker_test.rs

@@ -16,7 +16,7 @@ use core::time::Duration;
 use std::collections::HashMap;
 use std::env;
 use std::ffi::OsString;
-use std::io::Write;
+use std::io::{Write};
 #[cfg(target_family = "unix")]
 use std::os::unix::fs::OpenOptionsExt;
 use std::path::PathBuf;
@@ -29,7 +29,7 @@ mod utils {
 }
 
 use hyper::body::Frame;
-use nativelink_config::cas_server::{LocalWorkerConfig, WorkerProperty};
+use nativelink_config::cas_server::{ExecutionCompletionBehaviour, LocalWorkerConfig, WorkerProperty};
 use nativelink_config::stores::{
     FastSlowSpec, FilesystemSpec, MemorySpec, StoreDirection, StoreSpec,
 };
@@ -421,6 +421,125 @@ async fn simple_worker_start_action_test() -> Result<(), Error> {
     Ok(())
 }
 
+#[nativelink_test]
+async fn one_shot_shutdowns_worker_test() -> Result<(), Error> {
+    let config = LocalWorkerConfig {
+        execution_completion_behaviour: ExecutionCompletionBehaviour::OneShotAlways,
+        ..Default::default()
+    };
+    let mut test_context = setup_local_worker_with_config(config).await;
+    let streaming_response = test_context.maybe_streaming_response.take().unwrap();
+
+    {
+        let props = test_context
+            .client
+            .expect_connect_worker(Ok(streaming_response))
+            .await;
+        assert_eq!(props, ConnectWorkerRequest::default());
+    }
+
+    let expected_worker_id = "foobar".to_string();
+
+    let tx_stream = test_context.maybe_tx_stream.take().unwrap();
+    {
+        // First initialize our worker by sending the response to the connection request.
+        tx_stream
+            .send(Frame::data(
+                encode_stream_proto(&UpdateForWorker {
+                    update: Some(Update::ConnectionResult(ConnectionResult {
+                        worker_id: expected_worker_id.clone(),
+                    })),
+                })
+                    .unwrap(),
+            ))
+            .await
+            .map_err(|e| make_input_err!("Could not send : {:?}", e))?;
+    }
+
+    let action_digest = DigestInfo::new([3u8; 32], 10);
+    let action_info = ActionInfo {
+        command_digest: DigestInfo::new([1u8; 32], 10),
+        input_root_digest: DigestInfo::new([2u8; 32], 10),
+        timeout: Duration::from_secs(1),
+        platform_properties: HashMap::new(),
+        priority: 0,
+        load_timestamp: SystemTime::UNIX_EPOCH,
+        insert_timestamp: SystemTime::UNIX_EPOCH,
+        unique_qualifier: ActionUniqueQualifier::Uncacheable(ActionUniqueKey {
+            instance_name: INSTANCE_NAME.to_string(),
+            digest_function: DigestHasherFunc::Sha256,
+            digest: action_digest,
+        }),
+    };
+
+    {
+        // Send execution request.
+        tx_stream
+            .send(Frame::data(
+                encode_stream_proto(&UpdateForWorker {
+                    update: Some(Update::StartAction(StartExecute {
+                        execute_request: Some((&action_info).into()),
+                        operation_id: String::new(),
+                        queued_timestamp: None,
+                        platform: Some(Platform::default()),
+                        worker_id: expected_worker_id.clone(),
+                    })),
+                })
+                    .unwrap(),
+            ))
+            .await
+            .map_err(|e| make_input_err!("Could not send : {:?}", e))?;
+    }
+
+    let running_action = Arc::new(MockRunningAction::new());
+
+    let action_result = ActionResult {
+        output_files: vec![],
+        output_folders: vec![],
+        output_file_symlinks: vec![],
+        output_directory_symlinks: vec![],
+        exit_code: 5,
+        stdout_digest: DigestInfo::new([21u8; 32], 10),
+        stderr_digest: DigestInfo::new([22u8; 32], 10),
+        execution_metadata: ExecutionMetadata {
+            worker: expected_worker_id.clone(),
+            queued_timestamp: SystemTime::UNIX_EPOCH,
+            worker_start_timestamp: SystemTime::UNIX_EPOCH,
+            worker_completed_timestamp: SystemTime::UNIX_EPOCH,
+            input_fetch_start_timestamp: SystemTime::UNIX_EPOCH,
+            input_fetch_completed_timestamp: SystemTime::UNIX_EPOCH,
+            execution_start_timestamp: SystemTime::UNIX_EPOCH,
+            execution_completed_timestamp: SystemTime::UNIX_EPOCH,
+            output_upload_start_timestamp: SystemTime::UNIX_EPOCH,
+            output_upload_completed_timestamp: SystemTime::UNIX_EPOCH,
+        },
+        server_logs: HashMap::new(),
+        error: None,
+        message: String::new(),
+    };
+
+    // Send and wait for response from create_and_add_action to RunningActionsManager.
+    test_context
+        .actions_manager
+        .expect_create_and_add_action(Ok(running_action.clone()))
+        .await;
+
+
+    // Now the RunningAction needs to send a series of state updates. This shortcuts them
+    // into a single call (shortcut for prepare, execute, upload, collect_results, cleanup).
+    running_action
+        .simple_expect_get_finished_result(Ok(action_result.clone()))
+        .await?;
+
+    test_context.client.expect_execution_response(Ok(())).await;
+
+    test_context.client
+        .expect_going_away(Ok(()))
+        .await;
+
+    Ok(())
+}
+
 #[nativelink_test]
 async fn new_local_worker_creates_work_directory_test() -> Result<(), Error> {
     let cas_store = Store::new(FastSlowStore::new(