rustfs/crates/heal/tests/heal_integration_test.rs

// Copyright 2024 RustFS Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use http::HeaderMap;
use rustfs_common::heal_channel::{HealOpts, HealScanMode};
use rustfs_heal::heal::{
    manager::{HealConfig, HealManager},
    storage::{ECStoreHealStorage, HealObjectOptions as ObjectOptions, HealPutObjReader as PutObjReader, HealStorageAPI},
    task::{HealOptions, HealPriority, HealRequest, HealTaskStatus, HealType},
};
use serial_test::serial;
use std::{
    path::{Path, PathBuf},
    sync::{Arc, Once},
    time::Duration,
};
use tokio::fs;
use tokio_util::sync::CancellationToken;
use tracing::info;
use walkdir::WalkDir;

mod storage_api;

use storage_api::integration::{
    BucketOperations, BucketOptions, ECStore, Endpoint, EndpointServerPools, Endpoints, ObjectIO as _, ObjectOperations as _,
    PoolEndpoints, init_bucket_metadata_sys, init_local_disks,
};

const HEAL_FORMAT_WAIT_TIMEOUT: Duration = Duration::from_secs(25);
const HEAL_FORMAT_WAIT_INTERVAL: Duration = Duration::from_millis(250);
const NON_INLINE_TEST_DATA_SIZE: usize = 256 * 1024 + 137;

fn non_inline_test_data() -> Vec<u8> {
    (0..NON_INLINE_TEST_DATA_SIZE).map(|idx| (idx % 251) as u8).collect()
}

async fn wait_for_path_exists(path: &Path, timeout: Duration, interval: Duration) -> bool {
    let deadline = tokio::time::Instant::now() + timeout;
    loop {
        if path.exists() {
            return true;
        }
        if tokio::time::Instant::now() >= deadline {
            return false;
        }
        tokio::time::sleep(interval).await;
    }
}

static INIT: Once = Once::new();

pub fn init_tracing() {
    INIT.call_once(|| {
        let _ = tracing_subscriber::fmt()
            .with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
            .with_timer(tracing_subscriber::fmt::time::UtcTime::rfc_3339())
            .with_thread_names(true)
            .try_init();
    });
}

/// Test helper: Create test environment with ECStore
async fn setup_test_env() -> (Vec<PathBuf>, Arc<ECStore>, Arc<ECStoreHealStorage>) {
    init_tracing();

    // create temp dir as 4 disks with unique base dir
    let test_base_dir = format!("/tmp/rustfs_heal_heal_test_{}", uuid::Uuid::new_v4());
    let temp_dir = std::path::PathBuf::from(&test_base_dir);
    if temp_dir.exists() {
        fs::remove_dir_all(&temp_dir).await.ok();
    }
    fs::create_dir_all(&temp_dir).await.unwrap();

    // create 4 disk dirs
    let disk_paths = vec![
        temp_dir.join("disk1"),
        temp_dir.join("disk2"),
        temp_dir.join("disk3"),
        temp_dir.join("disk4"),
    ];

    for disk_path in &disk_paths {
        fs::create_dir_all(disk_path).await.unwrap();
    }

    // create EndpointServerPools
    let mut endpoints = Vec::new();
    for (i, disk_path) in disk_paths.iter().enumerate() {
        let mut endpoint = Endpoint::try_from(disk_path.to_str().unwrap()).unwrap();
        // set correct index
        endpoint.set_pool_index(0);
        endpoint.set_set_index(0);
        endpoint.set_disk_index(i);
        endpoints.push(endpoint);
    }

    let pool_endpoints = PoolEndpoints {
        legacy: false,
        set_count: 1,
        drives_per_set: 4,
        endpoints: Endpoints::from(endpoints),
        cmd_line: "test".to_string(),
        platform: format!("OS: {} | Arch: {}", std::env::consts::OS, std::env::consts::ARCH),
    };

    let endpoint_pools = EndpointServerPools::from(vec![pool_endpoints]);

    // format disks (only first time)
    init_local_disks(endpoint_pools.clone()).await.unwrap();

    // Use port 0 so nextest can run this integration binary in parallel
    // with other ECStore-backed tests without sharing a fixed peer port.
    let server_addr: std::net::SocketAddr = "127.0.0.1:0".parse().unwrap();
    let ecstore = ECStore::new(server_addr, endpoint_pools, CancellationToken::new())
        .await
        .unwrap();

    // init bucket metadata system
    let buckets_list = ecstore
        .list_bucket(&BucketOptions {
            no_metadata: true,
            ..Default::default()
        })
        .await
        .unwrap();
    let buckets = buckets_list.into_iter().map(|v| v.name).collect();
    init_bucket_metadata_sys(ecstore.clone(), buckets).await;

    // Create heal storage layer
    let heal_storage = Arc::new(ECStoreHealStorage::new(ecstore.clone()));

    (disk_paths, ecstore, heal_storage)
}

