fix(ecstore): harden tier reader and restore cleanup races (#5035)

* fix(tier): hold generation lease through readers

Refs rustfs/backlog#1354

Co-Authored-By: heihutu <heihutu@gmail.com>

* fix(restore): fence failed cleanup by source identity

Refs rustfs/backlog#1356

Co-Authored-By: heihutu <heihutu@gmail.com>

---------

Co-authored-by: heihutu <heihutu@gmail.com>

crates/ecstore/src/bucket/lifecycle/bucket_lifecycle_ops.rs

@@ -39,7 +39,7 @@ use crate::error::StorageError;
 use crate::error::{error_resp_to_object_err, is_err_object_not_found, is_err_version_not_found, is_network_or_host_down};
 use crate::object_api::{GetObjectReader, ObjectInfo, ObjectOptions};
 use crate::services::tier::{
-    tier::{TierConfigMgr, tier_destination_id_from_metadata},
+    tier::{TierConfigMgr, TierOperationLease, tier_destination_id_from_metadata},
     warm_backend::WarmBackendGetOpts,
 };
 use crate::set_disk::{MAX_PARTS_COUNT, RUSTFS_MULTIPART_BUCKET_KEY, RUSTFS_MULTIPART_OBJECT_KEY, SetDisks};
@@ -77,8 +77,10 @@ use std::env;
 use std::pin::Pin;
 use std::sync::atomic::{AtomicI64, Ordering};
 use std::sync::{Arc, Mutex, OnceLock, Weak};
+use std::task::{Context, Poll};
 use std::time::Duration as StdDuration;
 use time::OffsetDateTime;
+use tokio::io::{AsyncRead, ReadBuf};
 use tokio::select;
 use tokio::sync::mpsc::{Receiver, Sender};
 use tokio::sync::{RwLock, mpsc};
@@ -2539,7 +2541,34 @@ pub(crate) async fn get_transitioned_object_reader_with_tier_manager(
             );
             e
         })?;
