// Copyright 2025 PingCAP, Inc.

use std::{
    sync::{
        Arc, Condvar, Mutex,
        atomic::{AtomicUsize, Ordering},
    },
    thread,
    time::Duration,
};

use bytes::Bytes;
use chrono::{TimeZone, Utc};
use cloud_encryption::KeyspaceEncryptionConfig;
use cloud_worker::native_br::{
    BackupItem, RestoreState,
    test_utils::NativeBrSvcClient,
    v1x::{Backup, BackupStatusResponse, PackedBackupView, S3Override, TaskState},
};
use collections::HashMap;
use futures::executor::block_on;
use kvengine::dfs::new_dfs_from_config;
use kvproto::{
    kvrpcpb::Op,
    metapb,
    metapb::StoreState::{Offline, Tombstone, Up},
};
use native_br::{
    archive, backup,
    common::send_request_to_store,
    online_snapshot::restore_keyspace_from_online_snapshot_with_cfg,
    packing::CopiedPackedBackup,
    restore::RestoreConfig,
    restore_keyspace,
    restore_keyspace::{ReportRestoreStepTrait, RestoreStep},
};
use pd_client::{PdClient, util::get_all_stores_except_tiflash_async};
use rand::Rng;
use rfengine::service_worker::DfsWorkerStatus;
use security::{GetSecurityManager, SecurityConfig, SecurityManager};
use test_cloud_server::{
    ServerCluster, ServerClusterBuilder, TikvWorkerOptions, TryWaiter, alloc_node_id,
    alloc_node_id_vec,
    client::{
        ClusterClientOptions, CommitAction, MutateOptions, PrewriteExt, RequestOptions,
        TxnMutations, TxnWriteMethod,
    },
    keyspace::CreateKeyspaceOptions,
    oss::prepare_dfs,
    util::Mutation,
};
use test_pd_client::{PdWrapper, TestPdControlServer};
use tikv::{config::TikvConfig, storage::mvcc::TimeStamp};
use tikv_util::{codec::bytes::encode_bytes, config::ReadableDuration, info};
use tokio::runtime::Runtime;
use txn_types::Key;

use crate::cases::{
    FailpointGuard, cfg_notify_after, cfg_notify_once, i_to_keyspace_key, new_security_config,
    random_value,
};

const TEST_CMEK_ID: &str = "random";
const TEST_CMEK_VENDOR: &str = "test";

async fn set_lightweight_backup(cluster: &ServerCluster, enabled: bool) {
    let pd_client = cluster.get_pd_client();
    let security_mgr = pd_client.get_security_mgr();
    let http_client = security_mgr.http_client(hyper::Client::builder()).unwrap();
    let change = serde_json::json!({"rfengine.lightweight-backup": enabled});
    for store in get_all_stores_except_tiflash_async(pd_client.as_ref(), true)
        .await
        .unwrap()
    {
        let uri = security_mgr
            .build_uri(format!("{}/config", store.status_address))
            .unwrap();
        let request = hyper::Request::post(uri)
            .body(hyper::Body::from(serde_json::to_vec(&change).unwrap()))
            .unwrap();
        send_request_to_store(request, &store, &http_client, Duration::from_secs(30))
            .await
            .unwrap();
    }
}

async fn all_dfs_workers_match(
    cluster: &ServerCluster,
    expected: impl Fn(DfsWorkerStatus) -> bool,
) -> bool {
    let pd_client = cluster.get_pd_client();
    let security_mgr = pd_client.get_security_mgr();
    let http_client = security_mgr.http_client(hyper::Client::builder()).unwrap();
    let stores = get_all_stores_except_tiflash_async(pd_client.as_ref(), true)
        .await
        .unwrap();
    if stores.is_empty() {
        return false;
    }
    for store in stores {
        let uri = security_mgr
            .build_uri(format!("{}/rfengine/dfs_worker", store.status_address))
            .unwrap();
        let request = hyper::Request::get(uri).body(hyper::Body::empty()).unwrap();
        let Ok((_, body)) =
            send_request_to_store(request, &store, &http_client, Duration::from_secs(30)).await
        else {
            return false;
        };
        let Ok(status) = serde_json::from_slice(&body) else {
            return false;
        };
        if !expected(status) {
            return false;
        }
    }
    true
}

