fix(ecstore): repair lifecycle transition and restore flows (#2240)

Co-authored-by: heihutu <heihutu@gmail.com>
Co-authored-by: 安正超 <anzhengchao@gmail.com>
Co-authored-by: houseme <housemecn@gmail.com>
Co-authored-by: loverustfs <hello@rustfs.com>
Co-authored-by: weisd <im@weisd.in>

crates/scanner/tests/lifecycle_integration_test.rs

@@ -13,28 +13,42 @@
 // limitations under the License.
 
 use rustfs_ecstore::{
+    bucket::lifecycle::lifecycle::TransitionOptions,
     bucket::metadata::BUCKET_LIFECYCLE_CONFIG,
-    bucket::metadata_sys,
+    bucket::{lifecycle::bucket_lifecycle_ops::enqueue_transition_for_existing_objects, metadata_sys},
+    client::transition_api::{ReadCloser, ReaderImpl},
     disk::endpoint::Endpoint,
     endpoints::{EndpointServerPools, Endpoints, PoolEndpoints},
     global::GLOBAL_TierConfigMgr,
     store::ECStore,
-    store_api::{BucketOperations, MakeBucketOptions, ObjectIO, ObjectOperations, ObjectOptions, PutObjReader},
-    tier::tier_config::{TierConfig, TierMinIO, TierType},
+    store_api::{
+        BucketOperations, MakeBucketOptions, MultipartOperations, ObjectIO, ObjectOperations, ObjectOptions, PutObjReader,
+    },
+    tier::{
+        tier_config::{TierConfig, TierMinIO, TierType},
+        warm_backend::{WarmBackend, WarmBackendGetOpts},
+    },
 };
 use rustfs_scanner::scanner::init_data_scanner;
+use s3s::dto::RestoreRequest;
 use serial_test::serial;
 use std::{
+    collections::HashMap,
+    io::Cursor,
     path::PathBuf,
     sync::{Arc, Once, OnceLock},
     time::Duration,
 };
 use tokio::fs;
+use tokio::io::AsyncReadExt;
+use tokio::sync::Mutex;
 use tokio_util::sync::CancellationToken;
 use tracing::info;
+use uuid::Uuid;
 
 static GLOBAL_ENV: OnceLock<(Vec<PathBuf>, Arc<ECStore>)> = OnceLock::new();
 static INIT: Once = Once::new();
+const TRANSITION_WAIT_TIMEOUT: Duration = Duration::from_secs(15);
 
 fn init_tracing() {
     INIT.call_once(|| {
@@ -209,8 +223,17 @@ async fn set_bucket_lifecycle_deletemarker(bucket_name: &str) -> Result<(), Box<
 
 #[allow(dead_code)]
 async fn set_bucket_lifecycle_transition(bucket_name: &str) -> Result<(), Box<dyn std::error::Error>> {
+    set_bucket_lifecycle_transition_with_tier(bucket_name, "COLDTIER44").await
+}
+
+#[allow(dead_code)]
+async fn set_bucket_lifecycle_transition_with_tier(
+    bucket_name: &str,
+    storage_class: &str,
+) -> Result<(), Box<dyn std::error::Error>> {
     // Create a simple lifecycle configuration XML with 0 days expiry for immediate testing
-    let lifecycle_xml = r#"<?xml version="1.0" encoding="UTF-8"?>
+    let lifecycle_xml = format!(
+        r#"<?xml version="1.0" encoding="UTF-8"?>
 <LifecycleConfiguration>
     <Rule>
         <ID>test-rule</ID>
@@ -220,7 +243,7 @@ async fn set_bucket_lifecycle_transition(bucket_name: &str) -> Result<(), Box<dy
         </Filter>
         <Transition>
           <Days>0</Days>
-          <StorageClass>COLDTIER44</StorageClass>
+          <StorageClass>{storage_class}</StorageClass>
         </Transition>
     </Rule>
     <Rule>
@@ -231,12 +254,13 @@ async fn set_bucket_lifecycle_transition(bucket_name: &str) -> Result<(), Box<dy
         </Filter>
         <NoncurrentVersionTransition>
           <NoncurrentDays>0</NoncurrentDays>
-          <StorageClass>COLDTIER44</StorageClass>
+          <StorageClass>{storage_class}</StorageClass>
         </NoncurrentVersionTransition>
     </Rule>
-</LifecycleConfiguration>"#;
+</LifecycleConfiguration>"#
+    );
 
-    metadata_sys::update(bucket_name, BUCKET_LIFECYCLE_CONFIG, lifecycle_xml.as_bytes().to_vec()).await?;
+    metadata_sys::update(bucket_name, BUCKET_LIFECYCLE_CONFIG, lifecycle_xml.into_bytes()).await?;
 
     Ok(())
 }
@@ -350,6 +374,114 @@ async fn wait_for_object_absence(ecstore: &Arc<ECStore>, bucket: &str, object: &
     }
 }
 
+#[derive(Clone, Default)]
+struct MockWarmBackend {
+    objects: Arc<Mutex<HashMap<String, Vec<u8>>>>,
+}
+
+impl MockWarmBackend {
+    async fn put_bytes(&self, object: &str, bytes: Vec<u8>) -> String {
+        self.objects.lock().await.insert(object.to_string(), bytes);
+        Uuid::new_v4().to_string()
+    }
+
+    async fn read_bytes(&self, reader: ReaderImpl) -> Result<Vec<u8>, std::io::Error> {
+        match reader {
+            ReaderImpl::Body(bytes) => Ok(bytes.to_vec()),
+            ReaderImpl::ObjectBody(mut reader) => {
+                let mut buf = Vec::new();
+                reader.stream.read_to_end(&mut buf).await?;
+                Ok(buf)
+            }
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl WarmBackend for MockWarmBackend {
+    async fn put(&self, object: &str, r: ReaderImpl, _length: i64) -> Result<String, std::io::Error> {
+        let bytes = self.read_bytes(r).await?;
+        Ok(self.put_bytes(object, bytes).await)
+    }
+
+    async fn put_with_meta(
+        &self,
+        object: &str,
+        r: ReaderImpl,
+        _length: i64,
+        _meta: HashMap<String, String>,
+    ) -> Result<String, std::io::Error> {
+        let bytes = self.read_bytes(r).await?;
+        Ok(self.put_bytes(object, bytes).await)
+    }
+
+    async fn get(&self, object: &str, _rv: &str, opts: WarmBackendGetOpts) -> Result<ReadCloser, std::io::Error> {
+        let objects = self.objects.lock().await;
+        let Some(bytes) = objects.get(object) else {
+            return Err(std::io::Error::new(std::io::ErrorKind::NotFound, "mock object not found"));
+        };
+
+        let start = opts.start_offset.max(0) as usize;
+        let end = if opts.length > 0 {
+            start.saturating_add(opts.length as usize).min(bytes.len())
+        } else {
+            bytes.len()
+        };
+
+        Ok(tokio::io::BufReader::new(Cursor::new(bytes[start.min(bytes.len())..end].to_vec())))
+    }
+
+    async fn remove(&self, object: &str, _rv: &str) -> Result<(), std::io::Error> {
+        self.objects.lock().await.remove(object);
+        Ok(())
+    }
+
+    async fn in_use(&self) -> Result<bool, std::io::Error> {
+        Ok(false)
+    }
+}
+
+async fn register_mock_tier(tier_name: &str) -> MockWarmBackend {
+    let backend = MockWarmBackend::default();
+    let mut tier_config_mgr = GLOBAL_TierConfigMgr.write().await;
+    tier_config_mgr.tiers.insert(
+        tier_name.to_string(),
+        TierConfig {
+            version: "v1".to_string(),
+            tier_type: TierType::MinIO,
+            name: tier_name.to_string(),
+            ..Default::default()
+        },
+    );
+    tier_config_mgr
+        .driver_cache
+        .insert(tier_name.to_string(), Box::new(backend.clone()));
+    backend
+}
+
+async fn wait_for_transition(
+    ecstore: &Arc<ECStore>,
+    bucket: &str,
+    object: &str,
+    timeout: Duration,
+) -> Option<rustfs_ecstore::store_api::ObjectInfo> {
+    let deadline = tokio::time::Instant::now() + timeout;
+
+    loop {
+        if let Ok(info) = (**ecstore).get_object_info(bucket, object, &ObjectOptions::default()).await
+            && info.transitioned_object.status == "complete"
+        {
+            return Some(info);
+        }
+
+        if tokio::time::Instant::now() >= deadline {
+            return None;
+        }
+
+        tokio::time::sleep(Duration::from_millis(50)).await;
+    }
+}
+
 mod serial_tests {
     use super::*;
 
@@ -443,4 +575,283 @@ mod serial_tests {
 
         println!("Lifecycle transition basic test completed");
     }
+
+    #[tokio::test(flavor = "multi_thread", worker_threads = 1)]
+    #[serial]
+    async fn test_transition_and_restore_flows() {
+        let (_disk_paths, ecstore) = setup_test_env().await;
+
+        let tier_name = format!("COLDTIER{}", &Uuid::new_v4().simple().to_string()[..8]).to_uppercase();
+        let backend = register_mock_tier(&tier_name).await;
+
+        let put_bucket = format!("test-immediate-put-{}", &Uuid::new_v4().simple().to_string()[..8]);
+        let put_object = "test/object.txt";
+        let put_payload = b"Hello, immediate transition!";
+
+        create_test_bucket(&ecstore, put_bucket.as_str()).await;
+        set_bucket_lifecycle_transition_with_tier(put_bucket.as_str(), &tier_name)
+            .await
+            .expect("Failed to set lifecycle configuration");
+
+        upload_test_object(&ecstore, put_bucket.as_str(), put_object, put_payload).await;
+
+        enqueue_transition_for_existing_objects(ecstore.clone(), put_bucket.as_str())
+            .await
+            .expect("Failed to enqueue transitioned put object");
+
+        let put_info = wait_for_transition(&ecstore, put_bucket.as_str(), put_object, TRANSITION_WAIT_TIMEOUT)
+            .await
+            .expect("object should transition after enqueueing existing objects");
+
+        assert_eq!(put_info.transitioned_object.status, "complete");
+        assert_eq!(put_info.transitioned_object.tier, tier_name);
+        assert!(backend.objects.lock().await.contains_key(&put_info.transitioned_object.name));
+
+        let multipart_bucket = format!("test-immediate-mpu-{}", &Uuid::new_v4().simple().to_string()[..8]);
+        let multipart_object = "test/multipart.txt";
+
+        create_test_bucket(&ecstore, multipart_bucket.as_str()).await;
+        set_bucket_lifecycle_transition_with_tier(multipart_bucket.as_str(), &tier_name)
+            .await
+            .expect("Failed to set lifecycle configuration");
+
+        let upload = ecstore
+            .new_multipart_upload(multipart_bucket.as_str(), multipart_object, &ObjectOptions::default())
+            .await
+            .expect("Failed to create multipart upload");
+
+        let part_data = b"multipart immediate transition";
+        let mut reader = PutObjReader::from_vec(part_data.to_vec());
+        let part = ecstore
+            .put_object_part(
+                multipart_bucket.as_str(),
+                multipart_object,
+                &upload.upload_id,
+                1,
+                &mut reader,
+                &ObjectOptions::default(),
+            )
+            .await
+            .expect("Failed to upload multipart part");
+
+        ecstore
+            .clone()
+            .complete_multipart_upload(
+                multipart_bucket.as_str(),
+                multipart_object,
+                &upload.upload_id,
+                vec![rustfs_ecstore::store_api::CompletePart {
+                    part_num: 1,
+                    etag: part.etag.clone(),
+                    ..Default::default()
+                }],
+                &ObjectOptions::default(),
+            )
+            .await
+            .expect("Failed to complete multipart upload");
+
+        enqueue_transition_for_existing_objects(ecstore.clone(), multipart_bucket.as_str())
+            .await
+            .expect("Failed to enqueue transitioned multipart object");
+
+        let multipart_info = wait_for_transition(&ecstore, multipart_bucket.as_str(), multipart_object, TRANSITION_WAIT_TIMEOUT)
+            .await
+            .expect("object should transition after enqueueing existing objects");
+
+        assert_eq!(multipart_info.transitioned_object.status, "complete");
+        assert_eq!(multipart_info.transitioned_object.tier, tier_name);
+        assert!(
+            backend
+                .objects
+                .lock()
+                .await
+                .contains_key(&multipart_info.transitioned_object.name)
+        );
+
+        let src_bucket = format!("test-immediate-copy-src-{}", &Uuid::new_v4().simple().to_string()[..8]);
+        let dst_bucket = format!("test-immediate-copy-dst-{}", &Uuid::new_v4().simple().to_string()[..8]);
+        let src_object = "test/source.txt";
+        let dst_object = "test/copied.txt";
+        let payload = b"copy object immediate transition";
+
+        create_test_bucket(&ecstore, src_bucket.as_str()).await;
+        create_test_bucket(&ecstore, dst_bucket.as_str()).await;
+        set_bucket_lifecycle_transition_with_tier(dst_bucket.as_str(), &tier_name)
+            .await
+            .expect("Failed to set destination lifecycle configuration");
+
+        upload_test_object(&ecstore, src_bucket.as_str(), src_object, payload).await;
+
+        let mut src_info = ecstore
+            .get_object_info(src_bucket.as_str(), src_object, &ObjectOptions::default())
+            .await
+            .expect("Failed to load source object info");
+        src_info.put_object_reader = Some(PutObjReader::from_vec(payload.to_vec()));
+
+        ecstore
+            .copy_object(
+                src_bucket.as_str(),
+                src_object,
+                dst_bucket.as_str(),
+                dst_object,
+                &mut src_info,
+                &ObjectOptions::default(),
+                &ObjectOptions::default(),
+            )
+            .await
+            .expect("Failed to copy object");
+
+        enqueue_transition_for_existing_objects(ecstore.clone(), dst_bucket.as_str())
+            .await
+            .expect("Failed to enqueue transitioned copied object");
+
+        let copy_info = wait_for_transition(&ecstore, dst_bucket.as_str(), dst_object, TRANSITION_WAIT_TIMEOUT)
+            .await
+            .expect("copied object should transition after enqueueing existing objects");
+
+        assert_eq!(copy_info.transitioned_object.status, "complete");
+        assert_eq!(copy_info.transitioned_object.tier, tier_name);
+        assert!(backend.objects.lock().await.contains_key(&copy_info.transitioned_object.name));
+
+        let bucket_name = format!("test-lifecycle-update-{}", &Uuid::new_v4().simple().to_string()[..8]);
+        let object_name = "test/existing.txt";
+        let payload = b"existing object before lifecycle";
+
+        create_test_bucket(&ecstore, bucket_name.as_str()).await;
+        upload_test_object(&ecstore, bucket_name.as_str(), object_name, payload).await;
+
+        set_bucket_lifecycle_transition_with_tier(bucket_name.as_str(), &tier_name)
+            .await
+            .expect("Failed to set lifecycle configuration");
+
+        enqueue_transition_for_existing_objects(ecstore.clone(), bucket_name.as_str())
+            .await
+            .expect("Failed to enqueue transition for existing objects");
+
+        let info = wait_for_transition(&ecstore, bucket_name.as_str(), object_name, TRANSITION_WAIT_TIMEOUT)
+            .await
+            .expect("existing object should transition after lifecycle update");
+
+        assert_eq!(info.transitioned_object.status, "complete");
+        assert_eq!(info.transitioned_object.tier, tier_name);
+        assert!(backend.objects.lock().await.contains_key(&info.transitioned_object.name));
+
+        let bucket_name = format!("test-restore-mpu-{}", &Uuid::new_v4().simple().to_string()[..8]);
+        let object_name = "test/restore.txt";
+        let part1 = vec![b'a'; 5 * 1024 * 1024];
+        let part2 = b"restored-tail".to_vec();
+        let expected = [part1.clone(), part2.clone()].concat();
+
+        create_test_bucket(&ecstore, bucket_name.as_str()).await;
+        set_bucket_lifecycle_transition_with_tier(bucket_name.as_str(), &tier_name)
+            .await
+            .expect("Failed to set lifecycle configuration");
+
+        let upload = ecstore
+            .new_multipart_upload(bucket_name.as_str(), object_name, &ObjectOptions::default())
+            .await
+            .expect("Failed to create multipart upload");
+
+        let mut part1_reader = PutObjReader::from_vec(part1);
+        let uploaded_part1 = ecstore
+            .put_object_part(
+                bucket_name.as_str(),
+                object_name,
+                &upload.upload_id,
+                1,
+                &mut part1_reader,
+                &ObjectOptions::default(),
+            )
+            .await
+            .expect("Failed to upload first multipart part");
+
+        let mut part2_reader = PutObjReader::from_vec(part2);
+        let uploaded_part2 = ecstore
+            .put_object_part(
+                bucket_name.as_str(),
+                object_name,
+                &upload.upload_id,
+                2,
+                &mut part2_reader,
+                &ObjectOptions::default(),
+            )
+            .await
+            .expect("Failed to upload second multipart part");
+
+        ecstore
+            .clone()
+            .complete_multipart_upload(
+                bucket_name.as_str(),
+                object_name,
+                &upload.upload_id,
+                vec![
+                    rustfs_ecstore::store_api::CompletePart {
+                        part_num: 1,
+                        etag: uploaded_part1.etag.clone(),
+                        ..Default::default()
+                    },
+                    rustfs_ecstore::store_api::CompletePart {
+                        part_num: 2,
+                        etag: uploaded_part2.etag.clone(),
+                        ..Default::default()
+                    },
+                ],
+                &ObjectOptions::default(),
+            )
+            .await
+            .expect("Failed to complete multipart upload");
+
+        enqueue_transition_for_existing_objects(ecstore.clone(), bucket_name.as_str())
+            .await
+            .expect("Failed to enqueue transitioned restore object");
+
+        let transitioned = wait_for_transition(&ecstore, bucket_name.as_str(), object_name, TRANSITION_WAIT_TIMEOUT)
+            .await
+            .expect("multipart object should transition after enqueueing existing objects");
+        assert_eq!(transitioned.parts.len(), 2);
+
+        ecstore
+            .clone()
+            .restore_transitioned_object(
+                bucket_name.as_str(),
+                object_name,
+                &ObjectOptions {
+                    transition: TransitionOptions {
+                        restore_request: RestoreRequest {
+                            days: Some(1),
+                            description: None,
+                            glacier_job_parameters: None,
+                            output_location: None,
+                            select_parameters: None,
+                            tier: None,
+                            type_: None,
+                        },
+                        ..Default::default()
+                    },
+                    ..Default::default()
+                },
+            )
+            .await
+            .expect("Failed to restore transitioned multipart object");
+
+        let restored = ecstore
+            .get_object_info(bucket_name.as_str(), object_name, &ObjectOptions::default())
+            .await
+            .expect("Failed to load restored object info");
+        assert_eq!(restored.parts.len(), 2);
+        assert!(restored.restore_expires.is_some());
+        assert!(!restored.restore_ongoing);
+
+        let mut reader = ecstore
+            .get_object_reader(bucket_name.as_str(), object_name, None, http::HeaderMap::new(), &ObjectOptions::default())
+            .await
+            .expect("Failed to read restored object");
+        let mut data = Vec::new();
+        reader
+            .stream
+            .read_to_end(&mut data)
+            .await
+            .expect("Failed to consume restored object stream");
+        assert_eq!(data, expected);
+    }
 }