fix(put): reap cancelled eager commit owners (#6569)

crates/ecstore/src/object_api/types.rs

@@ -379,6 +379,11 @@ pub struct ObjectOptions {
     pub skip_rebalancing: bool,
     pub skip_free_version: bool,
 
+    /// Cooperative cancellation for an owned PutObject before authoritative
+    /// rename begins. Storage ignores it after entering the durable commit.
+    #[doc(hidden)]
+    pub put_object_cancellation: Option<tokio_util::sync::CancellationToken>,
+
     pub data_movement: bool,
     pub raw_data_movement_read: bool,
     /// Materialize the data-movement per-part checksum sidecar for APIs that

crates/ecstore/src/set_disk/ops/object.rs

@@ -149,6 +149,23 @@ impl Drop for PutObjectCommitCancellation {
     }
 }
 
+async fn wait_for_put_object_commit_cancellation(
+    owner_cancellation: Option<&CancellationToken>,
+    request_cancellation: Option<&CancellationToken>,
+) {
+    match (owner_cancellation, request_cancellation) {
+        (Some(owner), Some(request)) => {
+            tokio::select! {
+                _ = owner.cancelled() => {}
+                _ = request.cancelled() => {}
+            }
+        }
+        (Some(owner), None) => owner.cancelled().await,
+        (None, Some(request)) => request.cancelled().await,
+        (None, None) => std::future::pending().await,
+    }
+}
+
 #[inline]
 fn duration_millis_f64(duration: std::time::Duration) -> f64 {
     duration.as_secs_f64() * 1000.0
@@ -2314,8 +2331,15 @@ impl SetDisks {
 
         let tmp_object = format!("{}/{}/part.1", tmp_dir, fi.data_dir.unwrap());
 
+        let (operation_cancellation, mut commit_cancellation_handoff_tx, cancellation_wait) =
+            if let Some(cancellation) = opts.put_object_cancellation.clone() {
+                let (handoff_tx, handoff_rx) = tokio::sync::oneshot::channel();
+                (Some(cancellation.clone()), Some(handoff_tx), Some((cancellation, handoff_rx)))
+            } else {
+                (None, None, None)
+            };
         let mut tmp_cleanup_owned = false;
-        let result: Result<(ObjectInfo, Option<OldCurrentSize>)> = async {
+        let operation = async {
             let erasure = Arc::new(erasure_from_file_info(&fi, false)?);
 
             let put_object_size = known_put_object_storage_size(data.size());
@@ -2981,6 +3005,7 @@ impl SetDisks {
             let commit_capacity_scope_token = opts.capacity_scope_token;
             let commit_replication_state = replication_state_to_filemeta(&opts.put_replication_state());
             let commit_scanner_publication_lease_tokens = scanner_publication_lease_tokens;
+            let request_cancellation = operation_cancellation.clone();
             tmp_cleanup_owned = true;
 
             let commit = move |cancellation: Option<CancellationToken>| async move {
@@ -3065,10 +3090,12 @@ impl SetDisks {
                     }
                     Ok(())
                 };
-                let pre_rename_result = if let Some(cancellation) = cancellation {
+                let pre_rename_result = if cancellation.is_some() || request_cancellation.is_some() {
                     tokio::select! {
                         biased;
-                        _ = cancellation.cancelled() => Err(StorageError::OperationCanceled),
+                        _ = wait_for_put_object_commit_cancellation(cancellation.as_ref(), request_cancellation.as_ref()) => {
+                            Err(StorageError::OperationCanceled)
+                        },
                         result = pre_rename => result,
                     }
                 } else {
@@ -3422,6 +3449,9 @@ impl SetDisks {
                 ))
             };
 
+            if let Some(handoff) = commit_cancellation_handoff_tx.take() {
+                let _ = handoff.send(());
+            }
             if detach_commit_owner {
                 let mut cancellation = PutObjectCommitCancellation::new();
                 let child_token = cancellation.child_token();
@@ -3433,8 +3463,19 @@ impl SetDisks {
             } else {
                 Box::pin(commit(None)).await
             }
-        }
-        .await;
+        };
+        let result: Result<(ObjectInfo, Option<OldCurrentSize>)> = if let Some((cancellation, mut handoff_rx)) = cancellation_wait
+        {
+            tokio::pin!(operation);
+            tokio::select! {
+                biased;
+                _ = &mut handoff_rx => operation.await,
+                result = &mut operation => result,
+                _ = cancellation.cancelled() => Err(StorageError::OperationCanceled),
+            }
+        } else {
+            operation.await
+        };
 
         if issue3031_diag_enabled()
             && let Err(err) = &result
@@ -14269,6 +14310,71 @@ mod put_object_tmp_cleanup_tests {
         drop(temp_dirs);
     }
 
+    #[tokio::test]
+    async fn cooperative_cancellation_while_waiting_for_namespace_lock_cleans_tmp_workspace() {
+        let (temp_dirs, disk_stores, set_disks) = hermetic_set_disks(4).await;
+        let bucket = "tmp-clean-namespace-cancel-bucket";
+        let object = "contended-object";
+        for disk in &disk_stores {
+            disk.make_volume(bucket).await.expect("bucket volume should be created");
+        }
+
+        let first_barrier = PutObjectCommitBarrier::install(bucket, object, PutObjectCommitPause::AfterNamespace);
+        let first_set = set_disks.clone();
+        let first = tokio::spawn(async move {
+            let mut reader = PutObjReader::from_vec(vec![b'1'; TEST_OBJECT_SIZE]);
+            first_set
+                .put_object(bucket, object, &mut reader, &ObjectOptions::default())
+                .await
+        });
+        first_barrier.wait_until_paused().await;
+        let first_workspace = non_trash_tmp_entries(&temp_dirs).await.into_iter().collect::<HashSet<_>>();
+        assert!(!first_workspace.is_empty(), "the lock holder should own a staged tmp workspace");
+
+        let second_barrier = PutObjectCommitBarrier::install(bucket, object, PutObjectCommitPause::BeforeNamespace);
+        let cancellation = CancellationToken::new();
+        let second_cancellation = cancellation.clone();
+        let second_set = set_disks.clone();
+        let second = tokio::spawn(async move {
+            let mut reader = PutObjReader::from_vec(vec![b'2'; TEST_OBJECT_SIZE]);
+            let opts = ObjectOptions {
+                put_object_cancellation: Some(second_cancellation),
+                ..Default::default()
+            };
+            second_set.put_object(bucket, object, &mut reader, &opts).await
+        });
+        second_barrier.release_and_wait_until_namespace_pending().await;
+        let both_workspaces = non_trash_tmp_entries(&temp_dirs).await.into_iter().collect::<HashSet<_>>();
+        assert!(
+            both_workspaces.len() > first_workspace.len() && first_workspace.is_subset(&both_workspaces),
+            "the contending PUT should stage its own tmp workspace before cancellation"
+        );
+
+        cancellation.cancel();
+        let error = tokio::time::timeout(Duration::from_secs(10), second)
+            .await
+            .expect("cooperative cancellation should finish storage cleanup")
+            .expect("the cancelled PUT task should join")
+            .expect_err("the cancelled PUT must not commit");
+        assert!(matches!(error, StorageError::OperationCanceled));
+        let remaining = non_trash_tmp_entries(&temp_dirs).await.into_iter().collect::<HashSet<_>>();
+        assert_eq!(
+            remaining, first_workspace,
+            "the cancelled contender must clean only its own tmp workspace"
+        );
+
+        drop(second_barrier);
+        first_barrier.release();
+        first
+            .await
+            .expect("the lock-holding PUT task should join")
+            .expect("the lock-holding PUT should commit");
+        wait_for_tmp_workspace_to_drain(&temp_dirs, "the committed lock holder should eventually drain its tmp workspace").await;
+
+        drop(first_barrier);
+        drop(temp_dirs);
+    }
+
     #[tokio::test]
     async fn put_object_failure_cleans_tmp_workspace_inline() {
         let (temp_dirs, _disk_stores, set_disks) = hermetic_set_disks(4).await;

rustfs/src/app/object_usecase.rs

@@ -396,6 +396,7 @@ const LOG_COMPONENT_APP: &str = "app";
 const LOG_SUBSYSTEM_OBJECT: &str = "object";
 const EVENT_PUT_OBJECT_STORE_INFLIGHT_SLOW: &str = "put_object_store_inflight_slow";
 const EVENT_PUT_OBJECT_STORE_RETURNED: &str = "put_object_store_returned";
+const EVENT_PUT_OBJECT_COMMIT_OWNER_DEADLINE: &str = "put_object_commit_owner_deadline";
 const EVENT_GET_OBJECT_STREAM_BODY: &str = "get_object_stream_body";
 const EVENT_PUT_OBJECT_BODY_READ_STALLED: &str = "put_object_body_read_stalled";
 const GET_OBJECT_STAGE_PATH_S3_HANDLER: &str = "s3_handler";
@@ -411,6 +412,12 @@ const GET_OBJECT_STAGE_CHECKSUM_HEADERS: &str = "checksum_headers";
 const GET_OBJECT_STAGE_LIFECYCLE_EXPIRATION: &str = "lifecycle_expiration";
 const GET_OBJECT_STAGE_METADATA_FILTER: &str = "metadata_filter";
 const PUT_OBJECT_STORE_WARN_THRESHOLD: Duration = Duration::from_secs(5);
+// Eager PUT bodies are fully materialized before the storage owner starts. On
+// request cancellation, keep the commit/publication tail alive briefly, then
+// request pre-commit rollback and await cleanup so its write-health guard is
+// reaped without abandoning staged shards.
+const EAGER_PUT_COMMIT_CANCELLATION_GRACE: Duration =
+    Duration::from_secs(rustfs_config::DEFAULT_DRIVE_MAX_TIMEOUT_DURATION_SECS * 4);
 const GET_OBJECT_STREAM_WARN_THRESHOLD: Duration = Duration::from_secs(5);
 static GET_OBJECT_BUFFER_THRESHOLD_WARNED: AtomicBool = AtomicBool::new(false);
 
@@ -3203,6 +3210,62 @@ struct PutObjectCommitResult {
     put_versioned: bool,
 }
 
+struct EagerPutCommitOwner<T: Send + 'static> {
+    task: Option<tokio::task::JoinHandle<T>>,
+    cancellation: tokio_util::sync::CancellationToken,
+    cancellation_grace: Duration,
+}
+
+impl<T: Send + 'static> EagerPutCommitOwner<T> {
+    fn new(
+        task: tokio::task::JoinHandle<T>,
+        cancellation: tokio_util::sync::CancellationToken,
+        cancellation_grace: Duration,
+    ) -> Self {
+        Self {
+            task: Some(task),
+            cancellation,
+            cancellation_grace,
+        }
+    }
+
+    async fn join(mut self) -> Result<T, tokio::task::JoinError> {
+        let result = self.task.as_mut().expect("eager PUT commit owner task must be present").await;
+        self.task = None;
+        result
+    }
+}
+
+impl<T: Send + 'static> Drop for EagerPutCommitOwner<T> {
+    fn drop(&mut self) {
+        let Some(mut task) = self.task.take() else {
+            return;
+        };
+        if tokio::runtime::Handle::try_current().is_err() {
+            task.abort();
+            return;
+        }
+        let cancellation = self.cancellation.clone();
+        let cancellation_grace = self.cancellation_grace;
+        spawn_traced(async move {
+            if tokio::time::timeout(cancellation_grace, &mut task).await.is_err() {
+                cancellation.cancel();
+                metrics::counter!("rustfs_put_commit_owner_deadline_total", "put_path" => "eager").increment(1);
+                warn!(
+                    target: "rustfs::app::object_usecase",
+                    event = EVENT_PUT_OBJECT_COMMIT_OWNER_DEADLINE,
+                    component = LOG_COMPONENT_APP,
+                    subsystem = LOG_SUBSYSTEM_OBJECT,
+                    state = "cancellation_requested",
+                    cancellation_grace_ms = cancellation_grace.as_millis() as u64,
+                    "cancelled eager PutObject commit owner exceeded its grace period and requested storage cleanup"
+                );
+                let _ = task.await;
+            }
+        });
+    }
+}
+
 fn normalize_delete_objects_version_id(
     version_id: Option<String>,
 ) -> std::result::Result<(Option<String>, Option<Uuid>), String> {
@@ -6106,6 +6169,9 @@ impl DefaultObjectUsecase {
             object_lock_retain_until_date,
             &mut opts,
         )?;
+        let eager_put_commit_cancellation =
+            (use_zero_copy_eager_put_path || use_empty_or_small_eager_put_path).then(tokio_util::sync::CancellationToken::new);
+        opts.put_object_cancellation = eager_put_commit_cancellation.clone();
 
         // rustfs/backlog#1009: the pre-PUT lookup has exactly two consumers —
         // the existing-object WORM validation and usage accounting's
@@ -6488,7 +6554,14 @@ impl DefaultObjectUsecase {
                 Ok::<_, S3Error>(PutObjectCommitResult { obj_info, put_versioned })
             }
         });
-        let PutObjectCommitResult { obj_info, put_versioned } = match put_commit.await {
+        let put_commit_result = if let Some(cancellation) = eager_put_commit_cancellation {
+            EagerPutCommitOwner::new(put_commit, cancellation, EAGER_PUT_COMMIT_CANCELLATION_GRACE)
+                .join()
+                .await
+        } else {
+            put_commit.await
+        };
+        let PutObjectCommitResult { obj_info, put_versioned } = match put_commit_result {
             Ok(Ok(result)) => result,
             Ok(Err(err)) => {
                 let result: S3Result<S3Response<PutObjectOutput>> = Err(err);
@@ -10274,6 +10347,38 @@ mod tests {
     use tokio::io::{AsyncRead, ReadBuf};
     use tokio_tar::{Builder, EntryType, Header};
 
+    #[tokio::test]
+    async fn cancelled_eager_put_commit_owner_reaps_stalled_storage_task() {
+        let health = Arc::new(ObjectTrafficHealth::enabled_for_test(Duration::ZERO));
+        let task_health = Arc::clone(&health);
+        let cancellation = tokio_util::sync::CancellationToken::new();
+        let task_cancellation = cancellation.clone();
+        let task = spawn_traced_join(async move {
+            let _progress = task_health.track_write_storage().expect("write tracking must be enabled");
+            task_cancellation.cancelled().await;
+        });
+        let owner = EagerPutCommitOwner::new(task, cancellation, Duration::from_millis(10));
+        let request = spawn_traced_join(owner.join());
+
+        tokio::time::timeout(Duration::from_secs(2), async {
+            while !health.snapshot().write_stalled {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .expect("stalled owner must publish write-storage progress");
+
+        request.abort();
+        let _ = request.await;
+        tokio::time::timeout(Duration::from_secs(2), async {
+            while health.snapshot().write_stalled {
+                tokio::task::yield_now().await;
+            }
+        })
+        .await
+        .expect("cancelled owner must abort and reap the stalled storage task");
+    }
+
     #[test]
     fn delete_response_version_id_preserves_null_and_synthetic_semantics() {
         let version_id = Uuid::new_v4();