#[test]
fn test_backup_on_scaling_up() {
    const KEYSPACE_ID: u32 = 1;
    const DATA_LEN: usize = 100;
    const VALUE_SIZE: usize = 64;

    let get_all_stores_fp = "test_pd::get_all_stores";

    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("t_");
    let s3fs = new_dfs_from_config(dfs_config.clone());
    let reporter = Arc::new(DummyStepReporter::default());
    let runtime = Runtime::new().unwrap();

    let nodes = alloc_node_id_vec(6);
    let mut cluster = ServerCluster::new(nodes.clone(), |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = true;
    });
    cluster.wait_region_replicated(&[], 3);
    let pd_client = cluster.get_pd_client();
    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);

    let stores = nodes
        .iter()
        .map(|&node_id| {
            let store = pd_client.get_store(cluster.get_store_id(node_id)).unwrap();
            (node_id, store)
        })
        .collect::<HashMap<_, _>>();

    for idx in [3, 4, 5] {
        stop_node(&mut cluster, nodes[idx], &stores, Tombstone);
    }

    let i_to_key = i_to_keyspace_key(KEYSPACE_ID);
    client.put_kv(0..DATA_LEN, &i_to_key, random_value::<VALUE_SIZE>);
    client.verify_data_with_ref_store();
    let origin_ref_store = client.dump_ref_store();

    let backup_config = backup::BackupConfig {
        dfs: dfs_config.clone(),
        backup_delay: ReadableDuration::secs(1),
        tolerate_err: 1,
        skip_keyspace_meta: true,
        ..Default::default()
    };

    // Perform backup.
    let backup_name = generate_backup_name();
    thread::scope(|s| {
        let stopped_store_id = stop_node(&mut cluster, nodes[0], &stores, Offline);

        let backup_config = backup_config.clone();
        let backup_name = backup_name.clone();
        let pd_client = pd_client.clone();
        let backup_ts = client.get_ts().into_inner();
        // Let the first store-list query establish the backup topology, then
        // pause exactly at the refresh that used to be synchronized by sleep.
        let (stores_refresh, _stores_refresh_guard) = cfg_notify_after(get_all_stores_fp, 1);
        let h = s.spawn(move || {
            backup::backup_cluster_with_ts(
                backup_config,
                backup::BackupType::Lightweight,
                backup_name,
                pd_client.as_ref(),
                backup_ts,
                None,
            )
        });

        let resume_stores_refresh = stores_refresh.blocking_recv().unwrap();

        // This topology change will be tolerated.
        start_node(&mut cluster, nodes[3], &stores, Up);
        let _ = resume_stores_refresh.send(());

        let (_, backup_meta) = h.join().unwrap().expect("backup");
        info!("backup cluster result: {}", backup_meta);
        assert_eq!(backup_meta.get_stores().len(), 3);
        assert_eq!(backup_meta.tolerated_err, 1);
        assert_eq!(backup_meta.tolerated_err_stores, vec![stopped_store_id]);
    });

    // Put another data.
    client.put_kv(
        DATA_LEN / 4..DATA_LEN / 2,
        &i_to_key,
        random_value::<VALUE_SIZE>,
    );
    client.verify_data_with_ref_store();
    let origin_ref_store_2 = client.dump_ref_store();

    // Perform backup.
    let backup_name_2 = generate_backup_name();
    thread::scope(|s| {
        start_node(&mut cluster, nodes[0], &stores, Up);

        let backup_name = backup_name_2.clone();
        let backup_ts = client.get_ts().into_inner();
        let (stores_refresh, _stores_refresh_guard) = cfg_notify_after(get_all_stores_fp, 1);
        let h = s.spawn(move || {
            backup::backup_cluster_with_ts(
                backup_config,
                backup::BackupType::Lightweight,
                backup_name,
                pd_client.as_ref(),
                backup_ts,
                None,
            )
        });

        let resume_stores_refresh = stores_refresh.blocking_recv().unwrap();

        // This topology change can not be tolerated, then retry to succeed.
        for idx in [4, 5] {
            start_node(&mut cluster, nodes[idx], &stores, Up);
        }
        let _ = resume_stores_refresh.send(());

        let (_, backup_meta) = h.join().unwrap().expect("backup");
        info!("backup cluster result: {}", backup_meta);
        assert_eq!(backup_meta.get_stores().len(), 6);
        assert_eq!(backup_meta.tolerated_err, 0);
    });

    // More data at last.
    client.put_kv(
        DATA_LEN / 3..DATA_LEN * 2 / 3,
        &i_to_key,
        random_value::<VALUE_SIZE>,
    );
    client.verify_data_with_ref_store();

    let object_cache = cluster.create_object_cache_randomly();
    let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone());
    // Perform restore to verify backup.
    for (bk_name, ref_store) in [
        (backup_name, origin_ref_store),
        (backup_name_2, origin_ref_store_2),
    ] {
        let restore_config = RestoreConfig {
            tolerate_err: 0,
            strict_tolerate: true,
            lower_memory: RestoreConfig::use_lower_memory(),
            ..Default::default()
        };
        let res = restore_keyspace::restore_keyspace(
            KEYSPACE_ID,
            KEYSPACE_ID,
            &bk_name,
            None,
            s3fs.clone(),
            restore_config.clone(),
            cluster.get_pd_client(),
            None,
            &runtime,
            None,
            None,
            reporter.clone(),
            object_cache.clone(),
            limiter.clone(),
        )
        .expect("restore");
        info!("restore keyspace result for backup {}: {:?}", bk_name, res);
        let (existed, deleted) = client
            .verify_data_with_given_ref_store(&ref_store, None, &RequestOptions::default())
            .unwrap();
        assert_eq!(existed, DATA_LEN);
        assert_eq!(deleted, 0);
    }

    fail::remove(get_all_stores_fp);
    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_restore_on_disk_full() {
    const KEYSPACE_ID: u32 = 1;
    const DATA_LEN: usize = 10;
    const VALUE_SIZE: usize = 64;

    let low_space_fp = if rand::thread_rng().gen_bool(0.5) {
        "engine_is_low_space" // leader is low space.
    } else {
        "is_store_low_space" // follower is low space.
    };

    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("t_");
    let s3fs = new_dfs_from_config(dfs_config.clone());
    let reporter = Arc::new(DummyStepReporter::default());
    let runtime = Runtime::new().unwrap();

    let nodes = alloc_node_id_vec(3);
    let mut cluster = ServerCluster::new(nodes.clone(), |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = true;
    });
    cluster.wait_region_replicated(&[], 3);
    let pd_client = cluster.get_pd_client();
    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);

    let i_to_key = i_to_keyspace_key(KEYSPACE_ID);
    client.put_kv(0..DATA_LEN, &i_to_key, random_value::<VALUE_SIZE>);
    client.verify_data_with_ref_store();
    let origin_ref_store = client.dump_ref_store();

    let backup_config = backup::BackupConfig {
        dfs: dfs_config.clone(),
        backup_delay: ReadableDuration::secs(1),
        skip_keyspace_meta: true,
        ..Default::default()
    };

    // Perform backup.
    let backup_name = generate_backup_name();
    let pd_client = pd_client.clone();
    let backup_ts = client.get_ts().into_inner();
    backup::backup_cluster_with_ts(
        backup_config,
        backup::BackupType::Lightweight,
        backup_name.clone(),
        pd_client.as_ref(),
        backup_ts,
        None,
    )
    .expect("backup");

    // Put another data.
    client.put_kv(
        DATA_LEN / 4..DATA_LEN / 2,
        &i_to_key,
        random_value::<VALUE_SIZE>,
    );
    client.verify_data_with_ref_store();

    fail::cfg(low_space_fp, "return").unwrap();

    let restore_config = RestoreConfig::default_for_test();
    let object_cache = cluster.create_object_cache_randomly();
    let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone());
    thread::scope(|s| {
        let pd_client = cluster.get_pd_client();
        let h = s.spawn(move || {
            restore_keyspace::restore_keyspace(
                KEYSPACE_ID,
                KEYSPACE_ID,
                &backup_name,
                None,
                s3fs.clone(),
                restore_config,
                pd_client,
                None,
                &runtime,
                None,
                None,
                reporter,
                object_cache,
                limiter,
            )
        });

        thread::sleep(Duration::from_secs(1));
        fail::cfg(low_space_fp, "off").unwrap();

        let res = h.join().unwrap().expect("restore");
        info!("restore keyspace result: {:?}", res);
    });
    let (existed, deleted) = client
        .verify_data_with_given_ref_store(&origin_ref_store, None, &RequestOptions::default())
        .unwrap();
    assert_eq!(existed, DATA_LEN);
    assert_eq!(deleted, 0);

    fail::remove(low_space_fp);
    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_native_br_service() {
    const KEYSPACE_A_ID: u32 = 1;
    const KEYSPACE_B_ID: u32 = 42;
    const DATA_LEN: usize = 100;
    const VALUE_SIZE: usize = 64;

    let mock_get_keyspace_offline_pd = "offline_pd::mock_get_keyspace_by_name";
    let mock_no_tiflash = "pd_ctl::mock_no_tiflash_placement_rule_group";
    let fail_load_keyspace = "test_pd::load_keyspace";
    fail::cfg(mock_get_keyspace_offline_pd, "return").unwrap();
    fail::cfg(fail_load_keyspace, "return").unwrap();

    let runtime = Runtime::new().unwrap();
    let _enter = runtime.enter();

    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("t_");
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let mut cluster = ServerClusterBuilder::new(alloc_node_id_vec(3), |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = true;
    })
    .pd(pd_wrapper)
    .build();
    let tikv_worker_id = alloc_node_id();
    let tikv_worker_opts = TikvWorkerOptions {
        backup_interval: Duration::from_millis(1010),
        backup_skip_keyspace_meta: true,
        restore_timeout_pd_control: Duration::from_secs(3),
        ..Default::default()
    };
    cluster.start_tikv_workers(vec![tikv_worker_id], tikv_worker_opts.clone());
    cluster.wait_region_replicated(&[], 3);

    let pd_client = cluster.get_pd_client();
    let cluster_id = pd_client.get_cluster_id().unwrap();

    // Register both fixed IDs in mock PD, not just their Region boundaries.
    // Cross-keyspace restore synchronizes the target encryption config through
    // PD after retaining the restored files.
    assert_eq!(
        cluster.keyspace_manager().new_keyspace_id(KEYSPACE_B_ID),
        KEYSPACE_B_ID
    );
    cluster.keyspace_manager().create_keyspaces(
        &[KEYSPACE_A_ID, KEYSPACE_B_ID],
        vec![format!("ks{KEYSPACE_A_ID}"), format!("ks{KEYSPACE_B_ID}")],
        &CreateKeyspaceOptions::default(),
        None,
    );

    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_A_ID);
    client.split_keyspace(KEYSPACE_B_ID);

    let i_to_key = i_to_keyspace_key(KEYSPACE_A_ID);
    client.put_kv(0..DATA_LEN, &i_to_key, random_value::<VALUE_SIZE>);
    client.verify_data_with_ref_store();

    let datetime0 = Utc::now();
    let ref_store0 = client.dump_ref_store();

    let security_mgr = Arc::new(SecurityManager::default());
    let br_cli =
        NativeBrSvcClient::new(cluster_id, cluster.tikv_worker_endpoints(), security_mgr).unwrap();
    let backup_v1 = TryWaiter::timeout(10)
        .interval_dur(Duration::from_millis(100))
        .try_wait_result(|| {
            let mut backups = block_on(br_cli.list_backups(&datetime0)).unwrap();
            backups
                .items
                .pop()
                .ok_or(Err::<BackupItem, String>("no backup found".to_string()))
        })
        .unwrap();
    info!("backup: {:?}", backup_v1);

    // Both listing APIs expose the same backup. Validate the compatibility
    // representation without repeating the expensive restore scenario.
    let backup_x = TryWaiter::timeout(10)
        .interval_dur(Duration::from_millis(100))
        .try_wait_result(|| {
            let mut backups = block_on(
                br_cli.list_backups_x(
                    // add 1s to fetch the backup after written.
                    (datetime0 + chrono::Duration::seconds(1))
                        .format("%Y%m%d/%H%M%S")
                        .to_string(),
                    true,
                    1000,
                ),
            )
            .unwrap();
            backups
                .data
                .pop()
                .ok_or(Err::<Backup, String>("no backup found".to_string()))
        })
        .unwrap();
    info!("backup v1x: {:?}", backup_x);

    let backup_from_v1x = backup_x.try_to_backup_item().unwrap();
    assert_eq!(Some(backup_from_v1x.id), backup_x.v1_compat_id);
    assert_eq!(backup_v1.id, backup_from_v1x.id);
    assert!(
        backup_x.details.as_ref().is_some_and(|v| v.backup_ts > 0),
        "{:?}",
        backup_x
    );
    assert!(
        backup_x
            .details
            .as_ref()
            .is_some_and(|v| { !v.keyspaces.is_empty() }),
        "{:?}",
        backup_x
    );
    let backup = backup_from_v1x;
    // More data.
    client.put_kv(
        DATA_LEN / 3..DATA_LEN * 2 / 3,
        &i_to_key,
        random_value::<VALUE_SIZE>,
    );
    client.verify_data_with_ref_store();

    // Inplace restore.
    let keyspace_a_name = format!("ks{KEYSPACE_A_ID}");
    let progress =
        block_on(br_cli.restore_keyspace_to_backup(1, keyspace_a_name.clone(), None, &backup))
            .unwrap();
    info!("create restore: {:?}", progress);

    // Must fail due to get keyspace name error.
    TryWaiter::timeout(30)
        .interval_dur(Duration::from_millis(100))
        .must_wait(
            || {
                let progress =
                    block_on(br_cli.get_restore_progress(1, keyspace_a_name.clone())).unwrap();
                info!("restore progress: {:?}", progress);
                progress.status == RestoreState::Error
            },
            || "restore not error".to_string(),
        );

    // Restore PD keyspace lookup and retry.
    fail::remove(fail_load_keyspace);
    fail::cfg(mock_no_tiflash, "return").unwrap();

    // Restore tikv worker.
    cluster.stop_tikv_workers();
    cluster.start_tikv_workers(vec![tikv_worker_id], tikv_worker_opts.clone());

    // Check task persistence.
    let progress = block_on(br_cli.get_restore_progress(1, keyspace_a_name.clone())).unwrap();
    info!("restore progress (after restart): {:?}", progress);
    assert_eq!(progress.status, RestoreState::Error);
    assert_eq!(progress.error, "interrupted");

    // Retry with incorrect restore params.
    let err =
        block_on(br_cli.restore_keyspace_to_point_in_time(1, keyspace_a_name.clone(), datetime0))
            .unwrap_err();
    info!("retry restore error: {:?}", err);
    assert!(err.to_string().contains("restore params not match"));

    // Retry restore.
    let progress =
        block_on(br_cli.restore_keyspace_to_backup(1, keyspace_a_name.clone(), None, &backup))
            .unwrap();
    info!("retry restore: {:?}", progress);

    // Wait succeed.
    TryWaiter::timeout(30)
        .interval_dur(Duration::from_millis(100))
        .must_wait(
            || {
                let progress =
                    block_on(br_cli.get_restore_progress(1, keyspace_a_name.clone())).unwrap();
                info!("restore progress: {:?}", progress);
                assert_ne!(
                    progress.status,
                    RestoreState::Error,
                    "progress: {:?}",
                    progress
                );
                progress.status == RestoreState::Succeed
            },
            || "restore not succeed".to_string(),
        );

    // Verify data.
    let (existed, deleted) = client
        .verify_data_with_given_ref_store(&ref_store0, None, &RequestOptions::default())
        .unwrap();
    assert_eq!(existed, DATA_LEN);
    assert_eq!(deleted, 0);

    // Restore from keyspace A to B and verify data.
    {
        let keyspace_b_name = format!("ks{KEYSPACE_B_ID}");
        let progress = block_on(br_cli.restore_keyspace_to_backup(
            2,
            keyspace_b_name.clone(),
            Some(keyspace_a_name.clone()),
            &backup_v1,
        ))
        .unwrap();
        info!("create restore (A->B): {:?}", progress);

        TryWaiter::timeout(60)
            .interval_dur(Duration::from_millis(100))
            .must_wait(
                || {
                    let progress =
                        block_on(br_cli.get_restore_progress(2, keyspace_b_name.clone())).unwrap();
                    info!("restore progress (A->B): {:?}", progress);
                    progress.status == RestoreState::Succeed
                },
                || "restore (A->B) not succeed".to_string(),
            );

        let mut expected_ref_store = ref_store0.clone();
        expected_ref_store.rewrite_keyspace_prefix(KEYSPACE_B_ID);
        let (existed, deleted) = client
            .verify_data_with_given_ref_store(&expected_ref_store, None, &RequestOptions::default())
            .unwrap();
        assert_eq!(existed, DATA_LEN);
        assert_eq!(deleted, 0);
    }

    fail::remove(mock_no_tiflash);
    fail::remove(mock_get_keyspace_offline_pd);

    cluster.stop();
    oss.shutdown();
}