/// Test helper: Create a test bucket
async fn create_test_bucket(ecstore: &Arc<ECStore>, bucket_name: &str) {
    (**ecstore)
        .make_bucket(bucket_name, &Default::default())
        .await
        .expect("Failed to create test bucket");
    info!("Created test bucket: {}", bucket_name);
}

/// Test helper: Upload test object
async fn upload_test_object(ecstore: &Arc<ECStore>, bucket: &str, object: &str, data: &[u8]) {
    let mut reader = PutObjReader::from_vec(data.to_vec());
    let object_info = (**ecstore)
        .put_object(bucket, object, &mut reader, &ObjectOptions::default())
        .await
        .expect("Failed to upload test object");

    info!("Uploaded test object: {}/{} ({} bytes)", bucket, object, object_info.size);
}

mod serial_tests {
    use super::*;

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    #[serial]
    async fn test_heal_object_basic() {
        let (disk_paths, ecstore, heal_storage) = setup_test_env().await;

        // Create test bucket and object
        let bucket_name = "test-heal-object-basic";
        let object_name = "test-object.txt";
        let test_data = non_inline_test_data();

        create_test_bucket(&ecstore, bucket_name).await;
        upload_test_object(&ecstore, bucket_name, object_name, &test_data).await;
        let _obj_dir = disk_paths[0].join(bucket_name).join(object_name);
        // ─── 1️⃣ delete single data shard file ─────────────────────────────────────
        let obj_dir = disk_paths[0].join(bucket_name).join(object_name);
        // find part file at depth 2, e.g. .../<uuid>/part.1
        let target_part = WalkDir::new(&obj_dir)
            .min_depth(2)
            .max_depth(2)
            .into_iter()
            .filter_map(Result::ok)
            .find(|e| e.file_type().is_file() && e.file_name().to_str().map(|n| n.starts_with("part.")).unwrap_or(false))
            .map(|e| e.into_path())
            .expect("Failed to locate part file to delete");

        std::fs::remove_file(&target_part).expect("failed to delete part file");
        assert!(!target_part.exists());
        println!("✅ Deleted shard part file: {target_part:?}");

        let heal_opts = HealOpts {
            recreate: true,
            remove: false,
            update_parity: true,
            ..Default::default()
        };
        let (object_result, object_error) = heal_storage
            .heal_object(bucket_name, object_name, None, &heal_opts)
            .await
            .expect("failed to heal object");
        info!("heal_object result: {:?}, error: {:?}", object_result, object_error);
        assert!(object_error.is_none(), "heal_object returned error: {object_error:?}");

        // `test_heal_format_with_data` covers on-disk shard restoration. Here we
        // focus on the object-level healing contract: the object must remain
        // readable with intact contents after healing.
        let mut reader = ecstore
            .get_object_reader(bucket_name, object_name, None, HeaderMap::new(), &ObjectOptions::default())
            .await
            .expect("Failed to get object reader after heal");

        let mut downloaded_data = Vec::new();
        tokio::io::copy(&mut reader, &mut downloaded_data)
            .await
            .expect("Failed to read healed object data");

        assert_eq!(downloaded_data, test_data, "Healed object data does not match original");

        info!("Heal object basic test passed");
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    #[serial]
    async fn test_heal_bucket_basic() {
        let (disk_paths, ecstore, heal_storage) = setup_test_env().await;

        // Create test bucket
        let bucket_name = "test-heal-bucket-basic";
        create_test_bucket(&ecstore, bucket_name).await;

        // ─── 1️⃣ delete bucket dir on disk ──────────────
        let broken_bucket_path = disk_paths[0].join(bucket_name);
        assert!(broken_bucket_path.exists(), "bucket dir does not exist on disk");
        std::fs::remove_dir_all(&broken_bucket_path).expect("failed to delete bucket dir on disk");
        assert!(!broken_bucket_path.exists(), "bucket dir still exists after deletion");
        println!("✅ Deleted bucket directory on disk: {broken_bucket_path:?}");

        // Create heal manager with faster interval
        let cfg = HealConfig {
            heal_interval: Duration::from_millis(1),
            ..Default::default()
        };
        let heal_manager = HealManager::new(heal_storage.clone(), Some(cfg));
        heal_manager.start().await.unwrap();

        // Submit heal request for the bucket
        let heal_request = HealRequest::new(
            HealType::Bucket {
                bucket: bucket_name.to_string(),
            },
            HealOptions {
                dry_run: false,
                recursive: true,
                remove_corrupted: false,
                recreate_missing: false,
                scan_mode: HealScanMode::Normal,
                update_parity: false,
                timeout: Some(Duration::from_secs(300)),
                pool_index: None,
                set_index: None,
            },
            HealPriority::Normal,
        );

        let task_id = heal_request.id.clone();
        let admission = heal_manager
            .submit_heal_request(heal_request)
            .await
            .expect("Failed to submit bucket heal request");
        assert!(admission.is_admitted(), "bucket heal request should be admitted");

        info!("Submitted bucket heal request with task ID: {}", task_id);

        // Wait for task completion
        tokio::time::sleep(tokio::time::Duration::from_secs(5)).await;

        // Attempt to fetch task status (optional)
        if let Ok(status) = heal_manager.get_task_status(&task_id).await {
            if status == HealTaskStatus::Completed {
                info!("Bucket heal task status: {:?}", status);
            } else {
                panic!("Bucket heal task status: {status:?}");
            }
        }

        // ─── 3️⃣ Verify bucket directory is restored on every disk ───────
        assert!(broken_bucket_path.exists(), "bucket dir does not exist on disk");

        info!("Heal bucket basic test passed");
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    #[serial]
    async fn test_heal_format_basic() {
        let (disk_paths, _ecstore, heal_storage) = setup_test_env().await;

        // ─── 1️⃣ delete format.json on one disk ──────────────
        let format_path = disk_paths[0].join(".rustfs.sys").join("format.json");
        assert!(format_path.exists(), "format.json does not exist on disk");
        std::fs::remove_file(&format_path).expect("failed to delete format.json on disk");
        assert!(!format_path.exists(), "format.json still exists after deletion");
        println!("✅ Deleted format.json on disk: {format_path:?}");

        let (_format_result, format_error) = heal_storage.heal_format(false).await.expect("failed to run heal_format");
        if let Some(err) = format_error {
            info!("heal_format returned error: {:?}", err);
        }

        let restored = wait_for_path_exists(&format_path, HEAL_FORMAT_WAIT_TIMEOUT, HEAL_FORMAT_WAIT_INTERVAL).await;
        assert!(restored, "format.json does not exist on disk after heal");

        // ─── 2️⃣ verify format.json is restored ───────
        assert!(format_path.exists(), "format.json does not exist on disk after heal");

        info!("Heal format basic test passed");
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    #[serial]
    async fn test_heal_format_with_data() {
        let (disk_paths, ecstore, heal_storage) = setup_test_env().await;

        // Create test bucket and object
        let bucket_name = "test-heal-format-with-data";
        let object_name = "test-object.txt";
        let test_data = non_inline_test_data();

        create_test_bucket(&ecstore, bucket_name).await;
        upload_test_object(&ecstore, bucket_name, object_name, &test_data).await;
        let obj_dir = disk_paths[0].join(bucket_name).join(object_name);
        let target_part = WalkDir::new(&obj_dir)
            .min_depth(2)
            .max_depth(2)
            .into_iter()
            .filter_map(Result::ok)
            .find(|e| e.file_type().is_file() && e.file_name().to_str().map(|n| n.starts_with("part.")).unwrap_or(false))
            .map(|e| e.into_path())
            .expect("Failed to locate part file to delete");

        // ─── 1️⃣ delete format.json on one disk ──────────────
        let format_path = disk_paths[0].join(".rustfs.sys").join("format.json");
        std::fs::remove_dir_all(&disk_paths[0]).expect("failed to delete all contents under disk_paths[0]");
        std::fs::create_dir_all(&disk_paths[0]).expect("failed to recreate disk_paths[0] directory");
        println!("✅ Deleted format.json on disk: {:?}", disk_paths[0]);

        let (_format_result, format_error) = heal_storage.heal_format(false).await.expect("failed to run heal_format");
        if let Some(err) = format_error {
            info!("heal_format returned warning/error: {:?}", err);
        }

        let bucket_heal_opts = HealOpts {
            recursive: true,
            recreate: true,
            ..Default::default()
        };
        heal_storage
            .heal_bucket(bucket_name, &bucket_heal_opts)
            .await
            .expect("failed to heal bucket");

        let heal_opts = HealOpts {
            recreate: true,
            remove: false,
            ..Default::default()
        };
        let (object_result, object_error) = heal_storage
            .heal_object(bucket_name, object_name, None, &heal_opts)
            .await
            .expect("failed to heal object");
        info!("heal_object result: {:?}, error: {:?}", object_result, object_error);
        assert!(object_error.is_none(), "heal_object returned error: {object_error:?}");

        let format_restored = wait_for_path_exists(&format_path, HEAL_FORMAT_WAIT_TIMEOUT, HEAL_FORMAT_WAIT_INTERVAL).await;
        assert!(format_restored, "format.json does not exist on disk after heal");
        let target_restored = wait_for_path_exists(&target_part, HEAL_FORMAT_WAIT_TIMEOUT, HEAL_FORMAT_WAIT_INTERVAL).await;

        // ─── 3️⃣ verify format.json is restored ───────
        assert!(format_path.exists(), "format.json does not exist on disk after heal");
        assert!(target_restored, "part file was not restored after heal");

        // ─── 3️⃣ verify each part file is restored ───────
        assert!(target_part.exists());

        // Verify object metadata is accessible
        let obj_info = ecstore
            .get_object_info(bucket_name, object_name, &ObjectOptions::default())
            .await
            .expect("Expected object to be readable after heal");
        assert_eq!(obj_info.size as usize, test_data.len());

        // Actually read the object data to verify integrity
        let mut reader = ecstore
            .get_object_reader(bucket_name, object_name, None, HeaderMap::new(), &ObjectOptions::default())
            .await
            .expect("Failed to get object reader");

        let mut downloaded_data = Vec::new();
        tokio::io::copy(&mut reader, &mut downloaded_data)
            .await
            .expect("Failed to read object data");

        assert_eq!(downloaded_data, test_data, "Healed object data does not match original");

        info!("Heal format with data test passed");
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    #[serial]
    async fn test_heal_storage_api_direct() {
        let (_disk_paths, ecstore, heal_storage) = setup_test_env().await;

        // Test direct heal storage API calls

        // Test heal_format
        let format_result = heal_storage.heal_format(true).await; // dry run
        assert!(format_result.is_ok());
        info!("Direct heal_format test passed");

        // Test heal_bucket
        let bucket_name = "test-bucket-direct";
        create_test_bucket(&ecstore, bucket_name).await;

        let heal_opts = HealOpts {
            recursive: true,
            dry_run: true,
            remove: false,
            recreate: false,
            scan_mode: HealScanMode::Normal,
            update_parity: false,
            no_lock: false,
            pool: None,
            set: None,
        };

        let bucket_result = heal_storage.heal_bucket(bucket_name, &heal_opts).await;
        assert!(bucket_result.is_ok());
        info!("Direct heal_bucket test passed");

        // Test heal_object
        let object_name = "test-object-direct.txt";
        let test_data = b"Test data for direct heal API";
        upload_test_object(&ecstore, bucket_name, object_name, test_data).await;

        let object_heal_opts = HealOpts {
            recursive: false,
            dry_run: true,
            remove: false,
            recreate: false,
            scan_mode: HealScanMode::Normal,
            update_parity: false,
            no_lock: false,
            pool: None,
            set: None,
        };

        let object_result = heal_storage
            .heal_object(bucket_name, object_name, None, &object_heal_opts)
            .await;
        assert!(object_result.is_ok());
        info!("Direct heal_object test passed");

        info!("Direct heal storage API test passed");
    }
}