Harden worker disconnect (TraceMachina#1972)

There are a number of issues with worker disconnect that cause jobs to be unable to be scheduled or stick permanently in Executing.

The first issue is that killing a process does not wait for the process to die which causes the cleanup of the work tree and associated file system store to fail when the system is heavily loaded or the file system is slow.

The second issue is that cleanup is required even if prepare_action hasn't been called yet, separate the concerns of cleanup to filesystem cleanup and manager cleanup as two separate AtomicBools and handle these separate cases.

The third issue is that apply_filter_predicate doesn't handle update_awaited_action failing due to a version mismatch (e.g. client keep-alive), so that can cause the worker state to be skipped.

The fourth and most important is if a worker doesn't exist when worker_notify_run_action is called, the action remains forever in Executing state as nothing updates it, this is trivially resolved by calling update_operation in this case which would otherwise be called by immediate_evict_worker if it actually existed.

Finally, there are a couple of places that log errors which are just noise, notably updating an action on worker disconnect or keep-alive that's already completed and a warning about killing a process that's already dead.

Author: chrisstaite-menlo

nativelink-scheduler/src/api_worker_scheduler.rs

@@ -340,15 +340,23 @@ impl ApiWorkerSchedulerImpl {
                         .await,
                 );
             }
+            Ok(())
         } else {
             warn!(
                 ?worker_id,
                 ?operation_id,
                 ?action_info,
                 "Worker not found in worker map in worker_notify_run_action"
             );
+            // Ensure the operation is put back to queued state.
+            self.worker_state_manager
+                .update_operation(
+                    &operation_id,
+                    &worker_id,
+                    UpdateOperationType::UpdateWithDisconnect,
+                )
+                .await
         }
-        Ok(())
     }
 
     /// Evicts the worker from the pool and puts items back into the queue if anything was being executed on it.

nativelink-scheduler/src/simple_scheduler_state_manager.rs

@@ -314,13 +314,17 @@ where
     async fn apply_filter_predicate(
         &self,
         awaited_action: &AwaitedAction,
+        subscriber: &T::Subscriber,
         filter: &OperationFilter,
     ) -> bool {
         // Note: The caller must filter `client_operation_id`.
 
+        let mut maybe_reloaded_awaited_action: Option<AwaitedAction> = None;
         if awaited_action.last_client_keepalive_timestamp() + self.client_action_timeout
             < (self.now_fn)().now()
         {
+            // This may change if the version is out of date.
+            let mut timed_out = true;
             if !awaited_action.state().stage.is_finished() {
                 let mut state = awaited_action.state().as_ref().clone();
                 state.stage = ActionStage::Completed(ActionResult {
@@ -331,20 +335,66 @@ where
                     )),
                     ..ActionResult::default()
                 });
-                let mut new_awaited_action = awaited_action.clone();
-                new_awaited_action.worker_set_state(Arc::new(state), (self.now_fn)().now());
-                if let Err(err) = self
-                    .action_db
-                    .update_awaited_action(new_awaited_action)
-                    .await
-                {
-                    warn!(
-                        "Failed to update action to timed out state after client keepalive timeout. This is ok if multiple schedulers tried to set the state at the same time: {err}",
-                    );
+                let state = Arc::new(state);
+                // We may be competing with an client timestamp update, so try
+                // this a few times.
+                for attempt in 1..=MAX_UPDATE_RETRIES {
+                    let mut new_awaited_action = match &maybe_reloaded_awaited_action {
+                        None => awaited_action.clone(),
+                        Some(reloaded_awaited_action) => reloaded_awaited_action.clone(),
+                    };
+                    new_awaited_action.worker_set_state(state.clone(), (self.now_fn)().now());
+                    let err = match self
+                        .action_db
+                        .update_awaited_action(new_awaited_action)
+                        .await
+                    {
+                        Ok(()) => break,
+                        Err(err) => err,
+                    };
+                    // Reload from the database if the action was outdated.
+                    let maybe_awaited_action =
+                        if attempt == MAX_UPDATE_RETRIES || err.code != Code::Aborted {
+                            None
+                        } else {
+                            subscriber.borrow().await.ok()
+                        };
+                    if let Some(reloaded_awaited_action) = maybe_awaited_action {
+                        maybe_reloaded_awaited_action = Some(reloaded_awaited_action);
+                    } else {
+                        warn!(
+                            "Failed to update action to timed out state after client keepalive timeout. This is ok if multiple schedulers tried to set the state at the same time: {err}",
+                        );
+                        break;
+                    }
+                    // Re-check the predicate after reload.
+                    if maybe_reloaded_awaited_action
+                        .as_ref()
+                        .is_some_and(|awaited_action| {
+                            awaited_action.last_client_keepalive_timestamp()
+                                + self.client_action_timeout
+                                >= (self.now_fn)().now()
+                        })
+                    {
+                        timed_out = false;
+                        break;
+                    } else if maybe_reloaded_awaited_action
+                        .as_ref()
+                        .is_some_and(|awaited_action| awaited_action.state().stage.is_finished())
+                    {
+                        break;
+                    }
                 }
             }
-            return false;
+            if timed_out {
+                return false;
+            }
         }