#[derive(Clone, Copy, PartialEq, Eq)]
enum OnlineSnapshotFinalCheckFailure {
    ConcurrentRestoreShard,
    SourceFileRefs,
}

#[test]
fn test_online_snapshot_reports_concurrent_restore_shard_at_final_check() {
    run_online_snapshot_final_check_failure(
        OnlineSnapshotFinalCheckFailure::ConcurrentRestoreShard,
    );
}

#[test]
fn test_online_snapshot_propagates_file_refs_error_before_retention() {
    run_online_snapshot_final_check_failure(OnlineSnapshotFinalCheckFailure::SourceFileRefs);
}

fn run_online_snapshot_final_check_failure(failure: OnlineSnapshotFinalCheckFailure) {
    const SOURCE_KEYSPACE_ID: u32 = 1;
    const TARGET_KEYSPACE_ID: u32 = 2;
    const DATA_LEN: usize = 100;
    const VALUE_SIZE: usize = 64;
    const SOURCE_CAPTURE_FAILPOINT: &str = "online_snapshot_after_source_capture";

    native_br::common::step_to_stdout();
    let runtime = Runtime::new().unwrap();
    let _enter = runtime.enter();
    let (temp_dir, mut oss, dfs_config) = prepare_dfs("t_");
    let dfs = new_dfs_from_config(dfs_config.clone());
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let mut cluster = ServerClusterBuilder::new(alloc_node_id_vec(3), |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = true;
    })
    .pd(pd_wrapper)
    .build();
    cluster.wait_region_replicated(&[], 3);

    runtime.block_on(cluster.keyspace_manager().create_single_keyspace(
        SOURCE_KEYSPACE_ID,
        format!("ks{SOURCE_KEYSPACE_ID}"),
        &Default::default(),
        false,
    ));
    runtime.block_on(cluster.keyspace_manager().create_single_keyspace(
        TARGET_KEYSPACE_ID,
        format!("ks{TARGET_KEYSPACE_ID}"),
        &Default::default(),
        false,
    ));
    let mut client = cluster.new_client();
    client.split_keyspace(SOURCE_KEYSPACE_ID);
    client.split_keyspace(TARGET_KEYSPACE_ID);
    client.put_kv(
        0..DATA_LEN,
        i_to_keyspace_key(SOURCE_KEYSPACE_ID),
        random_value::<VALUE_SIZE>,
    );

    let backup_name = generate_backup_name();
    backup::backup_cluster_with_ts(
        backup::BackupConfig {
            dfs: dfs_config.clone(),
            backup_delay: ReadableDuration::secs(1),
            skip_keyspace_meta: true,
            ..Default::default()
        },
        backup::BackupType::Lightweight,
        backup_name.clone(),
        cluster.get_pd_client().as_ref(),
        client.get_ts().into_inner(),
        None,
    )
    .expect("backup source before concurrent RestoreShard");

    let keyspace_manager = cluster.keyspace_manager().clone();
    let pd_control_server =
        TestPdControlServer::start(&runtime, cluster.get_pd_client(), move |keyspace_name| {
            keyspace_manager
                .get_all_keyspaces()
                .into_iter()
                .find_map(|keyspace_id| {
                    let keyspace = keyspace_manager.get_keyspace_meta(keyspace_id)?;
                    (keyspace.name() == keyspace_name).then_some(keyspace_id)
                })
        });
    let pd_config = pd_client::Config::new(vec![pd_control_server.endpoint().to_string()]);
    let restore_config = RestoreConfig {
        pd: pd_config,
        dfs: dfs_config.clone(),
        timeout_pd_control: ReadableDuration::secs(3),
        online_snapshot_attempt_timeout: ReadableDuration::secs(60),
        ..RestoreConfig::default_for_test()
    };
    let reporter = Arc::new(DummyStepReporter::default());
    let gate = Arc::new((Mutex::new(0_u8), Condvar::new()));
    let callback_gate = gate.clone();
    let _source_capture_guard =
        FailpointGuard::new().cfg_callback(SOURCE_CAPTURE_FAILPOINT, move || {
            let (state, changed) = &*callback_gate;
            let mut state = state.lock().unwrap();
            if *state != 0 {
                return;
            }
            *state = 1;
            changed.notify_all();
            while *state == 1 {
                state = changed.wait(state).unwrap();
            }
        });
    let _no_tiflash_guard =
        FailpointGuard::new().cfg("pd_ctl::mock_no_tiflash_placement_rule_group", "return");
    let retain_calls = Arc::new(AtomicUsize::new(0));
    let (_file_refs_guard, _retain_guard) = if failure
        == OnlineSnapshotFinalCheckFailure::SourceFileRefs
    {
        let callback_retain_calls = retain_calls.clone();
        (
            Some(FailpointGuard::new().cfg("online_snapshot_source_file_refs_request", "return")),
            Some(FailpointGuard::new().cfg_callback(
                "restore_keyspace_before_retain_sst_files",
                move || {
                    callback_retain_calls.fetch_add(1, Ordering::SeqCst);
                },
            )),
        )
    } else {
        (None, None)
    };

    thread::scope(|scope| {
        let branch_dfs = dfs.clone();
        let branch_pd_client = cluster.get_pd_client();
        let branch_reporter = reporter.clone();
        let branch_config = restore_config.clone();
        let working_path = temp_dir.path().join("online-snapshot-conflict");
        let branch = scope.spawn(|| {
            restore_keyspace_from_online_snapshot_with_cfg(
                branch_config,
                &format!("ks{SOURCE_KEYSPACE_ID}"),
                &format!("ks{TARGET_KEYSPACE_ID}"),
                working_path,
                branch_dfs,
                branch_pd_client,
                &runtime,
                None,
                branch_reporter,
                None,
            )
        });

        let (state, changed) = &*gate;
        let state = state.lock().unwrap();
        let (mut state, wait) = changed
            .wait_timeout_while(state, Duration::from_secs(30), |state| *state == 0)
            .unwrap();
        if wait.timed_out() {
            *state = 2;
            changed.notify_all();
            drop(state);
            let branch_result = branch.join().unwrap();
            panic!(
                "online snapshot did not reach the post-capture failpoint: {:?}",
                branch_result
            );
        }
        drop(state);

        let restore_result = (failure == OnlineSnapshotFinalCheckFailure::ConcurrentRestoreShard)
            .then(|| {
                // Apply a real in-place RestoreShard after source capture but
                // before the mandatory leader-local file_refs observation.
                restore_keyspace::restore_keyspace(
                    SOURCE_KEYSPACE_ID,
                    SOURCE_KEYSPACE_ID,
                    &backup_name,
                    None,
                    dfs.clone(),
                    restore_config.clone(),
                    cluster.get_pd_client(),
                    None,
                    &runtime,
                    None,
                    None,
                    reporter.clone(),
                    None,
                    None,
                )
            });

        let mut state = gate.0.lock().unwrap();
        *state = 2;
        gate.1.notify_all();
        drop(state);
        if let Some(result) = restore_result {
            result.expect("concurrent in-place RestoreShard");
        }

        let err = branch
            .join()
            .unwrap()
            .expect_err("branch must reject the injected final-check failure");
        let err = err.to_string();
        let expected = match failure {
            OnlineSnapshotFinalCheckFailure::ConcurrentRestoreShard => {
                "source Region version did not follow a split-only transition"
            }
            OnlineSnapshotFinalCheckFailure::SourceFileRefs => {
                "injected source file_refs request failure"
            }
        };
        assert!(err.contains(expected), "{}", err);
        assert!(
            failure != OnlineSnapshotFinalCheckFailure::SourceFileRefs
                || retain_calls.load(Ordering::SeqCst) == 0,
            "retain_sst_files ran after file_refs failed"
        );
    });

    cluster.stop();
    oss.shutdown();
}

