diff --git a/nativelink-worker/src/running_actions_manager.rs b/nativelink-worker/src/running_actions_manager.rs index 9e8e77208..742e267d4 100644 --- a/nativelink-worker/src/running_actions_manager.rs +++ b/nativelink-worker/src/running_actions_manager.rs @@ -951,12 +951,25 @@ async fn do_cleanup( action_directory: &str, ) -> Result<(), Error> { // Mark this operation as being cleaned up - let Some(_cleaning_guard) = running_actions_manager.perform_cleanup(operation_id.clone()) - else { + let Some(cleaning_guard) = running_actions_manager.perform_cleanup(operation_id.clone()) else { // Cleanup is already happening elsewhere. return Ok(()); }; + do_cleanup_with_guard( + cleaning_guard, + running_actions_manager, + operation_id, + action_directory, + ) + .await +} +async fn do_cleanup_with_guard( + _cleaning_guard: CleanupGuard, + running_actions_manager: &Arc, + operation_id: &OperationId, + action_directory: &str, +) -> Result<(), Error> { debug!("Worker cleaning up"); // Note: We need to be careful to keep trying to cleanup even if one of the steps fails. let remove_dir_result = fs::remove_dir_all(action_directory) @@ -2101,9 +2114,24 @@ impl Drop for RunningActionImpl { ); let running_actions_manager = self.running_actions_manager.clone(); let action_directory = self.action_directory.clone(); + // Register the cleanup synchronously, before the background task gets + // its first poll. A retry of this operation on this worker consults + // the registration in wait_for_cleanup_if_needed; registering it + // lazily leaves a window where the retry recreates the directory and + // this cleanup then deletes the retry's live files. + let Some(cleaning_guard) = running_actions_manager.perform_cleanup(operation_id.clone()) + else { + // Cleanup is already happening elsewhere. + return; + }; background_spawn!("running_action_impl_drop", async move { - let Err(err) = - do_cleanup(&running_actions_manager, &operation_id, &action_directory).await + let Err(err) = do_cleanup_with_guard( + cleaning_guard, + &running_actions_manager, + &operation_id, + &action_directory, + ) + .await else { return; }; @@ -2864,6 +2892,12 @@ impl RunningActionsManagerImpl { operation_id, }) } + + /// Whether a cleanup of `operation_id`'s working directory is currently + /// registered. Exposed for tests. + pub fn is_cleaning_up(&self, operation_id: &OperationId) -> bool { + self.cleaning_up_operations.lock().contains(operation_id) + } } impl RunningActionsManager for RunningActionsManagerImpl { diff --git a/nativelink-worker/tests/running_actions_manager_test.rs b/nativelink-worker/tests/running_actions_manager_test.rs index e2007c150..e5cdb2d9d 100644 --- a/nativelink-worker/tests/running_actions_manager_test.rs +++ b/nativelink-worker/tests/running_actions_manager_test.rs @@ -5298,4 +5298,132 @@ done assert_eq!(action_result.output_folders[0].path, "out"); Ok(()) } + + /// Regression test: dropping an uncleaned action must register its + /// working-directory cleanup *before* returning, not when the spawned + /// background task first runs. Otherwise a retry of the same operation + /// slips past wait_for_cleanup_if_needed, recreates the directory, and + /// the late background cleanup deletes the retry's files mid-run. + #[nativelink_test] + async fn dropped_action_registers_cleanup_before_yielding_test() + -> Result<(), Box> { + const WORKER_ID: &str = "foo_worker_id"; + + fn test_monotonic_clock() -> SystemTime { + static CLOCK: AtomicU64 = AtomicU64::new(0); + monotonic_clock(&CLOCK) + } + + let (_, _, cas_store, ac_store) = setup_stores().await?; + let root_action_directory = make_temp_path("root_action_directory"); + fs::create_dir_all(&root_action_directory).await?; + + let running_actions_manager = Arc::new(RunningActionsManagerImpl::new_with_callbacks( + RunningActionsManagerArgs { + root_action_directory: root_action_directory.clone(), + execution_configuration: ExecutionConfiguration::default(), + cas_store: cas_store.clone(), + ac_store: Some(Store::new(ac_store.clone())), + historical_store: Store::new(cas_store.clone()), + upload_action_result_config: &UploadActionResultConfig { + upload_ac_results_strategy: UploadCacheResultsStrategy::Never, + ..Default::default() + }, + max_action_timeout: Duration::MAX, + max_upload_timeout: Duration::from_secs(DEFAULT_MAX_UPLOAD_TIMEOUT), + max_cleanup_wait: Duration::from_secs(DEFAULT_MAX_CLEANUP_WAIT), + max_cleanup_backoff: Duration::from_millis(DEFAULT_MAX_CLEANUP_BACKOFF), + timeout_handled_externally: false, + directory_cache: None, + #[cfg(target_os = "linux")] + use_namespaces: use_namespaces(), + }, + Callbacks { + now_fn: test_monotonic_clock, + sleep_fn: |_duration| Box::pin(future::pending()), + }, + )?); + + let command = Command { + arguments: vec!["true".to_string()], + ..Default::default() + }; + let command_digest = serialize_and_upload_message( + &command, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let input_root_digest = serialize_and_upload_message( + &Directory::default(), + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + let action = Action { + command_digest: Some(command_digest.into()), + input_root_digest: Some(input_root_digest.into()), + ..Default::default() + }; + let action_digest = serialize_and_upload_message( + &action, + cas_store.as_pin(), + &mut DigestHasherFunc::Sha256.hasher(), + ) + .await?; + + let operation_id = OperationId::default(); + let start_execute = StartExecute { + execute_request: Some(ExecuteRequest { + action_digest: Some(action_digest.into()), + ..Default::default() + }), + operation_id: operation_id.to_string(), + queued_timestamp: None, + platform: action.platform.clone(), + worker_id: WORKER_ID.to_string(), + }; + + // First attempt: create and prepare, then get aborted (dropped + // without cleanup), as happens when the worker disconnects or the + // operation is cancelled. + let running_action = running_actions_manager + .create_and_add_action(WORKER_ID.to_string(), start_execute.clone()) + .await? + .prepare_action() + .await?; + let work_directory = running_action.get_work_directory().clone(); + assert!(fs::metadata(&work_directory).await.is_ok()); + assert!(!running_actions_manager.is_cleaning_up(&operation_id)); + + drop(running_action); + + // The cleanup must be registered synchronously by the drop: a retry + // of the same operation checks this registration, and anything later + // (e.g. when the background task first runs) lets the retry recreate + // the directory only to have this cleanup delete it mid-run. + assert!( + running_actions_manager.is_cleaning_up(&operation_id), + "drop must register the working-directory cleanup before yielding" + ); + + // Retry of the same operation: must wait out the cleanup and end up + // with an intact working directory. + let retry_action = running_actions_manager + .create_and_add_action(WORKER_ID.to_string(), start_execute) + .await? + .prepare_action() + .await?; + for _ in 0..20 { + tokio::task::yield_now().await; + } + assert!( + fs::metadata(retry_action.get_work_directory()) + .await + .is_ok(), + "retry working directory was deleted by the previous attempt's cleanup" + ); + retry_action.cleanup().await?; + Ok(()) + } }