+        // If the action was reloaded, then use that for the rest of the checks
+        // instead of the input parameter.
+        let awaited_action = maybe_reloaded_awaited_action
+            .as_ref()
+            .unwrap_or(awaited_action);
 
         if let Some(operation_id) = &filter.operation_id {
             if operation_id != awaited_action.operation_id() {
@@ -518,22 +568,38 @@ where
 
             // Make sure we don't update an action that is already completed.
             if awaited_action.state().stage.is_finished() {
-                return Err(make_err!(
-                    Code::Internal,
-                    "Action {operation_id:?} is already completed with state {:?} - maybe_worker_id: {:?}",
-                    awaited_action.state().stage,
-                    maybe_worker_id,
-                ));
+                match &update {
+                    UpdateOperationType::UpdateWithDisconnect | UpdateOperationType::KeepAlive => {
+                        // No need to error a keep-alive when it's completed, it's just
+                        // unnecessary log noise.
+                        return Ok(());
+                    }
+                    _ => {
+                        return Err(make_err!(
+                            Code::Internal,
+                            "Action {operation_id:?} is already completed with state {:?} - maybe_worker_id: {:?}",
+                            awaited_action.state().stage,
+                            maybe_worker_id,
+                        ));
+                    }
+                }
             }
 
             let stage = match &update {
                 UpdateOperationType::KeepAlive => {
                     awaited_action.worker_keep_alive((self.now_fn)().now());
-                    return self
+                    match self
                         .action_db
                         .update_awaited_action(awaited_action)
                         .await
-                        .err_tip(|| "Failed to send KeepAlive in SimpleSchedulerStateManager::update_operation");
+                        .err_tip(|| "Failed to send KeepAlive in SimpleSchedulerStateManager::update_operation") {
+                        // Try again if there was a version mismatch.
+                        Err(err) if err.code == Code::Aborted => {
+                            last_err = Some(err);
+                            continue;
+                        }
+                        result => return result,
+                    }
                 }
                 UpdateOperationType::UpdateWithActionStage(stage) => stage.clone(),
                 UpdateOperationType::UpdateWithError(err) => {
@@ -658,7 +724,10 @@ where
                 .borrow()
                 .await
                 .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?;
-            if !self.apply_filter_predicate(&awaited_action, &filter).await {
+            if !self
+                .apply_filter_predicate(&awaited_action, &subscriber, &filter)
+                .await
+            {
                 return Ok(Box::pin(stream::empty()));
             }
             return Ok(Box::pin(stream::once(async move {
@@ -678,7 +747,10 @@ where
                 .borrow()
                 .await
                 .err_tip(|| "In SimpleSchedulerStateManager::filter_operations")?;
-            if !self.apply_filter_predicate(&awaited_action, &filter).await {
+            if !self
+                .apply_filter_predicate(&awaited_action, &subscriber, &filter)
+                .await
+            {
                 return Ok(Box::pin(stream::empty()));
             }
             return Ok(Box::pin(stream::once(async move {
@@ -704,11 +776,10 @@ where
                 .try_filter_map(|(subscriber, awaited_action)| {
                     let filter = filter.clone();
                     async move {
-                        if self.apply_filter_predicate(&awaited_action, &filter).await {
-                            Ok(Some((subscriber, awaited_action.sort_key())))
-                        } else {
-                            Ok(None)
-                        }
+                        Ok(self
+                            .apply_filter_predicate(&awaited_action, &subscriber, &filter)
+                            .await
+                            .then_some((subscriber, awaited_action.sort_key())))
                     }
                 })
                 .try_collect()
@@ -750,11 +821,10 @@ where
             .try_filter_map(move |(subscriber, awaited_action)| {
                 let filter = filter.clone();
                 async move {
-                    if self.apply_filter_predicate(&awaited_action, &filter).await {
-                        Ok(Some(subscriber))
-                    } else {
-                        Ok(None)
-                    }
+                    Ok(self
+                        .apply_filter_predicate(&awaited_action, &subscriber, &filter)
+                        .await
+                        .then_some(subscriber))
                 }
             })
             .map(move |result| -> Box<dyn ActionStateResult> {

nativelink-worker/src/running_actions_manager.rs

@@ -661,6 +661,7 @@ pub struct RunningActionImpl {
     timeout: Duration,
     running_actions_manager: Arc<RunningActionsManagerImpl>,
     state: Mutex<RunningActionImplState>,
+    has_manager_entry: AtomicBool,
     did_cleanup: AtomicBool,
 }
 
@@ -691,7 +692,10 @@ impl RunningActionImpl {
                 execution_metadata,
                 error: None,
             }),
-            did_cleanup: AtomicBool::new(false),
+            // Always need to ensure that we're removed from the manager on Drop.
+            has_manager_entry: AtomicBool::new(true),
+            // Only needs to be cleaned up after a prepare_action call, set there.
+            did_cleanup: AtomicBool::new(true),
         }
     }
 
@@ -732,6 +736,8 @@ impl RunningActionImpl {
                 fs::create_dir(&self.work_directory)
                     .await
                     .err_tip(|| format!("Error creating work directory {}", self.work_directory))?;
+                // Now the work directory has been created, we have to clean up.
+                self.did_cleanup.store(false, Ordering::Release);
                 // Download the input files/folder and place them into the temp directory.
                 self.metrics()
                     .download_to_directory
@@ -923,6 +929,10 @@ impl RunningActionImpl {
             .err_tip(|| "Expected stderr to exist on command this should never happen")?;
 
         let mut child_process_guard = guard(child_process, |mut child_process| {
+            if child_process.try_wait().is_ok_and(|res| res.is_some()) {
+                // The child already exited, probably a timeout or kill operation.
+                return;
+            }
             error!(
                 "Child process was not cleaned up before dropping the call to execute(), killing in background spawn."
             );
@@ -966,7 +976,7 @@ impl RunningActionImpl {
                 () = &mut sleep_fut => {
                     self.running_actions_manager.metrics.task_timeouts.inc();
                     killed_action = true;
-                    if let Err(err) = child_process_guard.start_kill() {
+                    if let Err(err) = child_process_guard.kill().await {
                         error!(
                             ?err,
                             "Could not kill process in RunningActionsManager for action timeout",
@@ -1035,7 +1045,7 @@ impl RunningActionImpl {
                 },
                 _ = &mut kill_channel_rx => {
                     killed_action = true;
-                    if let Err(err) = child_process_guard.start_kill() {
+                    if let Err(err) = child_process_guard.kill().await {
                         error!(
                             operation_id = ?self.operation_id,
                             ?err,
@@ -1305,6 +1315,12 @@ impl RunningActionImpl {
 impl Drop for RunningActionImpl {
     fn drop(&mut self) {
         if self.did_cleanup.load(Ordering::Acquire) {
+            if self.has_manager_entry.load(Ordering::Acquire) {
+                drop(
+                    self.running_actions_manager
+                        .cleanup_action(&self.operation_id),
+                );
+            }
             return;
         }
         let operation_id = self.operation_id.clone();
@@ -1370,6 +1386,7 @@ impl RunningAction for RunningActionImpl {
                     &self.action_directory,
                 )
                 .await;
+                self.has_manager_entry.store(false, Ordering::Release);
                 self.did_cleanup.store(true, Ordering::Release);
                 result.map(move |()| self)
             })
@@ -2104,9 +2121,11 @@ impl RunningActionsManager for RunningActionsManagerImpl {
                         .filter_map(|(_operation_id, action)| action.upgrade())
                         .collect()
                 };
-                for action in kill_operations {
-                    Self::kill_operation(action).await;
-                }
+                let mut kill_futures: FuturesUnordered<_> = kill_operations
+                    .into_iter()
+                    .map(Self::kill_operation)
+                    .collect();
+                while kill_futures.next().await.is_some() {}
             })
             .await;
         // Ignore error. If error happens it means there's no sender, which is not a problem.