#[rstest::rstest]
#[case::with_override_pack(true, false, false, false, false)]
#[case::with_moved_point_in_time(false, true, true, false, false)]
#[case::trivial_lightweight_with_cmek(false, false, false, true, true)]
fn test_native_br_service_x(
    #[case] override_pack: bool,
    #[case] copy_to_another_pfx: bool,
    #[case] with_point_in_time: bool,
    #[case] use_lightweight_backup_api: bool,
    #[case] enable_cmek: bool,
) {
    native_br::common::step_to_stdout();

    const KEYSPACE_ID: u32 = 1;
    const DATA_LEN: usize = 100;
    const VALUE_SIZE: usize = 64;

    let mock_get_keyspace_fp = "pd_ctl::mock_get_keyspace_by_name";
    let mock_get_keyspace_fp2 = "offline_pd::mock_get_keyspace_by_name";
    fail::cfg(mock_get_keyspace_fp, "return").unwrap();
    fail::cfg(mock_get_keyspace_fp2, "return").unwrap();

    let runtime = Runtime::new().unwrap();
    let _enter = runtime.enter();

    let (_temp_dir, _oss, dfs_config) = prepare_dfs("t_");
    let security_conf = if enable_cmek {
        new_security_config()
    } else {
        SecurityConfig::default()
    };
    let pd_wrapper = PdWrapper::new_test(1, &security_conf, None);
    let mut cluster = ServerClusterBuilder::new(alloc_node_id_vec(3), |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = true;
        conf.raft_store.enable_inner_key_offset = true;
        conf.security = security_conf.clone();
    })
    .pd(pd_wrapper)
    .build();
    let tikv_worker_id = alloc_node_id();
    let tikv_worker_opts = TikvWorkerOptions {
        backup_interval: Duration::from_millis(1010),
        backup_skip_keyspace_meta: true,
        restore_timeout_pd_control: Duration::from_secs(3),
        ..Default::default()
    };
    cluster.start_tikv_workers(vec![tikv_worker_id], tikv_worker_opts.clone());
    cluster.wait_region_replicated(&[], 3);

    let pd_client = cluster.get_pd_client();
    let cluster_id = pd_client.get_cluster_id().unwrap();

    let mut client = cluster.new_client();
    runtime.block_on(cluster.keyspace_manager().create_single_keyspace(
        KEYSPACE_ID,
        format!("ks{KEYSPACE_ID}"),
        &Default::default(),
        false,
    ));
    runtime.block_on(cluster.keyspace_manager().create_single_keyspace(
        2,
        "ks2".to_owned(),
        &Default::default(),
        false,
    ));
    if enable_cmek {
        enable_test_cmek(pd_client.clone(), KEYSPACE_ID);
    }
    client.split_keyspace(KEYSPACE_ID);
    client.split_keyspace(2);

    let i_to_key = i_to_keyspace_key(KEYSPACE_ID);
    let baseline_commit_ts = client.put_kv(0..DATA_LEN, &i_to_key, random_value::<VALUE_SIZE>);
    client.verify_data_with_ref_store();

    // IMPORTANT: `point_in_time` is converted to a TSO by taking its physical
    // milliseconds and composing `TimeStamp::compose(physical_ms, 0)`.
    //
    // That means the truncation boundary is aligned to millisecond granularity
    // (with logical=0). To make the test stable, we craft a boundary *between*
    // the baseline commit and the later overwrite by using the baseline
    // commit's physical time + 1ms.
    let baseline_physical_ms = baseline_commit_ts.physical();
    let datetime0 = Utc
        .timestamp_millis_opt((baseline_physical_ms + 1) as i64)
        .single()
        .expect("valid timestamp from baseline TSO physical part");
    let mut ref_store0 = client.dump_ref_store();

    // When `point_in_time` is specified, we want to verify that the restore will
    // truncate the later writes (in the same selected backup) and only restore
    // data as-of `datetime0`.
    //
    // NOTE: We intentionally keep the later writes in the same backup archive.
    // `point_in_time` is not used to select a different backup here, it only
    // affects truncation.
    if with_point_in_time {
        // Ensure the overwrite commits strictly after `datetime0` in physical
        // time. Otherwise, if baseline + overwrite land in the same physical
        // millisecond, truncation boundary (logical=0) may not sit between
        // them.
        let wait_until_ms = baseline_physical_ms + 2;
        let start_wait = std::time::Instant::now();
        while client.get_ts().physical() < wait_until_ms {
            assert!(
                start_wait.elapsed() < Duration::from_secs(5),
                "waited too long for PD TSO physical time to advance"
            );
            thread::sleep(Duration::from_millis(1));
        }

        // Overwrite some keys with new values so that restore-with-truncate can
        // be observed behaviorally.
        client.put_kv(0..(DATA_LEN / 2), &i_to_key, |_| b"boo!".into());
        client.verify_data_with_ref_store();
    }

    let security_mgr = Arc::new(SecurityManager::default());
    let br_cli =
        NativeBrSvcClient::new(cluster_id, cluster.tikv_worker_endpoints(), security_mgr).unwrap();

    let target_datetime = datetime0 + chrono::Duration::seconds(1);
    let backup_x = TryWaiter::timeout(10)
        .interval_dur(Duration::from_millis(100))
        .try_wait_result(|| {
            runtime.block_on(async {
                if override_pack {
                    let s3_override = S3Override {
                        prefix: Some(dfs_config.prefix.clone()),
                        bucket: Some(dfs_config.s3_bucket.clone()),
                        use_endpoint: None,
                    };
                    br_cli
                        .get_backup_for_pitr_with_override(target_datetime, s3_override)
                        .await
                } else {
                    br_cli.get_backup_for_pitr(target_datetime).await
                }
            })
        })
        .unwrap();
    info!("backup v1x: {:?}", backup_x);

    let backup_name = if use_lightweight_backup_api {
        let resp = br_cli.create_lightweight_backup().unwrap();
        info!("create lightweight backup via v1x: {:?}", resp);
        assert!(!resp.backup_name.is_empty());
        resp.backup_name
    } else {
        // Keep the old behavior unchanged.
        br_cli.flush_wals().unwrap();
        backup_x.name.clone()
    };

    let pack_backup = if override_pack {
        let s3_override = S3Override {
            prefix: Some(dfs_config.prefix.clone()),
            bucket: Some(dfs_config.s3_bucket.clone()),
            use_endpoint: None,
        };
        block_on(br_cli.pack_backup_with_override(
            1,
            &backup_name,
            &format!("ks{KEYSPACE_ID}"),
            s3_override,
        ))
    } else {
        block_on(br_cli.pack_backup(1, &backup_name, &format!("ks{KEYSPACE_ID}")))
    }
    .unwrap();
    println!(">>> {pack_backup:?}");

    assert!(
        TryWaiter::timeout(10)
            .interval_dur(Duration::from_millis(100))
            .try_wait(|| {
                let task = block_on(br_cli.get_task_status(1)).unwrap();
                task.state().is_done()
            })
    );

    // pack backup for another keyspace to verify that the meta won't be overridden.
    block_on(br_cli.pack_backup(100, &backup_name, "ks2")).unwrap();
    assert!(
        TryWaiter::timeout(10)
            .interval_dur(Duration::from_millis(100))
            .try_wait(|| {
                let task = block_on(br_cli.get_task_status(100)).unwrap();
                task.state().is_done()
            })
    );

    let task = block_on(br_cli.get_task_status(1)).unwrap();
    let TaskState::Done { info } = task.state() else {
        unreachable!()
    };
    let packed_backup = serde_json::from_value::<PackedBackupView>(info.clone()).unwrap();
    let restore_path = if copy_to_another_pfx {
        let copy =
            block_on(br_cli.copy_packed_backup(3, &packed_backup.path, "somewhere-else")).unwrap();
        println!(">>> {copy:?}");
        assert!(
            TryWaiter::timeout(10)
                .interval_dur(Duration::from_millis(100))
                .try_wait(|| {
                    let task = block_on(br_cli.get_task_status(3)).unwrap();
                    task.state().is_done()
                })
        );
        let task = block_on(br_cli.get_task_status(3)).unwrap();
        let TaskState::Done { info } = task.state() else {
            unreachable!()
        };
        let copied_backup = serde_json::from_value::<CopiedPackedBackup>(info.clone()).unwrap();
        copied_backup.new_exotic_path
    } else {
        packed_backup.path
    };

    let restore = block_on(br_cli.restore_packed_backup(
        2,
        &restore_path,
        "ks2",
        with_point_in_time.then_some(datetime0),
    ))
    .unwrap();
    println!(">>> {restore:?}");

    assert!(
        TryWaiter::timeout(100)
            .interval_dur(Duration::from_millis(100))
            .try_wait(|| {
                let task = block_on(br_cli.get_task_status(2)).unwrap();
                task.state().is_done()
            })
    );
    ref_store0.rewrite_keyspace_prefix(2);
    client
        .verify_data_with_given_ref_store(&ref_store0, None, &Default::default())
        .unwrap();
    if enable_cmek {
        assert_test_cmek(pd_client.as_ref(), 2);
    }

    fail::remove(mock_get_keyspace_fp);
    fail::remove(mock_get_keyspace_fp2);
}