-    Ok(get_fn(reader, h.clone()))
+    Ok(attach_tier_operation_lease(get_fn(reader, h.clone()), tgt_client))
+}
+
+struct TierOperationLeaseReader {
+    inner: Box<dyn AsyncRead + Unpin + Send + Sync>,
+    lease: Option<TierOperationLease>,
+}
+
+impl AsyncRead for TierOperationLeaseReader {
+    fn poll_read(mut self: Pin<&mut Self>, cx: &mut Context<'_>, buf: &mut ReadBuf<'_>) -> Poll<std::io::Result<()>> {
+        let had_capacity = buf.remaining() > 0;
+        let filled_before = buf.filled().len();
+        let poll = Pin::new(&mut self.inner).poll_read(cx, buf);
+        if matches!(poll, Poll::Ready(Err(_)))
+            || (had_capacity && matches!(poll, Poll::Ready(Ok(()))) && buf.filled().len() == filled_before)
+        {
+            self.lease.take();
+        }
+        poll
+    }
+}
+
+fn attach_tier_operation_lease(mut reader: GetObjectReader, lease: TierOperationLease) -> GetObjectReader {
+    reader.stream = Box::new(TierOperationLeaseReader {
+        inner: reader.stream,
+        lease: Some(lease),
+    });
+    reader
 }
 
 pub async fn post_restore_opts(version_id: &str, bucket: &str, object: &str) -> Result<ObjectOptions, std::io::Error> {
@@ -3139,11 +3168,19 @@ mod tests {
     use crate::bucket::lifecycle::tier_sweeper::Jentry;
     use crate::bucket::metadata::BUCKET_LIFECYCLE_CONFIG;
     use crate::bucket::metadata_sys;
+    #[cfg(feature = "test-util")]
+    use crate::client::transition_api::ReaderImpl;
     use crate::disk::RUSTFS_META_MULTIPART_BUCKET;
     use crate::disk::endpoint::Endpoint;
     use crate::error::is_err_invalid_upload_id;
     use crate::layout::endpoints::{EndpointServerPools, Endpoints, PoolEndpoints};
     use crate::object_api::{ObjectInfo, ObjectOptions, PutObjReader};
+    #[cfg(feature = "test-util")]
+    use crate::services::tier::test_util::register_mock_tier;
+    #[cfg(feature = "test-util")]
+    use crate::services::tier::tier::TierConfigMgr;
+    #[cfg(feature = "test-util")]
+    use crate::services::tier::warm_backend::WarmBackend as _;
     use crate::set_disk::{RUSTFS_MULTIPART_BUCKET_KEY, RUSTFS_MULTIPART_OBJECT_KEY};
     use crate::storage_api_contracts::{
         bucket::{BucketOperations, BucketOptions, MakeBucketOptions},
@@ -3151,7 +3188,11 @@ mod tests {
         multipart::MultipartOperations as _,
     };
     use crate::store::ECStore;
+    #[cfg(feature = "test-util")]
+    use bytes::Bytes;
     use futures::FutureExt;
+    #[cfg(feature = "test-util")]
+    use http::HeaderMap;
     use rustfs_common::metrics::{IlmAction, global_metrics};
     use rustfs_config::ENV_TRANSITION_WORKERS_ABSOLUTE_MAX;
     use s3s::dto::{
@@ -3168,9 +3209,73 @@ mod tests {
     use std::time::Duration as StdDuration;
     use time::OffsetDateTime;
     use tokio::fs;
+    #[cfg(feature = "test-util")]
+    use tokio::io::AsyncReadExt;
     use tokio_util::sync::CancellationToken;
     use uuid::Uuid;
 
+    #[cfg(feature = "test-util")]
+    #[tokio::test]
+    async fn transitioned_get_reader_holds_tier_operation_lease_until_stream_finishes() {
+        let manager = TierConfigMgr::new();
+        let tier = format!("COLDTIER{}", &Uuid::new_v4().simple().to_string()[..8]).to_uppercase();
+        let backend = register_mock_tier(&manager, &tier).await;
+        let remote_object = format!("remote/{}", Uuid::new_v4());
+        let body = Bytes::from_static(b"transitioned object body");
+        let remote_version = backend
+            .put(
+                &remote_object,
+                ReaderImpl::Body(body.clone()),
+                i64::try_from(body.len()).expect("body length should fit"),
+            )
+            .await
+            .expect("mock remote object should be stored");
+        let object_info = ObjectInfo {
+            bucket: "bucket".to_string(),
+            name: "object".to_string(),
+            size: i64::try_from(body.len()).expect("body length should fit"),
+            transitioned_object: TransitionedObject {
+                name: remote_object,
+                version_id: remote_version,
+                status: crate::bucket::lifecycle::lifecycle::TRANSITION_COMPLETE.to_string(),
+                tier: tier.clone(),
+                ..Default::default()
+            },
+            ..Default::default()
+        };
+
+        let mut reader = get_transitioned_object_reader_with_tier_manager(
+            &object_info.bucket,
+            &object_info.name,
+            &None,
+            &HeaderMap::new(),
+            &object_info,
+            &ObjectOptions::default(),
+            &manager,
+        )
+        .await
+        .expect("transitioned reader should open");
+
+        assert_eq!(
+            TierConfigMgr::active_operation_lease_count(&manager, &tier).await,
+            1,
+            "returned reader must keep the tier generation leased"
+        );
+
+        let mut got = Vec::new();
+        reader
+            .stream
+            .read_to_end(&mut got)
+            .await
+            .expect("transitioned reader should drain");
+        assert_eq!(got, body.as_ref());
+        assert_eq!(
+            TierConfigMgr::active_operation_lease_count(&manager, &tier).await,
+            0,
+            "tier generation lease should release after EOF"
+        );
+    }
+
     #[cfg(feature = "test-util")]
     #[tokio::test]
     async fn free_version_remote_delete_requires_persisted_destination_identity() {

crates/ecstore/src/services/tier/test_util.rs

@@ -146,6 +146,7 @@ struct MockWarmBackendInner {
     put_versions: Mutex<Vec<(String, String)>>,
     remove_versions: Mutex<Vec<(String, String)>>,
     put_barrier: Mutex<Option<Arc<MockPutBarrierState>>>,
+    get_barrier: Mutex<Option<Arc<MockGetBarrierState>>>,
 }
 
 #[derive(Default)]
@@ -154,6 +155,13 @@ struct MockPutBarrierState {
     release: Notify,
 }
 
+#[derive(Default)]
+struct MockGetBarrierState {
+    arrived: Notify,
+    release: Notify,
+    fail_after_release: bool,
+}
+
 /// One-shot barrier that pauses a mock tier PUT after storing its remote body.
 pub struct MockPutBarrier {
     state: Arc<MockPutBarrierState>,
@@ -179,6 +187,31 @@ impl Drop for MockPutBarrier {
     }
 }
 
+/// One-shot barrier that pauses a mock tier GET before it reads remote bytes.
+pub struct MockGetBarrier {
+    state: Arc<MockGetBarrierState>,
+}
+
+impl MockGetBarrier {
+    /// Wait until the GET has reached the deterministic pause point.
+    pub async fn wait_until_paused(&self) {
+        tokio::time::timeout(Duration::from_secs(30), self.state.arrived.notified())
+            .await
+            .expect("mock tier GET should reach the deterministic barrier");
+    }
+
+    /// Release the paused GET.
+    pub fn release(&self) {
+        self.state.release.notify_one();
+    }
+}
+
+impl Drop for MockGetBarrier {
+    fn drop(&mut self) {
+        self.state.release.notify_one();
+    }
+}
+
 /// In-memory [`WarmBackend`] for lifecycle / tiering integration tests.
 ///
 /// Cloning shares the same underlying storage, fault configuration, and
@@ -202,6 +235,17 @@ impl MockWarmBackend {
         MockPutBarrier { state }
     }
 
+    /// Arm a one-shot pause before the next tier GET, then return an error
+    /// after the test releases it.
+    pub async fn arm_failing_get_barrier(&self) -> MockGetBarrier {
+        let state = Arc::new(MockGetBarrierState {
+            fail_after_release: true,
+            ..Default::default()
+        });
+        *self.inner.get_barrier.lock().await = Some(Arc::clone(&state));
+        MockGetBarrier { state }
+    }
+
     // ---- fault injection -------------------------------------------------
 
     /// Replace the entire fault configuration.
@@ -500,6 +544,14 @@ impl WarmBackend for MockWarmBackend {
 
     async fn get(&self, object: &str, _rv: &str, opts: WarmBackendGetOpts) -> Result<ReadCloser, std::io::Error> {
         self.precondition().await?;
+        let barrier = self.inner.get_barrier.lock().await.take();
+        if let Some(barrier) = barrier {
+            barrier.arrived.notify_one();
+            barrier.release.notified().await;
+            if barrier.fail_after_release {
+                return Err(std::io::Error::other("mock warm backend GET failed after barrier"));
+            }
+        }
         self.record(MockWarmOp::Get {
             object: object.to_string(),
         })

crates/ecstore/src/services/tier/tier.rs

@@ -2220,6 +2220,19 @@ impl TierConfigMgr {
         Ok(lease)
     }
 
+    #[cfg(test)]
+    pub(crate) async fn active_operation_lease_count(handle: &Arc<RwLock<Self>>, tier_name: &str) -> usize {
+        let manager = handle.read().await;
+        let Some(runtime) = registered_tier_driver_runtime(&manager) else {
+            return 0;
+        };
+        lock_unpoisoned(&runtime)
+            .generations
+            .get(tier_name)
+            .map(|generation| generation.active_leases.load(Ordering::Acquire))
+            .unwrap_or(0)
+    }
+
     fn replace_driver(&mut self, tier_name: &str, driver: WarmBackendImpl) -> std::result::Result<(), AdminError> {
         let Some(runtime) = registered_tier_driver_runtime(self) else {
             self.driver_cache.insert(tier_name.to_string(), driver);

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

@@ -3598,7 +3598,7 @@ mod metadata_mutation_generation_tests {
 mod transition_commit_failure_tests {
     use super::hermetic_set_disks_support::hermetic_set_disks;
     use super::*;
-    use crate::bucket::lifecycle::lifecycle::{TRANSITION_PENDING, TransitionOptions};
+    use crate::bucket::lifecycle::lifecycle::{TRANSITION_COMPLETE, TRANSITION_PENDING, TransitionOptions};
     use crate::disk::DiskAPI as _;
     use crate::services::tier::test_util::{MockWarmBackend, register_mock_tier};
     use crate::services::tier::tier::TierConfigMgr;
@@ -3813,6 +3813,89 @@ mod transition_commit_failure_tests {
         assert_eq!(old_backend.object_count().await, 0);
     }
 
+    #[tokio::test]
+    #[serial_test::serial]
+    async fn failed_restore_cleanup_does_not_overwrite_concurrent_unversioned_put() {
+        let (_temp_dirs, disk_stores, set_disks) = hermetic_set_disks(4).await;
+        let bucket = "restore-cleanup-cas-bucket";
+        let object = "object.bin";
+        let original_payload = b"old transitioned body".repeat(1024);
+        let replacement_payload = b"new visible unversioned body".repeat(1024);
+        for disk in &disk_stores {
+            disk.make_volume(bucket).await.expect("bucket volume should be created");
+        }
+
+        let mut reader = PutObjReader::from_vec(original_payload.clone());
+        let original = set_disks
+            .put_object(bucket, object, &mut reader, &ObjectOptions::default())
+            .await
+            .expect("source object should be written");
+        let tier_name = format!("COLDTIER{}", &Uuid::new_v4().simple().to_string()[..8]).to_uppercase();
+        let backend = register_mock_tier(&runtime_sources::global_tier_config_mgr(), &tier_name).await;
+        set_disks
+            .transition_object(
+                bucket,
+                object,
+                &ObjectOptions {
+                    no_lock: true,
+                    transition: TransitionOptions {
+                        status: TRANSITION_PENDING.to_string(),
+                        tier: tier_name,
+                        etag: original.etag.clone().unwrap_or_default(),
+                        ..Default::default()
+                    },
+                    version_id: original.version_id.map(|version| version.to_string()),
+                    mod_time: original.mod_time,
+                    ..Default::default()
+                },
+            )
+            .await
+            .expect("source object should transition before restore");
+
+        let get_barrier = backend.arm_failing_get_barrier().await;
+        let restore_set = Arc::clone(&set_disks);
+        let restore = tokio::spawn(async move {
+            let mut opts = ObjectOptions::default();
+            opts.transition.restore_request.days = Some(1);
+            restore_set.restore_transitioned_object(bucket, object, &opts).await
+        });
+        get_barrier.wait_until_paused().await;
+
+        let mut replacement_reader = PutObjReader::from_vec(replacement_payload.clone());
+        let replacement = set_disks
+            .put_object(bucket, object, &mut replacement_reader, &ObjectOptions::default())
+            .await
+            .expect("concurrent unversioned PUT should commit while restore GET is paused");
+        get_barrier.release();
+        restore
+            .await
+            .expect("restore task should join")
+            .expect_err("injected tier GET failure should surface");
+
+        let visible = set_disks
+            .get_object_info(bucket, object, &ObjectOptions::default())
+            .await
+            .expect("replacement metadata should remain visible");
+        assert_eq!(visible.etag, replacement.etag, "stale restore cleanup must not republish the old ETag");
+        assert_ne!(
+            visible.transitioned_object.status, TRANSITION_COMPLETE,
+            "replacement object must not regain stale transition metadata"
+        );
+        let mut body = Vec::new();
+        set_disks
+            .get_object_reader(bucket, object, None, HeaderMap::new(), &ObjectOptions::default())
+            .await
+            .expect("replacement object should be readable")
+            .stream
+            .read_to_end(&mut body)
+            .await
+            .expect("replacement body should drain");
+        assert_eq!(
+            body, replacement_payload,
+            "stale restore cleanup must not make the old remote body current again"
+        );
+    }
+
     #[tokio::test]
     #[serial_test::serial]
     async fn legacy_reload_rejects_route_change_after_local_transition_commit() {

crates/ecstore/src/set_disk/replication.rs

@@ -13,6 +13,34 @@
 // limitations under the License.
 
 use super::*;
+use rustfs_utils::http::headers::{AMZ_RESTORE_EXPIRY_DAYS, AMZ_RESTORE_REQUEST_DATE};
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+struct RestoreCleanupIdentity {
+    version_id: Option<Uuid>,
+    data_dir: Option<Uuid>,
+    mod_time: Option<OffsetDateTime>,
+    size: i64,
+}
+
+impl RestoreCleanupIdentity {
+    fn from_object_info(obj_info: &ObjectInfo) -> Self {
+        Self {
+            version_id: obj_info.version_id,
+            data_dir: obj_info.data_dir,
+            mod_time: obj_info.mod_time,
+            size: obj_info.size,
+        }
+    }
+
+    fn matches_file_info(&self, fi: &FileInfo, expected_etag: &str) -> bool {
+        self.version_id == fi.version_id
+            && self.data_dir == fi.data_dir
+            && self.mod_time == fi.mod_time
+            && self.size == fi.size
+            && expected_etag == get_raw_etag(&fi.metadata)
+    }
+}
 
 impl SetDisks {
     pub async fn update_restore_metadata(
@@ -20,29 +48,46 @@ impl SetDisks {
         bucket: &str,
         object: &str,
         obj_info: &ObjectInfo,
-        _opts: &ObjectOptions,
+        opts: &ObjectOptions,
     ) -> Result<()> {
-        let mut oi = obj_info.clone();
-        oi.metadata_only = true;
-        Arc::make_mut(&mut oi.user_defined).remove(X_AMZ_RESTORE.as_str());
-        let version_id = oi.version_id.map(|v| v.to_string());
-        let _obj = self
-            .copy_object(
-                bucket,
-                object,
-                bucket,
-                object,
-                &mut oi,
-                &ObjectOptions {
-                    version_id: version_id.clone(),
-                    ..Default::default()
-                },
-                &ObjectOptions {
-                    version_id,
-                    ..Default::default()
-                },
+        if obj_info.bucket.is_empty() || obj_info.name.is_empty() {
+            return Ok(());
+        }
+        let expected = RestoreCleanupIdentity::from_object_info(obj_info);
+        let expected_etag = obj_info
+            .etag
+            .clone()
+            .unwrap_or_else(|| get_raw_etag(obj_info.user_defined.as_ref()));
+        let version_id = expected.version_id.map(|v| v.to_string());
+        let lock_guard = if !opts.no_lock {
+            Some(
+                self.acquire_write_lock_diag("restore_cleanup_metadata", bucket, object)
+                    .await?,
             )
+        } else {
+            None
+        };
+        let read_opts = ObjectOptions {
+            version_id,
+            versioned: opts.versioned,
+            version_suspended: opts.version_suspended,
+            ..Default::default()
+        };
+        let (mut fi, _, disks) = self
+            .get_object_fileinfo_gated(bucket, object, &read_opts, false, false)
             .await?;
+        if !expected.matches_file_info(&fi, &expected_etag) {
+            return Ok(());
+        }
+        fi.metadata.remove(X_AMZ_RESTORE.as_str());
+        fi.metadata.remove(AMZ_RESTORE_EXPIRY_DAYS);
+        fi.metadata.remove(AMZ_RESTORE_REQUEST_DATE);
+        if lock_guard.as_ref().is_some_and(|guard| guard.is_lock_lost()) {
+            return Err(Error::other("restore cleanup lock lost before metadata update".to_string()));
+        }
+        self.invalidate_get_object_metadata_cache(bucket, object).await;
+        self.update_object_meta(bucket, object, fi, disks.as_slice()).await?;
+        self.invalidate_get_object_metadata_cache(bucket, object).await;
         Ok(())
     }
 }