#[test]
fn test_on_demand_backup_restores_disabled_period_writes() {
    const SOURCE_KEYSPACE_ID: u32 = 1;
    const RESTORE_KEYSPACE_ID: u32 = 2;
    const PACK_TASK_ID: u64 = 1;
    const RESTORE_TASK_ID: u64 = 2;
    const DATA_LEN: usize = 100;
    const VALUE_SIZE: usize = 64;

    let mock_get_keyspace = "pd_ctl::mock_get_keyspace_by_name";
    let mock_no_tiflash = "pd_ctl::mock_no_tiflash_placement_rule_group";
    let mock_offline_keyspace = "offline_pd::mock_get_keyspace_by_name";
    fail::cfg(mock_get_keyspace, "return").unwrap();
    fail::cfg(mock_no_tiflash, "return").unwrap();
    fail::cfg(mock_offline_keyspace, "return").unwrap();

    let runtime = Runtime::new().unwrap();
    let _enter = runtime.enter();

    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("on_demand_backup_");
    let nodes = alloc_node_id_vec(3);
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let mut cluster = ServerClusterBuilder::new(nodes, |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = false;
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(vec![alloc_node_id()], TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 3);

    let source_keyspace = format!("ks{SOURCE_KEYSPACE_ID}");
    let restore_keyspace = format!("ks{RESTORE_KEYSPACE_ID}");
    runtime.block_on(cluster.keyspace_manager().create_single_keyspace(
        SOURCE_KEYSPACE_ID,
        source_keyspace.clone(),
        &Default::default(),
        false,
    ));
    runtime.block_on(cluster.keyspace_manager().create_single_keyspace(
        RESTORE_KEYSPACE_ID,
        restore_keyspace.clone(),
        &Default::default(),
        false,
    ));

    let cluster_id = cluster.get_pd_client().get_cluster_id().unwrap();
    let mut br_client = NativeBrSvcClient::new(
        cluster_id,
        cluster.tikv_worker_endpoints(),
        Arc::new(SecurityManager::default()),
    )
    .unwrap();
    br_client.set_timeout(Duration::from_secs(30), Duration::from_secs(30));
    assert!(block_on(all_dfs_workers_match(&cluster, |status| {
        !status.enabled && !status.running
    })));

    // Write the only copy of the test data while WAL upload is disabled.
    let mut client = cluster.new_client();
    client.split_keyspace(SOURCE_KEYSPACE_ID);
    client.split_keyspace(RESTORE_KEYSPACE_ID);
    let source_key = i_to_keyspace_key(SOURCE_KEYSPACE_ID);
    client.put_kv(0..DATA_LEN, &source_key, random_value::<VALUE_SIZE>);
    client.verify_data_with_ref_store();
    let mut expected = client.dump_ref_store();
    expected.rewrite_keyspace_prefix(RESTORE_KEYSPACE_ID);

    // Compose the on-demand operation through each store's config/status API
    // and the ordinary instant-backup endpoint.
    block_on(set_lightweight_backup(&cluster, true));
    TryWaiter::timeout(30).interval(1).must_wait(
        || {
            block_on(all_dfs_workers_match(&cluster, |status| {
                status.enabled && status.running && status.healthy
            }))
        },
        || "DFS workers did not become ready".to_string(),
    );
    let backup = block_on(br_client.create_lightweight_backup_async()).unwrap();
    block_on(set_lightweight_backup(&cluster, false));
    assert!(block_on(all_dfs_workers_match(&cluster, |status| {
        !status.enabled && !status.running
    })));

    // Pack the backup before restore. Offline packing requires complete WAL
    // chunks in DFS, so the test cannot pass via the online WAL fallback.
    block_on(br_client.pack_backup(PACK_TASK_ID, &backup.backup_name, &source_keyspace)).unwrap();
    TryWaiter::timeout(60).interval(1).must_wait(
        || {
            let task = block_on(br_client.get_task_status(PACK_TASK_ID)).unwrap();
            assert!(
                !matches!(task.state(), TaskState::Error { .. }),
                "packing on-demand backup failed: {:?}",
                task.state()
            );
            task.state().is_done()
        },
        || "packing on-demand backup did not finish".to_string(),
    );
    let pack_task = block_on(br_client.get_task_status(PACK_TASK_ID)).unwrap();
    let TaskState::Done { info } = pack_task.state() else {
        unreachable!()
    };
    let packed_backup = serde_json::from_value::<PackedBackupView>(info.clone()).unwrap();

    // Restore the packed backup into a separate keyspace. This path only uses
    // the WAL materialized by packing and cannot read the source stores' WAL.
    block_on(br_client.restore_packed_backup(
        RESTORE_TASK_ID,
        &packed_backup.path,
        &restore_keyspace,
        None,
    ))
    .unwrap();
    TryWaiter::timeout(60).interval(1).must_wait(
        || {
            let task = block_on(br_client.get_task_status(RESTORE_TASK_ID)).unwrap();
            assert!(
                !matches!(task.state(), TaskState::Error { .. }),
                "restore from packed on-demand backup failed: {:?}",
                task.state()
            );
            task.state().is_done()
        },
        || "restore from packed on-demand backup did not finish".to_string(),
    );

    let (existed, deleted) = client
        .verify_data_with_given_ref_store(&expected, None, &RequestOptions::default())
        .unwrap();
    assert_eq!(existed, DATA_LEN);
    assert_eq!(deleted, 0);
    assert!(block_on(all_dfs_workers_match(&cluster, |status| {
        !status.enabled && !status.running
    })));

    fail::remove(mock_get_keyspace);
    fail::remove(mock_no_tiflash);
    fail::remove(mock_offline_keyspace);
    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_backup_status_api_reports_auto_and_archived_backups() {
    const KEYSPACE_ID: u32 = 1;
    const DATA_LEN: usize = 32;
    const VALUE_SIZE: usize = 64;
    const RFENGINE_NOW_FP: &str = "rfengine::now_unix_millis";

    let runtime = Runtime::new().unwrap();
    let _enter = runtime.enter();

    let day0 = Utc.with_ymd_and_hms(2030, 1, 1, 0, 0, 0).unwrap();
    let day1 = day0 + chrono::Duration::days(1);
    fail::cfg(
        RFENGINE_NOW_FP,
        &format!("return({})", day0.timestamp_millis()),
    )
    .unwrap();

    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("t_");
    oss.set_as_if_date(Some(day0));
    let s3fs = new_dfs_from_config(dfs_config.clone());
    let pd_wrapper = PdWrapper::new_test(1, &SecurityConfig::default(), None);
    let mut cluster = ServerClusterBuilder::new(alloc_node_id_vec(3), |_, conf| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = true;
    })
    .pd(pd_wrapper)
    .build();
    let tikv_worker_id = alloc_node_id();
    cluster.start_tikv_workers(vec![tikv_worker_id], TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 3);

    let pd_client = cluster.get_pd_client();
    let cluster_id = pd_client.get_cluster_id().unwrap();
    let security_mgr = Arc::new(SecurityManager::default());
    let br_cli =
        NativeBrSvcClient::new(cluster_id, cluster.tikv_worker_endpoints(), security_mgr).unwrap();

    let empty_status = block_on(br_cli.get_backup_status(None)).unwrap();
    assert_eq!(
        empty_status,
        BackupStatusResponse {
            last_archived_date: None,
            first_auto_backup_time: None,
            last_auto_backup_time: None,
        }
    );

    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);
    let i_to_key = i_to_keyspace_key(KEYSPACE_ID);
    let keyspace_region_id = client.get_region_by_key(&i_to_key(0)).id;
    let backup_config = backup::BackupConfig {
        dfs: dfs_config.clone(),
        skip_keyspace_meta: true,
        ..Default::default()
    };

    pd_client.set_tso(TimeStamp::compose(day0.timestamp_millis() as u64, 0));
    client.put_kv(0..DATA_LEN, &i_to_key, random_value::<VALUE_SIZE>);
    client.verify_data_with_ref_store();
    // Archive discovers persisted shard files, so flush before backup to avoid
    // treating an in-memory-only write as an empty backup.
    cluster.flush_memtable(keyspace_region_id).unwrap();
    assert!(
        cluster.wait_for_memtable_flushed(keyspace_region_id, Duration::from_secs(10)),
        "memtable not flushed for keyspace region {}",
        keyspace_region_id
    );
    let backup_ts0 = client.get_ts().into_inner();
    let (backup_key0, backup_meta0) = backup::backup_cluster_with_ts(
        backup_config.clone(),
        backup::BackupType::Lightweight,
        String::default(),
        pd_client.as_ref(),
        backup_ts0,
        None,
    )
    .expect("backup day0");
    assert_eq!(backup_meta0.backup_ts, backup_ts0);
    for node_id in cluster.get_nodes() {
        block_on(cluster.get_rfengine(node_id).upload_wal_chunk());
    }

    fail::cfg(
        RFENGINE_NOW_FP,
        &format!("return({})", day1.timestamp_millis()),
    )
    .unwrap();
    pd_client.set_tso(TimeStamp::compose(day1.timestamp_millis() as u64, 0));
    oss.set_as_if_date(Some(day1));
    client.put_kv(
        DATA_LEN..DATA_LEN * 2,
        &i_to_key,
        random_value::<VALUE_SIZE>,
    );
    client.verify_data_with_ref_store();
    // Archive discovers persisted shard files, so flush before backup to avoid
    // treating an in-memory-only write as an empty backup.
    cluster.flush_memtable(keyspace_region_id).unwrap();
    assert!(
        cluster.wait_for_memtable_flushed(keyspace_region_id, Duration::from_secs(10)),
        "memtable not flushed for keyspace region {}",
        keyspace_region_id
    );
    let backup_ts1 = client.get_ts().into_inner();
    let (_, backup_meta1) = backup::backup_cluster_with_ts(
        backup_config,
        backup::BackupType::Lightweight,
        String::default(),
        pd_client.as_ref(),
        backup_ts1,
        None,
    )
    .expect("backup day1");
    assert_eq!(backup_meta1.backup_ts, backup_ts1);
    for node_id in cluster.get_nodes() {
        block_on(cluster.get_rfengine(node_id).upload_wal_chunk());
    }

    let mut archive_config = archive::ArchiveConfig {
        dfs: dfs_config.clone(),
        max_archive_file_size: 1024 * 32,
        dry_run: false,
        ..Default::default()
    };
    archive_config.check_data_dir();
    archive::archive_cluster_backup(
        archive_config,
        cluster.get_pd_client(),
        s3fs.clone(),
        day0.date_naive(),
        day1.date_naive(),
    )
    .expect("archive day0 backup");
    block_on(s3fs.delete_object(backup_key0.clone(), backup_key0.clone()))
        .expect("delete archived day0 backup meta");

    let backup_file_time = |backup_ts| {
        backup::IncrementalBackupFile::from_backup_ts(backup_ts)
            .created_at()
            .to_owned()
    };
    let status = block_on(br_cli.get_backup_status(None)).unwrap();
    assert_eq!(
        status.first_auto_backup_time,
        Some(backup_file_time(backup_ts1))
    );
    assert_eq!(status.last_archived_date, Some(day0.date_naive()));
    assert_eq!(
        status.last_auto_backup_time,
        Some(backup_file_time(backup_ts1))
    );

    fail::remove(RFENGINE_NOW_FP);
    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_backup_pessimistic_lock() {
    const KEYSPACE_ID: u32 = 1;
    const DATA_LEN: usize = 100;
    const VALUE_SIZE: usize = 64;

    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("t_");
    let s3fs = new_dfs_from_config(dfs_config.clone());
    let reporter = Arc::new(DummyStepReporter::default());
    let runtime = Runtime::new().unwrap();

    let nodes = alloc_node_id_vec(2);
    let mut cluster = ServerCluster::new(nodes.clone(), |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config.clone();
        conf.rfengine.lightweight_backup = true;
    });
    cluster.wait_region_replicated_ext(&[], 2, true);

    let pd_client = cluster.get_pd_client();
    let mut client = cluster.new_client();

    pd_client.disable_default_operator();
    cluster.remove_node_peers(nodes[1]);
    info!("ConfChange removed node peers {}", nodes[1]);

    client.split_keyspace(KEYSPACE_ID);
    client.set_async_commit();

    let i_to_key = i_to_keyspace_key(KEYSPACE_ID);

    client.put_kv(1..DATA_LEN, &i_to_key, random_value::<VALUE_SIZE>);
    client.verify_data_with_ref_store();

    // split at DATA_LEN / 2
    let split_key = i_to_key(DATA_LEN / 2);
    client.split(&split_key);
    let enc = |k| Key::from_raw(k).into_encoded();

    let pk = i_to_key(DATA_LEN);
    let sk = i_to_key(0);

    let region = client.pd_client().get_region(&enc(&pk)).unwrap();
    assert_eq!(region.peers.len(), 1, "{:?}", region);
    cluster.evict_peer(region.peers[0].id);
    let region = client.pd_client().get_region(&enc(&sk)).unwrap();
    let region2 = client.pd_client().get_region(&enc(&pk)).unwrap();
    assert_eq!(region.peers.len(), 1, "{:?}", region);
    assert_eq!(region2.peers.len(), 1, "{:?}", region);
    assert_ne!(
        region.peers[0].store_id, region2.peers[0].store_id,
        "{:?}\n{:?}",
        region, region2
    );

    let start_ts = client.get_ts();
    client
        .kv_pessimistic_lock(
            pk.clone().into(),
            vec![pk.clone().into(), sk.clone().into()],
            start_ts,
            20000,
            start_ts,
            None,
        )
        .unwrap();
    let origin_ref_store = client.dump_ref_store();

    let backup_config = backup::BackupConfig {
        dfs: dfs_config.clone(),
        backup_delay: ReadableDuration::secs(1),
        tolerate_err: 1,
        skip_keyspace_meta: true,
        ..Default::default()
    };

    // Perform backup.
    let backup_name = generate_backup_name();
    thread::scope(|s| {
        let backup_config = backup_config.clone();
        let backup_name = backup_name.clone();
        let pd_client = pd_client.clone();
        let backup_ts = client.get_ts().into_inner();

        let (rx, _backup_store_guard) = cfg_notify_once("native_br::backup_store");
        let (rx_finish, _backup_store_ret_guard) = cfg_notify_once("native_br::backup_store::ret");

        let h = s.spawn(move || {
            backup::backup_cluster_with_ts(
                backup_config,
                backup::BackupType::Lightweight,
                backup_name,
                pd_client.as_ref(),
                backup_ts,
                None,
            )
        });

        let tx = rx.blocking_recv().unwrap();
        let muts = vec![
            put(pk.clone(), random_value::<VALUE_SIZE>(0)),
            put(sk.clone(), random_value::<VALUE_SIZE>(0)),
        ];
        let txn_muts = TxnMutations::from_normal(muts);
        let _ = rx_finish.blocking_recv();
        let mut txn = client.begin_transaction(Some(start_ts));
        client
            .kv_prewrite_ext(
                Bytes::copy_from_slice(&pk),
                Some(&vec![Bytes::copy_from_slice(&sk)]),
                txn_muts,
                &mut txn,
                PrewriteExt {
                    resolve_lock: false,
                    for_update_ts: start_ts,
                    ..Default::default()
                },
            )
            .unwrap();
        tx.send(()).unwrap();

        let (_, backup_meta) = h.join().unwrap().expect("backup");
        info!("backup cluster result: {}", backup_meta);
    });

    // Perform restore to verify backup.
    let restore_config = RestoreConfig {
        tolerate_err: 0,
        strict_tolerate: true,
        lower_memory: RestoreConfig::use_lower_memory(),
        ..Default::default()
    };
    let object_cache = cluster.create_object_cache_randomly();
    let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone());
    let res = restore_keyspace::restore_keyspace(
        KEYSPACE_ID,
        KEYSPACE_ID,
        &backup_name,
        None,
        s3fs.clone(),
        restore_config.clone(),
        cluster.get_pd_client(),
        None,
        &runtime,
        None,
        None,
        reporter.clone(),
        object_cache,
        limiter,
    )
    .expect("restore");
    info!(
        "restore keyspace result for backup {}: {:?}",
        backup_name, res
    );
    let (existed, deleted) = client
        .verify_data_with_given_ref_store(&origin_ref_store, None, &RequestOptions::default())
        .unwrap();
    assert_eq!(existed, DATA_LEN - 1);
    assert_eq!(deleted, 0);

    cluster.stop();
    oss.shutdown();
}

#[rstest::rstest]
#[case(TxnWriteMethod::Normal)]
#[case::txn_file(TxnWriteMethod::FileBased)]
fn test_check_backup_ts(#[case] write_method: TxnWriteMethod) {
    const KEYSPACE_ID: u32 = 1;
    const DATA_LEN: usize = 10;
    const VALUE_SIZE: usize = 64;

    let process_commit_finish_fp = if matches!(write_method, TxnWriteMethod::Normal) {
        "txn::commit_process_write_finish"
    } else {
        "txn::txn_file_process_commit_finish"
    };
    let client_before_kv_commit_fp = "client::before_kv_commit";

    // Set to `false` to reproduce https://github.com/tidbcloud/cloud-storage-engine/issues/2779.
    // i.e., Txn2 has `commit_ts` < `backup_ts`, but not in the backup.
    let check_backup_ts = true;

    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("t_");
    let s3fs = new_dfs_from_config(dfs_config.clone());
    let reporter = Arc::new(DummyStepReporter::default());
    let runtime = Runtime::new().unwrap();

    let nodes = alloc_node_id_vec(3);
    let security_config = new_security_config();
    let pd_wrapper = PdWrapper::new_test(1, &security_config, None);
    let mut cluster = ServerClusterBuilder::new(nodes, |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config.clone();
        conf.security = security_config.clone();
        conf.rfengine.lightweight_backup = true;
        conf.storage.check_backup_ts = check_backup_ts;
    })
    .pd(pd_wrapper)
    .build();
    cluster.start_tikv_workers(alloc_node_id_vec(1), TikvWorkerOptions::default());
    cluster.wait_region_replicated(&[], 3);
    let pd_client = cluster.get_pd_client();
    let mut client = cluster.new_client();
    client.split_keyspace(KEYSPACE_ID);
    let i_to_key = i_to_keyspace_key(KEYSPACE_ID);

    if matches!(write_method, TxnWriteMethod::FileBased) {
        // Split region to avoid txn file command blocking by latch.
        block_on(pd_client.split_regions_with_retry(
            vec![encode_bytes(&i_to_key(DATA_LEN))],
            Duration::from_secs(30),
        ))
        .unwrap();
    }

    // To work around that not existed keys are not checked.
    client.del_kv(0..DATA_LEN * 3, i_to_key);

    let backup_name = generate_backup_name();
    let backup_name_cp = backup_name.clone();
    let ref_store0 = thread::scope(|s| {
        // Txn1 pause at "commit_process_write_finish":
        let (rx1, _process_commit_finish_guard) = cfg_notify_once(process_commit_finish_fp);
        let mut cli = cluster.new_client_opt(ClusterClientOptions {
            txn_file_max_chunk_size: Some(1024),
            ..Default::default()
        });
        let h1 = s.spawn(move || {
            cli.try_put_kv(
                0..DATA_LEN,
                i_to_keyspace_key(KEYSPACE_ID),
                random_value::<VALUE_SIZE>,
                MutateOptions {
                    write_method,
                    ..Default::default()
                },
            )
            .unwrap()
        });
        let tx1 = rx1.blocking_recv().unwrap();

        // Txn2, pause at "before_kv_commit":
        let (rx2, _client_before_kv_commit_guard) = cfg_notify_once(client_before_kv_commit_fp);
        let mut cli = cluster.new_client_opt(ClusterClientOptions {
            txn_file_max_chunk_size: Some(1024),
            ..Default::default()
        });
        let h2 = s.spawn(move || {
            cli.try_put_kv(
                DATA_LEN..DATA_LEN * 3,
                i_to_keyspace_key(KEYSPACE_ID),
                random_value::<VALUE_SIZE>,
                MutateOptions {
                    write_method,
                    ..Default::default()
                },
            )
            .unwrap()
        });
        let tx2 = rx2.blocking_recv().unwrap();

        // Perform backup.
        let backup_config = backup::BackupConfig {
            dfs: dfs_config.clone(),
            backup_delay: ReadableDuration::secs(1),
            skip_keyspace_meta: true,
            ..Default::default()
        };
        let pd_client = pd_client.clone();
        let backup_ts = client.get_ts().into_inner();
        let h_backup = s.spawn(move || {
            backup::backup_cluster_with_ts(
                backup_config,
                backup::BackupType::Lightweight,
                backup_name_cp,
                pd_client.as_ref(),
                backup_ts,
                None,
            )
            .expect("backup")
        });

        // Resume txn1.
        tx1.send(()).unwrap();
        let commit_ts1 = h1.join().unwrap().into_inner();
        let ref_store0 = client.dump_ref_store();
        fail::cfg(process_commit_finish_fp, "off").unwrap();

        // Wait backup.
        let (name, backup_meta) = h_backup.join().unwrap();
        info!("backup: {} {}", name, backup_meta);

        // Resume txn2.
        tx2.send(()).unwrap();
        let commit_ts2 = h2.join().unwrap().into_inner();
        fail::cfg(client_before_kv_commit_fp, "off").unwrap();

        let (existed, deleted) = client.verify_data_with_ref_store();
        assert_eq!(existed, DATA_LEN * 3);
        assert_eq!(deleted, 0);

        assert!(commit_ts1 < backup_ts);
        if check_backup_ts {
            assert!(backup_ts < commit_ts2);
        } else {
            assert!(commit_ts2 < backup_ts);
        }

        ref_store0
    });

    // Perform restore.
    let restore_config = RestoreConfig::default_for_test();
    let object_cache = cluster.create_object_cache_randomly();
    let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone());
    let res = restore_keyspace::restore_keyspace(
        KEYSPACE_ID,
        KEYSPACE_ID,
        &backup_name,
        None,
        s3fs,
        restore_config,
        cluster.get_pd_client(),
        None,
        &runtime,
        None,
        None,
        reporter.clone(),
        object_cache,
        limiter,
    )
    .expect("restore");
    info!("restore: {:?}", res);

    // Verify that only txn1 is restored.
    let (existed, deleted) = client
        .verify_data_with_given_ref_store(&ref_store0, None, &RequestOptions::default())
        .unwrap();
    assert_eq!(existed, DATA_LEN);
    assert_eq!(deleted, DATA_LEN * 2);

    fail::remove(process_commit_finish_fp);
    fail::remove(client_before_kv_commit_fp);
    cluster.stop();
    oss.shutdown();
}

#[test]
fn test_check_backup_ts_with_async_commit() {
    const KEYSPACE_ID: u32 = 1;
    const DATA_LEN: usize = 10;
    const VALUE_SIZE: usize = 64;

    let process_prewrite_finish_fp = "txn::prewrite_process_write_finish";
    let client_before_kv_prewrite_fp = "client::before_kv_prewrite";
    let cm_before_update_max_ts_fp = "cm_before_update_max_ts";

    // Set to `false` to reproduce https://github.com/tidbcloud/cloud-storage-engine/issues/2779.
    // i.e., Txn2 has `commit_ts` < `backup_ts`, but not in the backup.
    let check_backup_ts = true;

    let (_temp_dir, mut oss, dfs_config) = prepare_dfs("t_");
    let s3fs = new_dfs_from_config(dfs_config.clone());
    let reporter = Arc::new(DummyStepReporter::default());
    let runtime = Runtime::new().unwrap();

    let nodes = alloc_node_id_vec(3);
    let security_config = new_security_config();
    let pd_wrapper = PdWrapper::new_test(1, &security_config, None);
    let mut cluster = ServerClusterBuilder::new(nodes, |_, conf: &mut TikvConfig| {
        conf.dfs = dfs_config.clone();
        conf.security = security_config.clone();
        conf.rfengine.lightweight_backup = true;
        conf.storage.check_backup_ts = check_backup_ts;
    })
    .pd(pd_wrapper)
    .build();
    cluster.wait_region_replicated(&[], 3);
    let pd_client = cluster.get_pd_client();
    let mut client = cluster.new_client();
    client.set_async_commit();
    client.split_keyspace(KEYSPACE_ID);
    let i_to_key = i_to_keyspace_key(KEYSPACE_ID);

    // Split region to write to multiple regions.
    block_on(pd_client.split_regions_with_retry(
        vec![encode_bytes(&i_to_key(DATA_LEN))],
        Duration::from_secs(30),
    ))
    .unwrap();

    // To work around that not existed keys are not checked.
    client.del_kv(0..DATA_LEN * 3, i_to_key);

    let backup_name = generate_backup_name();
    let backup_name_cp = backup_name.clone();
    let ref_store0 = thread::scope(|s| {
        // Skip update_max_ts to ensure that `min_commit_ts` is not advanced by
        // `max_ts`.
        fail::cfg(cm_before_update_max_ts_fp, "return").unwrap();

        // Txn1 pause at "prewrite_process_write_finish":
        let (rx1, _process_prewrite_finish_guard) = cfg_notify_once(process_prewrite_finish_fp);
        let mut cli = cluster.new_client();
        cli.set_async_commit();
        let h1 = s.spawn(move || {
            cli.try_put_kv(
                0..DATA_LEN,
                i_to_keyspace_key(KEYSPACE_ID),
                random_value::<VALUE_SIZE>,
                MutateOptions {
                    commit_action: CommitAction::AsyncCommit(Duration::from_secs(3)),
                    no_verify: true,
                    ..Default::default()
                },
            )
            .unwrap()
        });
        let tx1 = rx1.blocking_recv().unwrap();

        // Txn2, pause at "before_kv_prewrite":
        let (rx2, _client_before_kv_prewrite_guard) = cfg_notify_once(client_before_kv_prewrite_fp);
        let mut cli = cluster.new_client();
        cli.set_async_commit();
        let h2 = s.spawn(move || {
            cli.try_put_kv(
                DATA_LEN..DATA_LEN * 3,
                i_to_keyspace_key(KEYSPACE_ID),
                random_value::<VALUE_SIZE>,
                MutateOptions {
                    commit_action: CommitAction::AsyncCommit(Duration::from_secs(3)),
                    no_verify: true,
                    ..Default::default()
                },
            )
            .unwrap()
        });
        let tx2 = rx2.blocking_recv().unwrap();

        // Perform backup.
        let backup_config = backup::BackupConfig {
            dfs: dfs_config.clone(),
            backup_delay: ReadableDuration::secs(1),
            skip_keyspace_meta: true,
            ..Default::default()
        };
        let pd_client = pd_client.clone();
        // Get 2 TSO to ensure that backup_ts is greater than txn2.start_ts + 1;
        let backup_ts = block_on(pd_client.batch_get_tso(2)).unwrap().into_inner();
        let h_backup = s.spawn(move || {
            backup::backup_cluster_with_ts(
                backup_config,
                backup::BackupType::Lightweight,
                backup_name_cp,
                pd_client.as_ref(),
                backup_ts,
                None,
            )
            .expect("backup")
        });

        // Resume txn1.
        tx1.send(()).unwrap();
        let commit_ts1 = h1.join().unwrap().into_inner();
        let ref_store0 = client.dump_ref_store();
        fail::cfg(process_prewrite_finish_fp, "off").unwrap();

        // Wait backup.
        let (name, backup_meta) = h_backup.join().unwrap();
        info!("backup: {} {}", name, backup_meta);

        // Resume txn2.
        tx2.send(()).unwrap();
        let commit_ts2 = h2.join().unwrap().into_inner();
        fail::cfg(client_before_kv_prewrite_fp, "off").unwrap();

        fail::cfg(cm_before_update_max_ts_fp, "off").unwrap();

        let (existed, deleted) = client.verify_data_with_ref_store();
        assert_eq!(existed, DATA_LEN * 3);
        assert_eq!(deleted, 0);

        info!(
            "backup_ts: {}, commit_ts1: {}, commit_ts2: {}",
            backup_ts, commit_ts1, commit_ts2
        );
        assert!(commit_ts1 < backup_ts);
        if check_backup_ts {
            assert!(backup_ts <= commit_ts2); // Can be equal.
        } else {
            assert!(commit_ts2 < backup_ts);
        }

        ref_store0
    });

    let object_cache = cluster.create_object_cache_randomly();
    let limiter = cluster.create_restore_limiter_randomly(runtime.handle().clone());
    // Perform restore.
    let res = restore_keyspace::restore_keyspace(
        KEYSPACE_ID,
        KEYSPACE_ID,
        &backup_name,
        None,
        s3fs,
        RestoreConfig::default(),
        cluster.get_pd_client(),
        None,
        &runtime,
        None,
        None,
        reporter.clone(),
        object_cache,
        limiter,
    )
    .expect("restore");
    info!("restore: {:?}", res);

    // Verify that only txn1 is restored.
    let (existed, deleted) = client
        .verify_data_with_given_ref_store(&ref_store0, None, &RequestOptions::default())
        .unwrap();
    assert_eq!(existed, DATA_LEN);
    assert_eq!(deleted, DATA_LEN * 2);

    fail::remove(process_prewrite_finish_fp);
    fail::remove(client_before_kv_prewrite_fp);
    fail::remove(cm_before_update_max_ts_fp);
    cluster.stop();
    oss.shutdown();
}

#[derive(Default)]
struct DummyStepReporter {}

impl ReportRestoreStepTrait for DummyStepReporter {
    fn report_step(&self, _step: RestoreStep) {}
}

fn generate_backup_name() -> String {
    static BACKUP_ID: AtomicUsize = AtomicUsize::new(0);
    format!("{:04}", BACKUP_ID.fetch_add(1, Ordering::Relaxed))
}

fn stop_node(
    cluster: &mut ServerCluster,
    node_id: u16,
    stores: &HashMap<u16, metapb::Store>,
    state: metapb::StoreState,
) -> u64 /* store_id */ {
    cluster.stop_node(node_id);
    let store_id = set_store_state(cluster, node_id, stores, state);
    info!("node stopped"; "node" => node_id, "store" => store_id, "state" => ?state);
    store_id
}

fn start_node(
    cluster: &mut ServerCluster,
    node_id: u16,
    stores: &HashMap<u16, metapb::Store>,
    state: metapb::StoreState,
) -> u64 /* store_id */ {
    let store_id = set_store_state(cluster, node_id, stores, state);
    cluster.start_node(node_id, |_, _| {});
    info!("node started"; "node" => node_id, "store" => store_id, "state" => ?state);
    store_id
}

fn set_store_state(
    cluster: &ServerCluster,
    node_id: u16,
    stores: &HashMap<u16, metapb::Store>,
    state: metapb::StoreState,
) -> u64 /* store_id */ {
    let mut store = stores.get(&node_id).unwrap().clone();
    let store_id = store.id;
    store.state = state;
    cluster.get_pd_client().put_store(store).unwrap();
    store_id
}

fn enable_test_cmek(pd_client: Arc<dyn PdClient>, keyspace_id: u32) {
    pd_client
        .set_keyspace_encryption(
            keyspace_id,
            KeyspaceEncryptionConfig {
                enabled: true,
                cmek_id: Some(TEST_CMEK_ID.to_owned()),
                vendor: Some(TEST_CMEK_VENDOR.to_owned()),
                ..Default::default()
            },
        )
        .unwrap();
}

fn assert_test_cmek(pd_client: &dyn PdClient, keyspace_id: u32) {
    let cfg = pd_client.get_keyspace_encryption(keyspace_id).unwrap();
    assert!(cfg.enabled);
    assert_eq!(cfg.cmek_id.as_deref(), Some(TEST_CMEK_ID));
    assert_eq!(cfg.vendor.as_deref(), Some(TEST_CMEK_VENDOR));
}

fn put(k: Vec<u8>, v: Vec<u8>) -> Mutation {
    let mut m = Mutation::default();
    m.set_op(Op::Put);
    m.set_key(k);
    m.set_value(v);
    m
}