// Copyright 2023 PingCAP, Inc.

use std::{
    borrow::Cow,
    cmp,
    collections::hash_map::Entry,
    default::Default,
    fmt, fs,
    future::Future,
    mem, ops,
    ops::Deref,
    path::{Path, PathBuf},
    sync::{
        Arc, Mutex, RwLock,
        atomic::{AtomicBool, Ordering},
    },
    thread,
    time::{Duration, Instant as StdInstant},
};

use api_version::{ApiV2, api_v2::KEYSPACE_PREFIX_LEN};
use async_trait::async_trait;
use bytes::Bytes;
use cloud_encryption::{
    EncryptionKey, EncryptionKeyManager, EncryptionState, KeyspaceEncryptionConfig, MasterKey,
    ToEncryptionState, master_key_encryption_context,
};
use cloud_server::{RestoreShardResponse, TikvServer};
use collections::{HashMap, HashMapExt, HashSet};
pub use engine_traits::ObjectCache;
use engine_traits::ObjectCacheWithHook;
use file_system::{IoRateLimitMode, IoRateLimiter, IoType};
use futures::{
    StreamExt,
    stream::{self, FuturesUnordered},
};
use http::{Request, header};
use hyper::Body;
use itertools::Itertools;
use kvengine::{
    ENCRYPTION_KEY, ENCRYPTION_META_KEY, GLOBAL_SHARD_END_KEY, IdAllocator, IdVer,
    LoadTableFilterFn, PrepareOpts, ShardMeta, ShardRange, ShardStats, ShardTag, WRITE_CF,
    dfs::{self, Dfs, FileType},
    encryption_key_from_shard_properties,
    limiter::StoreLimiter,
    table::{
        BoundedDataSet, DataBound, InnerKey,
        file::InMemFile,
        schema_file::{SchemaFile, build_schema_file, schema_from_table_info},
    },
};
use kvenginepb as pb;
use kvproto::{metapb, metapb::PeerRole, pdpb, raft_serverpb::MergeState};
use pb::{PackedBackup, VectorIndex};
use pd_client::{
    PdClient,
    pd_control::{KEYSPACE_CONFIG_ENCRYPTION_KEY, PdControl, TransferRegionPeerRole},
    util::get_all_stores_except_tiflash,
};
use protobuf::Message;
use raft::eraftpb;
use rand::seq::SliceRandom;
use rfengine::{RfEngine, convert_to_decompressed_wal_chunk};
use rfenginepb::ClusterBackupMeta;
use rfstore::store::{
    ApplyMsgs, PdIdAllocator, PeerTag, PreprocessContext, PreprocessResult, PreprocessScene,
    RegionIdVer, StoreMsg, load_raft_engine_meta, parse_raft_cmd, rlog, state::RaftState,
};
use security::SecurityConfig;
use slog_global::{debug, error, info, warn};
use tempfile::Builder as TempDirBuilder;
use tikv::{config::TikvConfig, storage::mvcc::Key};
use tikv_util::{
    HandyRwLock, backoff::ExponentialBackoff, box_err, box_try, config::AbsoluteOrPercentSize,
    defer, http::CONTENT_TYPE_PROTOBUF, merge_range::MergeRanges, mpsc,
    retry::try_wait_result_async, sys::thread::StdThreadBuildWrapper, time::Instant,
};
use tokio::{
    runtime::Runtime,
    sync::{Semaphore, watch},
};
use txn_types::TimeStamp;

use crate::{
    archive::{
        ArchiveReader, StoreMeta, get_cluster_backup_file_and_meta,
        get_incremental_backup_with_name, get_not_found_files, get_packed_backup_meta,
    },
    common::{
        RawRegion, RegionMetaGetter, ReplayWalLogsContext, StorePeer, TableFile,
        collect_snapshot_meta_rlog_files, is_derived_file_type, load_peer_raft_state, now,
        replay_wal_logs, retain_sst_files, send_request_to_store,
    },
    error::{
        Error,
        Error::{MetaNotFound, RetryLimitExceeded},
        Result,
    },
    limiter::{SnapshotSize, ThroughputLimiter},
    lock::LockResolver,
    metrics::{
        self, NATIVE_BR_RESTORE_PENDING_DATA_SIZE, NATIVE_BR_RESTORED_DATA_SIZE,
        NATIVE_BR_RESTORED_KV_SIZE,
    },
    packing::RestorePackEnv,
    restore::RestoreConfig,
    step,
    tiflash::remove_tiflash_replica_of_keyspace,
};

const WORKING_PATH_PREFIX: &str = "keyspace-restore";
const ZSTD_COMPRESSION_LEVEL: &str = "5"; // The same as ZSTD_COMPRESSION_LEVEL_FOR_REMOTE.

// TODO: get replicas from keyspace config to handle 5 replicas in the future.
const REPLICAS: usize = 3; // Number of replicas for each region.

const RESOLVE_LOCKS_BATCH_SIZE: usize = 1024;

const SPLIT_REGIONS_BATCH_SIZE: usize = 256;
pub(crate) const SCAN_REGION_BATCH_SIZE: usize = 256;
const SPLIT_REGIONS_TIMEOUT_PER_KEY: Duration = Duration::from_millis(100);
/// Files younger than this interval cannot reach the production DFS GC cutoff,
/// which is at least thirty minutes, even with a fifteen-minute safety margin.
const RECENT_FILE_RETAIN_SKIP_INTERVAL: Duration = Duration::from_secs(15 * 60);
const RESTORE_SCHEDULE_GUARD_TTL: Duration = Duration::from_secs(60);
const RESTORE_SCHEDULE_GUARD_RENEW_INTERVAL: Duration = Duration::from_secs(20);

const SNAPSHOT_RESTORE_STORE_ID: u64 = 0;
const SNAPSHOT_RESTORE_PEER_ID: u64 = 0;

pub const RESTORE_RFENGINE_CONCURRENCY: usize = 3;

#[derive(Clone)]
struct SnapshotNoopRecovery;

impl kvengine::RecoverHandler for SnapshotNoopRecovery {
    fn recover(
        &self,
        _: &kvengine::Engine,
        _: &Arc<kvengine::Shard>,
        _: &ShardMeta,
        _: bool,
    ) -> kvengine::Result<()> {
        Ok(())
    }
}

fn build_packed_restore_table_filter(shards_need_truncate: HashSet<u64>) -> LoadTableFilterFn {
    Arc::new(move |shard_id, _: &_| shards_need_truncate.contains(&shard_id))
}

fn build_lock_resolving_restore_table_filter(
    shards_need_truncate: HashSet<u64>,
) -> LoadTableFilterFn {
    Arc::new(move |shard_id, table: &kvengine::FileMeta| {
        shards_need_truncate.contains(&shard_id) || table.has_locks()
    })
}

#[derive(Clone, Debug, Default)]
pub struct RestoredKeyspace {
    pub keyspace_id: u32,
    pub target_keyspace_id: u32,
    pub ts: u64,
    pub restore_bytes: u64,
    pub tolerated_err: usize,
    pub resolved_ts: TimeStamp,
}

#[derive(Clone, Debug)]
struct RestoreRegionFence {
    region_id: u64,
    version: u64,
    conf_ver: Option<u64>,
    start: Vec<u8>,
    end: Vec<u8>,
}

#[derive(Clone)]
struct RestoreSnapshotRequest {
    change_set: pb::ChangeSet,
    expected_conf_ver: Option<u64>,
}

impl RestoreSnapshotRequest {
    fn fence(&self) -> RestoreRegionFence {
        let snapshot = self.change_set.get_restore_shard();
        RestoreRegionFence {
            region_id: self.change_set.get_shard_id(),
            version: self.change_set.get_shard_ver(),
            conf_ver: self.expected_conf_ver,
            start: snapshot.get_outer_start().to_vec(),
            end: snapshot.get_outer_end().to_vec(),
        }
    }
}

#[derive(Clone, Debug, Default)]
pub struct RestoredSnapshots {
    pub count: u64,
    pub restore_bytes: u64,
    retry_regions: Vec<RestoreRegionFence>,
}

#[derive(Clone, Debug, Default)]
pub struct RestoredShard {
    pub shard_id: u64,
    pub restore_bytes: u64,
}

/// RestoreStep is the steps of restore process.
///
/// Some steps are performed in retry loop, and will be reported multiple times.
///
/// Must keep the order of enums the same as they happen, then reporters can
/// avoid the progress jumping back and forth.
#[derive(Clone, Debug, PartialEq, PartialOrd)]
pub enum RestoreStep {
    Pending,
    Init,
    InstantBackup,
    CaptureOnlineSnapshot,
    RemoveTiFlashReplicas,
    LoadBackupMeta,
    ExtractBackupShards,
    ResolveLocks,
    FlushShards,
    TruncateTs,
    SplitRegions,
    AlignRegions,
    RestoreSnapshotsToServers,
    RestoreSnapshotsFinished,
    RetainSstFiles,
    Finalize,
}

pub trait ReportRestoreStepTrait: Send + Sync {
    fn report_step(&self, step: RestoreStep);

    fn report_pending_restore_data_size(&self, _data_size: i64) {}
}

/// Selects the subset of restored files that must be retained explicitly.
///
/// Online snapshot restore uses this after the last target `RestoreShard` to
/// subtract files that are still referenced by the live source keyspace. An
/// error aborts the restore before file retention.
pub(crate) trait RetainFileFilter {
    fn filter_files(&self, target_files: Vec<TableFile>) -> Result<Vec<TableFile>>;
}

/// In-memory hard deadline for one restore attempt.
///
/// It is deliberately not serializable: an interrupted online snapshot task
/// restarts from a new branch timestamp and creates a new deadline.
#[derive(Clone, Copy, Debug)]
pub(crate) struct RestoreDeadline {
    expires_at: StdInstant,
}

impl RestoreDeadline {
    pub(crate) fn after(timeout: Duration) -> Result<Self> {
        if timeout.is_zero() {
            return Err(box_err!(
                "online snapshot attempt timeout must be greater than zero"
            ));
        }
        let expires_at = StdInstant::now()
            .checked_add(timeout)
            .ok_or_else(|| -> Error { box_err!("online snapshot attempt deadline overflowed") })?;
        Ok(Self { expires_at })
    }

    pub(crate) fn remaining(&self, phase: &str) -> Result<Duration> {
        let remaining = self.expires_at.saturating_duration_since(StdInstant::now());
        if remaining.is_zero() {
            Err(box_err!(
                "online snapshot attempt deadline exceeded before {}",
                phase
            ))
        } else {
            Ok(remaining)
        }
    }

    pub(crate) fn cap_timeout(&self, timeout: Duration, phase: &str) -> Result<Duration> {
        Ok(cmp::min(timeout, self.remaining(phase)?))
    }

    /// Returns a child deadline capped by both this deadline and `timeout`.
    pub(crate) fn capped(self, timeout: Duration) -> Result<Self> {
        let child = Self::after(timeout)?;
        Ok(Self {
            expires_at: cmp::min(self.expires_at, child.expires_at),
        })
    }
}

#[derive(Clone, Debug)]
pub(crate) struct RestoreScheduleGuardState {
    healthy: Arc<AtomicBool>,
    valid_until: Option<Arc<Mutex<StdInstant>>>,
}

impl RestoreScheduleGuardState {
    pub(crate) fn check_healthy(&self) -> Result<()> {
        let expired = self
            .valid_until
            .as_ref()
            .is_some_and(|valid_until| StdInstant::now() >= *valid_until.lock().unwrap());
        if !expired && self.healthy.load(Ordering::Acquire) {
            Ok(())
        } else {
            self.healthy.store(false, Ordering::Release);
            Err(box_err!(
                "online snapshot schedule=deny lease expired during restore"
            ))
        }
    }

    #[cfg(test)]
    pub(crate) fn with_valid_until_for_test(valid_until: StdInstant) -> Self {
        Self {
            healthy: Arc::new(AtomicBool::new(true)),
            valid_until: Some(Arc::new(Mutex::new(valid_until))),
        }
    }
}

/// Owns a renewable `schedule=deny` rule. A renewal failure is tolerated only
/// while the last conservatively acknowledged TTL is still valid. Once that
/// deadline passes the state is terminal because a later PUT cannot prove that
/// no scheduler ran during the lease gap.
pub(crate) struct RestoreScheduleGuard {
    state: RestoreScheduleGuardState,
    stop_tx: Option<watch::Sender<bool>>,
}

impl RestoreScheduleGuard {
    pub(crate) fn install(
        pd: PdControl,
        key_ranges: &[(&[u8], &[u8])],
        required: bool,
        runtime: &Runtime,
    ) -> Result<Self> {
        let rule = pd_client::pd_control::LabelRule {
            id: uuid::Uuid::new_v4().to_string(),
            labels: vec![pd_client::pd_control::RegionLabel {
                key: "schedule".into(),
                value: "deny".into(),
                ttl: "1m".into(),
            }],
            rule_type: "key-range".into(),
            data: key_ranges
                .iter()
                .map(|(start, end)| {
                    let encoded_start_key = Key::from_raw(start).into_encoded();
                    let encoded_end_key = Key::from_raw(end).into_encoded();
                    pd_client::pd_control::KeyRangeRule {
                        start_key: log_wrappers::hex_encode_upper(encoded_start_key),
                        end_key: log_wrappers::hex_encode_upper(encoded_end_key),
                    }
                })
                .collect(),
        };
        let rule = Arc::new(rule);
        if !required {
            let (stop_tx, mut stop_rx) = watch::channel(false);
            runtime.spawn(async move {
                info!("start optional restore scheduling guard");
                if let Err(err) = pd.set_region_label_rule(&rule).await {
                    warn!("install optional restore scheduling guard failed: {err:?}");
                }
                loop {
                    tokio::select! {
                        changed = stop_rx.changed() => {
                            if changed.is_err() || *stop_rx.borrow() {
                                info!("stop optional restore scheduling guard");
                                break;
                            }
                        }
                        _ = tokio::time::sleep(RESTORE_SCHEDULE_GUARD_RENEW_INTERVAL) => {
                            if let Err(err) = pd.set_region_label_rule(&rule).await {
                                warn!("renew optional restore scheduling guard failed: {err:?}");
                            }
                        }
                    }
                }
            });
            return Ok(Self {
                state: RestoreScheduleGuardState {
                    healthy: Arc::new(AtomicBool::new(true)),
                    valid_until: None,
                },
                stop_tx: Some(stop_tx),
            });
        }

        let install_started_at = StdInstant::now();
        if let Err(err) = runtime.block_on(pd.set_region_label_rule(&rule)) {
            return Err(err.into());
        }
        let Some(initial_deadline) = install_started_at.checked_add(RESTORE_SCHEDULE_GUARD_TTL)
        else {
            return Err(box_err!("restore scheduling guard deadline overflowed"));
        };
        if StdInstant::now() >= initial_deadline {
            return Err(box_err!(
                "initial schedule=deny installation exceeded its TTL"
            ));
        }

        let state = RestoreScheduleGuardState {
            healthy: Arc::new(AtomicBool::new(true)),
            valid_until: Some(Arc::new(Mutex::new(initial_deadline))),
        };
        let renew_state = state.clone();
        let renew_valid_until = state.valid_until.as_ref().unwrap().clone();
        let (stop_tx, mut stop_rx) = watch::channel(false);
        runtime.spawn(async move {
            let mut valid_until = initial_deadline;
            let mut renew_at = install_started_at
                .checked_add(RESTORE_SCHEDULE_GUARD_RENEW_INTERVAL)
                .unwrap_or(install_started_at);
            loop {
                let now = StdInstant::now();
                if now >= valid_until {
                    error!("online snapshot schedule=deny lease expired");
                    renew_state.healthy.store(false, Ordering::Release);
                    break;
                }

                let wake_at = cmp::min(renew_at, valid_until);
                tokio::select! {
                    changed = stop_rx.changed() => {
                        if changed.is_err() || *stop_rx.borrow() {
                            break;
                        }
                    }
                    _ = tokio::time::sleep_until(tokio::time::Instant::from_std(wake_at)) => {}
                }
                if *stop_rx.borrow() {
                    break;
                }

                let now = StdInstant::now();
                if now >= valid_until {
                    error!("online snapshot schedule=deny lease expired");
                    renew_state.healthy.store(false, Ordering::Release);
                    break;
                }
                if now < renew_at {
                    continue;
                }

                let renewal_started_at = now;
                let renewal_timeout = valid_until.saturating_duration_since(renewal_started_at);
                let renewal =
                    tokio::time::timeout(renewal_timeout, pd.set_region_label_rule(&rule));
                let renewal_result = tokio::select! {
                    changed = stop_rx.changed() => {
                        if changed.is_err() || *stop_rx.borrow() {
                            break;
                        }
                        continue;
                    }
                    result = renewal => result,
                };
                match renewal_result {
                    Ok(Ok(())) => {
                        let Some(next_deadline) =
                            renewal_started_at.checked_add(RESTORE_SCHEDULE_GUARD_TTL)
                        else {
                            error!("restore scheduling guard deadline overflowed");
                            renew_state.healthy.store(false, Ordering::Release);
                            break;
                        };
                        valid_until = next_deadline;
                        *renew_valid_until.lock().unwrap() = next_deadline;
                    }
                    Ok(Err(err)) => {
                        warn!("renew region schedule=deny rule failed: {:?}", err);
                    }
                    Err(_) => {
                        error!("renew region schedule=deny rule timed out at lease deadline");
                        renew_state.healthy.store(false, Ordering::Release);
                        break;
                    }
                }
                renew_at = StdInstant::now()
                    .checked_add(RESTORE_SCHEDULE_GUARD_RENEW_INTERVAL)
                    .unwrap_or(valid_until);
            }
        });

        Ok(Self {
            state,
            stop_tx: Some(stop_tx),
        })
    }

    pub(crate) fn state(&self) -> RestoreScheduleGuardState {
        self.state.clone()
    }

    pub(crate) fn check_healthy(&self) -> Result<()> {
        self.state.check_healthy()
    }
}

impl Drop for RestoreScheduleGuard {
    fn drop(&mut self) {
        if let Some(stop_tx) = self.stop_tx.take() {
            let _ = stop_tx.send(true);
        }
    }
}

fn check_restore_deadline(deadline: Option<RestoreDeadline>, phase: &str) -> Result<()> {
    if let Some(deadline) = deadline {
        deadline.remaining(phase)?;
    }
    Ok(())
}

pub fn restore_keyspace_with_cfg(
    config: RestoreConfig,
    keyspace_name: &str,
    target_keyspace_name: &str,
    backup_name: &str,
    working_path: Option<PathBuf>,
    dfs: Arc<dyn Dfs>,
    pd_client: Arc<dyn PdClient>,
    runtime: &Runtime,
    truncate_ts: Option<u64>,
    master_key: Option<MasterKey>,
    reporter: Arc<dyn ReportRestoreStepTrait>,
    object_cache: Option<ObjectCache>,
    limiter: Option<Arc<ThroughputLimiter>>,
) -> Result<RestoredKeyspace> {
    let load_keyspace_id = |name: &str| {
        pd_client
            .load_keyspace(name.to_owned())
            .map(|keyspace| keyspace.get_id())
    };

    let mut pd_control = PdControl::new(config.pd.clone(), pd_client.get_security_mgr())?;
    pd_control.set_timeout(config.timeout_pd_control.0);
    let keyspace_id = load_keyspace_id(keyspace_name)?;
    step!("Keyspace name {keyspace_name}'s id is {keyspace_id}");

    let target_keyspace_id = if keyspace_name != target_keyspace_name {
        let target_id = load_keyspace_id(target_keyspace_name)?;
        step!("Target keyspace name {target_keyspace_name}'s id is {target_id}");
        target_id
    } else {
        reporter.report_step(RestoreStep::RemoveTiFlashReplicas);
        if let Err(e) = runtime.block_on(remove_tiflash_replica_of_keyspace(
            keyspace_id,
            &pd_control,
            pd_client.clone(),
        )) {
            error!("Keyspace {keyspace_id} remove tiflash replica err {:?}", e);
            return Err(e);
        }
        step!("Removed keyspace {keyspace_name} tiflash replica");
        keyspace_id
    };
    restore_keyspace(
        keyspace_id,
        target_keyspace_id,
        backup_name,
        working_path,
        dfs,
        config,
        pd_client,
        Some(pd_control),
        runtime,
        truncate_ts,
        master_key,
        reporter,
        object_cache,
        limiter,
    )
}

pub fn load_norm_backup_meta(
    dfs: Arc<dyn Dfs>,
    backup_name: &str,
) -> Result<(ClusterBackupMeta, Option<ArchiveReader>)> {
    match get_cluster_backup_file_and_meta(dfs.as_ref(), backup_name.to_owned()) {
        Ok((_, backup_meta)) => Ok((backup_meta, None)),
        Err(dfs::Error::NoSuchKey(err)) => {
            warn!(
                "The cluster backup meta {} is not found: {:?}, try the archive reader",
                backup_name.to_owned(),
                err
            );
            let Some(backup_file) =
                get_incremental_backup_with_name(dfs.get_prefix(), backup_name.to_owned())
            else {
                // The key cannot be parsed as an periodical backup. Skip it.
                return Err(Error::DfsError(dfs::Error::NoSuchKey(err)));
            };
            let backup_date = backup_file.created_at().date_naive();
            let archive_reader = ArchiveReader::new(dfs.clone(), &backup_date)?;
            let backup_meta = archive_reader
                .read_meta_file()
                .map_err(|_| MetaNotFound(backup_file.id()))?;
            Ok((backup_meta, Some(archive_reader)))
        }
        Err(err) => {
            error!("get cluster backup meta {} failed: {:?}", backup_name, err);
            Err(Error::DfsError(err))
        }
    }
}

fn target_keyspace_encryption_config_patch(
    exported_encryption_meta: Option<&[u8]>,
    has_exported_encryption_key: bool,
) -> Result<(Option<String>, &'static str)> {
    if exported_encryption_meta.is_some() && has_exported_encryption_key {
        return Err(Error::Other(box_err!(
            "restored shard contains both ENCRYPTION_META_KEY and ENCRYPTION_KEY"
        )));
    }

    let (encryption_config, mode) = if let Some(exported_encryption_meta) = exported_encryption_meta
    {
        let mut meta = pb::EncryptionMeta::default();
        meta.merge_from_bytes(exported_encryption_meta)
            .map_err(|err| {
                Error::Other(box_err!(
                    "failed to parse exported ENCRYPTION_META_KEY: {}",
                    err
                ))
            })?;
        let master_key = meta.get_master_key();
        (
            Some(KeyspaceEncryptionConfig::enabled(
                (!master_key.get_cmek_id().is_empty()).then(|| master_key.get_cmek_id().into()),
                (!master_key.get_vendor().is_empty()).then(|| master_key.get_vendor().into()),
                (!master_key.get_region().is_empty()).then(|| master_key.get_region().into()),
                (!master_key.get_endpoint().is_empty()).then(|| master_key.get_endpoint().into()),
                master_key_encryption_context(master_key),
            )),
            "metadata",
        )
    } else if has_exported_encryption_key {
        (
            Some(KeyspaceEncryptionConfig::enabled(
                None,
                None,
                None,
                None,
                Default::default(),
            )),
            "legacy",
        )
    } else {
        (None, "disabled")
    };

    let config_value = encryption_config
        .map(|config| {
            config.to_pd_config_value().map_err(|err| {
                Error::Other(box_err!(
                    "failed to serialize keyspace encryption config: {}",
                    err
                ))
            })
        })
        .transpose()?;
    Ok((config_value, mode))
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum TargetEncryptionConfigSync {
    /// Preserve the established in-place restore behavior: only exported
    /// metadata encryption requires a target PD config update.
    ExportedMetadataOnly,
    /// Cross-keyspace restores must make future target files use the restored
    /// source mode, including legacy and plaintext.
    RestoredMode,
}

impl TargetEncryptionConfigSync {
    fn for_restore(is_in_place: bool) -> Self {
        if is_in_place {
            Self::ExportedMetadataOnly
        } else {
            Self::RestoredMode
        }
    }

    fn should_patch(self, exported_encryption_meta: Option<&[u8]>) -> bool {
        self == Self::RestoredMode || exported_encryption_meta.is_some()
    }
}

async fn patch_target_keyspace_encryption_config(
    pd_client: Arc<dyn PdClient>,
    source_keyspace_id: u32,
    target_keyspace_id: u32,
    exported_encryption_meta: Option<&[u8]>,
    has_exported_encryption_key: bool,
) -> Result<&'static str> {
    let (encryption_config, mode) = target_keyspace_encryption_config_patch(
        exported_encryption_meta,
        has_exported_encryption_key,
    )?;
    let mut config = std::collections::HashMap::new();
    config.insert(KEYSPACE_CONFIG_ENCRYPTION_KEY.to_owned(), encryption_config);
    pd_client
        .patch_keyspace_config(target_keyspace_id, config, None)
        .await?;
    info!(
        "patched target keyspace encryption config";
        "source_keyspace_id" => source_keyspace_id,
        "target_keyspace_id" => target_keyspace_id,
        "mode" => mode,
    );
    Ok(mode)
}

/// Republishes source encryption metadata under a newer file ID without
/// changing its key lineage or key material.
fn regenerate_encryption_meta(
    source_meta: &pb::EncryptionMeta,
    new_file_id: u64,
) -> Result<pb::EncryptionMeta> {
    let mut meta = source_meta.clone();
    let current = meta.get_current();
    if new_file_id <= current.get_file_id() {
        return Err(Error::Other(box_err!(
            "target PD allocated encryption meta file_id {} which is not newer than source file_id {}",
            new_file_id,
            current.get_file_id()
        )));
    }

    // This republishes the same encryption epoch; it is not a key rotation.
    // Keep the source keyspace ID and all key material unchanged so retries can
    // safely publish another, newer copy without introducing a new data key.
    meta.mut_current().set_file_id(new_file_id);
    Ok(meta)
}

fn sync_target_keyspace_encryption_config(
    runtime: &Runtime,
    pd_client: Arc<dyn PdClient>,
    cluster: &BackupCluster,
) -> Result<()> {
    // Cross-keyspace restores make future target files use the restored source
    // mode. In-place restores preserve their historical behavior and only sync
    // the republished ENCRYPTION_META_KEY.
    let exported_encryption_meta = cluster.get_restore_encryption_meta();
    let is_in_place = cluster.source_cluster_id == pd_client.get_cluster_id()?
        && cluster.keyspace_id == cluster.target_keyspace_id;
    let sync_policy = TargetEncryptionConfigSync::for_restore(is_in_place);
    if !sync_policy.should_patch(exported_encryption_meta.as_deref()) {
        info!(
            "skip patching target keyspace encryption config because in-place restore has no exported ENCRYPTION_META_KEY";
            "source_keyspace_id" => cluster.keyspace_id,
            "target_keyspace_id" => cluster.target_keyspace_id,
        );
        return Ok(());
    }
    let has_exported_encryption_key = cluster.get_keyspace_exported_encryption_key().is_some();
    let mode = runtime.block_on(patch_target_keyspace_encryption_config(
        pd_client,
        cluster.keyspace_id,
        cluster.target_keyspace_id,
        exported_encryption_meta.as_deref(),
        has_exported_encryption_key,
    ))?;
    step!(
        "Patched target keyspace {} encryption config to {} mode",
        cluster.target_keyspace_id,
        mode
    );
    Ok(())
}

pub fn restore_keyspace(
    keyspace_id: u32,
    target_keyspace_id: u32,
    backup_name: &str,
    working_path: Option<PathBuf>,
    dfs: Arc<dyn Dfs>,
    config: RestoreConfig,
    pd_client: Arc<dyn PdClient>,
    pd_control: Option<PdControl>,
    runtime: &Runtime,
    truncate_ts: Option<u64>,
    master_key: Option<MasterKey>,
    reporter: Arc<dyn ReportRestoreStepTrait>,
    object_cache: Option<ObjectCache>,
    limiter: Option<Arc<ThroughputLimiter>>,
) -> Result<RestoredKeyspace> {
    let keyspace_tag = make_keyspace_tag(keyspace_id, target_keyspace_id);

    let inplace_restore = target_keyspace_id == keyspace_id;
    if !inplace_restore
        && (ApiV2::is_default_keyspace(keyspace_id)
            || ApiV2::is_default_keyspace(target_keyspace_id))
    {
        error!(
            "{} restore other keyspace from/to default keyspace",
            keyspace_tag
        );
        return Err(Error::RestoreWithDefaultKeyspace);
    }

    let (working_path, _guard) = match working_path {
        Some(path) => (path, None),
        None => {
            let temp_dir = box_try!(TempDirBuilder::new().prefix(WORKING_PATH_PREFIX).tempdir());
            let path = temp_dir.path().to_path_buf();
            (path, Some(temp_dir))
        }
    };

    let (keyspace_start, keyspace_end) = ApiV2::get_keyspace_range_by_id(keyspace_id);
    let (target_keyspace_start, target_keyspace_end) =
        ApiV2::get_keyspace_range_by_id(target_keyspace_id);

    reporter.report_step(RestoreStep::LoadBackupMeta);
    if config.restore_packed_backup {
        let packed = get_packed_backup_meta(dfs.as_ref(), backup_name.to_owned())?;
        if let Some(tts) = truncate_ts {
            check_packed_backup_meta_ts(&packed, tts)?;
        }
        let mut env = RestorePackEnv {
            dfs: dfs.clone(),
            pd_client: pd_client.clone(),
            pd_control,
            target_keyspace: target_keyspace_id,
            reporter: reporter.clone(),
            restore_config: config.clone(),
            data_dir: &working_path,
            truncate_ts,
            master_key,
        };
        let restored = env.execute(packed)?;
        return Ok(restored);
    }
    let (cluster_backup, archive_reader) = load_norm_backup_meta(dfs.clone(), backup_name)?;
    let truncate_ts = if archive_reader.is_none() {
        truncate_ts
    } else {
        Some(cluster_backup.backup_ts)
    };
    step!(
        "Start restore keyspace {} from backup <{}>, lightweight:{} range:[{},{}), target_range:[{},{}), backup ts:{}, safe ts:{} truncate ts:{:?}",
        keyspace_tag,
        backup_name,
        cluster_backup.is_lightweight,
        log_wrappers::hex_encode_upper(keyspace_start),
        log_wrappers::hex_encode_upper(keyspace_end),
        log_wrappers::hex_encode_upper(target_keyspace_start),
        log_wrappers::hex_encode_upper(target_keyspace_end),
        cluster_backup.backup_ts,
        cluster_backup.safe_ts,
        truncate_ts,
    );
    info!("{} restore config: {:?}", keyspace_tag, config);
    debug!(
        "Keyspace {} get cluster backup meta: {:?}",
        keyspace_tag, cluster_backup
    );

    reporter.report_step(RestoreStep::ExtractBackupShards);
    if let Some(truncate_ts) = truncate_ts {
        check_backup_meta_ts(&cluster_backup, truncate_ts)?;
    }
    let truncate_ts = truncate_ts.unwrap_or(cluster_backup.backup_ts);
    let opt = BackupClusterOptions {
        cluster_meta: cluster_backup.clone(),
        path: working_path,
        pd_client: pd_client.clone(),
        pd_control,
        dfs: dfs.clone(),
        restore_conf: config.clone(),
        keyspace_id,
        target_keyspace_id,
        truncate_ts,
        archiving: false,
        archive_reader,
        load_all_tables: false,
        override_dfs: None,
        offline_packing: false,
        master_key,
        object_cache,
    };
    let mut cluster = BackupCluster::new_opt(opt)?;
    step!(
        "Keyspace {} restore {} shards from backup",
        keyspace_tag,
        cluster.shards_count()
    );

    resolve_flush_and_restore_cluster(
        &mut cluster,
        config,
        runtime,
        reporter,
        limiter,
        None,
        None,
        None,
    )
}

fn resolve_flush_and_restore_cluster(
    cluster: &mut BackupCluster,
    config: RestoreConfig,
    runtime: &Runtime,
    reporter: Arc<dyn ReportRestoreStepTrait>,
    limiter: Option<Arc<ThroughputLimiter>>,
    retain_file_filter: Option<&dyn RetainFileFilter>,
    deadline: Option<RestoreDeadline>,
    scheduling_guard: Option<RestoreScheduleGuard>,
) -> Result<RestoredKeyspace> {
    let keyspace_tag = cluster.tag.clone();
    let truncate_ts = cluster.truncate_ts;

    // Resolve locks
    // In order to avoid data inconsistency, it is necessary to commit all locks of
    // secondary keys of a committed transaction that are in the backup.
    // If this is not done, uncommitted secondary keys may be dropped when the
    // committed record of primary key is compacted after restoration.
    // This is likely to happen because the timestamp of restored data is generally
    // smaller than the GC safepoint at that time.
    check_restore_deadline(deadline, "lock resolution")?;
    reporter.report_step(RestoreStep::ResolveLocks);
    let lock_resolver = LockResolver::new(
        &keyspace_tag,
        cluster.get_kvengine(),
        truncate_ts,
        RESOLVE_LOCKS_BATCH_SIZE,
        cluster.get_shard_meta_getter(),
    );
    let resolved_locks = runtime.block_on(lock_resolver.resolve_locks())?;
    let resolved_ts = resolved_locks.resolved_ts;
    cluster.add_shards_need_flush(&resolved_locks.resolved_shards);
    step!(
        "Keyspace {keyspace_tag} resolve {} locks / {} lock_txn_files of shards {:?}",
        resolved_locks.total_normal_locks_cnt,
        resolved_locks.total_lock_txn_files_cnt,
        resolved_locks.resolved_shards
    );

    // Trigger initial flush & flush mem-table to S3
    // NOTE: `cluster.shards` is NOT available before `flush_shards` finished.
    check_restore_deadline(deadline, "memtable flush")?;
    reporter.report_step(RestoreStep::FlushShards);
    let flush_timeout = match deadline {
        Some(deadline) => deadline.cap_timeout(config.timeout_wait_flush.0, "memtable flush")?,
        None => config.timeout_wait_flush.0,
    };
    let flush_cnt = cluster.flush_shards(flush_timeout)?;
    step!("Keyspace {keyspace_tag} flush {flush_cnt} shards");

    let mut res = prepare_and_restore_cluster_with_schedule_guard(
        cluster,
        config,
        runtime,
        reporter,
        limiter,
        retain_file_filter,
        deadline,
        scheduling_guard,
    )?;
    res.resolved_ts = resolved_ts.unwrap_or_else(|| TimeStamp::from(truncate_ts));
    Ok(res)
}

/// Restores already validated online source snapshots through the common
/// native-BR pipeline while retaining the pre-capture target scheduling guard.
pub(crate) fn restore_captured_online_snapshots(
    keyspace_id: u32,
    target_keyspace_id: u32,
    snapshots: Vec<OnlineSnapshotShard>,
    working_path: PathBuf,
    dfs: Arc<dyn Dfs>,
    config: RestoreConfig,
    pd_client: Arc<dyn PdClient>,
    pd_control: PdControl,
    runtime: &Runtime,
    branch_ts: u64,
    master_key: Option<MasterKey>,
    reporter: Arc<dyn ReportRestoreStepTrait>,
    limiter: Option<Arc<ThroughputLimiter>>,
    retain_file_filter: Option<&dyn RetainFileFilter>,
    deadline: RestoreDeadline,
    scheduling_guard: RestoreScheduleGuard,
) -> Result<RestoredKeyspace> {
    if keyspace_id == target_keyspace_id {
        return Err(box_err!(
            "online snapshot branch requires different source and target keyspaces"
        ));
    }
    if ApiV2::is_default_keyspace(keyspace_id) || ApiV2::is_default_keyspace(target_keyspace_id) {
        return Err(Error::RestoreWithDefaultKeyspace);
    }

    deadline.remaining("online snapshot materialization")?;
    reporter.report_step(RestoreStep::ExtractBackupShards);
    let mut cluster = BackupCluster::of_online_snapshots(
        keyspace_id,
        snapshots,
        working_path,
        pd_client,
        Some(pd_control),
        dfs,
        config.clone(),
        target_keyspace_id,
        branch_ts,
        master_key,
    )?;
    step!(
        "Keyspace {} restore {} shards from online snapshots at {}",
        cluster.tag(),
        cluster.shards_count(),
        branch_ts
    );
    resolve_flush_and_restore_cluster(
        &mut cluster,
        config,
        runtime,
        reporter,
        limiter,
        retain_file_filter,
        Some(deadline),
        Some(scheduling_guard),
    )
}

/// This should be call after all data are loaded into the KVEngine.
/// Rfengine won't be involved in this function.
pub(crate) fn prepare_and_restore_cluster(
    cluster: &mut BackupCluster,
    config: RestoreConfig,
    runtime: &Runtime,
    reporter: Arc<dyn ReportRestoreStepTrait>,
    limiter: Option<Arc<ThroughputLimiter>>,
    retain_file_filter: Option<&dyn RetainFileFilter>,
    deadline: Option<RestoreDeadline>,
) -> Result<RestoredKeyspace> {
    prepare_and_restore_cluster_with_schedule_guard(
        cluster,
        config,
        runtime,
        reporter,
        limiter,
        retain_file_filter,
        deadline,
        None,
    )
}

#[allow(clippy::too_many_arguments)]
fn prepare_and_restore_cluster_with_schedule_guard(
    cluster: &mut BackupCluster,
    config: RestoreConfig,
    runtime: &Runtime,
    reporter: Arc<dyn ReportRestoreStepTrait>,
    limiter: Option<Arc<ThroughputLimiter>>,
    retain_file_filter: Option<&dyn RetainFileFilter>,
    deadline: Option<RestoreDeadline>,
    scheduling_guard: Option<RestoreScheduleGuard>,
) -> Result<RestoredKeyspace> {
    let keyspace_tag = cluster.tag.clone();
    check_restore_deadline(deadline, "schema preparation")?;
    cluster.prepare_schema_at_truncate_ts(runtime)?;
    cluster.prepare_derived_metadata_retention();
    check_restore_deadline(deadline, "timestamp truncation")?;
    reporter.report_step(RestoreStep::TruncateTs);
    let truncate_ts_cnt = runtime.block_on(cluster.truncate_ts())?;
    step!(
        "Keyspace {} truncate {} shards to ts {}",
        keyspace_tag,
        truncate_ts_cnt,
        cluster.truncate_ts
    );

    cluster.prepare_restore_encryption_meta(runtime)?;

    let (target_keyspace_start, target_keyspace_end) =
        ApiV2::get_keyspace_range_by_id(cluster.target_keyspace_id);
    if let Some(guard) = scheduling_guard.as_ref() {
        guard.check_healthy()?;
    }
    check_restore_deadline(deadline, "target Region split")?;
    reporter.report_step(RestoreStep::SplitRegions);
    cluster.pre_split_regions(&config, deadline, runtime)?;

    let key_range = vec![(
        target_keyspace_start.as_slice(),
        target_keyspace_end.as_slice(),
    )];
    let scheduling_guard = match scheduling_guard {
        Some(guard) => {
            guard.check_healthy()?;
            guard
        }
        None => cluster.pause_merge_and_scheduler_during_restore(&key_range, runtime)?,
    };
    let scheduling_guard_state =
        (!cluster.static_region_metas.is_empty()).then(|| scheduling_guard.state());

    if !cluster.static_region_metas.is_empty() {
        check_restore_deadline(deadline, "online snapshot placement")?;
        let placement_timeout = match deadline {
            Some(deadline) => deadline.cap_timeout(
                config.timeout_restore_snapshot.0,
                "online snapshot placement",
            )?,
            None => config.timeout_restore_snapshot.0,
        };
        reporter.report_step(RestoreStep::AlignRegions);
        let target_regions = runtime.block_on(get_target_regions_with_retry(
            cluster.pd_client.as_ref(),
            &target_keyspace_start,
            &target_keyspace_end,
            placement_timeout,
        ))?;
        cluster.align_online_snapshot_placement(
            target_regions,
            placement_timeout,
            config.online_snapshot_max_unaligned_voter_peers,
            config.online_snapshot_placement_concurrency,
            scheduling_guard_state.as_ref().unwrap(),
            runtime,
        )?;
    }

    // Restore snapshot.
    let mut success_ranges = MergeRanges::default();
    let mut bo = Backoff::new(ExponentialBackoff::new(
        Duration::from_secs(1),
        Duration::from_secs(10),
        config.max_retry,
    ));
    let mut restore_bytes = 0;
    let mut rng = rand::thread_rng();
    let mut restored_target_files = Vec::new();
    let mut restored_target_file_ids = HashSet::default();
    let mut retry_region_fences = Vec::new();
    // NOTE: must call `bo::on_error` before every "continue".
    loop {
        check_restore_deadline(deadline, "RestoreShard retry")?;
        let mut target_regions = match runtime.block_on(get_target_regions(
            cluster.pd_client.as_ref(),
            &target_keyspace_start,
            &target_keyspace_end,
        )) {
            Ok(regions) => regions,
            Err(err) => {
                warn!("get target regions failed: {:?}, retry", err);
                bo.on_error(err)?;
                continue;
            }
        };
        if let Some(state) = &scheduling_guard_state {
            state.check_healthy()?;
            validate_restore_retry_topology(&retry_region_fences, &target_regions)?;
            retry_region_fences.clear();
        }

        let target_regions_total_cnt = target_regions.len();
        if !success_ranges.is_empty() {
            target_regions.retain(|r| !success_ranges.covered(&r.raw_start, &r.raw_end));
        }
        step!(
            "Keyspace {} target regions: {} ({} in total)",
            keyspace_tag,
            target_regions.len(),
            target_regions_total_cnt
        );

        // Align target regions.
        reporter.report_step(RestoreStep::AlignRegions);
        let aligned_regions = cluster.align_target_regions(target_regions);
        let unaligned_voters = cluster.count_online_snapshot_unaligned_voters(&aligned_regions)?;
        if unaligned_voters > config.online_snapshot_max_unaligned_voter_peers {
            return Err(box_err!(
                "online snapshot placement changed before RestoreShard: {} unaligned voter peers, limit is {}",
                unaligned_voters,
                config.online_snapshot_max_unaligned_voter_peers
            ));
        }
        let trimmed_shards_cnt =
            runtime.block_on(cluster.trim_over_bound_shards(&aligned_regions))?;
        step!(
            "Keyspace {} align {} backup shards to {} target regions and trim {} over bound shards",
            keyspace_tag,
            cluster.shards_count(),
            aligned_regions.len(),
            trimmed_shards_cnt
        );

        // Gather sstables for target regions.
        let (target_shards, sstables_cnt) = cluster.gather_all_tables(&aligned_regions)?;
        step!(
            "Keyspace {} gather {} SSTables for {} target regions",
            keyspace_tag,
            sstables_cnt,
            target_shards.len(),
        );
        let target_shards_for_retain = retain_file_filter.is_some().then(|| target_shards.clone());

        reporter.report_step(RestoreStep::RestoreSnapshotsToServers);
        let expected_conf_vers = aligned_regions
            .iter()
            .map(|aligned| {
                (
                    aligned.target_region.id,
                    aligned.target_region.epoch.get_conf_ver(),
                )
            })
            .collect::<HashMap<_, _>>();
        let mut snapshots = cluster
            .generate_snapshots(target_shards)
            .into_iter()
            .map(|change_set| RestoreSnapshotRequest {
                expected_conf_ver: scheduling_guard_state.as_ref().map(|_| {
                    *expected_conf_vers
                        .get(&change_set.get_shard_id())
                        .expect("target Region epoch must match generated RestoreShard")
                }),
                change_set,
            })
            .collect::<Vec<_>>();
        let snapshots_count = snapshots.len();
        // Because we scatter regions then split it to finer region, continous regions
        // are probably placed to the same store.
        // So shuffle the sorted snapshots hence our requests can be scattered to
        // stores.
        snapshots.shuffle(&mut rng);
        let ret = restore_snapshots(
            &keyspace_tag,
            runtime,
            cluster.pd_client.clone(),
            &reporter,
            snapshots,
            &mut success_ranges,
            &config,
            &mut bo,
            &limiter,
            scheduling_guard_state.clone(),
            deadline,
        )?;
        if let Some(state) = &scheduling_guard_state {
            // A lease can expire while the final batch of HTTP requests is in
            // flight. Validate it after every request has completed before
            // accepting that batch, including the last successful one.
            state.check_healthy()?;
        }
        retry_region_fences = ret.retry_regions.clone();
        reporter.report_step(RestoreStep::RestoreSnapshotsFinished);
        restore_bytes += ret.restore_bytes;
        if let Some(target_shards) = target_shards_for_retain {
            for target_shard in target_shards {
                let start = target_shard.range.outer_start.to_vec();
                let end = target_shard.range.outer_end.to_vec();
                if success_ranges.covered(&start, &end) {
                    for file in BackupCluster::get_shard_files(target_shard.id, &target_shard) {
                        if restored_target_file_ids.insert((file.ftype, file.id)) {
                            restored_target_files.push(file);
                        }
                    }
                }
            }
        }
        step!(
            "Keyspace {keyspace_tag} restore {snapshots_count} regions, restored {} regions, {} bytes, retry {} regions",
            ret.count,
            ret.restore_bytes,
            ret.retry_regions.len(),
        );

        if success_ranges.covered(&target_keyspace_start, &target_keyspace_end) {
            break;
        }
        bo.on_error(Error::RestoreSnapshot)?;
        step!(
            "Keyspace {keyspace_tag} restore retry {}",
            bo.current_attempts()
        );
    }

    // The target-keyspace scheduling guard is needed only until the final
    // RestoreShard has completed. PD removes the TTL-backed rule naturally
    // after renewals stop.
    drop(scheduling_guard);

    check_restore_deadline(deadline, "file retention")?;
    reporter.report_step(RestoreStep::RetainSstFiles);
    let mut all_files = if retain_file_filter.is_some() {
        restored_target_files
    } else {
        cluster.get_all_shard_files(None, None)
    };
    append_restored_schema_file_to_retain(
        cluster
            .restored_schema
            .as_ref()
            .expect("schema must be prepared before retaining restored files"),
        &mut all_files,
    );
    let mut files = match retain_file_filter {
        Some(filter) => filter.filter_files(all_files)?,
        None => all_files,
    };
    let skipped_recent_files = skip_recent_files_from_retain(&mut files, TimeStamp::physical_now());
    if skipped_recent_files > 0 {
        info!(
            "skip retaining recently created files";
            "keyspace" => &keyspace_tag,
            "skipped_files" => skipped_recent_files,
        );
    }
    check_restore_deadline(deadline, "file retention")?;
    fail::fail_point!("restore_keyspace_before_retain_sst_files");
    match retain_sst_files(files, cluster.dfs.clone()) {
        Err(e) => {
            return Err(box_err!(
                "Keyspace {} fail to retain restored sst files in s3, {:?}",
                keyspace_tag,
                e
            ));
        }
        Ok(file_cnt) => {
            step!("Keyspace {keyspace_tag} retain {file_cnt} files successfully");
        }
    }
    check_restore_deadline(deadline, "file retention completion")?;

    sync_target_keyspace_encryption_config(runtime, cluster.pd_client.clone(), cluster)?;

    let restore_ret = RestoredKeyspace {
        keyspace_id: cluster.keyspace_id,
        target_keyspace_id: cluster.target_keyspace_id,
        ts: cluster.truncate_ts,
        restore_bytes,
        tolerated_err: cluster.tolerated_err(),
        resolved_ts: TimeStamp::zero(),
    };
    Ok(restore_ret)
}

fn skip_recent_files_from_retain(files: &mut Vec<TableFile>, now_physical_ms: u64) -> usize {
    let original_file_count = files.len();
    files.retain(|file| {
        // Physical file IDs come from TSO before the file is created. An ID
        // younger than the DFS GC grace period therefore proves the file is
        // not yet eligible for logical deletion.
        let created_at_ms = TimeStamp::new(file.id).physical();
        now_physical_ms.saturating_sub(created_at_ms)
            >= RECENT_FILE_RETAIN_SKIP_INTERVAL.as_millis() as u64
    });
    original_file_count - files.len()
}

fn append_restored_schema_file_to_retain(
    restored_schema: &RestoredSchema,
    retained_files: &mut Vec<TableFile>,
) {
    if let Some(schema_file) = restored_schema.target_file.as_ref() {
        retained_files.push(TableFile {
            id: schema_file.get_file_id(),
            ftype: FileType::Schema,
            shard_id: 0,
            size: 0,
        });
    }
}

fn check_backup_meta_ts(meta: &ClusterBackupMeta, truncate_ts: u64) -> Result<()> {
    if truncate_ts < meta.safe_ts || truncate_ts > meta.backup_ts {
        return Err(Error::PitrTsError(
            truncate_ts,
            meta.safe_ts,
            meta.backup_ts,
        ));
    }
    Ok(())
}

fn check_packed_backup_meta_ts(packed: &PackedBackup, truncate_ts: u64) -> Result<()> {
    if truncate_ts < packed.safe_ts || truncate_ts > packed.backup_ts {
        return Err(Error::PitrTsError(
            truncate_ts,
            packed.safe_ts,
            packed.backup_ts,
        ));
    }
    Ok(())
}

#[derive(Default, Clone)]
pub struct BackupShard {
    pub region_id: u64,
    pub store_id: u64,
    pub peer_id: u64,
    pub need_flush: bool,
    pub meta: ShardMeta,
    pub(crate) raw_meta: Option<pb::ChangeSet>, // used for kv engine recovery only.
    raft_state: RaftState,
}

/// A source leader snapshot captured by the online branch protocol.
#[derive(Clone)]
pub(crate) struct OnlineSnapshotShard {
    pub(crate) leader: metapb::Peer,
    pub(crate) region: metapb::Region,
    pub(crate) change_set: pb::ChangeSet,
    pub(crate) mem_table_data: Vec<u8>,
}

impl fmt::Debug for BackupShard {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("BackupShard")
            .field("tag", &format!("{}", self.tag()))
            .field("region_id", &self.region_id)
            .field("store_id", &self.store_id)
            .field("peer_id", &self.peer_id)
            .field("need_flush", &self.need_flush)
            .field("meta", &self.meta.to_change_set())
            .field("raw_meta", &self.raw_meta)
            .field("raft_state", &self.raft_state)
            .finish()
    }
}

#[cfg(test)]
impl PartialEq for BackupShard {
    fn eq(&self, other: &Self) -> bool {
        self.region_id == other.region_id
            && self.store_id == other.store_id
            && self.peer_id == other.peer_id
            && self.need_flush == other.need_flush
            && self.meta.to_change_set() == other.meta.to_change_set()
            && self.raw_meta == other.raw_meta
            && self.raft_state == other.raft_state
    }
}

impl BackupShard {
    pub fn tag(&self) -> ShardTag {
        ShardTag::new(self.store_id, IdVer::new(self.region_id, self.ver()))
    }

    pub fn ver(&self) -> u64 {
        self.meta.ver
    }

    pub fn id_ver(&self) -> IdVer {
        IdVer::new(self.region_id, self.ver())
    }

    pub fn start(&self) -> &[u8] {
        &self.meta.range.outer_start
    }

    pub fn end(&self) -> &[u8] {
        &self.meta.range.outer_end
    }

    pub fn inner_start(&self) -> InnerKey<'_> {
        self.meta.range.inner_start()
    }

    pub fn inner_end(&self) -> InnerKey<'_> {
        self.meta.range.inner_end()
    }

    pub fn inner_key_off(&self) -> usize {
        self.meta.inner_key_off
    }

    pub fn table_version(&self) -> u64 {
        self.meta.base_version + self.meta.data_sequence
    }

    pub fn reset_columnar(&mut self) {
        self.meta.clear_columnar_related_meta();
    }
}

impl BoundedDataSet for BackupShard {
    fn data_bound(&self) -> DataBound<'_> {
        self.meta.range.data_bound()
    }
}

#[derive(Clone)]
struct BackupSchemaShard {
    id: u64,
    version: u64,
    start_key: Bytes,
    end_key: Bytes,
}

#[derive(Clone)]
struct BackupSchemaReader {
    engine: kvengine::Engine,
    shards: Vec<BackupSchemaShard>,
}

impl BackupSchemaReader {
    fn new(engine: kvengine::Engine, shards: impl Iterator<Item = BackupSchemaShard>) -> Self {
        Self {
            engine,
            shards: shards.collect(),
        }
    }

    fn first_shard_at_or_after(&self, key: &[u8]) -> usize {
        self.shards
            .partition_point(|shard| shard.end_key.as_ref() <= key)
    }

    fn shard_for_key(&self, key: &[u8]) -> Option<&BackupSchemaShard> {
        self.shards
            .get(self.first_shard_at_or_after(key))
            .filter(|shard| shard.start_key.as_ref() <= key && key < shard.end_key.as_ref())
    }

    fn snapshot(
        &self,
        shard: &BackupSchemaShard,
    ) -> std::result::Result<kvengine::SnapAccess, String> {
        let load = || {
            self.engine
                .try_get_snap_access(shard.id, shard.version)
                .map_err(|err| format!("load backup shard {} snapshot failed: {err}", shard.id))?
                .ok_or_else(|| {
                    format!(
                        "backup shard {} version {} not found",
                        shard.id, shard.version
                    )
                })
        };
        let mut snapshot = load()?;
        if snapshot.has_unloaded_tables() {
            self.engine
                .load_unloaded_tables(shard.id, shard.version, false, IoType::Import)
                .map_err(|err| {
                    format!(
                        "load backup shard {} version {} tables failed: {err}",
                        shard.id, shard.version
                    )
                })?;
            snapshot = load()?;
        }
        Ok(snapshot)
    }
}

#[async_trait]
impl schema::KvGetter for BackupSchemaReader {
    async fn get(
        &self,
        key: &[u8],
        start_ts: TimeStamp,
    ) -> std::result::Result<Option<Vec<u8>>, String> {
        let Some(shard) = self.shard_for_key(key) else {
            return Ok(None);
        };
        let snapshot = self.snapshot(shard)?;
        let item = snapshot
            .get_async(WRITE_CF, key, start_ts.into_inner())
            .await;
        let value = item.get_value();
        // Committed MVCC deletes are stored as empty user values rather than table
        // tombstones. Hide them to match transactional snapshot semantics.
        if item.version == 0 || item.is_deleted() || value.is_empty() {
            Ok(None)
        } else {
            Ok(Some(value.to_vec()))
        }
    }

    async fn batch_get(
        &self,
        keys: &[Vec<u8>],
        start_ts: TimeStamp,
    ) -> std::result::Result<Vec<Option<Vec<u8>>>, String> {
        let mut values = Vec::with_capacity(keys.len());
        for key in keys {
            values.push(self.get(key, start_ts).await?);
        }
        Ok(values)
    }
}

#[async_trait]
impl schema::KvScanner for BackupSchemaReader {
    async fn scan(
        &self,
        start: &[u8],
        end: &[u8],
        start_ts: TimeStamp,
    ) -> std::result::Result<Vec<(Vec<u8>, Vec<u8>)>, String> {
        if start >= end {
            return Ok(Vec::new());
        }

        let mut pairs = Vec::with_capacity(schema::DEFAULT_SCAN_BATCH_SIZE as usize);
        for shard in &self.shards[self.first_shard_at_or_after(start)..] {
            if shard.start_key.as_ref() >= end {
                break;
            }
            let snapshot = self.snapshot(shard)?;
            let mut iterator = snapshot
                .new_iterator_async(WRITE_CF, false, false, Some(start_ts.into_inner()), true)
                .await;
            let scan_start = cmp::max(start, shard.start_key.as_ref());
            iterator.seek_async(scan_start).await;
            while iterator.valid() && iterator.key() < end {
                let value = iterator.val().to_vec();
                // See `get`: an empty value represents a committed MVCC delete.
                if !value.is_empty() {
                    pairs.push((iterator.key().to_vec(), value));
                    if pairs.len() == schema::DEFAULT_SCAN_BATCH_SIZE as usize {
                        return Ok(pairs);
                    }
                }
                iterator.next_async().await;
            }
        }
        Ok(pairs)
    }
}

#[derive(Clone, Debug, PartialEq)]
struct AlignedRegion {
    target_region: RawRegion,
    backup_shards_id: Vec<u64>,
}

struct OnlinePlacementTarget {
    region_id: u64,
    desired_peers: Vec<(u64, TransferRegionPeerRole)>,
    desired_stores: HashSet<u64>,
}

impl OnlinePlacementTarget {
    fn is_converged(&self, region: &metapb::Region) -> bool {
        region
            .get_peers()
            .iter()
            .filter(|peer| peer.get_role() == PeerRole::Voter)
            .map(|peer| peer.get_store_id())
            .collect::<HashSet<_>>()
            == self.desired_stores
    }
}

/// Validates that locality tolerance is applied only to a complete target
/// voter set. A missing or differently replicated target is a topology error,
/// not an unaligned peer that can be repaired by downloading files from DFS.
fn validate_online_snapshot_voter_stores(
    target_region_id: u64,
    actual_stores: Option<HashSet<u64>>,
    desired_stores: &HashSet<u64>,
) -> Result<(HashSet<u64>, usize)> {
    let Some(actual_stores) = actual_stores else {
        return Err(box_err!(
            "online snapshot target Region {} is missing or has no leader after placement",
            target_region_id
        ));
    };
    if actual_stores.len() != desired_stores.len() {
        return Err(box_err!(
            "online snapshot target Region {} voter count does not match source placement: actual {:?}, desired {:?}",
            target_region_id,
            actual_stores,
            desired_stores
        ));
    }
    let unaligned = actual_stores.difference(desired_stores).count();
    Ok((actual_stores, unaligned))
}

#[derive(Clone)]
struct OnlinePlacementSource {
    region: metapb::Region,
    leader: metapb::Peer,
    inner_start: Vec<u8>,
    inner_end: Vec<u8>,
}

/// Lightweight source topology needed to split and place the target before
/// the source snapshot data has been captured.
pub(crate) struct OnlineSnapshotPlacementPlan {
    sources: Vec<OnlinePlacementSource>,
}

impl OnlineSnapshotPlacementPlan {
    /// Freezes the source Region ranges and voter placement observed while the
    /// source scheduling guard is healthy.
    pub(crate) fn new(source_regions: &[(metapb::Region, metapb::Peer)]) -> Result<Self> {
        if source_regions.is_empty() {
            return Err(box_err!(
                "online snapshot placement requires at least one source Region"
            ));
        }

        let mut sources = Vec::with_capacity(source_regions.len());
        let mut previous_end = None;
        for (region, leader) in source_regions {
            let raw_region = RawRegion::from(region.clone());
            if raw_region.raw_start >= raw_region.raw_end {
                return Err(box_err!(
                    "source Region {} has an invalid range",
                    region.get_id()
                ));
            }
            if previous_end
                .as_ref()
                .is_some_and(|end| end != &raw_region.raw_start)
            {
                return Err(box_err!(
                    "source Regions are not a contiguous ordered range at Region {}",
                    region.get_id()
                ));
            }
            previous_end = Some(raw_region.raw_end.clone());
            sources.push(OnlinePlacementSource {
                region: region.clone(),
                leader: leader.clone(),
                inner_start: InnerKey::from_outer_key(&raw_region.raw_start).to_vec(),
                inner_end: InnerKey::from_outer_end_key(&raw_region.raw_end).to_vec(),
            });
        }
        Ok(Self { sources })
    }

    fn align_target_regions(&self, target_regions: Vec<RawRegion>) -> Vec<(RawRegion, Vec<usize>)> {
        let mut aligned_regions = Vec::with_capacity(target_regions.len());
        let mut source_idx = 0;
        for target_region in target_regions {
            let target_inner_start =
                InnerKey::from_outer_key(target_region.get_start_key()).to_vec();
            let target_inner_end =
                InnerKey::from_outer_end_key(target_region.get_end_key()).to_vec();
            if source_idx > 0 && self.sources[source_idx - 1].inner_end > target_inner_start {
                source_idx -= 1;
            }

            let mut aligned_source_indexes = Vec::new();
            while source_idx < self.sources.len()
                && self.sources[source_idx].inner_start < target_inner_end
            {
                if self.sources[source_idx].inner_end > target_inner_start {
                    aligned_source_indexes.push(source_idx);
                }
                source_idx += 1;
            }
            aligned_regions.push((target_region, aligned_source_indexes));
        }
        aligned_regions
    }

    fn split_keys(
        &self,
        target_regions: &[RawRegion],
        target_keyspace_start: &[u8],
        coarse_split_regions_factor: usize,
    ) -> (Vec<Vec<u8>>, Vec<Vec<u8>>) {
        let target_boundaries = target_regions
            .iter()
            .map(|region| InnerKey::from_outer_key(region.get_start_key()).to_vec())
            .collect::<HashSet<_>>();
        let mut coarse_split_keys =
            Vec::with_capacity(self.sources.len() / coarse_split_regions_factor);
        let mut split_keys = Vec::with_capacity(self.sources.len().saturating_sub(1));
        for (idx, source) in self.sources.iter().enumerate().skip(1) {
            if target_boundaries.contains(&source.inner_start) {
                continue;
            }
            let split_key = [target_keyspace_start, source.inner_start.as_slice()].concat();
            if idx % coarse_split_regions_factor == 0 {
                coarse_split_keys.push(split_key);
            } else {
                split_keys.push(split_key);
            }
        }
        (coarse_split_keys, split_keys)
    }

    fn placement_targets(
        &self,
        target_regions: Vec<RawRegion>,
    ) -> Result<HashMap<u64, OnlinePlacementTarget>> {
        let mut targets = HashMap::default();
        for (target_region, source_indexes) in self.align_target_regions(target_regions) {
            let [source_idx] = source_indexes.as_slice() else {
                return Err(box_err!(
                    "target Region {} is not aligned to exactly one online source Region: {:?}",
                    target_region.id,
                    source_indexes
                ));
            };
            let source = &self.sources[*source_idx];
            let desired_leader_store = source.leader.get_store_id();
            let mut desired_stores = HashSet::default();
            let mut desired_peers = Vec::new();
            for peer in source.region.get_peers() {
                match peer.get_role() {
                    PeerRole::Learner => continue,
                    PeerRole::Voter => {}
                    role => {
                        return Err(box_err!(
                            "source Region {} has an unstable voter role {:?}",
                            source.region.get_id(),
                            role
                        ));
                    }
                }
                let store_id = peer.get_store_id();
                if !desired_stores.insert(store_id) {
                    return Err(box_err!(
                        "source Region {} contains duplicate voter Store {}",
                        source.region.get_id(),
                        store_id
                    ));
                }
                let role = if store_id == desired_leader_store {
                    TransferRegionPeerRole::Leader
                } else {
                    TransferRegionPeerRole::Follower
                };
                desired_peers.push((store_id, role));
            }
            if desired_peers.is_empty() || !desired_stores.contains(&desired_leader_store) {
                return Err(box_err!(
                    "source Region {} has no valid voter placement",
                    source.region.get_id()
                ));
            }
            targets.insert(
                target_region.id,
                OnlinePlacementTarget {
                    region_id: target_region.id,
                    desired_peers,
                    desired_stores,
                },
            );
        }
        Ok(targets)
    }

    fn align_placement(
        &self,
        target_regions: Vec<RawRegion>,
        pd_client: Arc<dyn PdClient>,
        pd_control: PdControl,
        timeout: Duration,
        max_unaligned_voter_peers: usize,
        placement_concurrency: usize,
        scheduling_guards: &[RestoreScheduleGuardState],
        runtime: &Runtime,
        tag: &str,
    ) -> Result<usize> {
        if placement_concurrency == 0 {
            return Err(box_err!(
                "online snapshot placement concurrency must be greater than 0"
            ));
        }
        let targets = self.placement_targets(target_regions)?;
        let unaligned = align_online_snapshot_placement_targets(
            targets,
            pd_client,
            pd_control,
            timeout,
            placement_concurrency,
            scheduling_guards,
            runtime,
        )?;
        if unaligned > max_unaligned_voter_peers {
            return Err(box_err!(
                "online snapshot placement has {} unaligned voter peers, limit is {}",
                unaligned,
                max_unaligned_voter_peers
            ));
        }
        info!(
            "online snapshot placement phase finished";
            "keyspace" => tag,
            "unaligned_voter_peers" => unaligned,
            "limit" => max_unaligned_voter_peers,
        );
        Ok(unaligned)
    }

    /// Splits the empty target at missing source boundaries and aligns its
    /// voters. Every supplied scheduling guard is checked throughout the wait.
    #[allow(clippy::too_many_arguments)]
    pub(crate) fn prepare_target(
        &self,
        target_keyspace_id: u32,
        pd_client: Arc<dyn PdClient>,
        pd_control: PdControl,
        config: &RestoreConfig,
        scheduling_guards: &[RestoreScheduleGuardState],
        runtime: &Runtime,
        deadline: RestoreDeadline,
    ) -> Result<usize> {
        if config.coarse_split_regions_factor == 0 {
            return Err(box_err!(
                "coarse split Regions factor must be greater than 0"
            ));
        }
        check_schedule_guards(scheduling_guards)?;
        let (target_start, target_end) = ApiV2::get_keyspace_range_by_id(target_keyspace_id);
        let scan_timeout = deadline.cap_timeout(
            config.timeout_restore_snapshot.0,
            "online snapshot target Region scan",
        )?;
        let target_regions = runtime.block_on(get_target_regions_with_retry(
            pd_client.as_ref(),
            &target_start,
            &target_end,
            scan_timeout,
        ))?;
        let (coarse_split_keys, split_keys) = self.split_keys(
            &target_regions,
            &target_start,
            config.coarse_split_regions_factor,
        );
        let split_timeout = deadline.cap_timeout(
            config.timeout_split_regions.0,
            "online snapshot target Region split",
        )?;
        // Exact source-voter placement follows this split pass, so preserve
        // the inherited placement instead of creating scatter operators.
        let scatter_coarse_regions = false;
        split_regions_for_keys(
            &format!("target-keyspace-{target_keyspace_id}"),
            pd_client.as_ref(),
            coarse_split_keys,
            split_keys,
            scatter_coarse_regions,
            split_timeout,
            Some(deadline),
            runtime,
        )?;

        check_schedule_guards(scheduling_guards)?;
        let placement_timeout = deadline.cap_timeout(
            config.timeout_restore_snapshot.0,
            "online snapshot target placement",
        )?;
        let target_regions = runtime.block_on(get_target_regions_with_retry(
            pd_client.as_ref(),
            &target_start,
            &target_end,
            placement_timeout,
        ))?;
        self.align_placement(
            target_regions,
            pd_client,
            pd_control,
            placement_timeout,
            config.online_snapshot_max_unaligned_voter_peers,
            config.online_snapshot_placement_concurrency,
            scheduling_guards,
            runtime,
            &format!("target-keyspace-{target_keyspace_id}"),
        )
    }
}

pub(crate) fn check_schedule_guards(scheduling_guards: &[RestoreScheduleGuardState]) -> Result<()> {
    for guard in scheduling_guards {
        guard.check_healthy()?;
    }
    Ok(())
}

#[allow(clippy::too_many_arguments)]
fn align_online_snapshot_placement_targets(
    targets: HashMap<u64, OnlinePlacementTarget>,
    pd_client: Arc<dyn PdClient>,
    pd_control: PdControl,
    timeout: Duration,
    placement_concurrency: usize,
    scheduling_guards: &[RestoreScheduleGuardState],
    runtime: &Runtime,
) -> Result<usize> {
    runtime.block_on(async move {
        let begin = Instant::now_coarse();
        let mut pending = targets;
        let mut last_submit = HashMap::<u64, StdInstant>::default();
        while !pending.is_empty() {
            let remaining = timeout.saturating_sub(begin.saturating_elapsed());
            if remaining.is_zero() {
                break;
            }
            check_schedule_guards(scheduling_guards)?;
            let active_operators =
                match tokio::time::timeout(remaining, pd_control.get_operators()).await {
                    Ok(operators) => operators?,
                    Err(_) => break,
                }
                .into_iter()
                .map(|operator| operator.region_id)
                .collect::<HashSet<_>>();
            let remaining = timeout.saturating_sub(begin.saturating_elapsed());
            if remaining.is_zero() {
                break;
            }

            let round = stream::iter(pending.values())
                .map(|target| {
                    let pd_client = pd_client.clone();
                    let pd_control = pd_control.clone();
                    let operator_active = active_operators.contains(&target.region_id);
                    let should_submit = !operator_active
                        && last_submit
                            .get(&target.region_id)
                            .is_none_or(|last| last.elapsed() >= Duration::from_secs(2));
                    async move {
                        check_schedule_guards(scheduling_guards)?;
                        let current = pd_client.get_region_leader_by_id(target.region_id).await?;
                        check_schedule_guards(scheduling_guards)?;
                        if let Some((region, _)) = current {
                            // A move operator retains OpRegion while only its
                            // best-effort leader step may remain, so observed
                            // voters take precedence over the operator snapshot.
                            if target.is_converged(&region) {
                                return Ok::<_, Error>((target.region_id, true, false));
                            }
                        }
                        if should_submit {
                            check_schedule_guards(scheduling_guards)?;
                            if let Err(err) = pd_control
                                .transfer_region(target.region_id, &target.desired_peers)
                                .await
                            {
                                warn!(
                                    "failed to submit online snapshot placement operator";
                                    "region_id" => target.region_id,
                                    "err" => ?err,
                                );
                            }
                            return Ok((target.region_id, false, true));
                        }
                        Ok((target.region_id, false, false))
                    }
                })
                .buffer_unordered(placement_concurrency)
                .collect::<Vec<_>>();
            let round_results = match tokio::time::timeout(remaining, round).await {
                Ok(results) => results,
                Err(_) => break,
            };

            for result in round_results {
                let (region_id, converged, submitted) = result?;
                if converged {
                    pending.remove(&region_id);
                }
                if submitted {
                    last_submit.insert(region_id, StdInstant::now());
                }
            }
            if !pending.is_empty() {
                let remaining = timeout.saturating_sub(begin.saturating_elapsed());
                if !remaining.is_zero() {
                    tokio::time::sleep(cmp::min(Duration::from_millis(500), remaining)).await;
                }
            }
        }

        check_schedule_guards(scheduling_guards)?;
        if !pending.is_empty() {
            let active_operators = pd_control
                .get_operators()
                .await?
                .into_iter()
                .map(|operator| operator.region_id)
                .collect::<HashSet<_>>();
            check_schedule_guards(scheduling_guards)?;
            if let Some(region_id) = pending
                .keys()
                .find(|region_id| active_operators.contains(region_id))
            {
                return Err(box_err!(
                    "online snapshot placement operator for Region {} did not finish before the deadline",
                    region_id
                ));
            }
        }

        let placement_states = stream::iter(pending.values())
            .map(|target| {
                let pd_client = pd_client.clone();
                async move {
                    check_schedule_guards(scheduling_guards)?;
                    let current = pd_client.get_region_leader_by_id(target.region_id).await?;
                    check_schedule_guards(scheduling_guards)?;
                    let actual_stores = current.map(|(region, _)| {
                        region
                            .get_peers()
                            .iter()
                            .filter(|peer| peer.get_role() == PeerRole::Voter)
                            .map(|peer| peer.get_store_id())
                            .collect::<HashSet<_>>()
                    });
                    let (actual_stores, unaligned) = validate_online_snapshot_voter_stores(
                        target.region_id,
                        actual_stores,
                        &target.desired_stores,
                    )?;
                    Ok::<_, Error>((target.region_id, actual_stores, unaligned))
                }
            })
            .buffer_unordered(placement_concurrency)
            .collect::<Vec<_>>()
            .await;

        let mut unaligned_voter_peers = 0;
        for state in placement_states {
            let (region_id, actual_stores, unaligned) = state?;
            let target = pending
                .get(&region_id)
                .expect("placement state must belong to a pending Region");
            unaligned_voter_peers += unaligned;
            warn!(
                "online snapshot target placement did not converge";
                "region_id" => region_id,
                "desired_stores" => ?target.desired_stores,
                "actual_stores" => ?actual_stores,
                "unaligned_voter_peers" => unaligned,
            );
        }
        Ok(unaligned_voter_peers)
    })
}

fn split_regions_for_keys(
    tag: &str,
    pd_client: &dyn PdClient,
    coarse_split_keys: Vec<Vec<u8>>,
    split_keys: Vec<Vec<u8>>,
    scatter_coarse_regions: bool,
    timeout: Duration,
    deadline: Option<RestoreDeadline>,
    runtime: &Runtime,
) -> Result<()> {
    if !coarse_split_keys.is_empty() {
        debug!("{} coarse split regions", tag; "keys" => ?coarse_split_keys);
        split_regions_with_retry(
            tag,
            pd_client,
            coarse_split_keys,
            scatter_coarse_regions,
            timeout,
            deadline,
            runtime,
        )?;
    }
    if !split_keys.is_empty() {
        debug!("{} split regions", tag; "keys" => ?split_keys);
        split_regions_with_retry(
            tag, pd_client, split_keys, false, timeout, deadline, runtime,
        )?;
    }
    Ok(())
}

fn split_regions_with_retry(
    tag: &str,
    pd_client: &dyn PdClient,
    raw_keys: Vec<Vec<u8>>,
    scatter: bool,
    timeout: Duration,
    deadline: Option<RestoreDeadline>,
    runtime: &Runtime,
) -> Result<()> {
    for batch in raw_keys.chunks(SPLIT_REGIONS_BATCH_SIZE) {
        let encoded_keys = batch
            .iter()
            .map(|key| Key::from_raw(key).into_encoded())
            .collect::<Vec<_>>();
        let batch_timeout = cmp::max(
            SPLIT_REGIONS_TIMEOUT_PER_KEY * encoded_keys.len() as u32,
            timeout,
        );
        let batch_timeout = match deadline {
            Some(deadline) => {
                deadline.cap_timeout(batch_timeout, "online snapshot target Region split")?
            }
            None => batch_timeout,
        };
        step!(
            "Keyspace {} split regions for {} keys",
            tag,
            encoded_keys.len()
        );
        let regions = match deadline {
            Some(_) => runtime.block_on(async {
                match tokio::time::timeout(
                    batch_timeout,
                    pd_client.split_regions_with_retry(encoded_keys, batch_timeout),
                )
                .await
                {
                    Ok(result) => result.map_err(Error::SpitRegionsError),
                    Err(_) => Err(box_err!(
                        "online snapshot target Region split batch timed out after {:?}",
                        batch_timeout
                    )),
                }
            })?,
            None => runtime
                .block_on(pd_client.split_regions_with_retry(encoded_keys, batch_timeout))
                .map_err(Error::SpitRegionsError)?,
        };
        check_restore_deadline(deadline, "target Region split completion")?;
        if scatter {
            step!("Keyspace {} scatter {} regions", tag, regions.len());
            if let Err(err) = pd_client.scatter_regions_by_id(regions) {
                warn!("{} scatter regions failed: {:?}", tag, err);
            }
        }
    }
    Ok(())
}

struct FullShardCollection {
    shards: Vec<BackupShard>,
    tombstone_peer_count: usize,
}

pub struct BackupClusterOptions {
    pub cluster_meta: ClusterBackupMeta,
    pub path: PathBuf,
    pub pd_client: Arc<dyn PdClient>,
    pub pd_control: Option<PdControl>,
    pub dfs: Arc<dyn Dfs>,
    pub restore_conf: RestoreConfig,
    pub object_cache: Option<ObjectCache>,
    pub keyspace_id: u32,
    pub target_keyspace_id: u32,
    pub truncate_ts: u64,
    pub archiving: bool,
    pub archive_reader: Option<ArchiveReader>,
    pub load_all_tables: bool, // `true` for "check_table" ONLY.
    pub override_dfs: Option<Arc<dyn Dfs>>,
    pub offline_packing: bool,
    /// A runtime-resolved master key to reuse for restore setup.
    ///
    /// When this is `Some`, `BackupCluster` must use it directly instead of
    /// resolving a new key from `restore_conf.security`.
    pub master_key: Option<MasterKey>,
}

pub struct BackupCluster {
    tag: String,
    path: PathBuf,
    pd_client: Arc<dyn PdClient>,
    pd_control: Option<PdControl>,
    dfs: Arc<dyn Dfs>,
    security_conf: SecurityConfig,
    master_key: MasterKey,
    // Physical cluster recorded in the source backup.
    source_cluster_id: u64,
    keyspace_id: u32,
    target_keyspace_id: u32,
    keyspace_start: Vec<u8>,
    keyspace_end: Vec<u8>,
    truncate_ts: u64,
    archive_reader: Option<ArchiveReader>,

    // store_id -> rf_engine.
    raft_engines: HashMap<u64, RfEngine>,
    // store_id -> store_config.
    store_configs: HashMap<u64, TikvConfig>,
    // must be `Some` after `new()`.
    kv_engine: Option<kvengine::Engine>,

    // region_id -> shard.
    pub(crate) shards: HashMap<u64, BackupShard>,
    // region_id -> raw_meta, used for kv_engine recover.
    raw_metas: HashMap<u64, pb::ChangeSet>,
    // store_id -> Vec<shard_id>.
    store_shards: HashMap<u64, Vec<u64>>,
    // shard_id -> StorePeer{store_id, peer_id}. Used to resolve locks.
    shards_store_map: Option<Arc<HashMap<u64 /* shard_id */, StorePeer>>>,
    // Static source Region metadata used by online snapshot lock resolution.
    static_region_metas: HashMap<u64, (metapb::Region, metapb::Peer)>,
    // shard_id sorted by start_key.
    sorted_shards: Vec<u64>,
    // Set<shard_id>.
    shards_need_flush: HashSet<u64>,
    // Set<shard_id>.
    shards_need_truncate: HashSet<u64>,
    // Set<shard_id>.
    shards_need_reset_columnar: HashSet<u64>,

    // Schema at truncate_ts, prepared once and shared by all target Regions.
    restored_schema: Option<RestoredSchema>,
    // Derived metadata retention decided from source file IDs before truncate_ts
    // rewrites create new files.
    derived_metadata_retention: Option<HashMap<u64, DerivedMetadataRetention>>,
    // Source encryption meta republished with a target-PD file ID when the
    // restore is not same-cluster in-place.
    restore_encryption_meta: Option<Vec<u8>>,

    meta_applier: Option<Arc<MetaApplier>>,

    meta_sender: Option<mpsc::Sender<StoreMsg>>,

    tolerated_err: usize,

    // Wrap by `Option` to detect use before filled.
    txn_chunk_ids_in_wal: Option<Vec<u64>>,

    // `false` for restoration.
    // `true` for "check_table" to read data from backup directly.
    load_all_tables: bool,

    id_allocator: Arc<dyn IdAllocator>,

    /// Override the dfs used by kvengine.
    override_dfs: Option<Arc<dyn Dfs>>,
    // NOTE: New members need to check if need to clear in reset_keyspace.
}

struct RestoredSchema {
    version: i64,
    target_file: Option<SchemaFile>,
}

#[derive(Clone, Debug, Default)]
struct DerivedMetadataRetention {
    retain_columnar: bool,
    retain_fts: bool,
    retained_vector_indexes: HashSet<(i64, i64, i64)>,
}

impl RestoredSchema {
    fn target_file_id(&self) -> u64 {
        self.target_file
            .as_ref()
            .map(SchemaFile::get_file_id)
            .unwrap_or_default()
    }
}

impl Drop for BackupCluster {
    fn drop(&mut self) {
        self.stop_engines();
    }
}

fn resolve_master_key(
    dfs: &Arc<dyn Dfs>,
    security_conf: &SecurityConfig,
    injected: Option<MasterKey>,
) -> Result<MasterKey> {
    if let Some(master_key) = injected {
        return Ok(master_key);
    }
    Ok(box_try!(
        dfs.get_runtime().block_on(security_conf.new_master_key())
    ))
}

impl BackupCluster {
    pub fn reset_keyspace(&mut self, keyspace_id: u32, target_keyspace_id: u32) -> Result<()> {
        let (keyspace_start, keyspace_end) = ApiV2::get_keyspace_range_by_id(keyspace_id);
        self.tag = make_keyspace_tag(keyspace_id, target_keyspace_id);
        self.keyspace_id = keyspace_id;
        self.target_keyspace_id = target_keyspace_id;
        self.keyspace_start = keyspace_start;
        self.keyspace_end = keyspace_end;

        let kv_engine = self.kv_engine.take().unwrap();
        kv_engine.close();
        drop(kv_engine);

        if let Some(sender) = self.meta_sender.take() {
            sender.send(StoreMsg::Stop).unwrap();
        }
        self.meta_applier.take();
        self.sorted_shards.clear();
        self.raw_metas.clear();
        self.store_shards.clear();
        self.shards_store_map = None;
        self.static_region_metas.clear();
        self.shards.clear();
        self.shards_need_flush.clear();
        self.shards_need_truncate.clear();
        self.shards_need_reset_columnar.clear();
        self.restored_schema = None;
        self.derived_metadata_retention = None;
        self.restore_encryption_meta = None;
        self.tolerated_err = 0;
        self.txn_chunk_ids_in_wal = None;
        self.load_shards()?;
        self.collect_txn_chunks_in_wal()?;

        self.check_all_shard_files()?;

        for store_id in self.get_all_stores_id() {
            self.setup_kv_engine(store_id)?;
        }
        Ok(())
    }

    pub fn new(
        cluster_meta: &ClusterBackupMeta,
        path: PathBuf,
        pd_client: Arc<dyn PdClient>,
        dfs: Arc<dyn Dfs>,
        restore_conf: RestoreConfig,
        keyspace_id: u32,
        target_keyspace_id: u32,
        truncate_ts: u64,
        archiving: bool,
        archive_reader: Option<ArchiveReader>,
        load_all_tables: bool, // `true` for "check_table" ONLY.
        object_cache: Option<ObjectCache>,
    ) -> Result<BackupCluster> {
        Self::new_opt(BackupClusterOptions {
            cluster_meta: cluster_meta.clone(),
            path,
            pd_client,
            pd_control: None,
            dfs,
            restore_conf,
            object_cache,
            keyspace_id,
            target_keyspace_id,
            truncate_ts,
            archiving,
            archive_reader,
            load_all_tables,
            override_dfs: None,
            offline_packing: false,
            master_key: None,
        })
    }

    pub fn new_opt(options: BackupClusterOptions) -> Result<BackupCluster> {
        let BackupClusterOptions {
            cluster_meta,
            path,
            pd_client,
            pd_control,
            dfs,
            restore_conf,
            object_cache,
            keyspace_id,
            target_keyspace_id,
            truncate_ts,
            archiving,
            archive_reader,
            load_all_tables,
            override_dfs,
            offline_packing,
            master_key: injected_master_key,
        } = options;

        let (keyspace_start, keyspace_end) = if archiving {
            (Vec::default(), GLOBAL_SHARD_END_KEY.to_vec())
        } else {
            ApiV2::get_keyspace_range_by_id(keyspace_id)
        };
        let security_conf = restore_conf.security.clone();
        let master_key = resolve_master_key(&dfs, &security_conf, injected_master_key)?;
        let tag = make_keyspace_tag(keyspace_id, target_keyspace_id);
        let mut cluster = Self {
            tag: tag.clone(),
            path,
            pd_client: pd_client.clone(),
            pd_control,
            dfs,
            security_conf,
            master_key,
            source_cluster_id: cluster_meta.cluster_id,
            keyspace_id,
            target_keyspace_id,
            keyspace_start,
            keyspace_end,
            truncate_ts,
            archive_reader,
            shards: Default::default(),
            raw_metas: Default::default(),
            raft_engines: Default::default(),
            store_configs: Default::default(),
            kv_engine: Default::default(),
            store_shards: Default::default(),
            shards_store_map: Default::default(),
            static_region_metas: Default::default(),
            sorted_shards: Default::default(),
            shards_need_flush: Default::default(),
            shards_need_truncate: Default::default(),
            shards_need_reset_columnar: Default::default(),
            restored_schema: None,
            derived_metadata_retention: None,
            restore_encryption_meta: None,
            meta_applier: None,
            meta_sender: None,
            tolerated_err: 0,
            txn_chunk_ids_in_wal: None,
            load_all_tables,
            id_allocator: Arc::new(PdIdAllocator::new(pd_client)),
            override_dfs,
        };

        let mut store_configs = HashMap::with_capacity(cluster_meta.stores.len());

        // If the cluster backup already has tolerate error, we can't tolerate more
        // errors in restoration.
        let mut tolerate_err = if cluster_meta.get_tolerated_err() > 0 {
            0
        } else {
            restore_conf.tolerate_err
        };

        let is_error_can_tolerate = |err: &Error| -> bool {
            if matches!(
                err,
                Error::RfengineHttpRequestError(_)
                    | Error::PdError(pd_client::Error::StoreTombstone(_))
            ) {
                true
            } else if !restore_conf.strict_tolerate {
                crate::metrics::NATIVE_BR_RESTORE_ERROR
                    .with_label_values(&["setup_raft_engine"])
                    .inc();
                true
            } else {
                false
            }
        };
        let mut raft_engines: HashMap<u64, RfEngine> = Default::default();
        let mut errors = vec![];
        let mut cluster_tolerated_err = 0;
        let mut recv_restore_rfengine =
            |result_rx: &mpsc::Receiver<Result<(u64, TikvConfig, RfEngine)>>| {
                // We can tolerate one store failure for lightweight restoration during fetch
                // latest wal chunk from store.
                match result_rx.recv().unwrap() {
                    Err(err) => errors.push(err),
                    Ok((store_id, store_config, rf_engine)) => {
                        store_configs.insert(store_id, store_config);
                        raft_engines.insert(store_id, rf_engine);
                        info!(
                            "Keyspace {} setup raft engine for store {} done",
                            cluster.tag(),
                            store_id
                        );
                    }
                }
            };
        let (result_tx, result_rx) = tikv_util::mpsc::bounded(cluster_meta.stores.len());
        let mut msg_count = 0;

        for store in &cluster_meta.stores {
            let store_id = store.get_store_id();
            if !archiving
                && !store.keyspace_size.is_empty()
                && !store.keyspace_size.contains_key(&keyspace_id)
            {
                info!(
                    "Keyspace {} skip store {} which has no data",
                    cluster.tag(),
                    store_id
                );
                continue;
            }
            let tag = format!("{tag}:{store_id}");
            let store_config = cluster.generate_store_config(store_id);

            let tx = result_tx.clone();
            let cluster_backup_meta = cluster_meta.clone();
            let pd_client = cluster.pd_client.clone();
            let dfs = cluster.dfs.clone();
            let keyspace_tag = cluster.tag().to_string();
            let archive_store_meta = if let Some(archive_reader) = &cluster.archive_reader {
                let store_meta = archive_reader.get_store_wal_rlog_meta(store_id)?;
                Some((archive_reader.get_start_date().to_string(), store_meta))
            } else {
                None
            };
            let restore_conf_cp = restore_conf.clone();
            let object_cache = object_cache.clone();
            std::thread::spawn(move || {
                let res = BackupCluster::setup_raft_engine(
                    &tag,
                    store_id,
                    keyspace_id,
                    &cluster_backup_meta,
                    &store_config,
                    pd_client,
                    dfs,
                    archiving,
                    archive_store_meta,
                    &restore_conf_cp,
                    object_cache,
                    offline_packing,
                );
                if let Err(err) = &res {
                    warn!(
                        "Keyspace {} setup raft engine for store {} failed: {:?}",
                        keyspace_tag, store_id, err,
                    );
                }
                let _ = tx.send(res.map(|rf_engine| (store_id, store_config, rf_engine)));
            });

            if msg_count < restore_conf.store_concurrency {
                msg_count += 1;
            } else {
                // Do NOT return on error to avoid thread leaky.
                recv_restore_rfengine(&result_rx);
            }
        }
        for _ in 0..msg_count {
            recv_restore_rfengine(&result_rx);
        }

        for err in errors {
            if is_error_can_tolerate(&err) {
                if tolerate_err == 0 {
                    return Err(err);
                }
                warn!(
                    "Keyspace {} setup raft engine failed, tolerate it: {:?}",
                    cluster.tag(),
                    err
                );
                tolerate_err -= 1;
                cluster_tolerated_err += 1;
            } else {
                return Err(err);
            }
        }

        cluster.raft_engines = raft_engines;
        cluster.store_configs = store_configs;
        cluster.tolerated_err = cluster_tolerated_err;
        cluster.load_shards()?;
        cluster.collect_txn_chunks_in_wal()?;

        if archiving {
            return Ok(cluster);
        }
        cluster.check_all_shard_files()?;

        for store_id in cluster.get_all_stores_id() {
            cluster.setup_kv_engine(store_id)?;
        }
        Ok(cluster)
    }

    fn new_snapshot_restore_cluster(
        source_cluster_id: u64,
        keyspace_id: u32,
        path: PathBuf,
        pd_client: Arc<dyn PdClient>,
        pd_control: Option<PdControl>,
        dfs: Arc<dyn Dfs>,
        restore_conf: RestoreConfig,
        target_keyspace_id: u32,
        truncate_ts: u64,
        master_key: Option<MasterKey>,
    ) -> Result<Self> {
        let (keyspace_start, keyspace_end) = ApiV2::get_keyspace_range_by_id(keyspace_id);
        let tag = make_keyspace_tag(keyspace_id, target_keyspace_id);
        let security_conf = restore_conf.security.clone();
        let master_key = resolve_master_key(&dfs, &security_conf, master_key)?;
        Ok(Self {
            tag,
            path,
            pd_client: pd_client.clone(),
            pd_control,
            dfs,
            security_conf,
            master_key,
            source_cluster_id,
            keyspace_id,
            target_keyspace_id,
            keyspace_start,
            keyspace_end,
            truncate_ts,
            archive_reader: None,
            shards: Default::default(),
            raw_metas: Default::default(),
            raft_engines: Default::default(),
            store_configs: Default::default(),
            kv_engine: Default::default(),
            store_shards: Default::default(),
            shards_store_map: Default::default(),
            static_region_metas: Default::default(),
            sorted_shards: Default::default(),
            shards_need_flush: Default::default(),
            shards_need_truncate: Default::default(),
            shards_need_reset_columnar: Default::default(),
            restored_schema: None,
            derived_metadata_retention: None,
            restore_encryption_meta: None,
            meta_applier: None,
            meta_sender: None,
            tolerated_err: 0,
            txn_chunk_ids_in_wal: Some(vec![]),
            load_all_tables: false,
            id_allocator: Arc::new(PdIdAllocator::new(pd_client)),
            override_dfs: None,
        })
    }

    fn initialize_snapshot_shards<'a>(
        &mut self,
        change_sets: impl IntoIterator<Item = &'a pb::ChangeSet>,
    ) -> Result<()> {
        let shards = change_sets
            .into_iter()
            .map(|change_set| {
                (
                    change_set.shard_id,
                    BackupShard {
                        region_id: change_set.shard_id,
                        store_id: SNAPSHOT_RESTORE_STORE_ID,
                        peer_id: SNAPSHOT_RESTORE_PEER_ID,
                        need_flush: false,
                        meta: ShardMeta::new(SNAPSHOT_RESTORE_STORE_ID, change_set),
                        raw_meta: Some(change_set.clone()),
                        raft_state: RaftState::default(),
                    },
                )
            })
            .collect::<HashMap<_, _>>();
        if shards.is_empty() {
            return Err(Error::BackupEmptyForKeyspace(self.keyspace_id));
        }
        let shard_ids = shards
            .keys()
            .copied()
            .sorted_by_key(|shard_id| shards.get(shard_id).unwrap().lower_bound())
            .collect::<Vec<_>>();
        self.raw_metas = shards
            .iter()
            .map(|(shard_id, shard)| (*shard_id, shard.raw_meta.clone().unwrap()))
            .collect();
        self.shards = shards;
        self.store_shards
            .insert(SNAPSHOT_RESTORE_STORE_ID, shard_ids.clone());
        self.sorted_shards = shard_ids;
        self.shards_need_truncate.extend(
            self.shards
                .iter()
                .filter(|(_, shard)| shard.meta.max_ts >= self.truncate_ts)
                .map(|(&shard_id, _)| shard_id),
        );
        Ok(())
    }

    pub fn of_packed(
        packed_backup: &PackedBackup,
        path: PathBuf,
        pd_client: Arc<dyn PdClient>,
        pd_control: Option<PdControl>,
        dfs: Arc<dyn Dfs>,
        restore_conf: RestoreConfig,
        target_keyspace_id: u32,
        truncate_ts: u64,
        master_key: Option<MasterKey>,
    ) -> Result<Self> {
        let mut cluster = Self::new_snapshot_restore_cluster(
            packed_backup.cluster_id,
            packed_backup.keyspace_id,
            path,
            pd_client,
            pd_control,
            dfs,
            restore_conf,
            target_keyspace_id,
            truncate_ts,
            master_key,
        )?;
        cluster.initialize_snapshot_shards(packed_backup.get_shards())?;
        cluster.verify_snapshot_shards()?;
        let config = cluster.generate_store_config(SNAPSHOT_RESTORE_STORE_ID);
        let table_filter = cluster.packed_restore_table_filter();
        cluster.setup_kv_engine_with(
            SNAPSHOT_RESTORE_STORE_ID,
            &config,
            SnapshotNoopRecovery,
            table_filter,
        )?;
        cluster.shards = cluster
            .meta_applier
            .as_ref()
            .unwrap()
            .take_shards()
            .unwrap();
        Ok(cluster)
    }

    /// Builds the normal restore-side temporary KVEngine directly from online
    /// source snapshots, without reconstructing per-Store rfengines.
    fn of_online_snapshots(
        keyspace_id: u32,
        snapshots: Vec<OnlineSnapshotShard>,
        path: PathBuf,
        pd_client: Arc<dyn PdClient>,
        pd_control: Option<PdControl>,
        dfs: Arc<dyn Dfs>,
        restore_conf: RestoreConfig,
        target_keyspace_id: u32,
        truncate_ts: u64,
        master_key: Option<MasterKey>,
    ) -> Result<Self> {
        if snapshots.is_empty() {
            return Err(Error::BackupEmptyForKeyspace(keyspace_id));
        }

        let source_cluster_id = pd_client.get_cluster_id()?;
        let mut cluster = Self::new_snapshot_restore_cluster(
            source_cluster_id,
            keyspace_id,
            path,
            pd_client,
            pd_control,
            dfs,
            restore_conf,
            target_keyspace_id,
            truncate_ts,
            master_key,
        )?;
        cluster
            .initialize_snapshot_shards(snapshots.iter().map(|snapshot| &snapshot.change_set))?;
        cluster.verify_snapshot_shards()?;

        let mut shard_store_map = HashMap::with_capacity(snapshots.len());
        let mut static_region_metas = HashMap::with_capacity(snapshots.len());
        for snapshot in &snapshots {
            let shard_id = snapshot.change_set.get_shard_id();
            shard_store_map.insert(
                shard_id,
                StorePeer {
                    store_id: snapshot.leader.get_store_id(),
                    peer_id: snapshot.leader.get_id(),
                },
            );
            static_region_metas
                .insert(shard_id, (snapshot.region.clone(), snapshot.leader.clone()));
        }
        cluster.shards_store_map = Some(Arc::new(shard_store_map));
        cluster.static_region_metas = static_region_metas;

        let config = cluster.generate_store_config(SNAPSHOT_RESTORE_STORE_ID);
        let table_filter = cluster.lock_resolving_restore_table_filter();
        cluster.setup_kv_engine_with(
            SNAPSHOT_RESTORE_STORE_ID,
            &config,
            SnapshotNoopRecovery,
            table_filter,
        )?;

        let engine = cluster.get_kvengine();
        for snapshot in snapshots {
            let shard_id = snapshot.change_set.get_shard_id();
            let shard_ver = snapshot.change_set.get_shard_ver();
            let snap = snapshot.change_set.get_snapshot();
            let encryption_key = encryption_key_from_shard_properties(
                snap.get_properties(),
                engine.get_encryption_key_manager(),
                &engine.get_master_key(),
                Some(keyspace_id),
            )
            .map_err(|err| -> Error {
                box_err!(
                    "failed to resolve encryption key for online snapshot shard {}: {}",
                    shard_id,
                    err
                )
            })?;
            let mem_tbls =
                cluster
                    .dfs
                    .get_runtime()
                    .block_on(kvengine::SnapAccess::construct_memtables(
                        shard_id,
                        shard_ver,
                        &snapshot.mem_table_data,
                        &engine.get_txn_chunk_manager(),
                        encryption_key,
                    ))?;
            let needs_flush = mem_tbls.iter().any(|mem_tbl| !mem_tbl.is_empty());
            engine
                .get_shard_with_ver(shard_id, shard_ver)?
                .install_memtables_for_restore(mem_tbls)?;
            if needs_flush {
                cluster.shards_need_flush.insert(shard_id);
            }
        }
        Ok(cluster)
    }

    pub fn tag(&self) -> &str {
        &self.tag
    }

    fn setup_raft_engine_for_lightweight(
        tag: &str,
        store_id: u64,
        keyspace_id: u32,
        cluster_backup: &ClusterBackupMeta,
        conf: &TikvConfig,
        pd_client: Arc<dyn PdClient>,
        dfs: Arc<dyn Dfs>,
        archiving: bool,
        archive_store_meta: Option<(String, StoreMeta)>, // archive date, archive store meta
        restore_conf: &RestoreConfig,
        object_cache: Option<ObjectCache>,
        offline_packing: bool,
    ) -> Result<RfEngine> {
        let object_cache_no_hook: Option<ObjectCacheWithHook> = object_cache.map(Into::into);
        let rlog_files = if let Some((date, store_meta)) = &archive_store_meta {
            ArchiveReader::read_store_rlog_files(dfs.as_ref(), date, store_meta)?
        } else {
            collect_snapshot_meta_rlog_files(
                dfs.clone(),
                &dfs.get_prefix(),
                cluster_backup,
                store_id,
                object_cache_no_hook.as_ref(),
            )?
        };
        rfengine::lightweight_restore(
            store_id,
            (!archiving).then_some(keyspace_id),
            Path::new(&conf.raft_store.raftdb_path),
            rlog_files.snap_epoch,
            rlog_files.snap_meta,
            rlog_files.snap_rlog,
            conf.rfengine.epoch_rotate_len,
        )
        .map_err(|x| Error::RfEngine(x))?;

        let rf_engine = TikvServer::init_raft_engine(conf, None)?;

        // When archiving or offline packing, the dfs should have complete wal chunks.
        let complete_wal_chunks = archiving || offline_packing;
        let cache_dir = if restore_conf.lower_memory {
            let dir = Path::new(&conf.storage.data_dir).join("cache");
            box_try!(fs::create_dir_all(&dir));
            Some(dir)
        } else {
            None
        };
        let wal_chunks_cache = if restore_conf.cache_for_decompressed_wal_chunks {
            object_cache_no_hook.map(|cache| {
                let hook =
                    |chunk| convert_to_decompressed_wal_chunk(chunk).map_err(|e| e.to_string());
                cache.with_hook(Arc::new(hook))
            })
        } else {
            object_cache_no_hook
        };
        let ctx = ReplayWalLogsContext {
            pd_client,
            dfs,
            store_id,
            cluster_backup,
            rf_engine: &rf_engine,
            complete_wal_chunks,
            full_restore: false,
            fetch_wal_timeout: restore_conf.timeout_fetch_wal.0,
            cache_dir,
            wal_chunks_cache,
            from_archive: archive_store_meta.is_some(),
        };
        replay_wal_logs(tag, ctx, archive_store_meta, rlog_files.snap_epoch)?;
        Ok(rf_engine)
    }

    pub fn setup_raft_engine(
        tag: &str,
        store_id: u64,
        keyspace_id: u32,
        cluster_backup: &ClusterBackupMeta,
        conf: &TikvConfig,
        pd_client: Arc<dyn PdClient>,
        dfs: Arc<dyn Dfs>,
        archiving: bool,
        archive_store_meta: Option<(String, StoreMeta)>, // archive date, archive store meta
        restore_conf: &RestoreConfig,
        object_cache: Option<ObjectCache>,
        offline_packing: bool,
    ) -> Result<RfEngine> {
        if !cluster_backup.is_lightweight {
            return Err(Error::BackupError(
                "only support lightweight restore".to_string(),
            ));
        }
        Self::setup_raft_engine_for_lightweight(
            tag,
            store_id,
            keyspace_id,
            cluster_backup,
            conf,
            pd_client,
            dfs,
            archiving,
            archive_store_meta,
            restore_conf,
            object_cache,
            offline_packing,
        )
    }

    // Take all raw metas to release memory.
    fn take_raw_metas(&mut self, store_id: u64) -> HashMap<u64, pb::ChangeSet> {
        let shards = self.store_shards.get(&store_id).unwrap().to_owned();
        let mut raw_metas = HashMap::with_capacity(shards.len());
        for shard_id in shards {
            raw_metas.insert(shard_id, self.raw_metas.remove(&shard_id).unwrap());
        }
        raw_metas
    }

    fn get_shards_need_flush_and_truncate(&self) -> Vec<u64> {
        self.shards_need_flush
            .union(&self.shards_need_truncate)
            .cloned()
            .collect()
    }

    fn setup_kv_engine(&mut self, store_id: u64) -> Result<()> {
        let conf = self.store_configs.get(&store_id).unwrap().clone();
        let rf_engine = self.raft_engines.get(&store_id).unwrap();
        let recoverer = rfstore::store::RecoverHandler::new(rf_engine.clone());
        let table_filter = self.lock_resolving_restore_table_filter();
        self.setup_kv_engine_with(store_id, &conf, recoverer, table_filter)
    }

    fn packed_restore_table_filter(&self) -> Option<LoadTableFilterFn> {
        if self.load_all_tables {
            None
        } else {
            Some(build_packed_restore_table_filter(
                self.shards_need_truncate.clone(),
            ))
        }
    }

    fn lock_resolving_restore_table_filter(&self) -> Option<LoadTableFilterFn> {
        if self.load_all_tables {
            None
        } else {
            Some(build_lock_resolving_restore_table_filter(
                self.shards_need_truncate.clone(),
            ))
        }
    }

    fn setup_kv_engine_with(
        &mut self,
        store_id: u64,
        conf: &TikvConfig,
        recoverer: impl kvengine::RecoverHandler + 'static,
        table_filter: Option<LoadTableFilterFn>,
    ) -> Result<()> {
        if self.kv_engine.is_none() {
            let io_rate_limiter =
                Arc::new(IoRateLimiter::new(IoRateLimitMode::WriteOnly, true, true));
            io_rate_limiter
                .set_io_rate_limit(conf.storage.io_rate_limit.max_bytes_per_sec.0 as usize);
            // TODO: enable flow control to avoid OOM during recovery.
            let store_limiter = Arc::new(StoreLimiter::dummy());

            let mut meta_iter = MetaIterator::new(store_id, Vec::new(), HashMap::default());
            let dfs = if let Some(odfs) = self.override_dfs.as_ref() {
                odfs.clone()
            } else {
                self.dfs.clone()
            };
            let (kv_engine, sender, receiver) = TikvServer::init_kv_engine(
                self.pd_client.clone(),
                conf,
                dfs,
                io_rate_limiter,
                store_limiter,
                &mut meta_iter,
                recoverer.clone(),
                true,
                self.master_key.clone(),
                self.pd_client.get_security_mgr(),
            )?;

            let shard_meta = &self.get_sorted_shard(0).meta;
            let encryption_key = encryption_key_from_shard_properties(
                &shard_meta.get_properties_pb(),
                kv_engine.get_encryption_key_manager(),
                &kv_engine.get_master_key(),
                Some(shard_meta.range.keyspace_id),
            )
            .map_err(|err| -> Error {
                box_err!("failed to resolve encryption key for restore: {}", err)
            })?;
            // Move shards to meta_applier to apply meta change in `flush_mem_table`.
            // Move back after flush finished.
            let mut shards = mem::take(&mut self.shards);
            for shard in shards.values_mut() {
                if self.shards_need_reset_columnar.contains(&shard.region_id) {
                    info!(
                        "Keyspace {} shard {} reset columnar",
                        self.keyspace_id, shard.region_id
                    );
                    shard.reset_columnar();
                }
            }
            let meta_applier = Arc::new(MetaApplier::new(
                self.keyspace_id,
                kv_engine.clone(),
                encryption_key,
                shards,
                receiver,
            ));
            self.meta_applier = Some(meta_applier.clone());
            self.meta_sender = Some(sender);
            std::thread::Builder::new()
                .name("native-br-handler".into())
                .spawn_wrapper(move || {
                    meta_applier.run();
                })
                .unwrap();

            self.kv_engine = Some(kv_engine);
        }

        // Load immutable tables selected by the source path. Normal and online
        // restores include L0 and LOCK_CF tables for lock resolution; packed
        // backups have already resolved locks.
        // TODO: L0 & LOCK_CF with `max_ts` < `safe_ts` can be skipped.
        // (Currently there is no `max_ts` for LOCK_CF of L0)
        let store_shards = self.store_shards.get(&store_id).unwrap().to_owned();
        let raw_metas = self.take_raw_metas(store_id);
        let kv_engine = self.kv_engine.as_mut().unwrap();
        if !store_shards.is_empty() {
            let mut meta_iter =
                MetaIterator::new(kv_engine.get_engine_id(), store_shards, raw_metas);
            let (mut metas, _) = kvengine::EngineCore::read_meta(&mut meta_iter)?;
            info!(
                "Keyspace {} kv_engine load {} shards in restore keyspace",
                self.keyspace_id,
                metas.len()
            );

            // Clear columnar related meta in metas if the shard need reset columnar.
            for meta in metas.values_mut() {
                if self.shards_need_reset_columnar.contains(&meta.id) {
                    info!(
                        "Keyspace {} shard {} clear columnar related meta",
                        self.keyspace_id, meta.id
                    );
                    meta.clear_columnar_related_meta();
                }
            }

            tikv_util::init_task_local_sync(|| {
                kv_engine.load_shards(metas, recoverer, table_filter)
            })?;
        }
        Ok(())
    }

    fn stop_engines(&mut self) {
        let raft_engines = mem::take(&mut self.raft_engines);
        for rf in raft_engines.into_values() {
            rf.stop_worker(true);
            drop(rf);
        }
        if let Some(kv_engine) = self.kv_engine.take() {
            // `self.kv_engine` will be None when `BackupCluster` fail to initialize.
            kv_engine.close();
            drop(kv_engine);
        }
        // Receiver cannot get err msg even close_sender, send Stop msg instead.
        if let Some(sender) = &self.meta_sender {
            sender.send(StoreMsg::Stop).unwrap();
        }
    }

    fn generate_store_config(&self, store_id: u64) -> TikvConfig {
        let (store_path, rf_engine_path) = (
            self.path.join(store_id.to_string()),
            self.path.join(store_id.to_string()).join("raft"),
        );

        let mut config = TikvConfig::default();
        config.storage.data_dir = store_path.to_str().unwrap().to_string();
        config.raft_store.raftdb_path = rf_engine_path.to_str().unwrap().to_string();
        config.dfs.zstd_compression_level = ZSTD_COMPRESSION_LEVEL.to_string();
        config.rocksdb.max_background_jobs = 2;
        config.rocksdb.max_sub_compactions = 1;
        config.security = self.security_conf.clone();

        config.rfengine.lightweight_backup = false;
        config.kvengine.ia.force_ia = true;
        config.kvengine.ia.mem_cap = AbsoluteOrPercentSize::Percent(5.0);
        config.kvengine.ia.disk_cap = AbsoluteOrPercentSize::Percent(20.0);

        TikvServer::init_config(config).get_current()
    }

    fn collect_keyspace_shards(
        tag: &str,
        store_id: u64,
        rf: &RfEngine,
        keyspace_start: &[u8],
        keyspace_end: &[u8],
    ) -> Result<Vec<BackupShard>> {
        let region_peers = rf.get_region_peer_map();
        let mut prefix_shards = vec![];
        let mut no_meta_peers = vec![];
        for (region_id, peer_id) in region_peers {
            if region_id == 0 {
                continue;
            }
            tikv_util::set_current_region_thread_local(region_id);
            let meta = match load_raft_engine_meta(rf, peer_id) {
                Some(meta) => meta,
                None => {
                    debug!(
                        "region {} peer {} has no rf_engine meta",
                        region_id, peer_id
                    );
                    if no_meta_peers.len() < 5 {
                        no_meta_peers.push((region_id, peer_id));
                    }
                    continue;
                }
            };
            assert_eq!(region_id, meta.shard_id);

            let snap = meta.get_snapshot();
            if keyspace_start <= snap.get_outer_start() && snap.get_outer_end() <= keyspace_end {
                let shard = Self::create_backup_shard(rf, store_id, region_id, peer_id, meta)?;
                debug!("{} collect_keyspace_shards", shard.tag(); "inner_key_off" => shard.inner_key_off());
                prefix_shards.push(shard);
            }
        }
        if !no_meta_peers.is_empty() {
            warn!(
                "{} collect_keyspace_shards: peers has no meta: {:?}",
                tag, no_meta_peers
            );
        }
        Ok(prefix_shards)
    }

    fn collect_full_shards(store_id: u64, rf: &RfEngine) -> Result<FullShardCollection> {
        let region_peers = rf.get_region_peer_map();
        let mut prefix_shards = vec![];
        let mut tombstone_peer_count = 0;
        for (region_id, peer_id) in region_peers {
            if region_id == 0 {
                continue;
            }
            tikv_util::set_current_region_thread_local(region_id);
            let meta = match load_raft_engine_meta(rf, peer_id) {
                Some(meta) => meta,
                None => {
                    warn!(
                        "region {} peer {} has no rf_engine meta",
                        region_id, peer_id
                    );
                    continue;
                }
            };
            assert_eq!(region_id, meta.shard_id);

            if rf
                .load_region_state(peer_id, meta.shard_ver)
                .is_some_and(|state| {
                    state.get_state() == kvproto::raft_serverpb::PeerState::Tombstone
                })
            {
                tombstone_peer_count += 1;
                continue;
            }

            let shard = Self::create_backup_shard(rf, store_id, region_id, peer_id, meta)?;
            prefix_shards.push(shard);
        }
        Ok(FullShardCollection {
            shards: prefix_shards,
            tombstone_peer_count,
        })
    }

    pub fn is_full_range(&self) -> bool {
        self.keyspace_start.is_empty() && self.keyspace_end.eq(GLOBAL_SHARD_END_KEY)
    }

    fn create_backup_shard(
        rf: &RfEngine,
        store_id: u64,
        region_id: u64,
        peer_id: u64,
        meta: kvenginepb::ChangeSet,
    ) -> Result<BackupShard> {
        let tag = ShardTag::new(store_id, IdVer::from_change_set(&meta));
        let snap = meta.get_snapshot();
        let raft_last_index = rf.get_last_index(peer_id);
        let need_flush_mem_table = raft_last_index.is_some_and(|last_idx| {
            if last_idx < snap.get_data_sequence() {
                warn!("{} create_backup_shard: last_index < data_sequence", tag;
                    "peer_id" => peer_id, "last_idx" => last_idx, "meta" => ?meta,
                    "truncated_index" => rf.get_truncated_index(peer_id),
                );
            }
            last_idx > snap.get_data_sequence()
        });
        let need_initial_flush = meta.has_parent();

        let raft_state = load_peer_raft_state(rf, peer_id, meta.shard_ver).ok_or_else(|| -> Error {
            let states = rf.get_peer_all_states(peer_id, false);
            error!(
                "{} failed to load peer raft state, store_id: {}, region_id: {}, peer_id: {}, state: {:?}",
                tag, store_id, region_id, peer_id, states
            );

            box_err!("failed to load peer raft state, store_id: {}, region_id: {} peer_id: {}",
                store_id, region_id, peer_id)
        })?;

        let shard = BackupShard {
            region_id,
            store_id,
            peer_id,
            need_flush: need_flush_mem_table || need_initial_flush,
            meta: ShardMeta::new(rf.get_engine_id(), &meta),
            raw_meta: Some(meta),
            raft_state,
        };
        debug!(
            "{} create_backup_shard: raft_last_index {:?}, {:?}",
            tag, raft_last_index, shard
        );
        Ok(shard)
    }

    fn load_shards(&mut self) -> Result<()> {
        let tag = self.tag();
        // Will need to retry when `get_leader_shards_and_preprocess` produce new shards
        // after preprocess.
        let mut leader_shards = HashMap::default();
        let mut retry = 0_usize;
        let is_full_range = self.is_full_range();
        loop {
            let mut all_shards = HashMap::default();
            let mut tombstone_peer_count = 0;
            for (&store_id, rf_engine) in &self.raft_engines {
                let store_shards = if is_full_range {
                    let collection = Self::collect_full_shards(store_id, rf_engine)?;
                    tombstone_peer_count += collection.tombstone_peer_count;
                    collection.shards
                } else {
                    Self::collect_keyspace_shards(
                        tag,
                        store_id,
                        rf_engine,
                        &self.keyspace_start,
                        &self.keyspace_end,
                    )?
                };
                all_shards.insert(store_id, store_shards);
            }

            // Check whether backup is empty (for this keyspace).
            // Empty backup should be invalid or before the creation of this
            // keyspace, so the restore is not allowed.
            if all_shards.iter().all(|(_, shards)| shards.is_empty()) {
                if tombstone_peer_count > 0 {
                    return Err(Error::ShardRangeIntegrityError(format!(
                        "full-range reconstruction excluded {tombstone_peer_count} persisted tombstone peers and no eligible peers remain"
                    )));
                }
                return Err(Error::BackupEmptyForKeyspace(self.keyspace_id));
            }

            match self.get_leader_shards_and_preprocess(all_shards)? {
                (_, true) => {
                    if retry > 20 {
                        // Should never happen. Just a safe guard to avoid infinite loop.
                        return Err(box_err!(
                            "get_leader_shards_and_preprocess retry too many times"
                        ));
                    }
                    retry += 1;
                    continue;
                }
                (shards, false) => {
                    let _ = mem::replace(&mut leader_shards, shards);
                    break;
                }
            }
        }

        if is_full_range {
            // Reserve all the leader shards for archiving.
            let mut shards: HashMap<u64, BackupShard> = HashMap::default();
            let sorted_shards = leader_shards
                .values()
                .sorted_by(|a, b| a.start().cmp(b.start()))
                .map(|x| x.region_id)
                .collect();
            shards.extend(leader_shards.into_iter().map(|(k, v)| (k.id, v)));
            self.shards = shards;
            self.sorted_shards = sorted_shards;
        } else {
            // The `find_intact_shards` in the following methods can only handle the shards
            // of a single keyspace.
            let (intact_shards, sorted_shards_id) =
                Self::handle_overlapping_shards(self.tag(), leader_shards)?;
            self.shards = intact_shards;
            self.sorted_shards = sorted_shards_id;
        }
        info!("{}: shards cnt {}", self.tag(), self.shards.len());
        debug!(
            "Keyspace {} BackupCluster.load_shards: {:?}",
            self.tag(),
            self.shards
        );
        debug!(
            "Keyspace {} BackupCluster.sorted_shards: {:?}",
            self.tag(),
            self.sorted_shards
        );

        for (&shard_id, shard) in self.shards.iter_mut() {
            self.store_shards
                .entry(shard.store_id)
                .and_modify(|ids| ids.push(shard_id))
                .or_insert_with(|| vec![shard_id]);

            // Take `raw_meta`s out from `shards`, as `shards` will be moved into
            // `MetaApplier` for flush shards.
            self.raw_metas
                .insert(shard_id, shard.raw_meta.take().unwrap());
        }

        self.shards_store_map = Some(Arc::new(
            self.shards
                .iter()
                .map(|(&shard_id, shard)| {
                    (
                        shard_id,
                        StorePeer {
                            store_id: shard.store_id,
                            peer_id: shard.peer_id,
                        },
                    )
                })
                .collect(),
        ));
        self.shards_need_flush.extend(
            self.shards
                .iter()
                .filter(|(_, shard)| shard.need_flush)
                .map(|(&shard_id, _)| shard_id),
        );
        self.shards_need_truncate.extend(
            self.shards
                .iter()
                .filter(|(_, shard)| shard.meta.max_ts >= self.truncate_ts)
                .map(|(&shard_id, _)| shard_id),
        );

        if is_full_range {
            self.verify_full_shards()?;
        }
        Ok(())
    }

    fn get_shard_files(shard_id: u64, meta: &ShardMeta) -> Vec<TableFile> {
        let mut files = meta
            .all_files()
            .iter()
            .filter(|(_, file)| {
                !matches!(
                    file.file_type,
                    FileType::VectorIndex | FileType::FtsPackedFile | FileType::FtsDedicatedFile
                )
            })
            .map(|(&id, file)| {
                let size = if file.file_type == FileType::Blob {
                    file.table_meta_off as u64
                } else if file.l0_size > 0 {
                    file.l0_size as u64
                } else {
                    (file.table_meta_off as f64 * crate::limiter::TABLE_SIZE_MULTIPLIER) as u64
                };
                TableFile {
                    id,
                    ftype: file.file_type,
                    shard_id,
                    size,
                }
            })
            .collect::<Vec<_>>();
        files.extend(meta.all_txn_chunk_ids().into_iter().map(|id| TableFile {
            id,
            ftype: FileType::TxnChunk,
            shard_id,
            size: 0,
        }));
        if meta.schema.is_valid() {
            files.push(TableFile {
                id: meta.schema.file_id(),
                ftype: FileType::Schema,
                shard_id,
                size: 0,
            });
        }
        for vector_index in &meta.vector_indexes {
            files.extend(vector_index.get_files().iter().map(|file| TableFile {
                id: file.get_id(),
                ftype: FileType::VectorIndex,
                shard_id,
                size: (file.get_meta_offset() as f64 * crate::limiter::TABLE_SIZE_MULTIPLIER)
                    as u64,
            }));
        }
        files.extend(
            meta.fts_l0_files
                .iter()
                .chain(&meta.fts_l1_files)
                .map(|file| TableFile {
                    id: file.get_id(),
                    ftype: FileType::FtsPackedFile,
                    shard_id,
                    size: (file.get_meta_offset() as f64 * crate::limiter::TABLE_SIZE_MULTIPLIER)
                        as u64,
                }),
        );
        files.extend(meta.fts_l2_files.iter().map(|file| TableFile {
            id: file.get_id(),
            ftype: FileType::FtsDedicatedFile,
            shard_id,
            size: (file.get_meta_offset() as f64 * crate::limiter::TABLE_SIZE_MULTIPLIER) as u64,
        }));
        files
    }

    // Should be called after load_shard & collect_txn_chunks_in_wal.
    pub fn get_all_shard_files(
        &self,
        skip_shards: Option<HashSet<u64>>,
        skip_keyspace_ids: Option<HashSet<u32>>,
    ) -> Vec<TableFile> {
        let mut files = vec![];
        for shard in self.shards.values() {
            if let Some(skip_shards) = &skip_shards {
                if skip_shards.contains(&shard.region_id) {
                    info!(
                        "skip shard {} files {}",
                        shard.tag(),
                        shard.meta.all_files().len()
                    );
                    continue;
                }
            }
            if let Some(skip_keyspace_ids) = &skip_keyspace_ids {
                if let Some(keyspace_id) = ApiV2::get_u32_keyspace_id_by_key(shard.start()) {
                    if skip_keyspace_ids.contains(&keyspace_id) {
                        let keyspace_end_key = ApiV2::get_keyspace_end_by_id(keyspace_id);
                        if shard.end() <= keyspace_end_key.as_slice() {
                            info!(
                                "skip keyspace {} shard {} files {}",
                                keyspace_id,
                                shard.tag(),
                                shard.meta.all_files().len()
                            );
                            continue;
                        }
                    }
                }
            }
            files.extend(Self::get_shard_files(shard.region_id, &shard.meta));
        }
        files.extend(
            self.txn_chunk_ids_in_wal
                .as_ref()
                .unwrap()
                .iter()
                .map(|&id| TableFile {
                    id,
                    ftype: FileType::TxnChunk,
                    shard_id: 0,
                    size: 0,
                }),
        );
        files
    }

    fn plan_archive_file_restore(
        not_found_files: Vec<TableFile>,
    ) -> (HashSet<u64>, Vec<TableFile>) {
        let mut shards_need_reset_columnar = HashSet::default();
        let mut files_to_restore = Vec::with_capacity(not_found_files.len());
        for file in not_found_files {
            if !is_derived_file_type(file.ftype) {
                files_to_restore.push(file);
            } else {
                shards_need_reset_columnar.insert(file.shard_id);
            }
        }
        (shards_need_reset_columnar, files_to_restore)
    }

    // Should be called after load_shard and before setup_kv_engine.
    pub fn check_all_shard_files(&mut self) -> Result<()> {
        let files = self.get_all_shard_files(None, None);
        let mut shards_need_reset_columnar = HashSet::default();
        if let Some(archive_reader) = &self.archive_reader {
            let not_found_files = get_not_found_files(self.dfs.clone(), files)?;
            let (shards_to_reset, files_to_restore) =
                Self::plan_archive_file_restore(not_found_files);
            shards_need_reset_columnar = shards_to_reset;
            archive_reader.restore_files(files_to_restore)?;
        }
        step!(
            "Keyspace {} need reset columnar shards: {:?}",
            self.tag(),
            shards_need_reset_columnar
        );
        self.shards_need_reset_columnar = shards_need_reset_columnar;
        Ok(())
    }

    fn get_leader_shards_and_preprocess(
        &self,
        all_shards: HashMap<u64, Vec<BackupShard>>,
    ) -> Result<(
        HashMap<IdVer, BackupShard>, // leader_shards
        bool,                        // has_new_peer
    )> {
        let shards_cnt = all_shards.values().map(|x| x.len()).sum::<usize>() / REPLICAS;
        let mut shard_peers: HashMap<IdVer, Vec<BackupShard>> = HashMap::with_capacity(shards_cnt);

        for shard in all_shards.into_values().flatten() {
            shard_peers.entry(shard.id_ver()).or_default().push(shard);
        }
        let mut leader_shards = HashMap::with_capacity(shard_peers.len());
        for (&id_ver, peers) in shard_peers.iter_mut() {
            // Since we enabled raft pre-vote, the correct leader must have the max term and
            // then max last index.
            // Also sort by store_id (descending) to get peer with smallest store id as
            // leader, to make the result determined and improve efficiency for
            // retry.
            peers.sort_by(|a, b| {
                let term_last_a = (
                    a.raft_state.get_term(),
                    a.raft_state.get_last_index(),
                    !a.store_id,
                );
                let term_last_b = (
                    b.raft_state.get_term(),
                    b.raft_state.get_last_index(),
                    !b.store_id,
                );
                term_last_a.cmp(&term_last_b)
            });
            let mut leader = peers.pop().unwrap();
            if leader.raft_state.get_commit() < leader.raft_state.get_last_index() {
                // If we tolerate error, the actual leader maybe missing, it's possible that the
                // most up-to-date follower's last_index is already committed.
                // It's possible that the original cluster eventually not committed to
                // the last index, and the last index contains commit primary key timestamp
                // lower than backup ts. But the probability is very low and the
                // data is still consistent.
                let mut hard_state = leader.raft_state.get_hard_state();
                hard_state.set_commit(leader.raft_state.get_last_index());
                leader.raft_state.set_hard_state(&hard_state);
            }
            leader_shards.insert(id_ver, leader);
        }

        let mut has_new_peer = false;
        let mut ver_changed = vec![];
        for (&id_ver, shard) in leader_shards.iter_mut() {
            tikv_util::set_current_region_thread_local(shard.region_id);
            // Preprocess if needed.
            if shard.raft_state.get_commit() > shard.raft_state.get_last_preprocessed_index() {
                info!(
                    "Keyspace {} shard {} need preprocess, commit {}, last_preprocessed_index {}",
                    self.tag(),
                    shard.region_id,
                    shard.raft_state.get_commit(),
                    shard.raft_state.get_last_preprocessed_index()
                );
                let rf_engine = self.raft_engines.get(&shard.store_id).unwrap();
                if self.preprocess_shard(rf_engine, shard)? {
                    has_new_peer = true;
                }

                if shard.ver() != id_ver.ver {
                    ver_changed.push(id_ver);
                }
            }
        }

        for id_ver in ver_changed {
            let shard = leader_shards.remove(&id_ver).unwrap();
            leader_shards.insert(shard.id_ver(), shard);
        }

        Ok((leader_shards, has_new_peer))
    }

    /// Handle overlapping shards.
    /// Shards overlap will happen when some followers had not finished split or
    /// merge.
    fn handle_overlapping_shards(
        tag: &str,
        mut leader_shards: HashMap<IdVer, BackupShard>,
    ) -> Result<(
        HashMap<u64, BackupShard>, // intact_shards
        Vec<u64>,                  // Vec<shard_id> sorted by BackupShard.start()
    )> {
        let mut shards_map: HashMap<InnerKey<'_>, Vec<&BackupShard>> = HashMap::default();
        for shard in leader_shards.values() {
            shards_map
                .entry(shard.inner_start())
                .or_default()
                .push(shard);
        }
        // Increasingly sort as `find_intact_shards` will try the last shard first.
        for shards in shards_map.values_mut() {
            shards.sort_by(|a, b| a.ver().cmp(&b.ver()));
        }
        let Some(intact_shards) = Self::find_intact_shards(&shards_map) else {
            error!("{}: shards not intact for keyspace", tag; "shards" => ?shards_map);
            return Err(box_err!("shards not intact for keyspace"));
        };
        info!("{}: intact shards: {:?}", tag, intact_shards); // TODO: debug log.
        let mut shards: HashMap<u64, BackupShard> = HashMap::default();
        let mut sorted_shards: Vec<u64> = vec![];
        for id_ver in intact_shards {
            let shard = leader_shards.remove(&id_ver).unwrap();
            shards.insert(id_ver.id, shard);
            sorted_shards.push(id_ver.id);
        }

        Ok((shards, sorted_shards))
    }

    // Consider as a DAG, where the start/end keys are the nodes, and the shards are
    // the edges.
    fn find_intact_shards(
        shards_map: &HashMap<InnerKey<'_>, Vec<&BackupShard>>,
    ) -> Option<Vec<IdVer>> {
        let shards = shards_map.get(&InnerKey::default())?;

        let mut stack: Vec<&BackupShard> = Vec::with_capacity(shards_map.len());
        stack.extend(shards);
        let mut visited: HashMap<InnerKey<'_>, &BackupShard> = HashMap::default();

        while let Some(curr) = stack.pop() {
            let next_key = curr.inner_end();
            if next_key.deref() == GLOBAL_SHARD_END_KEY {
                let mut edges = vec![];
                let mut curr = curr;
                loop {
                    edges.push(curr.id_ver());
                    let prev_key = curr.inner_start();
                    if prev_key.is_empty() {
                        break;
                    } else {
                        curr = *visited.get(&prev_key).unwrap()
                    }
                }

                edges.reverse();
                return Some(edges);
            }

            if let Entry::Vacant(e) = visited.entry(next_key) {
                if let Some(shards) = shards_map.get(&next_key) {
                    stack.extend(shards);
                    e.insert(curr);
                    continue;
                }
            }
        }
        None
    }

    fn preprocess_shard(
        &self,
        rf_engine: &RfEngine,
        old_shard: &mut BackupShard,
    ) -> Result<bool /* has_new_peer */> {
        let mut preprocessor = PeerPreprocessor::new(rf_engine, old_shard);

        // Get raft logs.
        let low_idx = preprocessor.preprocessed_index + 1;
        let high_idx = old_shard.raft_state.get_commit() + 1;
        info!(
            "Keyspace {} shard {} preprocess Raft log [{}, {}), raft_state {:?}",
            self.tag(),
            old_shard.tag(),
            low_idx,
            high_idx,
            preprocessor.raft_state,
        );

        let mut entries = Vec::with_capacity((high_idx.saturating_sub(low_idx)) as usize);
        let peer_id = old_shard.peer_id;
        rf_engine
            .fetch_raft_entries_to(peer_id, low_idx, high_idx, None, &mut entries, None)
            .map_err(|e| -> Error {
                let stats = rf_engine.get_peer_stats(peer_id);
                let truncated_state = rfstore::store::load_raft_truncated_state(rf_engine, peer_id);
                box_err!(
                    "{} entries unavailable err: {:?}, stats {:?}, truncated_state: {:?}, low: {}, high {}",
                    old_shard.tag(),
                    e,
                    stats,
                    truncated_state,
                    low_idx,
                    high_idx
                )
            })?;

        // Prepare context.
        let mut raft_wb = rfengine::WriteBatch::new();
        let mut remove_dependents = Vec::new();
        let mut apply_msgs = ApplyMsgs::default();
        let raft_cfg = rfstore::store::Config::default();
        let mut destroying = std::collections::HashSet::default();
        let mut ctx = PreprocessContext {
            store_id: old_shard.store_id,
            raft: rf_engine,
            raft_wb: &mut raft_wb,
            remove_dependents: &mut remove_dependents,
            apply_msgs: &mut apply_msgs,
            cfg: &raft_cfg,
            destroying: &mut destroying,
            scene: PreprocessScene::Restore,
        };

        // Preprocess
        let mut preprocess_ref = preprocessor.as_ref();
        for entry in &entries {
            let res = preprocess_ref.preprocess_committed_entry(&mut ctx, entry);
            match res {
                PreprocessResult::Ok => {}
                PreprocessResult::Rejected { reason, report_err } => {
                    info!("Keyspace {} preprocess committed entry rejected", self.tag();
                        "region" => %old_shard.tag(),
                        "index" => entry.get_index(),
                        "reason" => ?reason,
                        "report_err" => ?report_err,
                    );
                    if let Some(rfstore::Error::PreprocessRestoreShardRejected {
                        region_id,
                        log_index,
                        log_term,
                    }) = report_err
                    {
                        error!(
                            "Keyspace {} reject replay restore for region {}",self.tag(), region_id;
                            "region" => %old_shard.tag(),
                            "log_index" => log_index,
                            "log_term" => log_term,
                        );
                        metrics::NATIVE_BR_RESTORE_ERROR
                            .with_label_values(&["reject_replay_restore"])
                            .inc();
                        return Err(Error::RejectReplayRestore);
                    }
                }
            }
        }

        // Apply to rf_engine
        preprocess_ref
            .raft_state
            .set_last_preprocessed_index(*preprocess_ref.preprocessed_index);
        preprocess_ref.write_raft_state(&mut ctx); // update `last_preprocessed_index` to rf_engine.
        rf_engine.apply(ctx.raft_wb); // Do not need persist.

        // Get new shard.
        // TODO: In-place update `old_shard`, other than get a new one from rf_engine.
        let meta = load_raft_engine_meta(rf_engine, peer_id).unwrap();
        let new_shard = Self::create_backup_shard(
            rf_engine,
            old_shard.store_id,
            old_shard.region_id,
            peer_id,
            meta,
        )?;
        info!(
            "Keyspace {} preprocessed shard {}, old {:?}, new {:?}",
            self.tag(),
            old_shard.tag(),
            old_shard,
            new_shard
        );

        // Verify.
        let old_commit = old_shard.raft_state.get_commit();
        let new_commit = new_shard.raft_state.get_commit();
        let new_last_preprocessed_index = new_shard.raft_state.get_last_preprocessed_index();
        assert!(
            old_commit == new_commit && new_commit == new_last_preprocessed_index,
            "Keyspace {} shard {} preprocess: commit_index(old:{}, new:{})/last_preprocessed_index({}) not match",
            self.tag(),
            old_shard.tag(),
            old_commit,
            new_commit,
            new_last_preprocessed_index,
        );

        // When new shard has smaller range, there must be new peer due to split.
        // Raise error to retry from `load_shards` for simplicity, as it is a very rare
        // case.
        // Do not need to care about region merge, and `handle_overlapping_shards` will
        // handle the overlapping caused by merge.
        let has_new_peer =
            old_shard.start() < new_shard.start() || new_shard.end() < old_shard.end();
        if has_new_peer {
            info!(
                "{} split detected, retry load_shards, {} old_shard {:?}, new_shard {:?}",
                self.tag,
                old_shard.tag(),
                old_shard,
                new_shard
            );
        }

        let _ = mem::replace(old_shard, new_shard);
        Ok(has_new_peer)
    }

    /// Strictly verifies that a canonical snapshot shard list covers the source
    /// keyspace without gaps or overlaps.
    fn verify_snapshot_shards(&self) -> Result<()> {
        debug_assert!(!self.sorted_shards.is_empty());
        if self.is_full_range() {
            return self.verify_full_shards();
        }

        let first_shard = self.get_sorted_shard(0);
        let last_shard = self.get_sorted_shard(self.sorted_shards.len() - 1);
        if first_shard.start() != self.keyspace_start {
            return Err(box_err!(
                "start key of first region not match: {:?}",
                first_shard
            ));
        }
        if last_shard.end() != self.keyspace_end {
            return Err(box_err!(
                "end key of last region not match: {:?}",
                last_shard
            ));
        }

        for shards_id in self.sorted_shards.windows(2) {
            let left = self.get_shard(shards_id[0]).unwrap();
            let right = self.get_shard(shards_id[1]).unwrap();
            if left.end() != right.start() {
                return Err(box_err!(
                    "region boundary not match: {:?}, {:?}",
                    left,
                    right
                ));
            }
        }

        Ok(())
    }

    fn verify_full_shards(&self) -> Result<()> {
        let first_shard = self.get_sorted_shard(0);
        let last_shard = self.get_sorted_shard(self.sorted_shards.len() - 1);
        if first_shard.start() != self.keyspace_start {
            return Err(Error::ShardRangeIntegrityError(format!(
                "start key of first region not match: {:?}",
                first_shard
            )));
        }
        if last_shard.end() != self.keyspace_end {
            return Err(Error::ShardRangeIntegrityError(format!(
                "end key of last region not match: {:?}",
                last_shard
            )));
        }

        for shards_id in self.sorted_shards.windows(2) {
            let left = self.get_shard(shards_id[0]).unwrap();
            let right = self.get_shard(shards_id[1]).unwrap();
            if left.end() != right.start() {
                return Err(Error::ShardRangeIntegrityError(format!(
                    "region boundary not match: {:?}, {:?}",
                    left, right
                )));
            }
        }
        Ok(())
    }

    fn get_sorted_shard(&self, sorted_idx: usize) -> &BackupShard {
        self.get_shard(self.sorted_shards[sorted_idx]).unwrap()
    }

    pub fn get_shard(&self, shard_id: u64) -> Option<&BackupShard> {
        self.shards.get(&shard_id)
    }

    pub fn get_shard_mut(&mut self, shard_id: u64) -> Option<&mut BackupShard> {
        self.shards.get_mut(&shard_id)
    }

    pub fn shards_count(&self) -> usize {
        self.sorted_shards.len()
    }

    pub fn get_shard_metas_before_flush(&self) -> Vec<ShardMeta> {
        let mut shards = Vec::with_capacity(self.sorted_shards.len());
        let guard = self.meta_applier.as_ref().unwrap().shards.read().unwrap();
        let backup_shards = guard.as_ref().unwrap();
        for &id in &self.sorted_shards {
            let backup_shard = backup_shards.get(&id).unwrap();
            shards.push(backup_shard.meta.clone());
        }
        shards
    }

    pub fn get_kvengine(&self) -> kvengine::Engine {
        self.kv_engine.as_ref().unwrap().clone()
    }

    pub fn get_rfengine(&self, store_id: u64) -> Option<RfEngine> {
        self.raft_engines.get(&store_id).cloned()
    }

    /// Returns the master key decrypted during cluster construction.
    pub fn get_master_key(&self) -> MasterKey {
        self.master_key.clone()
    }

    pub fn get_shard_meta_getter(&self) -> RegionMetaGetter {
        RegionMetaGetter {
            shard_store_map: self.shards_store_map.as_ref().unwrap().clone(),
            raft_engines: Arc::new(self.raft_engines.clone()),
            static_regions: Arc::new(self.static_region_metas.clone()),
        }
    }

    // The returned id of stores are sorted.
    // This will help to reuse downloaded table files when re-running the
    // restoration process.
    // See `BackupCluster::setup_kv_engine`.
    fn get_all_stores_id(&self) -> Vec<u64> {
        let mut stores_id: Vec<_> = self.store_shards.keys().copied().collect();
        stores_id.sort();
        stores_id
    }

    pub fn add_shards_need_flush(&mut self, shard_ids: &[u64]) {
        self.shards_need_flush.extend(shard_ids);
    }

    pub fn flush_shards(
        &mut self,
        timeout: Duration,
    ) -> Result<usize /* number of flushed shards */> {
        let kv_engine = self.kv_engine.as_ref().unwrap();
        // TODO: run in parallel.
        for &shard_id in &self.shards_need_flush {
            let engine_shard = kv_engine
                .get_shard(shard_id)
                .unwrap_or_else(|| panic!("shard not found: {}", shard_id));
            kv_engine.flush_shard_for_restore(&engine_shard)?;
        }

        self.check_flushed(timeout)?;

        let errors = self.meta_applier.as_ref().unwrap().take_errors();
        if !errors.is_empty() {
            return Err(box_err!(
                "{} meta_applier meet errors: {:?}",
                self.tag,
                errors
            ));
        }

        // Take back from meta_applier.
        let mut shards = self.meta_applier.as_ref().unwrap().take_shards().unwrap();
        self.shards = mem::take(&mut shards);

        Ok(self.shards_need_flush.len())
    }

    fn check_flushed(&self, timeout: Duration) -> Result<()> {
        let is_flushed = |stats: &ShardStats| {
            stats.flushed && stats.mem_table_size == 0 && stats.mem_table_count == 1
        };

        let kv_engine = self.kv_engine.as_ref().unwrap();
        let begin = Instant::now_coarse();
        while begin.saturating_elapsed() < timeout {
            let done = kv_engine
                .get_all_shard_stats()
                .into_iter()
                .all(|stats| is_flushed(&stats));
            if done {
                return Ok(());
            }
            thread::sleep(Duration::from_millis(500));
        }

        let stats: Vec<_> = kv_engine
            .get_all_shard_stats()
            .into_iter()
            .filter(|stats| !is_flushed(stats))
            .collect();
        error!(
            "Keyspace {} wait_for_mem_table_flush timeout, stats: {:?}",
            self.keyspace_id, stats
        );
        Err(box_err!("wait_for_mem_table_flush timeout"))
    }

    async fn truncate_ts(&mut self) -> Result<usize> {
        let mut truncate_cnt = 0_usize;
        for shard_id in self.get_shards_need_flush_and_truncate() {
            truncate_cnt += self.truncate_shard_ts(shard_id).await? as usize;
        }
        Ok(truncate_cnt)
    }

    // NOTE: `sorted_backup_shards_id` & `sorted_target_regions` must be sorted by
    // start_key.
    fn align_target_regions_impl(
        sorted_backup_shards_id: &[u64],
        backup_shards: &HashMap<u64, BackupShard>,
        sorted_target_regions: Vec<RawRegion>,
    ) -> Vec<AlignedRegion> {
        let mut aligned_regions = Vec::with_capacity(sorted_target_regions.len());
        let mut idx = 0;

        // NOTE: use inner_key to check overlapping, target_regions may belongs to
        // a new keyspace.
        let key_in_shard = |key: &[u8], shard: &BackupShard| {
            let inner_key = InnerKey::from_outer_key(key);
            shard.data_bound().overlap_key(inner_key)
        };

        let get_backup_shard =
            |idx: usize| backup_shards.get(&sorted_backup_shards_id[idx]).unwrap();

        for target_region in sorted_target_regions {
            // Check overlapping with previous backup_region.
            if idx > 0 && key_in_shard(target_region.get_start_key(), get_backup_shard(idx - 1)) {
                idx -= 1;
            }
            let target_inner_begin = InnerKey::from_outer_key(target_region.get_start_key());
            let target_inner_end = InnerKey::from_outer_end_key(target_region.get_end_key());
            let mut aligned_shards_id = vec![];
            while idx < sorted_backup_shards_id.len()
                && get_backup_shard(idx).inner_start() < target_inner_end
            {
                // When retrying, some of regions may be filtered out.
                // In this scenario, we need to advance the cursor of `backup` to the first
                // region actually overlapping with target.
                //
                // An example: without this `if`, the shard `a` will be mistakenly added to
                // `backup_shards_id` when processing the target region `b`.
                //
                // Backup Shards  |-------|--a--|----|--...
                // Target Regions |-----|----|     |----b---|
                //                            ^^^^^ filtered out by preivous restore.
                if get_backup_shard(idx).inner_end() > target_inner_begin {
                    aligned_shards_id.push(sorted_backup_shards_id[idx]);
                }
                idx += 1;
            }

            if aligned_shards_id.len() > 1 {
                warn!("unaligned backup shard and region.";
                    "region" => ?target_region,
                    "aligned_shards_id" => ?aligned_shards_id,
                );
            }

            aligned_regions.push(AlignedRegion {
                target_region,
                backup_shards_id: aligned_shards_id,
            });
        }

        aligned_regions
    }

    fn align_target_regions(&self, sorted_target_regions: Vec<RawRegion>) -> Vec<AlignedRegion> {
        Self::align_target_regions_impl(&self.sorted_shards, &self.shards, sorted_target_regions)
    }

    fn count_online_snapshot_unaligned_voters(
        &self,
        aligned_regions: &[AlignedRegion],
    ) -> Result<usize> {
        if self.static_region_metas.is_empty() {
            return Ok(0);
        }
        let mut unaligned = 0;
        for aligned in aligned_regions {
            let [source_shard_id] = aligned.backup_shards_id.as_slice() else {
                return Err(box_err!(
                    "target Region {} is not aligned to exactly one online source shard",
                    aligned.target_region.id
                ));
            };
            let Some((source_region, _)) = self.static_region_metas.get(source_shard_id) else {
                return Err(box_err!(
                    "source Region metadata is missing for shard {}",
                    source_shard_id
                ));
            };
            let source_stores = source_region
                .get_peers()
                .iter()
                .filter(|peer| peer.get_role() == PeerRole::Voter)
                .map(|peer| peer.get_store_id())
                .collect::<HashSet<_>>();
            let target_stores = aligned
                .target_region
                .peers
                .iter()
                .filter(|peer| peer.get_role() == PeerRole::Voter)
                .map(|peer| peer.get_store_id())
                .collect::<HashSet<_>>();
            let (_, region_unaligned) = validate_online_snapshot_voter_stores(
                aligned.target_region.id,
                Some(target_stores),
                &source_stores,
            )?;
            unaligned += region_unaligned;
        }
        Ok(unaligned)
    }

    fn align_online_snapshot_placement(
        &mut self,
        target_regions: Vec<RawRegion>,
        timeout: Duration,
        max_unaligned_voter_peers: usize,
        placement_concurrency: usize,
        scheduling_guard: &RestoreScheduleGuardState,
        runtime: &Runtime,
    ) -> Result<()> {
        if self.static_region_metas.is_empty() {
            return Ok(());
        }
        let Some(pd_control) = self.pd_control.clone() else {
            return Err(box_err!(
                "online snapshot placement requires PD HTTP control"
            ));
        };
        let source_regions = self
            .sorted_shards
            .iter()
            .map(|shard_id| {
                self.static_region_metas
                    .get(shard_id)
                    .cloned()
                    .ok_or_else(|| {
                        box_err!("source Region metadata is missing for shard {}", shard_id)
                    })
            })
            .collect::<Result<Vec<_>>>()?;
        let plan = OnlineSnapshotPlacementPlan::new(&source_regions)?;
        let unaligned = plan.align_placement(
            target_regions,
            self.pd_client.clone(),
            pd_control,
            timeout,
            max_unaligned_voter_peers,
            placement_concurrency,
            std::slice::from_ref(scheduling_guard),
            runtime,
            self.tag(),
        )?;
        self.tolerated_err += unaligned;
        Ok(())
    }
    pub(crate) fn pause_merge_and_scheduler_during_restore(
        &self,
        key_ranges: &[(&[u8], &[u8])],
        runtime: &Runtime,
    ) -> Result<RestoreScheduleGuard> {
        let required = !self.static_region_metas.is_empty();
        let Some(pd) = self.pd_control.clone() else {
            if required {
                return Err(box_err!(
                    "online snapshot target scheduling guard requires PD HTTP control"
                ));
            }
            return Ok(RestoreScheduleGuard {
                state: RestoreScheduleGuardState {
                    healthy: Arc::new(AtomicBool::new(true)),
                    valid_until: None,
                },
                stop_tx: None,
            });
        };
        RestoreScheduleGuard::install(pd, key_ranges, required, runtime)
    }

    fn pre_split_regions(
        &self,
        config: &RestoreConfig,
        deadline: Option<RestoreDeadline>,
        runtime: &Runtime,
    ) -> Result<()> {
        let (target_keyspace_start, target_keyspace_end) =
            ApiV2::get_keyspace_range_by_id(self.target_keyspace_id);
        let scan_timeout = match deadline {
            Some(deadline) => deadline.cap_timeout(
                config.timeout_restore_snapshot.0,
                "online snapshot target Region scan",
            )?,
            None => config.timeout_restore_snapshot.0,
        };
        let target_regions = runtime.block_on(get_target_regions_with_retry(
            self.pd_client.as_ref(),
            &target_keyspace_start,
            &target_keyspace_end,
            scan_timeout,
        ))?;
        let aligned_regions = self.align_target_regions(target_regions);

        let mut coarse_split_keys =
            Vec::with_capacity(self.sorted_shards.len() / config.coarse_split_regions_factor);
        let mut split_keys = Vec::with_capacity(self.sorted_shards.len() - 1);
        for r in aligned_regions {
            for (idx, &shard_id) in r.backup_shards_id.iter().enumerate().skip(1) {
                let shard = self.get_shard(shard_id).unwrap();
                let split_key =
                    [target_keyspace_start.clone(), shard.inner_start().to_vec()].concat();

                if idx % config.coarse_split_regions_factor == 0 {
                    coarse_split_keys.push(split_key);
                } else {
                    split_keys.push(split_key);
                }
            }
        }

        // Online snapshot restore follows every split reconciliation with an
        // exact source-voter placement pass. A best-effort scatter would only
        // create unrelated PD operators that the placement pass must wait for.
        let scatter_coarse_regions = self.static_region_metas.is_empty();
        self.split_regions_for_keys(
            coarse_split_keys,
            split_keys,
            scatter_coarse_regions,
            config.timeout_split_regions.0,
            deadline,
            runtime,
        )
    }

    fn split_regions_for_keys(
        &self,
        coarse_split_keys: Vec<Vec<u8>>,
        split_keys: Vec<Vec<u8>>,
        scatter_coarse_regions: bool,
        timeout: Duration,
        deadline: Option<RestoreDeadline>,
        runtime: &Runtime,
    ) -> Result<()> {
        split_regions_for_keys(
            self.tag(),
            self.pd_client.as_ref(),
            coarse_split_keys,
            split_keys,
            scatter_coarse_regions,
            timeout,
            deadline,
            runtime,
        )
    }
    fn select_schema_file(shards: &[&BackupShard]) -> Option<(u64, i64)> {
        shards
            .iter()
            .filter(|shard| shard.meta.schema.is_valid())
            .map(|shard| (shard.meta.schema.file_id(), shard.meta.schema.file_ver()))
            .max_by_key(|&(file_id, _)| file_id)
    }

    fn schema_reader(&self) -> Arc<BackupSchemaReader> {
        Arc::new(BackupSchemaReader::new(
            self.kv_engine.as_ref().unwrap().clone(),
            self.sorted_shards.iter().map(|shard_id| {
                let shard = self.get_shard(*shard_id).unwrap();
                BackupSchemaShard {
                    id: shard.meta.id,
                    version: shard.meta.ver,
                    start_key: shard.meta.range.outer_start.clone(),
                    end_key: shard.meta.range.outer_end.clone(),
                }
            }),
        ))
    }

    fn load_source_schema_file(
        &self,
        runtime: &Runtime,
        file_id: u64,
        file_version: i64,
    ) -> Result<SchemaFile> {
        let schema_file = runtime.block_on(
            self.kv_engine
                .as_ref()
                .unwrap()
                .load_schema_file(file_id, IoType::Import),
        )?;
        if schema_file.get_keyspace_id() != self.keyspace_id
            || schema_file.get_version() != file_version
        {
            return Err(box_err!(
                "schema file {} metadata mismatch: expected keyspace {} version {}, got keyspace {} version {}",
                file_id,
                self.keyspace_id,
                file_version,
                schema_file.get_keyspace_id(),
                schema_file.get_version()
            ));
        }
        Ok(schema_file)
    }

    fn write_restored_schema_file(
        &self,
        runtime: &Runtime,
        schema_data: Vec<u8>,
    ) -> Result<SchemaFile> {
        let file_id = *self.id_allocator.alloc_id(1)?.first().unwrap();
        let schema_data = Bytes::from(schema_data);
        runtime.block_on(self.dfs.create(
            file_id,
            schema_data.clone(),
            dfs::Options::default().with_type(dfs::FileType::Schema),
        ))?;
        SchemaFile::open(Arc::new(InMemFile::new(file_id, schema_data))).map_err(|err| {
            box_err!(
                "{} open restored schema file {} failed: {}",
                self.tag(),
                file_id,
                err
            )
        })
    }

    fn rewrite_schema_file(
        &self,
        runtime: &Runtime,
        schema_file: SchemaFile,
    ) -> Result<SchemaFile> {
        info!(
            "Keyspace {} set schema file {} restore version to {}",
            self.tag(),
            schema_file.get_file_id(),
            self.truncate_ts
        );
        let new_schema_file = schema_file
            .set_restore_version_and_keyspace_id(self.truncate_ts, self.target_keyspace_id);
        let source_file_id = new_schema_file.get_file_id();
        let new_schema_file =
            self.write_restored_schema_file(runtime, new_schema_file.to_bytes())?;
        info!(
            "Keyspace {} update schema file {} restore version to {} and keyspace_id to {}, new file id {}",
            self.tag(),
            source_file_id,
            self.truncate_ts,
            self.target_keyspace_id,
            new_schema_file.get_file_id(),
        );
        Ok(new_schema_file)
    }

    fn rebuild_schema_file(
        &self,
        runtime: &Runtime,
        reader: Arc<BackupSchemaReader>,
        effective_schema_version: i64,
    ) -> Result<Option<SchemaFile>> {
        let sync_result = runtime
            .block_on(schema::sync_schema(
                reader.clone(),
                reader,
                self.keyspace_id,
                None,
                TimeStamp::new(self.truncate_ts),
            ))
            .map_err(|err| -> Error {
                box_err!(
                    "{} rebuild schema at {} failed: {}",
                    self.tag(),
                    self.truncate_ts,
                    err
                )
            })?;
        if !sync_result.is_full_sync || sync_result.schema_version != effective_schema_version {
            return Err(box_err!(
                "{} inconsistent schema at {}: expected version {}, full sync returned version {} (full_sync={})",
                self.tag(),
                self.truncate_ts,
                effective_schema_version,
                sync_result.schema_version,
                sync_result.is_full_sync
            ));
        }

        let schemas = sync_result
            .table_infos
            .iter()
            .map(schema_from_table_info)
            .collect::<std::result::Result<Vec<_>, _>>()
            .map_err(|err| -> Error {
                box_err!(
                    "{} convert schema at {} failed: {}",
                    self.tag(),
                    self.truncate_ts,
                    err
                )
            })?;
        if schemas.is_empty() {
            return Ok(None);
        }
        self.write_restored_schema_file(
            runtime,
            build_schema_file(
                self.target_keyspace_id,
                effective_schema_version,
                schemas,
                self.truncate_ts,
            ),
        )
        .map(Some)
    }

    fn prepare_schema_at_truncate_ts(&mut self, runtime: &Runtime) -> Result<()> {
        if self.restored_schema.is_some() {
            return Ok(());
        }
        let source_schema = {
            let shards = self
                .sorted_shards
                .iter()
                .map(|shard_id| self.get_shard(*shard_id).unwrap())
                .collect::<Vec<_>>();
            Self::select_schema_file(&shards)
        };
        let reader = self.schema_reader();
        let effective_schema_version = runtime
            .block_on(schema::get_effective_schema_version(
                reader.clone(),
                self.keyspace_id,
                TimeStamp::new(self.truncate_ts),
            ))
            .map_err(|err| -> Error {
                box_err!(
                    "{} read schema version at {} failed: {}",
                    self.tag(),
                    self.truncate_ts,
                    err
                )
            })?;

        let target_file = match source_schema {
            Some((file_id, file_version)) if file_version == effective_schema_version => {
                let source_file = self.load_source_schema_file(runtime, file_id, file_version)?;
                Some(self.rewrite_schema_file(runtime, source_file)?)
            }
            _ if effective_schema_version == 0 => None,
            _ => self.rebuild_schema_file(runtime, reader, effective_schema_version)?,
        };
        info!(
            "{} prepared schema at truncate timestamp", self.tag();
            "truncate_ts" => self.truncate_ts,
            "schema_version" => effective_schema_version,
            "schema_file_id" => target_file.as_ref().map(SchemaFile::get_file_id).unwrap_or_default(),
            "source_schema" => ?source_schema,
        );
        self.restored_schema = Some(RestoredSchema {
            version: effective_schema_version,
            target_file,
        });
        Ok(())
    }

    fn expected_columnar_table_ids(
        restored_schema: &RestoredSchema,
        shard: &BackupShard,
    ) -> Vec<i64> {
        restored_schema
            .target_file
            .as_ref()
            .map(|schema_file| {
                schema_file.overlap_columnar_tables(shard.inner_start(), shard.inner_end())
            })
            .unwrap_or_default()
    }

    fn fts_metadata_matches_schema(
        restored_schema: &RestoredSchema,
        source: &BackupShard,
        expected_columnar_table_ids: &[i64],
        truncate_ts: u64,
    ) -> bool {
        let Some(schema_file) = restored_schema.target_file.as_ref() else {
            return source.meta.fts_indexes.is_empty()
                && source.meta.fts_index_metas.is_empty()
                && source.meta.fts_l0_files.is_empty()
                && source.meta.fts_l1_files.is_empty()
                && source.meta.fts_l2_files.is_empty()
                && source.meta.fts_pending_l0_ids.is_empty()
                && source.meta.fts_l0_snap_version.into_inner() == 0;
        };
        let expected_columnar_table_ids = expected_columnar_table_ids
            .iter()
            .copied()
            .collect::<HashSet<_>>();
        let expected_indexes = schema_file
            .iter_tables()
            .filter(|(table_id, _)| expected_columnar_table_ids.contains(table_id))
            .flat_map(|(&table_id, schema)| {
                schema
                    .fulltext_indexes
                    .iter()
                    .map(move |index| ((table_id, index.get_index_id()), index))
            })
            .collect::<HashMap<_, _>>();
        let expected_index_ids = expected_indexes.keys().copied().collect::<HashSet<_>>();
        if source.meta.fts_indexes.len() != expected_index_ids.len()
            || source
                .meta
                .fts_indexes
                .iter()
                .any(|index_id| !expected_index_ids.contains(index_id))
        {
            return false;
        }

        let actual_meta_ids = source
            .meta
            .fts_index_metas
            .iter()
            .flat_map(|(&table_id, indexes)| {
                indexes.keys().map(move |&index_id| (table_id, index_id))
            })
            .collect::<HashSet<_>>();
        if actual_meta_ids != expected_index_ids {
            return false;
        }
        for (&(table_id, index_id), index_def) in &expected_indexes {
            let tracked_meta = &source.meta.fts_index_metas[&table_id][&index_id];
            let expected_index_fields = kvengine::table::fts::tracked_index_fields_for_layout(
                index_def,
                tracked_meta.get_field_layout(),
            );
            let actual_index_fields =
                kvengine::table::fts::tracked_index_fields_from_meta_or_index_def(
                    tracked_meta,
                    index_def,
                );
            if actual_index_fields != expected_index_fields {
                return false;
            }
        }

        let fts_file_ids = source
            .meta
            .fts_l0_files
            .iter()
            .map(|file| file.get_id())
            .chain(source.meta.fts_l1_files.iter().map(|file| file.get_id()))
            .chain(source.meta.fts_l2_files.iter().map(|file| file.get_id()))
            .collect::<Vec<_>>();
        if fts_file_ids.iter().any(|&file_id| file_id > truncate_ts) {
            return false;
        }
        // Fileless watermarks cannot be placed on either side of truncate_ts.
        // Rebuilding an empty index is cheap and avoids treating such metadata as
        // proof of coverage.
        if !expected_indexes.is_empty() && fts_file_ids.is_empty() {
            return false;
        }
        if expected_indexes.is_empty()
            && (!fts_file_ids.is_empty()
                || !source.meta.fts_pending_l0_ids.is_empty()
                || source.meta.fts_l0_snap_version.into_inner() != 0)
        {
            return false;
        }

        let columnar_l0_ids = source
            .meta
            .all_files()
            .iter()
            .filter(|(_, file)| file.is_columnar_file() && file.get_level() == 0)
            .map(|(&file_id, _)| file_id)
            .collect::<HashSet<_>>();
        source
            .meta
            .fts_pending_l0_ids
            .iter()
            .all(|file_id| columnar_l0_ids.contains(file_id))
    }

    fn vector_index_matches_schema(
        restored_schema: &RestoredSchema,
        expected_columnar_table_ids: &HashSet<i64>,
        index: &VectorIndex,
        truncate_ts: u64,
    ) -> bool {
        if index.get_files().is_empty() {
            return false;
        }
        if index
            .get_files()
            .iter()
            .any(|file| file.get_id() > truncate_ts)
        {
            return false;
        }
        if !expected_columnar_table_ids.contains(&index.get_table_id()) {
            return false;
        }
        restored_schema
            .target_file
            .as_ref()
            .and_then(|schema_file| schema_file.get_table(index.get_table_id()))
            .is_some_and(|table_schema| {
                table_schema.vector_indexes.iter().any(|definition| {
                    definition.index_id == index.get_index_id()
                        && definition.col_id == index.get_col_id()
                })
            })
    }

    fn derive_metadata_retention(
        restored_schema: &RestoredSchema,
        shard: &BackupShard,
        truncate_ts: u64,
    ) -> DerivedMetadataRetention {
        let mut expected_columnar_table_ids =
            Self::expected_columnar_table_ids(restored_schema, shard);
        expected_columnar_table_ids.sort_unstable();
        expected_columnar_table_ids.dedup();
        let mut source_columnar_table_ids = shard.meta.columnar_table_ids.clone();
        source_columnar_table_ids.sort_unstable();
        source_columnar_table_ids.dedup();
        let columnar_file_ids = shard
            .meta
            .all_files()
            .iter()
            .filter(|(_, file)| file.is_columnar_file())
            .map(|(&file_id, _)| file_id)
            .collect::<Vec<_>>();

        // A physical file ID is allocated after the MVCC timestamps represented
        // by that file. Therefore ID <= truncate_ts proves that the file cannot
        // contain post-truncate data; larger IDs are conservatively rebuilt.
        let retain_columnar = source_columnar_table_ids == expected_columnar_table_ids
            && columnar_file_ids
                .iter()
                .all(|&file_id| file_id <= truncate_ts)
            && (expected_columnar_table_ids.is_empty()
                || (shard.meta.schema.is_valid() && !columnar_file_ids.is_empty()));
        let mut retention = DerivedMetadataRetention {
            retain_columnar,
            ..Default::default()
        };
        if retain_columnar {
            retention.retain_fts = Self::fts_metadata_matches_schema(
                restored_schema,
                shard,
                &expected_columnar_table_ids,
                truncate_ts,
            );
            let expected_columnar_table_ids_set = expected_columnar_table_ids
                .iter()
                .copied()
                .collect::<HashSet<_>>();
            let mut vector_index_reusable = HashMap::default();
            for index in &shard.meta.vector_indexes {
                let index_id = (
                    index.get_table_id(),
                    index.get_index_id(),
                    index.get_col_id(),
                );
                let reusable = Self::vector_index_matches_schema(
                    restored_schema,
                    &expected_columnar_table_ids_set,
                    index,
                    truncate_ts,
                );
                vector_index_reusable
                    .entry(index_id)
                    // Multiple entries for one logical index make its coverage
                    // ambiguous. Rebuild it instead of merging duplicate state.
                    .and_modify(|all_reusable| *all_reusable = false)
                    .or_insert(reusable);
            }
            retention.retained_vector_indexes.extend(
                vector_index_reusable
                    .into_iter()
                    .filter_map(|(index_id, reusable)| reusable.then_some(index_id)),
            );
        }
        retention
    }

    fn prepare_derived_metadata_retention(&mut self) {
        if self.derived_metadata_retention.is_some() {
            return;
        }
        let restored_schema = self
            .restored_schema
            .as_ref()
            .expect("schema must be prepared before derived metadata retention");
        let mut retention_by_shard = HashMap::with_capacity(self.shards.len());
        for (&shard_id, shard) in &self.shards {
            let retention =
                Self::derive_metadata_retention(restored_schema, shard, self.truncate_ts);
            info!(
                "{} prepared derived metadata retention", self.tag();
                "shard_id" => shard_id,
                "retain_columnar" => retention.retain_columnar,
                "retain_fts" => retention.retain_fts,
                "retained_vector_indexes" => retention.retained_vector_indexes.len(),
                "truncate_ts" => self.truncate_ts,
            );
            retention_by_shard.insert(shard_id, retention);
        }
        self.derived_metadata_retention = Some(retention_by_shard);
    }

    /// Copies FTS state already validated against schema-at-T and source file
    /// IDs, filtering physical files and pending columnar L0s to `target`.
    fn copy_retained_fts(source: &ShardMeta, target: &mut ShardMeta) -> usize {
        target.fts_indexes = source
            .fts_indexes
            .iter()
            .filter(|(table_id, _)| target.columnar_table_ids.contains(table_id))
            .copied()
            .collect();
        target.fts_index_metas = source
            .fts_index_metas
            .iter()
            .filter(|(table_id, _)| target.columnar_table_ids.contains(table_id))
            .map(|(&table_id, indexes)| (table_id, indexes.clone()))
            .collect();
        target.fts_l0_snap_version = source.fts_l0_snap_version;
        target.fts_l0_files = source
            .fts_l0_files
            .iter()
            .filter(|file| {
                target.overlap_bound(kvengine::FileMeta::from_fts_l0_file(file).data_bound())
            })
            .cloned()
            .collect();
        target.fts_l1_files = source
            .fts_l1_files
            .iter()
            .filter(|file| {
                target.overlap_bound(kvengine::FileMeta::from_fts_l1_file(file).data_bound())
            })
            .cloned()
            .collect();
        target.fts_l2_files = source
            .fts_l2_files
            .iter()
            .filter(|file| {
                target.overlap_bound(kvengine::FileMeta::from_fts_l2_file(file).data_bound())
            })
            .cloned()
            .collect();

        let file_count =
            target.fts_l0_files.len() + target.fts_l1_files.len() + target.fts_l2_files.len();
        if target.fts_indexes.is_empty() || file_count == 0 {
            // A split may filter all source FTS files out of this target range.
            // Do not carry a fileless watermark that could suppress rebuilding
            // historical columnar data.
            target.fts_indexes.clear();
            target.fts_index_metas.clear();
            target.fts_l0_snap_version = Default::default();
            target.fts_l0_files.clear();
            target.fts_l1_files.clear();
            target.fts_l2_files.clear();
            target.fts_pending_l0_ids.clear();
            return 0;
        }

        let columnar_l0_ids = target
            .all_files()
            .iter()
            .filter(|(_, file)| file.is_columnar_file() && file.get_level() == 0)
            .map(|(&file_id, _)| file_id)
            .collect::<HashSet<_>>();
        target.fts_pending_l0_ids = source
            .fts_pending_l0_ids
            .iter()
            .filter(|file_id| columnar_l0_ids.contains(file_id))
            .copied()
            .collect();
        file_count
    }

    /// Copies validated vector indexes and range-overlapping files to `target`.
    fn copy_retained_vector_indexes(
        source: &ShardMeta,
        target: &mut ShardMeta,
        retained_indexes: &HashSet<(i64, i64, i64)>,
    ) -> usize {
        let mut file_count = 0;
        for source_index in &source.vector_indexes {
            let index_id = (
                source_index.get_table_id(),
                source_index.get_index_id(),
                source_index.get_col_id(),
            );
            if !target
                .columnar_table_ids
                .contains(&source_index.get_table_id())
                || !retained_indexes.contains(&index_id)
            {
                continue;
            }
            let files = source_index
                .get_files()
                .iter()
                .filter(|file| {
                    target.overlap_bound(
                        kvengine::FileMeta::from_vector_index_file(file).data_bound(),
                    )
                })
                .cloned()
                .collect::<Vec<_>>();
            if files.is_empty() {
                continue;
            }
            for file in &files {
                target.add_file(
                    file.get_id(),
                    kvengine::FileMeta::from_vector_index_file(file),
                );
            }
            file_count += files.len();
            let mut index = source_index.clone();
            index.set_files(files.into());
            target.vector_indexes.push(index);
        }
        file_count
    }

    fn gather_all_tables(
        &self,
        aligned_regions: &[AlignedRegion],
    ) -> Result<(Vec<ShardMeta>, usize /* number of sstables */)> {
        let restored_schema = self
            .restored_schema
            .as_ref()
            .expect("schema must be prepared before gathering target tables");
        let retention_by_shard = self
            .derived_metadata_retention
            .as_ref()
            .expect("derived metadata retention must be prepared before gathering target tables");
        let mut sstables_cnt = 0;
        let mut target_shards = Vec::with_capacity(aligned_regions.len());
        for region in aligned_regions {
            let shards = region
                .backup_shards_id
                .iter()
                .map(|&shard_id| self.get_shard(shard_id).unwrap())
                .collect::<Vec<_>>();

            let mut meta = ShardMeta::default();
            meta.id = region.target_region.id;
            meta.ver = region.target_region.epoch.version;
            meta.range = ShardRange::new(
                region.target_region.get_start_key(),
                region.target_region.get_end_key(),
            );
            meta.inner_key_off = self.target_inner_key_off();
            let retention = if shards.len() == 1 {
                retention_by_shard
                    .get(&shards[0].region_id)
                    .cloned()
                    .unwrap_or_default()
            } else {
                DerivedMetadataRetention::default()
            };
            let target_columnar_table_ids = restored_schema
                .target_file
                .as_ref()
                .map(|schema_file| {
                    schema_file
                        .overlap_columnar_tables(meta.range.inner_start(), meta.range.inner_end())
                })
                .unwrap_or_default();
            let retain_target_columnar =
                retention.retain_columnar && !target_columnar_table_ids.is_empty();

            let shards_change_inner_key_off = shards
                .iter()
                .filter_map(|shard| {
                    (shard.inner_key_off() != meta.inner_key_off).then_some(shard.region_id)
                })
                .collect::<Vec<_>>();
            if !shards_change_inner_key_off.is_empty() {
                info!(
                    "{} shards change inner key offset after restore", self.tag();
                    "shards" => ?shards_change_inner_key_off, "target_inner_key_off" => meta.inner_key_off
                );
            }

            let encryption_meta = self.get_restore_encryption_meta();
            let encryption_key = self.get_keyspace_exported_encryption_key();
            if let Some(encryption_meta) = encryption_meta {
                meta.set_property(ENCRYPTION_META_KEY, encryption_meta.as_slice());
            } else if let Some(encryption_key) = encryption_key {
                meta.set_property(ENCRYPTION_KEY, encryption_key.as_slice());
            }
            let mut properties_helper =
                kvengine::util::PropertiesHelper::new_from_shard_meta(&meta);
            properties_helper.set_rewrite_range_prefix(!self.is_inplace_restore());
            for shard in &shards {
                let mut all_l0_ssts: HashSet<u64> = HashSet::default();
                for (&file_id, file_meta) in shard.meta.all_files() {
                    if file_meta.is_vector_index_file()
                        || file_meta.is_fts_file()
                        || (file_meta.is_columnar_file() && !retain_target_columnar)
                    {
                        continue;
                    }
                    if meta.overlap_bound(file_meta.data_bound()) {
                        meta.add_file(file_id, file_meta.clone());
                        sstables_cnt += 1;
                        if file_meta.get_level() == 0 && file_meta.is_sst_file() {
                            all_l0_ssts.insert(file_id);
                        }
                    }
                }
                // Only add the filtered l0s from ShardMeta.unconverted_l0s
                // to meta.unconverted_l0s
                meta.unconverted_l0s.extend(
                    shard
                        .meta
                        .unconverted_l0s
                        .iter()
                        .filter(|id| all_l0_ssts.contains(id)),
                );
                // Use `base_version` as table version, and `data_sequence` is 0.
                // And they will be adjusted at server side in `restore_shard` procedure.
                meta.base_version = cmp::max(meta.base_version, shard.table_version());
                properties_helper.merge_shard_meta(&shard.meta);
            }
            if retain_target_columnar {
                let source = shards[0];
                meta.columnar_table_ids = target_columnar_table_ids;
                meta.columnar_l2_snap_version = meta
                    .all_files()
                    .values()
                    .filter(|file| file.is_columnar_file() && file.get_level() == 2)
                    .map(|file| file.get_snap_version())
                    .max()
                    .unwrap_or_default();

                if retention.retain_fts {
                    sstables_cnt += Self::copy_retained_fts(&source.meta, &mut meta);
                }
                sstables_cnt += Self::copy_retained_vector_indexes(
                    &source.meta,
                    &mut meta,
                    &retention.retained_vector_indexes,
                );
            } else {
                meta.clear_columnar_related_meta();
                if !retention.retain_columnar
                    && shards.iter().any(|shard| {
                        !shard.meta.columnar_table_ids.is_empty()
                            || shard
                                .meta
                                .all_files()
                                .values()
                                .any(|file| file.is_columnar_file())
                    })
                {
                    warn!(
                        "{} reset derived metadata because columnar state is unsafe to reuse",
                        self.tag();
                        "target_region_id" => meta.id,
                        "source_shards" => ?region.backup_shards_id,
                        "truncate_ts" => self.truncate_ts,
                    );
                }
            }
            meta.schema.update_by_restore(
                restored_schema.target_file_id(),
                restored_schema.version,
                self.truncate_ts,
            );

            properties_helper.build_to_shard_meta(&mut meta);
            target_shards.push(meta);
        }
        Ok((target_shards, sstables_cnt))
    }

    async fn truncate_shard_ts(&mut self, shard_id: u64) -> Result<bool /* has_truncate_ts */> {
        let shard_meta = &self.get_shard(shard_id).unwrap().meta;
        let kvengine = self.kv_engine.as_ref().unwrap();
        let shard = kvengine
            .get_shard_with_ver(shard_meta.id, shard_meta.ver)
            .expect("Could not find shard with meta");
        let keyspace_id = self.keyspace_id;
        let mut res_cs = kvengine
            .truncate_with_ts(&shard, self.truncate_ts)
            .await?
            .unwrap();
        if res_cs.has_truncate_ts() && !ShardMeta::is_empty_table_change(res_cs.get_truncate_ts()) {
            let shard = self.get_shard_mut(shard_id).unwrap();
            res_cs.set_sequence(shard.meta.seq + 1);
            debug!(
                "Keyspace {} before truncate ts: {:?}, table change: {:?}",
                keyspace_id,
                shard,
                res_cs.get_truncate_ts()
            );
            shard.meta.apply_change_set(&res_cs);
            debug!(
                "Keyspace {} after apply_truncate_ts: {:?}",
                keyspace_id, shard
            );
            Ok(true)
        } else {
            Ok(false)
        }
    }

    async fn trim_over_bound(&mut self, shard_id: u64) -> Result<bool /* has_trim_over_bound */> {
        let shard = self.get_shard(shard_id).unwrap();
        let kvengine = self.kv_engine.as_ref().unwrap();

        let mut res_cs = kvengine.trim_over_bound_by_meta(&shard.meta).await?;
        let keyspace_id = self.keyspace_id;
        if res_cs.has_trim_over_bound() {
            let shard = self.get_shard_mut(shard_id).unwrap();
            res_cs.set_sequence(shard.meta.seq + 1);
            debug!(
                "Keyspace {} before trim_over_bound: {:?}, table change: {:?}",
                keyspace_id,
                shard,
                res_cs.get_trim_over_bound()
            );
            shard.meta.apply_change_set(&res_cs);
            debug!(
                "Keyspace {} after apply_trim_over_bound: {:?}",
                keyspace_id, shard
            );
            Ok(true)
        } else {
            Ok(false)
        }
    }

    async fn trim_over_bound_shards(
        &mut self,
        aligned_regions: &[AlignedRegion],
    ) -> Result<usize /* number of trimmed shards */> {
        let mut trim_shards_cnt = 0_usize;
        let mut unique_shards = HashSet::default();
        for region in aligned_regions
            .iter()
            .filter(|x| x.backup_shards_id.len() > 1)
        {
            for &shard_id in &region.backup_shards_id {
                if unique_shards.contains(&shard_id) {
                    continue;
                }
                unique_shards.insert(shard_id);

                // TODO: Run in parallel.
                if self.trim_over_bound(shard_id).await? {
                    trim_shards_cnt += 1;
                }
            }
        }

        Ok(trim_shards_cnt)
    }

    pub fn generate_snapshots(&mut self, target_shards: Vec<ShardMeta>) -> Vec<pb::ChangeSet> {
        target_shards
            .into_iter()
            .map(|meta| {
                let mut cs = meta.to_change_set();
                let snap = cs.take_snapshot();
                cs.set_restore_shard(snap);
                info!("{} generate_snapshot cs: {:?}", self.tag, cs);
                cs
            })
            .collect()
    }

    fn target_inner_key_off(&self) -> usize {
        if ApiV2::is_default_keyspace(self.target_keyspace_id) {
            0
        } else {
            KEYSPACE_PREFIX_LEN
        }
    }

    fn is_inplace_restore(&self) -> bool {
        self.keyspace_id == self.target_keyspace_id
    }

    fn get_keyspace_exported_encryption_key(&self) -> Option<Vec<u8>> {
        self.get_sorted_shard(0)
            .meta
            .get_property(ENCRYPTION_KEY)
            .map(|bytes| bytes.to_vec())
    }

    fn get_restore_encryption_meta(&self) -> Option<Vec<u8>> {
        self.restore_encryption_meta.clone().or_else(|| {
            self.get_sorted_shard(0)
                .meta
                .get_property(ENCRYPTION_META_KEY)
                .map(|bytes| bytes.to_vec())
        })
    }

    fn prepare_restore_encryption_meta(&mut self, runtime: &Runtime) -> Result<()> {
        if self.restore_encryption_meta.is_some() {
            return Ok(());
        }

        let target_cluster_id = self.pd_client.get_cluster_id()?;
        if self.source_cluster_id == target_cluster_id && self.is_inplace_restore() {
            return Ok(());
        }

        let source_meta = self
            .get_restore_encryption_meta()
            .map(|bytes| protobuf::parse_from_bytes(bytes.as_slice()))
            .transpose()
            .map_err(|err| {
                Error::Other(box_err!(
                    "failed to parse exported ENCRYPTION_META_KEY: {}",
                    err
                ))
            })?;
        let Some(source_meta) = source_meta else {
            return Ok(());
        };
        self.restore_encryption_meta = Some(self.generate_restore_encryption_meta(
            runtime,
            &source_meta,
            target_cluster_id,
        )?);
        Ok(())
    }

    fn generate_restore_encryption_meta(
        &self,
        runtime: &Runtime,
        source_meta: &pb::EncryptionMeta,
        target_cluster_id: u64,
    ) -> Result<Vec<u8>> {
        let new_file_id = runtime
            .block_on(self.pd_client.batch_get_tso(1))?
            .into_inner();
        let encryption_meta = regenerate_encryption_meta(source_meta, new_file_id)?;
        let encryption_meta = encryption_meta.write_to_bytes().map_err(|err| {
            Error::Other(box_err!(
                "failed to serialize regenerated ENCRYPTION_META_KEY: {}",
                err
            ))
        })?;
        let opts = dfs::Options::default().with_type(dfs::FileType::EncryptionDict);
        runtime.block_on(self.dfs.create(
            new_file_id,
            Bytes::from(encryption_meta.clone()),
            opts,
        ))?;
        info!(
            "regenerated source encryption meta for restore";
            "source_cluster_id" => self.source_cluster_id,
            "target_cluster_id" => target_cluster_id,
            "source_keyspace_id" => self.keyspace_id,
            "target_keyspace_id" => self.target_keyspace_id,
            "file_id" => new_file_id,
        );
        Ok(encryption_meta)
    }

    fn tolerated_err(&self) -> usize {
        self.tolerated_err
    }

    // Should be called after `load_shards`.
    fn collect_txn_chunks_in_wal(&mut self) -> Result<()> {
        let mut total_chunk_ids = Vec::new();

        let master_key = self.get_master_key();
        let encryption_manager = Arc::new(EncryptionKeyManager::new());
        let mut decryption_buf = Vec::new();

        for store_id in self.get_all_stores_id() {
            let rf_engine = self.raft_engines.get(&store_id).unwrap();
            for &shard_id in self.store_shards.get(&store_id).unwrap() {
                let shard = self.get_shard(shard_id).unwrap();

                let meta = self.raw_metas.get(&shard_id).unwrap();
                let snap = meta.get_snapshot();
                let encryption_key = encryption_key_from_shard_properties(
                    snap.get_properties(),
                    encryption_manager.clone(),
                    &master_key,
                    Some(shard.meta.range.keyspace_id),
                )
                .map_err(|err| -> Error {
                    box_err!("failed to resolve encryption key for restore: {}", err)
                })?;
                let tag = ShardTag::new(store_id, IdVer::from_change_set(meta));
                let peer_tag =
                    PeerTag::new(store_id, RegionIdVer::new(meta.shard_id, meta.shard_ver));

                let applied_index = snap.get_data_sequence();
                let commit_index = shard.raft_state.get_commit();
                let low_idx = applied_index + 1;
                let high_idx = commit_index + 1;
                debug_assert!(
                    low_idx <= high_idx,
                    "invalid entry index, low_idx {}, high_idx {}",
                    low_idx,
                    high_idx
                );
                if low_idx >= high_idx {
                    continue;
                }

                let mut entries = Vec::with_capacity((high_idx.saturating_sub(low_idx)) as usize);
                rf_engine
                    .fetch_raft_entries_to(
                        shard.peer_id,
                        low_idx,
                        high_idx,
                        None,
                        &mut entries,
                        None,
                    )
                    .map_err(|e| -> Error {
                        box_err!(
                            "{} entries unavailable err: {:?}, low: {}, high {}",
                            tag,
                            e,
                            low_idx,
                            high_idx
                        )
                    })?;

                let mut chunk_ids = Vec::new();
                for e in entries {
                    if e.data.is_empty() || e.entry_type != eraftpb::EntryType::EntryNormal {
                        continue;
                    }

                    let req =
                        parse_raft_cmd(peer_tag, &e, encryption_key.as_ref(), &mut decryption_buf);
                    if req.get_header().get_region_epoch().version != meta.shard_ver {
                        continue;
                    }

                    if let Some(custom) = rlog::get_custom_log(&req) {
                        if custom.is_txn_file_ref() {
                            let mut txn_file_ref = box_try!(custom.get_txn_file_ref());
                            chunk_ids.extend(txn_file_ref.take_chunk_ids());
                        }
                    }
                }

                if !chunk_ids.is_empty() {
                    info!("{} collect txn chunks in wal [{}, {})", tag, low_idx, high_idx;
                        "chunk_ids" => ?chunk_ids);
                    total_chunk_ids.extend(chunk_ids);
                }
            }
        }

        total_chunk_ids.sort();
        total_chunk_ids.dedup();
        self.txn_chunk_ids_in_wal = Some(total_chunk_ids);
        Ok(())
    }
}

struct MetaApplier {
    keyspace_id: u32,
    engine: kvengine::Engine,
    encryption_key: Option<EncryptionKey>,
    shards: RwLock<Option<HashMap<u64, BackupShard>>>,
    store_rx: mpsc::Receiver<StoreMsg>,
    errors: Mutex<Vec<Error>>,
}

impl Drop for MetaApplier {
    fn drop(&mut self) {
        info!(
            "Meta applier is dropped, engine {}",
            self.engine.get_engine_id()
        );
    }
}

impl MetaApplier {
    fn new(
        keyspace_id: u32,
        engine: kvengine::Engine,
        encryption_key: Option<EncryptionKey>,
        shards: HashMap<u64, BackupShard>,
        store_rx: mpsc::Receiver<StoreMsg>,
    ) -> Self {
        Self {
            keyspace_id,
            engine,
            encryption_key,
            shards: RwLock::new(Some(shards)),
            store_rx,
            errors: Default::default(),
        }
    }

    pub fn take_shards(&self) -> Option<HashMap<u64, BackupShard>> {
        self.shards.wl().take()
    }

    pub fn take_errors(&self) -> Vec<Error> {
        let mut errors = self.errors.lock().unwrap();
        mem::take(&mut *errors)
    }

    fn run(&self) {
        'outer: loop {
            let msg = match self.store_rx.recv() {
                Ok(msg) => msg,
                Err(err) => {
                    error!(
                        "Keyspace {} applier recv task error: {:?}",
                        self.keyspace_id, err
                    );
                    return;
                }
            };
            match msg {
                StoreMsg::GenerateEngineChangeSet(mut cs, _) => {
                    let tag =
                        ShardTag::new(self.engine.get_engine_id(), IdVer::from_change_set(&cs));
                    let engine_shard = self.engine.get_shard(cs.shard_id).unwrap();
                    let seq = cmp::max(
                        engine_shard.get_write_sequence(),
                        engine_shard.get_meta_sequence(),
                    ) + 1;
                    cs.set_sequence(seq);
                    debug!(
                        "Keyspace {} shard {} apply change set: {:?}",
                        self.keyspace_id, tag, cs
                    );

                    {
                        let mut shards = self.shards.wl();
                        if let Some(shards) = shards.as_mut() {
                            let shard = shards.get_mut(&cs.shard_id).unwrap();
                            shard.meta.apply_change_set(&cs);
                        } else {
                            if cs.has_compaction() {
                                debug!(
                                    "Keyspace {} MetaApplier: ignore changeset: {:?}",
                                    self.keyspace_id, cs
                                );
                                self.engine.meta_committed(&cs, true);
                                continue 'outer;
                            }
                            panic!(
                                "MetaApplier: changeset is lost due to shards map had be taken: {:?}",
                                cs
                            );
                        }
                    }

                    self.engine.meta_committed(&cs, false);
                    match self
                        .engine
                        .prepare_change_set(
                            cs,
                            PrepareOpts {
                                encryption_state: self.encryption_key.to_encryption_state(),
                                io_type: IoType::Import,
                                ..Default::default()
                            },
                        )
                        .and_then(|cs| self.engine.apply_change_set(&cs))
                    {
                        Ok(()) => debug!(
                            "Keyspace {} shard {} MetaApplier apply change set successfully",
                            self.keyspace_id, tag
                        ),
                        Err(err) => {
                            let msg = format!(
                                "Keyspace {} shard {} MetaApplier apply change set failed: {:?}",
                                self.keyspace_id, tag, err
                            );
                            error!("{}", msg);
                            self.errors
                                .lock()
                                .unwrap()
                                .push(Error::Other(box_err!("{}", msg)));
                        }
                    }
                }
                StoreMsg::Stop => {
                    info!(
                        "Engine {} restore keyspace {} meta applier receive stop msg and stop now",
                        self.engine.get_engine_id(),
                        self.keyspace_id
                    );
                    break 'outer;
                }
                _ => {
                    error!("unexpected msg");
                }
            }
        }
    }
}

pub struct MetaIterator {
    store_id: u64,
    shards: Vec<u64>,
    raw_metas: HashMap<u64, pb::ChangeSet>,
}

impl MetaIterator {
    pub fn new(store_id: u64, shards: Vec<u64>, raw_metas: HashMap<u64, pb::ChangeSet>) -> Self {
        Self {
            store_id,
            shards,
            raw_metas,
        }
    }
}

impl kvengine::MetaIterator for MetaIterator {
    fn iterate<F>(&mut self, mut f: F) -> kvengine::Result<()>
    where
        F: FnMut(kvenginepb::ChangeSet),
    {
        for shard_id in &self.shards {
            f(self.raw_metas.get(shard_id).unwrap().to_owned());
        }
        Ok(())
    }

    fn engine_id(&self) -> u64 {
        self.store_id
    }
}

/// Scans a bounded key range and requires the returned Regions to form an
/// exact, contiguous cover of that range.
pub(crate) async fn scan_regions(
    pd_client: &dyn PdClient,
    start_key: &[u8],
    end_key: &[u8],
) -> Result<Vec<pdpb::Region>> {
    let encoded_start_key = Key::from_raw(start_key).into_encoded();
    let encoded_end_key = Key::from_raw(end_key).into_encoded();
    let mut regions = Vec::new();
    let mut next_key = encoded_start_key;
    // PD uses an empty Region start key as the negative-infinity sentinel,
    // while a scan request encodes an empty raw start into a non-empty
    // memcomparable key.
    let mut expected_region_start = if start_key.is_empty() {
        Vec::new()
    } else {
        next_key.clone()
    };
    while next_key < encoded_end_key && !next_key.is_empty() {
        let scanned_regions = pd_client
            .scan_regions(
                next_key.clone(),
                encoded_end_key.clone(),
                SCAN_REGION_BATCH_SIZE,
            )
            .await?;
        if scanned_regions.is_empty() {
            break;
        }
        for scanned_region in scanned_regions {
            let region = scanned_region.get_region();
            if region.get_start_key() != expected_region_start.as_slice() {
                return Err(box_err!(
                    "Region {} leaves a gap or overlap in Region scan",
                    region.get_id()
                ));
            }
            let region_end = region.get_end_key();
            if region_end.is_empty() || region_end <= next_key.as_slice() {
                return Err(box_err!(
                    "Region {} does not advance Region scan",
                    region.get_id()
                ));
            }
            next_key = region_end.to_vec();
            expected_region_start.clone_from(&next_key);
            regions.push(scanned_region);
        }
    }
    if !regions.is_empty() && next_key != encoded_end_key {
        return Err(box_err!(
            "Regions do not align with the requested range boundaries"
        ));
    }
    Ok(regions)
}

async fn get_target_regions(
    pd_client: &dyn PdClient,
    start_key: &[u8],
    end_key: &[u8],
) -> Result<Vec<RawRegion>> {
    let scanned_regions = scan_regions(pd_client, start_key, end_key).await?;
    let regions = scanned_regions
        .into_iter()
        .map(|mut scanned_region| scanned_region.take_region().into())
        .collect::<Vec<_>>();
    debug!("target_regions: {:?}", regions);
    if regions.is_empty() {
        return Err(box_err!("no region"));
    }
    Ok(regions)
}

async fn get_target_regions_with_retry(
    pd_client: &dyn PdClient,
    start_key: &[u8],
    end_key: &[u8],
    timeout: Duration,
) -> Result<Vec<RawRegion>> {
    try_wait_result_async(
        || Box::pin(get_target_regions(pd_client, start_key, end_key)),
        timeout,
        || Duration::from_millis(500),
    )
    .await
}

/// Verifies that every failed online RestoreShard range has only been refined
/// by splits. A merge necessarily produces a current Region that crosses one
/// of the failed request's old boundaries. Split descendants inherit the
/// parent's conf version, so a different value is an independent membership
/// change and is terminal.
fn validate_restore_retry_topology(
    retry_regions: &[RestoreRegionFence],
    current_regions: &[RawRegion],
) -> Result<()> {
    for expected in retry_regions
        .iter()
        .filter(|expected| expected.conf_ver.is_some())
    {
        let descendants = current_regions
            .iter()
            .filter(|region| region.raw_end > expected.start && region.raw_start < expected.end)
            .collect::<Vec<_>>();
        let Some(first) = descendants.first() else {
            return Err(box_err!(
                "online RestoreShard Region {} range disappeared while retrying",
                expected.region_id
            ));
        };
        let mut cursor = expected.start.clone();
        for region in &descendants {
            if region.raw_start != cursor || region.raw_end > expected.end {
                return Err(box_err!(
                    "online RestoreShard detected a merge across Region {}:{} range while retrying",
                    expected.region_id,
                    expected.version
                ));
            }
            let actual_conf_ver = region.epoch.get_conf_ver();
            let expected_conf_ver = expected.conf_ver.unwrap();
            if actual_conf_ver != expected_conf_ver {
                return Err(Error::RegionConfVerNotMatch {
                    expected: expected_conf_ver,
                    actual: actual_conf_ver,
                });
            }
            cursor.clone_from(&region.raw_end);
        }
        if first.raw_start != expected.start || cursor != expected.end {
            return Err(box_err!(
                "online RestoreShard detected a merge across Region {}:{} range while retrying",
                expected.region_id,
                expected.version
            ));
        }
        if descendants.len() > 1
            || descendants[0].id != expected.region_id
            || descendants[0].epoch.get_version() != expected.version
        {
            info!(
                "online RestoreShard will rebuild snapshots for split descendants";
                "original_region_id" => expected.region_id,
                "original_region_version" => expected.version,
                "descendant_regions" => ?descendants.iter().map(|region| region.id).collect::<Vec<_>>(),
            );
        }
    }
    Ok(())
}

fn can_retry_restore_snapshot_batch(err: &Error) -> bool {
    matches!(
        err,
        Error::StoreDiskFull(_)
            | Error::RegionVerNotMatch { .. }
            | Error::RegionConfChangePending(_)
            | Error::RegionNotFoundOrNoLeader(_)
            | Error::PdError(_)
    )
}

enum RestoreSnapshotTaskOutcome {
    Finished(Result<RestoreShardResponse>),
    SkippedAfterTerminalError,
    JoinFailed(String),
}

/// Runs one RestoreShard batch to quiescence.
///
/// A terminal request error prevents queued tasks from starting, but requests
/// that may already have reached a Store are never aborted. The function waits
/// for every spawned task so its caller can safely publish a terminal task
/// state and stop renewing the scheduling guard after this future returns.
async fn run_restore_snapshot_batch<F>(
    requests: Vec<(RestoreRegionFence, F)>,
    concurrency: usize,
) -> Vec<(RestoreRegionFence, RestoreSnapshotTaskOutcome)>
where
    F: Future<Output = Result<RestoreShardResponse>> + Send + 'static,
{
    let semaphore = Arc::new(Semaphore::new(concurrency));
    let stop_queued_requests = Arc::new(AtomicBool::new(false));
    let mut handles = FuturesUnordered::new();

    for (fence, request) in requests {
        let semaphore = semaphore.clone();
        let stop_queued_requests = stop_queued_requests.clone();
        let handle = tokio::spawn(async move {
            let _permit = match semaphore.acquire_owned().await {
                Ok(permit) => permit,
                Err(err) => {
                    stop_queued_requests.store(true, Ordering::Release);
                    return RestoreSnapshotTaskOutcome::Finished(Err(box_err!(
                        "RestoreShard batch semaphore closed: {}",
                        err
                    )));
                }
            };
            if stop_queued_requests.load(Ordering::Acquire) {
                return RestoreSnapshotTaskOutcome::SkippedAfterTerminalError;
            }

            let result = request.await;
            if result
                .as_ref()
                .is_err_and(|err| !can_retry_restore_snapshot_batch(err))
            {
                stop_queued_requests.store(true, Ordering::Release);
            }
            RestoreSnapshotTaskOutcome::Finished(result)
        });
        handles.push(async move { (fence, handle.await) });
    }

    let mut results = Vec::with_capacity(handles.len());
    while let Some((fence, result)) = handles.next().await {
        let outcome = match result {
            Ok(outcome) => outcome,
            Err(err) => {
                stop_queued_requests.store(true, Ordering::Release);
                RestoreSnapshotTaskOutcome::JoinFailed(err.to_string())
            }
        };
        results.push((fence, outcome));
    }
    results
}

fn restore_snapshots(
    tag: &str,
    runtime: &Runtime,
    pd_client: Arc<dyn PdClient>,
    reporter: &Arc<dyn ReportRestoreStepTrait>,
    snapshots: Vec<RestoreSnapshotRequest>,
    success_ranges: &mut MergeRanges,
    config: &RestoreConfig,
    bo: &mut Backoff,
    limiter: &Option<Arc<ThroughputLimiter>>,
    scheduling_guard: Option<RestoreScheduleGuardState>,
    deadline: Option<RestoreDeadline>,
) -> Result<RestoredSnapshots> {
    let (snapshots_size, need_calibrate) = estimate_snapshot_size(
        snapshots.iter().map(|snapshot| &snapshot.change_set),
        limiter,
        REPLICAS,
    );
    debug_assert_eq!(snapshots.len(), snapshots_size.len());

    let security_mgr = pd_client.get_security_mgr();
    let client = security_mgr.http_client(hyper::Client::builder())?;
    let timeout = config.timeout_restore_snapshot.0;
    let store_count = get_all_stores_except_tiflash(pd_client.as_ref())?.len();
    let max_concurrency = config
        .restore_snapshot_concurrency_factor
        .saturating_mul(store_count)
        .max(1);
    let configured_concurrency = config.restore_snapshot_concurrency.max(1);
    let restore_snapshot_concurrency = configured_concurrency.min(max_concurrency);

    info!(
        "using restore snapshot concurrency";
        "restore_snapshot_concurrency" => restore_snapshot_concurrency,
        "configured_concurrency" => configured_concurrency,
        "restore_snapshot_concurrency_factor" => config.restore_snapshot_concurrency_factor,
        "store_count" => store_count,
        "max_by_factor" => max_concurrency
    );

    let mut requests = Vec::with_capacity(snapshots.len());
    for (snapshot, snapshot_size) in snapshots.into_iter().zip(snapshots_size.into_iter()) {
        let pd_client = pd_client.clone();
        let client = client.clone();
        let reporter = reporter.clone();
        let fence = snapshot.fence();
        let limiter = limiter.clone();
        let scheduling_guard = scheduling_guard.clone();
        let request = request_restore_snapshot(
            pd_client,
            client,
            reporter,
            snapshot.change_set,
            snapshot.expected_conf_ver,
            timeout,
            limiter,
            snapshot_size,
            need_calibrate,
            scheduling_guard,
            deadline,
        );
        requests.push((fence, request));
    }
    let results = runtime.block_on(run_restore_snapshot_batch(
        requests,
        restore_snapshot_concurrency,
    ));

    let mut disk_full_regions = vec![];
    let mut disk_full_stores: HashSet<u64> = HashSet::default();
    let mut restored = RestoredSnapshots::default();
    let mut terminal_error = None;
    for (fence, outcome) in results {
        let result = match outcome {
            RestoreSnapshotTaskOutcome::Finished(result) => result,
            RestoreSnapshotTaskOutcome::SkippedAfterTerminalError => continue,
            RestoreSnapshotTaskOutcome::JoinFailed(err) => {
                let err = box_err!(
                    "RestoreShard task for Region {} failed to join: {}",
                    fence.region_id,
                    err
                );
                error!("{} request_restore_snapshot task failed: {:?}", tag, err);
                if terminal_error.is_none() {
                    terminal_error = Some(err);
                }
                continue;
            }
        };
        match result {
            Ok(resp) => {
                success_ranges.insert(fence.start, fence.end);
                restored.count += 1;
                restored.restore_bytes += resp.restore_bytes;
            }
            Err(Error::StoreDiskFull(stores_id)) => {
                disk_full_regions.push(fence.region_id);
                disk_full_stores.extend(stores_id);
                restored.retry_regions.push(fence);
            }
            Err(e) if can_retry_restore_snapshot_batch(&e) => {
                let range = (
                    InnerKey::from_outer_key(&fence.start),
                    InnerKey::from_outer_end_key(&fence.end),
                );
                info!("{} request_restore_snapshot error: {:?}, retry", tag, e; "range" => ?range);
                restored.retry_regions.push(fence);
            }
            Err(e) => {
                let range = (
                    InnerKey::from_outer_key(&fence.start),
                    InnerKey::from_outer_end_key(&fence.end),
                );
                error!("{} request_restore_snapshot error: {:?}", tag, e; "range" => ?range);
                if terminal_error.is_none() {
                    terminal_error = Some(e);
                }
            }
        }
    }

    if let Some(err) = terminal_error {
        return Err(err);
    }

    if !disk_full_regions.is_empty() {
        let stores = Vec::from_iter(disk_full_stores);
        warn!("{} request_restore_snapshot error: disk full", tag;
            "regions" => ?disk_full_regions, "stores" => ?stores);
        bo.on_error(Error::StoreDiskFull(stores))?;

        if scheduling_guard.is_none() {
            if let Err(err) = pd_client.scatter_regions_by_id(disk_full_regions.clone()) {
                warn!("{} scatter regions failed", tag;
                    "regions" => ?disk_full_regions, "err" => ?err);
            }
        } else {
            warn!(
                "{} keep online RestoreShard placement unchanged after disk-full error",
                tag;
                "regions" => ?disk_full_regions,
            );
        }
    }

    Ok(restored)
}

async fn request_restore_snapshot(
    pd_client: Arc<dyn PdClient>,
    client: security::HttpClient,
    reporter: Arc<dyn ReportRestoreStepTrait>,
    cs: pb::ChangeSet,
    expected_conf_ver: Option<u64>,
    timeout: Duration,
    limiter: Option<Arc<ThroughputLimiter>>,
    snapshot_size: SnapshotSize,
    restore_size_need_calibrate: bool,
    scheduling_guard: Option<RestoreScheduleGuardState>,
    deadline: Option<RestoreDeadline>,
) -> Result<RestoreShardResponse> {
    let security_mgr = pd_client.get_security_mgr();
    let post_data = Cow::from(cs.write_to_bytes().unwrap());
    let mut last_err = None;
    let mut retry_cnt = 0;
    let start_time = Instant::now();
    defer! { metrics::NATIVE_BR_RESTORE_SNAPSHOT_REQUEST_SECS.observe(start_time.saturating_elapsed_secs()) }
    'retry: while start_time.saturating_elapsed() < timeout {
        check_restore_deadline(deadline, "RestoreShard request")?;
        if let Some(guard) = &scheduling_guard {
            guard.check_healthy()?;
        }
        retry_cnt += 1;
        let (region, leader) = match pd_client.get_region_leader_by_id(cs.shard_id).await? {
            Some((region, leader)) => (region, leader),
            None => {
                let e = Err(Error::RegionNotFoundOrNoLeader(cs.shard_id));
                warn!("{}:{}: {:?}", cs.shard_id, cs.shard_ver, e);
                last_err = Some(e);
                tokio::time::sleep(Duration::from_millis(200)).await;
                metrics::NATIVE_BR_RESTORE_ERROR
                    .with_label_values(&["restore_snapshot::pd_region_not_found_or_no_leader"])
                    .inc();
                continue 'retry;
            }
        };

        let region_epoch = region.get_region_epoch();
        if let Some(expected_conf_ver) = expected_conf_ver {
            let actual_conf_ver = region_epoch.get_conf_ver();
            if actual_conf_ver != expected_conf_ver {
                return Err(Error::RegionConfVerNotMatch {
                    expected: expected_conf_ver,
                    actual: actual_conf_ver,
                });
            }
        }
        let region_ver = region_epoch.get_version();
        let shard_ver = cs.get_shard_ver();
        if region_ver != shard_ver {
            metrics::NATIVE_BR_RESTORE_ERROR
                .with_label_values(&["restore_snapshot::pd_epoch_not_match"])
                .inc();
            return Err(Error::RegionVerNotMatch {
                expected: shard_ver,
                actual: region_ver,
            });
        }

        let store = pd_client.get_store_async(leader.get_store_id()).await?;
        let tag = ShardTag::new(store.get_id(), IdVer::new(cs.shard_id, cs.shard_ver));

        let replicas = region
            .get_peers()
            .iter()
            .filter(|p| p.role == PeerRole::Voter)
            .count();
        let mut estimated_restore_size = snapshot_size.estimated_size * replicas as u64;
        if let Some(limiter) = limiter.as_ref() {
            if restore_size_need_calibrate && !snapshot_size.files.is_empty() {
                match limiter
                    .calibrate_restore_size_from_tikv(&snapshot_size, &region)
                    .await
                {
                    Ok(calibrated_restore_size) => {
                        debug!(
                            "{} request_restore_snapshot: calibrated restore size: {} bytes",
                            tag, calibrated_restore_size
                        );
                        estimated_restore_size = calibrated_restore_size;
                    }
                    Err(err) => {
                        warn!(
                            "{} request_restore_snapshot: calibrate restore size error: {:?}",
                            tag, err
                        );
                        debug_assert!(false);
                    }
                }
            }
        }
        if estimated_restore_size > 0 {
            reporter.report_pending_restore_data_size(estimated_restore_size as i64);
            NATIVE_BR_RESTORE_PENDING_DATA_SIZE.add(estimated_restore_size as i64);
        }
        tikv_util::defer!({
            if estimated_restore_size > 0 {
                reporter.report_pending_restore_data_size(-(estimated_restore_size as i64));
                NATIVE_BR_RESTORE_PENDING_DATA_SIZE.sub(estimated_restore_size as i64);
            }
        });
        let restore_size_charge = match limiter.as_ref() {
            Some(limiter) if estimated_restore_size > 0 => {
                Some(limiter.consume_restore_size(estimated_restore_size).await)
            }
            _ => None,
        };

        if let Some(guard) = &scheduling_guard {
            guard.check_healthy()?;
        }
        let path = match expected_conf_ver {
            Some(expected_conf_ver) => format!(
                "{}/restore-shard?expected_conf_ver={}",
                &store.status_address, expected_conf_ver
            ),
            None => format!("{}/restore-shard", &store.status_address),
        };
        let uri = security_mgr.build_uri(path)?;
        let req = Request::post(uri)
            .header(header::ACCEPT, CONTENT_TYPE_PROTOBUF)
            .body(Body::from(post_data.clone()))
            .unwrap();
        let request_timeout = match deadline {
            Some(deadline) => deadline.cap_timeout(timeout / 2, "RestoreShard request")?,
            None => timeout / 2,
        };
        match send_request_to_store(req, &store, &client, request_timeout).await {
            Ok((_, resp)) => {
                if let Some(charge) = restore_size_charge {
                    charge.commit();
                }
                let resp: RestoreShardResponse = serde_json::from_slice(&resp).unwrap();
                metrics::NATIVE_BR_RESTORED_SNAPSHOT_KV_BYTES.inc_by(resp.restore_bytes);
                debug!("{} request_restore_snapshot succeed", tag);
                NATIVE_BR_RESTORED_DATA_SIZE.inc_by(estimated_restore_size);
                NATIVE_BR_RESTORED_KV_SIZE.inc_by(resp.restore_bytes);
                return Ok(resp);
            }
            Err(Error::HttpPbError(status, mut err)) => {
                warn!("{} request_restore_snapshot failed: {:?}", tag, err);
                let sleep_dur = if err.has_disk_full() {
                    metrics::NATIVE_BR_RESTORE_ERROR
                        .with_label_values(&["restore_snapshot::store_disk_full"]);
                    return Err(Error::StoreDiskFull(err.take_disk_full().take_store_id()));
                } else if err.has_epoch_not_match() {
                    metrics::NATIVE_BR_RESTORE_ERROR
                        .with_label_values(&["restore_snapshot::epoch_not_match"]);
                    let current_epoch = err
                        .get_epoch_not_match()
                        .get_current_regions()
                        .first()
                        .map(|region| region.get_region_epoch());
                    if let Some(expected_conf_ver) = expected_conf_ver {
                        let actual_conf_ver = current_epoch.map_or(0, |epoch| epoch.get_conf_ver());
                        if actual_conf_ver != expected_conf_ver {
                            return Err(Error::RegionConfVerNotMatch {
                                expected: expected_conf_ver,
                                actual: actual_conf_ver,
                            });
                        }
                    }
                    let actual_version = current_epoch.map_or(0, |epoch| epoch.get_version());
                    if actual_version != shard_ver {
                        return Err(Error::RegionVerNotMatch {
                            expected: shard_ver,
                            actual: actual_version,
                        });
                    }
                    return Err(Error::RegionConfChangePending(cs.get_shard_id()));
                } else if err.has_not_leader() && err.get_not_leader().has_leader() {
                    metrics::NATIVE_BR_RESTORE_ERROR
                        .with_label_values(&["restore_snapshot::not_leader"]);
                    Duration::from_millis(50)
                } else {
                    metrics::NATIVE_BR_RESTORE_ERROR
                        .with_label_values(&["restore_snapshot::unknown_http"]);
                    warn!("restore snapshot encounters generic error."; "err" => ?err);
                    Duration::from_millis(500)
                };
                drop(restore_size_charge);
                tokio::time::sleep(sleep_dur).await;
                last_err = Some(Err(Error::HttpPbError(status, err)));
                continue 'retry;
            }
            Err(e) => {
                let err_msg = format!(
                    "{} request_restore_snapshot #{retry_cnt} error: {:?}, region: {:?}, leader: {:?}",
                    tag, e, region, leader
                );
                warn!("{}", err_msg);
                metrics::NATIVE_BR_RESTORE_ERROR
                    .with_label_values(&["restore_snapshot::unknown_internal"]);
                last_err = Some(Err(box_err!(err_msg)));
                drop(restore_size_charge);
                tokio::time::sleep(Duration::from_millis(500)).await;
                continue 'retry;
            }
        }
    }
    last_err.expect("there must be error")
}

fn estimate_snapshot_size<'a>(
    snapshots: impl IntoIterator<Item = &'a pb::ChangeSet>,
    limiter: &Option<Arc<ThroughputLimiter>>,
    replicas: usize,
) -> (Vec<SnapshotSize>, bool /* need_calibrate */) {
    let mut total_snapshot_size = 0;
    let snapshots_size = snapshots
        .into_iter()
        .map(|cs| {
            let snapshot_size = ThroughputLimiter::estimate_snapshot_size_locally(cs);
            total_snapshot_size += snapshot_size.estimated_size;
            snapshot_size
        })
        .collect();
    let need_calibrate = limiter
        .as_ref()
        .is_some_and(|x| total_snapshot_size * replicas as u64 >= x.calibrate_threshold());

    (snapshots_size, need_calibrate)
}

fn make_keyspace_tag(source_keyspace_id: u32, target_keyspace_id: u32) -> String {
    format!("{}->{}", source_keyspace_id, target_keyspace_id)
}

struct PeerPreprocessor {
    preprocessed_index: u64,
    region: metapb::Region,
    preprocessed_region: Option<metapb::Region>,
    peer: metapb::Peer,
    shard_meta: Option<ShardMeta>,
    last_committed_split_idx: u64,
    pending_truncate: Option<(u64 /* term */, u64 /* index */)>,
    raft_hard_state: eraftpb::HardState,
    raft_state: RaftState,
    pending_merge_state: Option<MergeState>,
    want_rollback_merge_peers: HashSet<u64>,
    learner_skip_idx: u64,
    encryption_state: EncryptionState,
}

impl PeerPreprocessor {
    fn new(rf_engine: &RfEngine, shard: &mut BackupShard) -> Self {
        let peer = metapb::Peer {
            id: shard.peer_id,
            store_id: shard.store_id,
            role: PeerRole::Voter,
            ..Default::default()
        };
        let mut region_state = rf_engine
            .load_region_state(shard.peer_id, shard.ver())
            .unwrap_or_else(|| {
                panic!("failed to load region state for peer {}", shard.peer_id);
            });
        let region = region_state.take_region();
        let merge_state = region_state
            .has_merge_state()
            .then(|| region_state.take_merge_state());

        // Set `preprocessed_index` to the `data_sequence` to replay from last flush and
        // restore memory-based fields of `ShardMeta`.
        // See https://github.com/tidbcloud/cloud-storage-engine/issues/1680.
        //
        // Also note that when the shard is not yet initial flushed, in theory we should
        // preprocess from parent. But the only memory field `txn_file_locks` so far is
        // used for ignore split/merge, which will not happen when the shard is not
        // initial flushed. So currently we do not preprocess from parent for
        // simplicity.
        // See https://github.com/tidbcloud/cloud-storage-engine/issues/1862.
        let preprocessed_index = shard.meta.data_sequence;

        Self {
            preprocessed_index,
            region,
            preprocessed_region: None,
            peer,
            shard_meta: Some(shard.meta.clone()), // TODO: mem::take()
            last_committed_split_idx: 0,
            pending_truncate: None,
            raft_hard_state: shard.raft_state.get_hard_state(),
            raft_state: shard.raft_state,
            pending_merge_state: merge_state,
            want_rollback_merge_peers: HashSet::default(),
            learner_skip_idx: 0,
            encryption_state: EncryptionState::Disabled,
        }
    }

    fn as_ref(&mut self) -> rfstore::store::peer::PreprocessRef<'_> {
        rfstore::store::peer::PreprocessRef {
            preprocessed_index: &mut self.preprocessed_index,
            region: &mut self.region,
            preprocessed_region: &mut self.preprocessed_region,
            peer: &mut self.peer,
            shard_meta: &mut self.shard_meta,
            last_committed_split_idx: &mut self.last_committed_split_idx,
            pending_truncate: &mut self.pending_truncate,
            raft_hard_state: self.raft_hard_state.clone(),
            raft_state: &mut self.raft_state,
            is_leader: false,
            pending_merge_state: &mut self.pending_merge_state,
            want_rollback_merge_peers: &mut self.want_rollback_merge_peers,
            learner_skip_idx: &mut self.learner_skip_idx,
            encryption_state: &mut self.encryption_state,
        }
    }
}

struct Backoff {
    inner: ExponentialBackoff,
}

impl ops::Deref for Backoff {
    type Target = ExponentialBackoff;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

impl ops::DerefMut for Backoff {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.inner
    }
}

impl Backoff {
    fn new(bo: ExponentialBackoff) -> Self {
        Self { inner: bo }
    }

    fn on_error(&mut self, err: Error) -> Result<usize /* current_attempts */> {
        match self.next_delay() {
            Ok(delay) => {
                thread::sleep(delay);
                Ok(self.current_attempts())
            }
            Err(_) => Err(RetryLimitExceeded(Box::new(err))),
        }
    }
}

#[cfg(test)]
pub(crate) use tests::reconstruct_full_range_with_peer_states;

#[cfg(test)]
mod tests {
    use std::{
        collections::BTreeMap, convert::Infallible, net::TcpListener, sync::atomic::AtomicUsize,
    };

    use api_version::api_v2::KEYSPACE_PREFIX_LEN;
    use hyper::{
        Response, Server,
        service::{make_service_fn, service_fn},
    };
    use kvengine::{
        UserMeta, WriteBatch,
        dfs::test_util::new_test_s3fs,
        test_engine::{ApplyTask, new_test_engine_api_v2},
    };
    use schema::{
        KvGetter,
        schema::{DbInfo, STATE_PUBLIC},
    };
    use security::{GetSecurityManager, SecurityConfig};
    use tempfile::Builder as TempDirBuilder;
    use tikv_util::codec::{bytes::encode_bytes, number::NumberEncoder};

    use super::*;

    struct MockPdClient;

    impl PdClient for MockPdClient {}

    impl GetSecurityManager for MockPdClient {}

    struct SlowSplitPdClient {
        calls: Arc<AtomicUsize>,
    }

    impl PdClient for SlowSplitPdClient {
        fn split_regions_with_retry(
            &self,
            _keys: Vec<Vec<u8>>,
            _timeout: Duration,
        ) -> futures::future::BoxFuture<'_, pd_client::Result<Vec<u64>>> {
            self.calls.fetch_add(1, Ordering::AcqRel);
            Box::pin(std::future::pending())
        }
    }

    impl GetSecurityManager for SlowSplitPdClient {}

    struct RecordingSplitPdClient {
        split_calls: Arc<AtomicUsize>,
        scatter_calls: Arc<AtomicUsize>,
    }

    impl PdClient for RecordingSplitPdClient {
        fn split_regions_with_retry(
            &self,
            _keys: Vec<Vec<u8>>,
            _timeout: Duration,
        ) -> futures::future::BoxFuture<'_, pd_client::Result<Vec<u64>>> {
            self.split_calls.fetch_add(1, Ordering::AcqRel);
            Box::pin(async { Ok(vec![1]) })
        }

        fn scatter_regions_by_id(&self, _regions_id: Vec<u64>) -> pd_client::Result<()> {
            self.scatter_calls.fetch_add(1, Ordering::AcqRel);
            Ok(())
        }
    }

    impl GetSecurityManager for RecordingSplitPdClient {}

    #[derive(Clone, Copy)]
    struct FullRangePeer {
        store_id: u64,
        region_id: u64,
        peer_id: u64,
        state: kvproto::raft_serverpb::PeerState,
    }

    fn new_full_range_rfengine(peer: FullRangePeer) -> (tempfile::TempDir, RfEngine) {
        const REGION_VERSION: u64 = 4;

        let dir = TempDirBuilder::new()
            .prefix("full-range-rfengine")
            .tempdir()
            .unwrap();
        let rf =
            RfEngine::open(dir.path(), &rfengine::RfEngineConfig::default(), None, None).unwrap();
        let metapb_peer = metapb::Peer {
            id: peer.peer_id,
            store_id: peer.store_id,
            ..Default::default()
        };
        let region = metapb::Region {
            id: peer.region_id,
            end_key: GLOBAL_SHARD_END_KEY.to_vec(),
            region_epoch: Some(metapb::RegionEpoch {
                version: REGION_VERSION,
                ..Default::default()
            })
            .into(),
            peers: vec![metapb_peer].into(),
            ..Default::default()
        };
        let mut change_set =
            rfstore::store::initial_change_set(peer.region_id, REGION_VERSION).change_set;
        change_set
            .mut_snapshot()
            .set_outer_end(GLOBAL_SHARD_END_KEY.to_vec());
        let mut wb = rfengine::WriteBatch::new();
        rfstore::store::write_peer_state(&mut wb, peer.peer_id, &region, peer.state, None);
        rfstore::store::write_initial_raft_state(
            &mut wb,
            peer.peer_id,
            peer.region_id,
            REGION_VERSION,
            0,
        );
        wb.set_state(
            peer.peer_id,
            peer.region_id,
            0,
            rfengine::KV_ENGINE_META_KEY,
            &change_set.write_to_bytes().unwrap(),
        );
        rf.write(wb).unwrap();
        (dir, rf)
    }

    fn new_full_range_test_cluster(prefix: &str) -> (tempfile::TempDir, BackupCluster) {
        let dir = TempDirBuilder::new().prefix(prefix).tempdir().unwrap();
        let dfs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(b""));
        let mut cluster = BackupCluster::new_snapshot_restore_cluster(
            1,
            0,
            dir.path().to_path_buf(),
            Arc::new(MockPdClient),
            None,
            dfs,
            RestoreConfig::default(),
            0,
            0,
            None,
        )
        .unwrap();
        cluster.keyspace_start.clear();
        cluster.keyspace_end = GLOBAL_SHARD_END_KEY.to_vec();
        (dir, cluster)
    }

    pub(crate) fn reconstruct_full_range_with_peer_states(
        states: &[kvproto::raft_serverpb::PeerState],
    ) -> Result<()> {
        let (_cluster_dir, mut cluster) = new_full_range_test_cluster("load-full-shards");
        let mut rf_dirs = Vec::with_capacity(states.len());
        for (index, state) in states.iter().copied().enumerate() {
            let store_id = index as u64 + 1;
            let (rf_dir, rf) = new_full_range_rfengine(FullRangePeer {
                store_id,
                region_id: 100 + store_id,
                peer_id: 200 + store_id,
                state,
            });
            rf_dirs.push(rf_dir);
            cluster.raft_engines.insert(store_id, rf);
        }

        cluster.load_shards()
    }

    #[test]
    fn test_collect_full_shards_excludes_only_tombstone_peer() {
        const STORE_ID: u64 = 1;
        const REGION_ID: u64 = 2;
        const PEER_ID: u64 = 3;

        for (state, expected_count) in [
            (kvproto::raft_serverpb::PeerState::Tombstone, 0),
            (kvproto::raft_serverpb::PeerState::Applying, 1),
            (kvproto::raft_serverpb::PeerState::Merging, 1),
        ] {
            let (_dir, rf) = new_full_range_rfengine(FullRangePeer {
                store_id: STORE_ID,
                region_id: REGION_ID,
                peer_id: PEER_ID,
                state,
            });

            let collection = BackupCluster::collect_full_shards(STORE_ID, &rf).unwrap();

            assert_eq!(
                collection.shards.len(),
                expected_count,
                "only persisted tombstone peers may be excluded, state: {state:?}"
            );
            assert_eq!(
                collection.tombstone_peer_count,
                usize::from(state == kvproto::raft_serverpb::PeerState::Tombstone),
                "tombstone provenance must match the persisted peer state"
            );
        }
    }

    #[test]
    fn test_load_full_shards_distinguishes_tombstones_empty_and_live_coverage() {
        let all_tombstone_error = reconstruct_full_range_with_peer_states(&[
            kvproto::raft_serverpb::PeerState::Tombstone,
            kvproto::raft_serverpb::PeerState::Tombstone,
        ])
        .unwrap_err();
        assert!(matches!(
            all_tombstone_error,
            Error::ShardRangeIntegrityError(_)
        ));

        let genuinely_empty_error = reconstruct_full_range_with_peer_states(&[]).unwrap_err();
        assert!(matches!(
            genuinely_empty_error,
            Error::BackupEmptyForKeyspace(0)
        ));

        reconstruct_full_range_with_peer_states(&[
            kvproto::raft_serverpb::PeerState::Tombstone,
            kvproto::raft_serverpb::PeerState::Normal,
        ])
        .unwrap();
    }

    #[test]
    fn test_verify_full_shards_rejects_gaps_and_overlaps() {
        let dir = TempDirBuilder::new()
            .prefix("verify-full-shards-noncontiguous")
            .tempdir()
            .unwrap();
        let dfs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(b""));
        for second_start in [&b"n"[..], &b"l"[..]] {
            let mut cluster = BackupCluster::new_snapshot_restore_cluster(
                1,
                0,
                dir.path().to_path_buf(),
                Arc::new(MockPdClient),
                None,
                dfs.clone(),
                RestoreConfig::default(),
                0,
                0,
                None,
            )
            .unwrap();
            cluster.keyspace_start.clear();
            cluster.keyspace_end = GLOBAL_SHARD_END_KEY.to_vec();
            for (region_id, start, end) in [
                (1, &[][..], &b"m"[..]),
                (2, second_start, GLOBAL_SHARD_END_KEY),
            ] {
                let mut shard = BackupShard::default();
                shard.region_id = region_id;
                shard.meta.range = ShardRange::new(start, end);
                cluster.shards.insert(region_id, shard);
                cluster.sorted_shards.push(region_id);
            }

            let error = cluster.verify_full_shards().unwrap_err();

            assert!(matches!(error, Error::ShardRangeIntegrityError(_)));
        }
    }

    struct RegionScanPdClient {
        region: pdpb::Region,
    }

    impl PdClient for RegionScanPdClient {
        fn scan_regions(
            &self,
            _key: Vec<u8>,
            _end_key: Vec<u8>,
            _limit: usize,
        ) -> pd_client::PdFuture<Vec<pdpb::Region>> {
            let region = self.region.clone();
            Box::pin(async move { Ok(vec![region]) })
        }
    }

    impl GetSecurityManager for RegionScanPdClient {}

    fn retry_region_fence(start: &[u8], end: &[u8], conf_ver: u64) -> RestoreRegionFence {
        RestoreRegionFence {
            region_id: 10,
            version: 20,
            conf_ver: Some(conf_ver),
            start: start.to_vec(),
            end: end.to_vec(),
        }
    }

    fn raw_region(id: u64, start: &[u8], end: &[u8], version: u64, conf_ver: u64) -> RawRegion {
        let mut epoch = metapb::RegionEpoch::default();
        epoch.set_version(version);
        epoch.set_conf_ver(conf_ver);
        RawRegion {
            id,
            raw_start: start.to_vec(),
            raw_end: end.to_vec(),
            epoch,
            peers: Vec::new(),
        }
    }

    #[test]
    fn test_restore_retry_accepts_split_descendants() {
        let expected = retry_region_fence(b"a", b"z", 7);
        let current = vec![
            raw_region(10, b"a", b"m", 21, 7),
            raw_region(11, b"m", b"z", 21, 7),
        ];

        validate_restore_retry_topology(&[expected], &current).unwrap();
    }

    #[test]
    fn test_restore_retry_rejects_merge_across_old_boundary() {
        let expected = retry_region_fence(b"m", b"z", 7);
        let current = vec![raw_region(9, b"a", b"z", 21, 7)];

        let err = validate_restore_retry_topology(&[expected], &current).unwrap_err();
        assert!(err.to_string().contains("merge"), "{err:?}");
    }

    #[test]
    fn test_restore_retry_rejects_conf_ver_change() {
        let expected = retry_region_fence(b"a", b"z", 7);
        let current = vec![raw_region(10, b"a", b"z", 20, 8)];

        assert!(matches!(
            validate_restore_retry_topology(&[expected], &current),
            Err(Error::RegionConfVerNotMatch {
                expected: 7,
                actual: 8
            })
        ));
    }

    #[test]
    fn test_restore_snapshot_batch_drains_in_flight_and_skips_queued_after_terminal_error() {
        async fn request(
            active_requests_ready: Arc<tokio::sync::Barrier>,
            release_in_flight: Arc<tokio::sync::Notify>,
            started_requests: Arc<AtomicUsize>,
            terminal_finished: Arc<AtomicBool>,
            in_flight_finished: Arc<AtomicBool>,
        ) -> Result<RestoreShardResponse> {
            match started_requests.fetch_add(1, Ordering::AcqRel) {
                0 => {
                    active_requests_ready.wait().await;
                    terminal_finished.store(true, Ordering::Release);
                    Err(Error::RegionConfVerNotMatch {
                        expected: 1,
                        actual: 2,
                    })
                }
                1 => {
                    active_requests_ready.wait().await;
                    release_in_flight.notified().await;
                    in_flight_finished.store(true, Ordering::Release);
                    Ok(RestoreShardResponse::default())
                }
                request_index => panic!("queued request {request_index} unexpectedly started"),
            }
        }

        let runtime = Runtime::new().unwrap();
        runtime.block_on(async {
            let active_requests_ready = Arc::new(tokio::sync::Barrier::new(2));
            let release_in_flight = Arc::new(tokio::sync::Notify::new());
            let started_requests = Arc::new(AtomicUsize::new(0));
            let terminal_finished = Arc::new(AtomicBool::new(false));
            let in_flight_finished = Arc::new(AtomicBool::new(false));
            let requests = (0..3)
                .map(|kind| {
                    (
                        retry_region_fence(&[kind], &[kind + 1], 1),
                        request(
                            active_requests_ready.clone(),
                            release_in_flight.clone(),
                            started_requests.clone(),
                            terminal_finished.clone(),
                            in_flight_finished.clone(),
                        ),
                    )
                })
                .collect();

            let batch = tokio::spawn(run_restore_snapshot_batch(requests, 2));
            tokio::time::timeout(Duration::from_secs(5), async {
                while !terminal_finished.load(Ordering::Acquire) {
                    tokio::task::yield_now().await;
                }
            })
            .await
            .unwrap();

            assert!(!batch.is_finished());
            assert!(!in_flight_finished.load(Ordering::Acquire));
            assert_eq!(started_requests.load(Ordering::Acquire), 2);

            release_in_flight.notify_one();
            let results = tokio::time::timeout(Duration::from_secs(5), batch)
                .await
                .unwrap()
                .unwrap();
            assert!(in_flight_finished.load(Ordering::Acquire));
            assert_eq!(started_requests.load(Ordering::Acquire), 2);

            let mut succeeded = 0;
            let mut terminal = 0;
            let mut skipped = 0;
            for (_, outcome) in results {
                match outcome {
                    RestoreSnapshotTaskOutcome::Finished(Ok(_)) => succeeded += 1,
                    RestoreSnapshotTaskOutcome::Finished(Err(Error::RegionConfVerNotMatch {
                        ..
                    })) => terminal += 1,
                    RestoreSnapshotTaskOutcome::SkippedAfterTerminalError => skipped += 1,
                    RestoreSnapshotTaskOutcome::Finished(Err(err)) => {
                        panic!("unexpected RestoreShard error: {err:?}")
                    }
                    RestoreSnapshotTaskOutcome::JoinFailed(err) => {
                        panic!("unexpected RestoreShard join error: {err}")
                    }
                }
            }
            assert_eq!((succeeded, terminal, skipped), (1, 1, 1));
        });
    }

    #[test]
    fn test_restore_schedule_guard_state_rejects_expired_lease() {
        let healthy = Arc::new(AtomicBool::new(true));
        let state = RestoreScheduleGuardState {
            healthy: healthy.clone(),
            valid_until: Some(Arc::new(Mutex::new(StdInstant::now()))),
        };

        let err = state.check_healthy().unwrap_err();
        assert!(err.to_string().contains("lease expired"), "{err:?}");
        assert!(!healthy.load(Ordering::Acquire));
    }

    #[test]
    fn test_schedule_guard_collection_rejects_any_expired_lease() {
        let healthy = RestoreScheduleGuardState {
            healthy: Arc::new(AtomicBool::new(true)),
            valid_until: None,
        };
        let expired = RestoreScheduleGuardState {
            healthy: Arc::new(AtomicBool::new(true)),
            valid_until: Some(Arc::new(Mutex::new(StdInstant::now()))),
        };

        let err = check_schedule_guards(&[healthy, expired]).unwrap_err();
        assert!(err.to_string().contains("lease expired"), "{err:?}");
    }

    #[test]
    fn test_online_snapshot_split_obeys_attempt_deadline() {
        let runtime = Runtime::new().unwrap();
        let calls = Arc::new(AtomicUsize::new(0));
        let pd_client = SlowSplitPdClient {
            calls: calls.clone(),
        };
        let deadline = RestoreDeadline::after(Duration::from_millis(50)).unwrap();
        let started_at = StdInstant::now();

        let err = split_regions_with_retry(
            "deadline-test",
            &pd_client,
            vec![b"split-key".to_vec()],
            false,
            Duration::from_secs(30),
            Some(deadline),
            &runtime,
        )
        .unwrap_err();

        assert!(err.to_string().contains("split batch timed out"), "{err:?}");
        assert_eq!(calls.load(Ordering::Acquire), 1);
        assert!(started_at.elapsed() < Duration::from_secs(2));
    }

    #[test]
    fn test_split_regions_can_preserve_exact_placement() {
        let runtime = Runtime::new().unwrap();
        let split_calls = Arc::new(AtomicUsize::new(0));
        let scatter_calls = Arc::new(AtomicUsize::new(0));
        let pd_client = RecordingSplitPdClient {
            split_calls: split_calls.clone(),
            scatter_calls: scatter_calls.clone(),
        };

        split_regions_for_keys(
            "online-target",
            &pd_client,
            vec![b"coarse".to_vec()],
            vec![b"fine".to_vec()],
            false,
            Duration::from_secs(30),
            None,
            &runtime,
        )
        .unwrap();
        assert_eq!(split_calls.load(Ordering::Acquire), 2);
        assert_eq!(scatter_calls.load(Ordering::Acquire), 0);

        split_regions_for_keys(
            "ordinary-target",
            &pd_client,
            vec![b"coarse".to_vec()],
            Vec::new(),
            true,
            Duration::from_secs(30),
            None,
            &runtime,
        )
        .unwrap();
        assert_eq!(split_calls.load(Ordering::Acquire), 3);
        assert_eq!(scatter_calls.load(Ordering::Acquire), 1);
    }

    #[test]
    fn test_optional_restore_schedule_guard_install_does_not_wait_for_pd() {
        let runtime = Runtime::new().unwrap();
        let listener = TcpListener::bind("127.0.0.1:0").unwrap();
        let endpoint = listener.local_addr().unwrap().to_string();
        let (request_started_tx, request_started_rx) = std::sync::mpsc::sync_channel(1);
        let release_request = Arc::new(tokio::sync::Notify::new());
        let server_release_request = release_request.clone();
        let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel();
        let server_handle = {
            let _runtime_guard = runtime.enter();
            runtime.spawn(
                Server::from_tcp(listener)
                    .unwrap()
                    .serve(make_service_fn(move |_| {
                        let request_started_tx = request_started_tx.clone();
                        let release_request = server_release_request.clone();
                        async move {
                            Ok::<_, Infallible>(service_fn(move |_| {
                                let request_started_tx = request_started_tx.clone();
                                let release_request = release_request.clone();
                                async move {
                                    let _ = request_started_tx.send(());
                                    release_request.notified().await;
                                    Ok::<_, Infallible>(Response::new(Body::from("\"ok\"")))
                                }
                            }))
                        }
                    }))
                    .with_graceful_shutdown(async {
                        let _ = shutdown_rx.await;
                    }),
            )
        };

        let mut pd =
            PdControl::new(pd_client::Config::new(vec![endpoint]), Default::default()).unwrap();
        pd.set_timeout(Duration::from_secs(5));
        let started_at = StdInstant::now();
        let guard = RestoreScheduleGuard::install(pd, &[(b"a", b"z")], false, &runtime).unwrap();
        let elapsed = started_at.elapsed();

        request_started_rx
            .recv_timeout(Duration::from_secs(5))
            .unwrap();
        drop(guard);
        release_request.notify_one();
        let _ = shutdown_tx.send(());
        runtime.block_on(server_handle).unwrap().unwrap();

        assert!(elapsed < Duration::from_secs(1), "elapsed: {elapsed:?}");
    }

    #[test]
    fn test_scan_regions_accepts_empty_start_sentinel() {
        let raw_end = ApiV2::get_keyspace_end_by_id(0);
        let mut region = metapb::Region::default();
        region.set_id(1);
        region.set_end_key(Key::from_raw(&raw_end).into_encoded());
        let mut scanned_region = pdpb::Region::default();
        scanned_region.set_region(region);
        let pd_client = RegionScanPdClient {
            region: scanned_region,
        };

        let runtime = Runtime::new().unwrap();
        let regions = runtime
            .block_on(scan_regions(&pd_client, b"", &raw_end))
            .unwrap();

        assert_eq!(regions.len(), 1);
        assert!(regions[0].get_region().get_start_key().is_empty());
    }

    fn test_db_key(keyspace_id: u32, db_id: i64) -> Vec<u8> {
        let mut key = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
        key.push(b'm');
        key.extend_from_slice(&encode_bytes(b"DBs"));
        key.encode_u64(b'h' as u64).unwrap();
        key.extend_from_slice(&encode_bytes(format!("DB:{db_id}").as_bytes()));
        key
    }

    #[test]
    fn test_backup_schema_reader_skips_deleted_values_across_batch_boundary() {
        const KEYSPACE_ID: u32 = 1;
        const READ_TS: u64 = 10;
        const LIVE_DB_ID: i64 = 9_999;

        let (engine, apply_tx) = new_test_engine_api_v2();
        let shard = engine.get_shard(1).unwrap();
        let mut write_batch = WriteBatch::default();
        write_batch.reset(shard.id);
        let user_meta = UserMeta::new(READ_TS - 1, READ_TS).to_array();
        let first_deleted_key = test_db_key(KEYSPACE_ID, 1_000);
        for db_id in 1_000..1_000 + schema::DEFAULT_SCAN_BATCH_SIZE as i64 {
            let key = test_db_key(KEYSPACE_ID, db_id);
            write_batch.put(WRITE_CF, &key, &[], 0, &user_meta, READ_TS);
        }
        let live_db = DbInfo {
            id: LIVE_DB_ID,
            state: STATE_PUBLIC,
            ..Default::default()
        };
        write_batch.put(
            WRITE_CF,
            &test_db_key(KEYSPACE_ID, LIVE_DB_ID),
            &serde_json::to_vec(&live_db).unwrap(),
            0,
            &user_meta,
            READ_TS,
        );
        let (result_tx, result_rx) = mpsc::bounded(1);
        apply_tx
            .send(ApplyTask::new_wb(write_batch, result_tx))
            .unwrap();
        result_rx.recv().unwrap().unwrap();

        let reader = Arc::new(BackupSchemaReader::new(
            engine.engine.clone(),
            [BackupSchemaShard {
                id: shard.id,
                version: shard.ver,
                start_key: shard.range.outer_start.clone(),
                end_key: shard.range.outer_end.clone(),
            }]
            .into_iter(),
        ));
        let runtime = Runtime::new().unwrap();
        let read_ts = TimeStamp::new(READ_TS);

        assert_eq!(
            runtime
                .block_on(reader.get(&first_deleted_key, read_ts))
                .unwrap(),
            None
        );
        let keyspace_prefix = ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID);
        let dbs = runtime
            .block_on(schema::load_schema(reader, &keyspace_prefix, read_ts))
            .unwrap();
        assert_eq!(dbs.len(), 1);
        assert_eq!(dbs[0].id, LIVE_DB_ID);
    }

    fn new_test_packed_backup(keyspace_id: u32) -> PackedBackup {
        let (keyspace_start, keyspace_end) = ApiV2::get_keyspace_range_by_id(keyspace_id);
        let mut snapshot = pb::Snapshot::new();
        snapshot.set_outer_start(keyspace_start);
        snapshot.set_outer_end(keyspace_end);

        let mut shard = pb::ChangeSet::new();
        shard.set_shard_id(1);
        shard.set_shard_ver(1);
        shard.set_snapshot(snapshot);

        let mut packed = PackedBackup::new();
        packed.set_backup_ts(123);
        packed.set_safe_ts(120);
        packed.set_keyspace_id(keyspace_id);
        packed.set_unpacked(true);
        packed.mut_shards().push(shard);
        packed
    }

    fn new_test_encryption_meta(
        keyspace_id: u32,
        file_id: u64,
        data_key_id: u32,
    ) -> pb::EncryptionMeta {
        let mut current = pb::EncryptionEpoch::default();
        current.set_file_id(file_id);
        current.set_data_key_id(data_key_id);
        current.set_created_at(123);

        let mut master_key = pb::MasterKey::default();
        master_key.set_vendor("test".to_owned());
        master_key.set_cmek_id("cmek".to_owned());
        master_key.set_ciphertext(vec![1, 2, 3]);

        let mut data_key = pb::DataKey::default();
        data_key.set_ciphertext(vec![cloud_encryption::KEY_TYPE_AES_256_CTR, 4, 5, 6]);

        let mut meta = pb::EncryptionMeta::default();
        meta.set_keyspace_id(keyspace_id);
        meta.set_current(current);
        meta.set_master_key(master_key);
        meta.mut_data_keys().insert(data_key_id, data_key);
        meta
    }

    #[test]
    fn test_online_restore_loads_lock_files_but_packed_restore_skips_them() {
        let shards_need_truncate = HashSet::from_iter([1]);
        let packed_filter = build_packed_restore_table_filter(shards_need_truncate.clone());
        let online_filter = build_lock_resolving_restore_table_filter(shards_need_truncate);

        let lock_file = kvengine::FileMeta::from_l0_table(&pb::L0Create::default());
        let mut regular_table = pb::TableCreate::default();
        regular_table.set_cf(WRITE_CF as i32);
        regular_table.set_level(1);
        let regular_file = kvengine::FileMeta::from_table(&regular_table);
        let mut lock_table = pb::TableCreate::default();
        lock_table.set_cf(kvengine::LOCK_CF as i32);
        lock_table.set_level(1);
        let lock_cf_file = kvengine::FileMeta::from_table(&lock_table);

        assert!(packed_filter(1, &regular_file));
        assert!(online_filter(1, &regular_file));
        assert!(!packed_filter(2, &lock_file));
        assert!(online_filter(2, &lock_file));
        assert!(!packed_filter(2, &lock_cf_file));
        assert!(online_filter(2, &lock_cf_file));
        assert!(!packed_filter(2, &regular_file));
        assert!(!online_filter(2, &regular_file));
    }

    #[test]
    fn test_regenerate_encryption_meta_only_advances_file_id() {
        let source = new_test_encryption_meta(1, 100, 7);
        let encryption_meta = regenerate_encryption_meta(&source, 200).unwrap();

        assert_eq!(encryption_meta.get_keyspace_id(), source.get_keyspace_id());
        assert_eq!(encryption_meta.get_master_key(), source.get_master_key());
        assert_eq!(encryption_meta.get_data_keys(), source.get_data_keys());
        assert_eq!(encryption_meta.get_history(), source.get_history());
        assert_eq!(encryption_meta.get_current().get_file_id(), 200);
        assert_eq!(
            encryption_meta.get_current().get_data_key_id(),
            source.get_current().get_data_key_id()
        );
        assert_eq!(
            encryption_meta.get_current().get_created_at(),
            source.get_current().get_created_at()
        );
    }

    #[test]
    fn test_regenerate_encryption_meta_requires_newer_file_id() {
        let source = new_test_encryption_meta(1, 100, 7);

        regenerate_encryption_meta(&source, 100).unwrap_err();
        regenerate_encryption_meta(&source, 99).unwrap_err();
    }

    #[test]
    fn test_target_keyspace_encryption_patch_matches_restored_shard_mode() {
        let (config, mode) = target_keyspace_encryption_config_patch(None, false).unwrap();
        assert_eq!(config, None);
        assert_eq!(mode, "disabled");

        let (config, mode) = target_keyspace_encryption_config_patch(None, true).unwrap();
        let config = KeyspaceEncryptionConfig::from_pd_config_value(&config.unwrap()).unwrap();
        assert_eq!(
            config.intended_encryption_method().unwrap(),
            Some(cloud_encryption::EncryptionMethod::Legacy)
        );
        assert_eq!(mode, "legacy");

        let mut meta = pb::EncryptionMeta::default();
        let mut master_key = pb::MasterKey::default();
        master_key.set_cmek_id("test-cmek".to_owned());
        master_key.set_vendor("test-kms".to_owned());
        master_key.set_region("test-region".to_owned());
        meta.set_master_key(master_key);
        let meta = meta.write_to_bytes().unwrap();

        let (config, mode) = target_keyspace_encryption_config_patch(Some(&meta), false).unwrap();
        let config = KeyspaceEncryptionConfig::from_pd_config_value(&config.unwrap()).unwrap();
        assert_eq!(
            config.intended_encryption_method().unwrap(),
            Some(cloud_encryption::EncryptionMethod::Metadata)
        );
        assert_eq!(config.cmek_id.as_deref(), Some("test-cmek"));
        assert_eq!(config.vendor.as_deref(), Some("test-kms"));
        assert_eq!(config.region.as_deref(), Some("test-region"));
        assert_eq!(mode, "metadata");

        target_keyspace_encryption_config_patch(Some(&meta), true).unwrap_err();
    }

    #[test]
    fn test_target_encryption_sync_policy_distinguishes_in_place_and_branch() {
        assert_eq!(
            TargetEncryptionConfigSync::for_restore(true),
            TargetEncryptionConfigSync::ExportedMetadataOnly
        );
        assert_eq!(
            TargetEncryptionConfigSync::for_restore(false),
            TargetEncryptionConfigSync::RestoredMode
        );
        assert!(
            !TargetEncryptionConfigSync::ExportedMetadataOnly.should_patch(None),
            "in-place plaintext or legacy restores must not add a new PD patch"
        );
        assert!(TargetEncryptionConfigSync::ExportedMetadataOnly.should_patch(Some(b"meta")));
        assert!(TargetEncryptionConfigSync::RestoredMode.should_patch(None));
    }

    #[test]
    fn test_online_placement_convergence_uses_only_voter_stores() {
        let peer = |peer_id, store_id| {
            let mut peer = metapb::Peer::default();
            peer.set_id(peer_id);
            peer.set_store_id(store_id);
            peer
        };
        let leader = peer(1, 1);
        let mut region = metapb::Region::default();
        region.set_peers(vec![leader.clone(), peer(2, 2), peer(3, 3)].into());
        let target = OnlinePlacementTarget {
            region_id: 1,
            desired_peers: vec![
                (2, TransferRegionPeerRole::Leader),
                (1, TransferRegionPeerRole::Follower),
                (3, TransferRegionPeerRole::Follower),
            ],
            desired_stores: HashSet::from_iter([1, 2, 3]),
        };

        assert!(target.is_converged(&region));
    }

    #[test]
    fn test_online_placement_rejects_missing_or_differently_replicated_target() {
        let desired_stores = HashSet::from_iter([1, 2, 3]);

        assert!(
            validate_online_snapshot_voter_stores(1, None, &desired_stores)
                .unwrap_err()
                .to_string()
                .contains("missing or has no leader")
        );
        assert!(
            validate_online_snapshot_voter_stores(
                1,
                Some(HashSet::from_iter([1, 2])),
                &desired_stores,
            )
            .unwrap_err()
            .to_string()
            .contains("voter count does not match")
        );
        assert!(
            validate_online_snapshot_voter_stores(
                1,
                Some(HashSet::from_iter([1, 2, 3, 4])),
                &desired_stores,
            )
            .unwrap_err()
            .to_string()
            .contains("voter count does not match")
        );
    }

    #[test]
    fn test_online_placement_counts_only_same_size_store_mismatch() {
        let desired_stores = HashSet::from_iter([1, 2, 3]);
        let (_, unaligned) = validate_online_snapshot_voter_stores(
            1,
            Some(HashSet::from_iter([1, 2, 4])),
            &desired_stores,
        )
        .unwrap();

        assert_eq!(unaligned, 1);
    }

    #[test]
    fn test_resolve_master_key_prefers_injected_key() {
        let dfs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(b""));
        let injected = MasterKey::new(&[7; 32]);
        let security_conf = SecurityConfig {
            master_key: cloud_encryption::MasterKeyConfig {
                vendor: "aws".to_owned(),
                key_id: "test-key".to_owned(),
                cipher_text: "%%%not-base64%%%".to_owned(),
                ..Default::default()
            },
            ..Default::default()
        };

        let resolved = resolve_master_key(&dfs, &security_conf, Some(injected.clone())).unwrap();

        assert_eq!(resolved.export(), injected.export());
    }

    #[test]
    fn test_resolve_master_key_falls_back_to_security_config() {
        let dfs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(b""));
        let key_id = "resolve-master-key";
        let security_conf = SecurityConfig {
            master_key: cloud_encryption::MasterKeyConfig {
                vendor: "test".to_owned(),
                key_id: key_id.to_owned(),
                ..Default::default()
            },
            ..Default::default()
        };

        let resolved = resolve_master_key(&dfs, &security_conf, None).unwrap();
        let mut expected = key_id.as_bytes().to_vec();
        expected.resize(32, 11);

        assert_eq!(resolved.export(), MasterKey::new(&expected).export());
    }

    #[test]
    fn test_of_packed_prefers_injected_master_key() {
        const KEYSPACE_ID: u32 = 42;

        let dfs: Arc<dyn Dfs> = Arc::new(new_test_s3fs(b""));
        let injected = MasterKey::new(&[9; 32]);
        let restore_conf = RestoreConfig {
            security: SecurityConfig {
                master_key: cloud_encryption::MasterKeyConfig {
                    vendor: "aws".to_owned(),
                    key_id: "test-key".to_owned(),
                    cipher_text: "%%%not-base64%%%".to_owned(),
                    ..Default::default()
                },
                ..Default::default()
            },
            ..Default::default()
        };
        let packed = new_test_packed_backup(KEYSPACE_ID);
        let dir = TempDirBuilder::new()
            .prefix("packed-master-key")
            .tempdir()
            .unwrap();

        let cluster = BackupCluster::of_packed(
            &packed,
            dir.path().to_path_buf(),
            Arc::new(MockPdClient) as Arc<dyn PdClient>,
            None,
            dfs,
            restore_conf,
            KEYSPACE_ID,
            packed.backup_ts,
            Some(injected.clone()),
        )
        .unwrap();

        assert_eq!(cluster.get_master_key().export(), injected.export());
    }

    #[test]
    fn test_select_schema_file_uses_largest_file_id() {
        let mut higher_version = BackupShard::default();
        higher_version.meta.schema.update_by_restore(10, 21, 0);
        let mut higher_file_id = BackupShard::default();
        higher_file_id.meta.schema.update_by_restore(11, 20, 0);

        assert_eq!(
            BackupCluster::select_schema_file(&[&higher_version, &higher_file_id]),
            Some((11, 20)),
        );
    }

    #[test]
    fn test_append_restored_schema_file_to_retain() {
        let schema_data = build_schema_file(1, 20, Vec::new(), 30);
        let restored_schema = RestoredSchema {
            version: 20,
            target_file: Some(
                SchemaFile::open(Arc::new(InMemFile::new(10, schema_data.into()))).unwrap(),
            ),
        };
        let mut retained_files = Vec::new();

        append_restored_schema_file_to_retain(&restored_schema, &mut retained_files);

        assert_eq!(retained_files.len(), 1);
        assert_eq!(retained_files[0].id, 10);
        assert_eq!(retained_files[0].ftype, FileType::Schema);
    }

    #[test]
    fn test_skip_recent_files_from_retain() {
        const NOW_PHYSICAL_MS: u64 = 1_000_000;
        let file = |created_at_ms| TableFile {
            id: TimeStamp::compose(created_at_ms, 1).into_inner(),
            ftype: FileType::Sst,
            shard_id: 1,
            size: 1,
        };
        let min_age_ms = RECENT_FILE_RETAIN_SKIP_INTERVAL.as_millis() as u64;
        let mut files = vec![
            file(NOW_PHYSICAL_MS - min_age_ms - 1),
            file(NOW_PHYSICAL_MS - min_age_ms),
            file(NOW_PHYSICAL_MS - 1),
            file(NOW_PHYSICAL_MS + 1),
        ];

        let skipped = skip_recent_files_from_retain(&mut files, NOW_PHYSICAL_MS);

        assert_eq!(skipped, 2);
        assert_eq!(files.len(), 2);
        assert_eq!(
            TimeStamp::new(files[0].id).physical(),
            NOW_PHYSICAL_MS - min_age_ms - 1
        );
        assert_eq!(
            TimeStamp::new(files[1].id).physical(),
            NOW_PHYSICAL_MS - min_age_ms
        );
    }

    #[test]
    fn test_derived_file_inventory_and_archive_plan() {
        let start = key_with_suffix(1, b"_r");
        let end = key_with_suffix(1, b"_s");
        let mut meta = ShardMeta::default();
        meta.id = 10;
        let vector_index = new_vector_index(3, 700, &start, &end);
        let vector_file = &vector_index.get_files()[0];
        meta.add_file(
            vector_file.get_id(),
            kvengine::FileMeta::from_vector_index_file(vector_file),
        );
        meta.vector_indexes.push(vector_index);
        meta.fts_l0_files
            .push(new_fts_packed_file(800, &start, &end));

        let files = BackupCluster::get_shard_files(10, &meta);
        assert_eq!(
            files
                .iter()
                .filter(|file| file.ftype == FileType::VectorIndex && file.id == 700)
                .count(),
            1
        );
        assert!(
            files
                .iter()
                .any(|file| { file.ftype == FileType::FtsPackedFile && file.id == 800 })
        );

        let missing_files = vec![
            TableFile {
                id: 1,
                ftype: FileType::Sst,
                shard_id: 10,
                size: 100,
            },
            TableFile {
                id: 800,
                ftype: FileType::FtsPackedFile,
                shard_id: 10,
                size: 100,
            },
        ];
        let (shards_to_reset, files_to_restore) =
            BackupCluster::plan_archive_file_restore(missing_files);
        assert_eq!(shards_to_reset, HashSet::from_iter([10]));
        assert_eq!(files_to_restore.len(), 1);
        assert_eq!(files_to_restore[0].id, 1);
        assert_eq!(files_to_restore[0].ftype, FileType::Sst);
    }

    #[test]
    fn test_derived_metadata_retention_uses_schema_and_source_file_ids() {
        const TRUNCATE_TS: u64 = 1_000;

        let start = key_with_suffix(1, b"_r");
        let end = key_with_suffix(1, b"_s");
        let restored_schema = RestoredSchema {
            version: 20,
            target_file: Some(new_derived_test_schema_file()),
        };
        let mut shard = BackupShard::default();
        shard.region_id = 10;
        shard.meta.range = ShardRange::new(&start, &end);
        shard.meta.schema.update_by_restore(100, 20, 0);
        shard.meta.columnar_table_ids = vec![1];
        shard.meta.add_file(
            900,
            new_columnar_file(900, 10, start.as_slice(), end.as_slice()),
        );

        install_tracked_fts_index(&mut shard.meta, 1, 7);
        shard
            .meta
            .fts_l0_files
            .push(new_fts_packed_file(800, &start, &end));
        shard.meta.vector_indexes.extend([
            // A file ID equal to truncate_ts is safe because file IDs are
            // allocated after the data timestamps represented by the file.
            new_vector_index(3, TRUNCATE_TS, &start, &end),
            new_vector_index(4, 1_001, &start, &end),
            new_vector_index(5, 600, &start, &end),
        ]);

        let retention =
            BackupCluster::derive_metadata_retention(&restored_schema, &shard, TRUNCATE_TS);
        assert!(retention.retain_columnar);
        assert!(retention.retain_fts);
        assert_eq!(
            retention.retained_vector_indexes,
            HashSet::from_iter([(1, 3, 2)])
        );

        let mut target = ShardMeta::default();
        target.range = ShardRange::new(&start, &end);
        target.columnar_table_ids = vec![1];
        assert_eq!(
            BackupCluster::copy_retained_fts(&shard.meta, &mut target),
            1
        );
        assert_eq!(
            BackupCluster::copy_retained_vector_indexes(
                &shard.meta,
                &mut target,
                &retention.retained_vector_indexes,
            ),
            1
        );
        assert_eq!(target.vector_indexes.len(), 1);

        shard
            .meta
            .fts_index_metas
            .get_mut(&1)
            .unwrap()
            .get_mut(&7)
            .unwrap()
            .mut_index_fields()[0]
            .mut_index_field_option()
            .set_analyzer_type(pb::fts::FullTextAnalyzerType::FtsAnalyzerExactValueV1);
        let retention =
            BackupCluster::derive_metadata_retention(&restored_schema, &shard, TRUNCATE_TS);
        assert!(retention.retain_columnar);
        assert!(!retention.retain_fts);
        shard
            .meta
            .fts_index_metas
            .get_mut(&1)
            .unwrap()
            .get_mut(&7)
            .unwrap()
            .mut_index_fields()[0]
            .mut_index_field_option()
            .set_analyzer_type(pb::fts::FullTextAnalyzerType::FtsAnalyzerStandardV1);

        let mut non_overlapping_target = ShardMeta::default();
        non_overlapping_target.range =
            ShardRange::new(&key_with_suffix(1, b"_a"), &key_with_suffix(1, b"_b"));
        non_overlapping_target.columnar_table_ids = vec![1];
        assert_eq!(
            BackupCluster::copy_retained_fts(&shard.meta, &mut non_overlapping_target),
            0
        );
        assert!(non_overlapping_target.fts_indexes.is_empty());

        shard.meta.fts_l0_files[0].set_id(TRUNCATE_TS + 1);
        let retention =
            BackupCluster::derive_metadata_retention(&restored_schema, &shard, TRUNCATE_TS);
        assert!(retention.retain_columnar);
        assert!(!retention.retain_fts);
        assert_eq!(
            retention.retained_vector_indexes,
            HashSet::from_iter([(1, 3, 2)])
        );

        shard.meta.add_file(
            TRUNCATE_TS + 1,
            new_columnar_file(TRUNCATE_TS + 1, 11, &start, &end),
        );
        let retention =
            BackupCluster::derive_metadata_retention(&restored_schema, &shard, TRUNCATE_TS);
        assert!(!retention.retain_columnar);
        assert!(!retention.retain_fts);
        assert!(retention.retained_vector_indexes.is_empty());
    }

    #[test]
    fn test_derived_metadata_retention_rebuilds_columnar_on_table_id_mismatch() {
        let start = key_with_suffix(1, b"_r");
        let end = key_with_suffix(1, b"_s");
        let restored_schema = RestoredSchema {
            version: 20,
            target_file: Some(new_derived_test_schema_file()),
        };
        let mut shard = BackupShard::default();
        shard.meta.range = ShardRange::new(&start, &end);
        shard.meta.schema.update_by_restore(100, 20, 0);
        shard.meta.columnar_table_ids.clear();
        shard.meta.add_file(
            900,
            new_columnar_file(900, 10, start.as_slice(), end.as_slice()),
        );

        let retention = BackupCluster::derive_metadata_retention(&restored_schema, &shard, 1_000);

        assert!(!retention.retain_columnar);
    }

    fn key_with_suffix(table_id: i64, suffix: &[u8]) -> Vec<u8> {
        let mut key = kvengine::table_id::encode_table_prefix_key(table_id).to_vec();
        key.extend_from_slice(suffix);
        key
    }

    fn new_derived_test_schema_file() -> SchemaFile {
        let mut column = kvengine::table::columnar::new_int_handle_column_info();
        column.set_column_id(2);

        let vector_indexes = [3, 4]
            .into_iter()
            .map(|index_id| kvengine::table::columnar::VectorIndexDef {
                index_id,
                col_id: 2,
                dimension: 3,
                ..Default::default()
            })
            .collect();
        let mut fulltext_index = pb::fts::FullTextIndexDef::new();
        fulltext_index.set_index_id(7);
        fulltext_index.set_index_fields(vec![new_test_fts_index_field(2)].into());

        let mut builder = kvengine::table::schema_file::SchemaBufBuilder::new(1);
        builder.columns(
            kvengine::table::columnar::new_int_handle_column_info(),
            kvengine::table::columnar::new_version_column_info(),
            vec![column],
            Vec::new(),
            2,
            vector_indexes,
            vec![fulltext_index],
        );
        let schema = kvengine::table::schema_file::Schema::new(builder.build());
        let data = build_schema_file(2, 20, vec![schema], 1_000);
        SchemaFile::open(Arc::new(InMemFile::new(200, data.into()))).unwrap()
    }

    fn new_columnar_file(
        id: u64,
        snap_version: u64,
        smallest: &[u8],
        biggest: &[u8],
    ) -> kvengine::FileMeta {
        let mut file = pb::ColumnarCreate::new();
        file.set_id(id);
        file.set_level(0);
        file.set_snap_version(snap_version);
        file.set_smallest(smallest.to_vec());
        file.set_biggest(biggest.to_vec());
        kvengine::FileMeta::from_columnar_table(&file)
    }

    fn new_vector_index(
        index_id: i64,
        file_id: u64,
        smallest: &[u8],
        biggest: &[u8],
    ) -> VectorIndex {
        let mut file = pb::VectorIndexFile::new();
        file.set_id(file_id);
        file.set_snap_version(10);
        file.set_smallest(smallest.to_vec());
        file.set_biggest(biggest.to_vec());

        let mut index = VectorIndex::new();
        index.set_table_id(1);
        index.set_index_id(index_id);
        index.set_col_id(2);
        index.set_snap_version(10);
        index.set_files(vec![file].into());
        index
    }

    fn new_fts_packed_file(id: u64, smallest: &[u8], biggest: &[u8]) -> pb::FtsPackedFileInfo {
        let mut file = pb::FtsPackedFileInfo::new();
        file.set_id(id);
        file.set_smallest(smallest.to_vec());
        file.set_biggest(biggest.to_vec());
        file
    }

    fn install_tracked_fts_index(meta: &mut ShardMeta, table_id: i64, index_id: i64) {
        let tracked_meta = kvengine::table::fts::new_tracked_index_meta(
            table_id,
            index_id,
            vec![new_test_fts_index_field(2)],
            pb::fts::FtsFieldLayout::FtsFieldLayoutIndexFieldOrdinalString,
        );
        meta.fts_indexes.insert((table_id, index_id));
        meta.fts_index_metas
            .entry(table_id)
            .or_default()
            .insert(index_id, tracked_meta);
    }

    fn new_test_fts_index_field(col_id: i64) -> pb::fts::FullTextIndexFieldDef {
        let mut option = pb::fts::FullTextIndexFieldOption::new();
        option.set_analyzer_type(pb::fts::FullTextAnalyzerType::FtsAnalyzerStandardV1);
        option.set_enable_bm25(true);
        let mut field = pb::fts::FullTextIndexFieldDef::new();
        field.set_col_id(col_id);
        field.set_index_field_option(option);
        field
    }

    #[test]
    fn test_handle_overlapping_shards() {
        const KEYSPACE_ID: u32 = 42;
        let make_backup_shard = |tuple: (
            u64,  // shard_id
            u64,  // ver
            &str, // start, "00" for start of keyspace
            &str, // end, "99" for end of keyspace
        )|
         -> BackupShard {
            let mut shard = BackupShard::default();
            shard.region_id = tuple.0;
            shard.meta.ver = tuple.1;
            let outer_start = if tuple.2 == "00" {
                ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID)
            } else {
                [
                    ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID),
                    tuple.2.as_bytes().to_vec(),
                ]
                .concat()
            };
            let outer_end = if tuple.3 == "99" {
                ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID + 1)
            } else {
                [
                    ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID),
                    tuple.3.as_bytes().to_vec(),
                ]
                .concat()
            };
            shard.meta.range = ShardRange::new(&outer_start, &outer_end);
            shard.meta.inner_key_off = 4;
            shard
        };

        let add_leader_shard =
            |leader_shards: &mut HashMap<IdVer, BackupShard>,
             _store_id: u64,
             shards: Vec<(u64, u64, &str, &str)>| {
                for shard_tuple in shards {
                    let shard = make_backup_shard(shard_tuple);
                    leader_shards.insert(shard.id_ver(), shard);
                }
            };

        let cases = vec![
            (
                vec![(1, 100, "00", "99")], // Vec<(shard_id, ver, start, end)>
                vec![(1, 100, "00", "99")],
                vec![(1, 100, "00", "99")],
                // expected: Option<Vec<(shard_id, ver, start,end)>>, None means error
                Some(vec![(1, 100, "00", "99")]),
            ),
            (
                vec![(1, 100, "00", "01"), (2, 200, "01", "99")],
                vec![(1, 100, "00", "01"), (2, 200, "01", "99")],
                vec![(1, 100, "00", "01"), (2, 200, "01", "99")],
                Some(vec![(1, 100, "00", "01"), (2, 200, "01", "99")]),
            ),
            (
                vec![(1, 100, "00", "01"), (2, 200, "01", "99")],
                vec![(1, 201, "00", "99")], // merge from shard 1 & 2
                vec![(1, 100, "00", "01"), (2, 200, "01", "99")],
                Some(vec![(1, 201, "00", "99")]),
            ),
            (
                vec![(1, 100, "00", "01"), (2, 200, "01", "99")],
                vec![(2, 201, "00", "99")], // merge from shard 1 & 2
                vec![(1, 100, "00", "01"), (2, 200, "01", "99")],
                Some(vec![(2, 201, "00", "99")]),
            ),
            (
                vec![(1, 101, "00", "01"), (2, 101, "01", "99")], // split from shard 1
                vec![(1, 100, "00", "99")],
                vec![(1, 100, "00", "99")],
                Some(vec![(1, 101, "00", "01"), (2, 101, "01", "99")]),
            ),
            (
                vec![(2, 100, "01", "02")],
                vec![(1, 100, "00", "01")],
                vec![(3, 100, "02", "99")],
                Some(vec![
                    (1, 100, "00", "01"),
                    (2, 100, "01", "02"),
                    (3, 100, "02", "99"),
                ]),
            ),
            (
                vec![(1, 100, "00", "99")],
                vec![(1, 100, "00", "99")],
                vec![(1, 101, "01", "99")],
                Some(vec![(1, 100, "00", "99")]),
            ),
            (
                vec![(1, 100, "00", "99")],
                vec![(1, 100, "00", "99")],
                vec![(1, 101, "00", "01")],
                Some(vec![(1, 100, "00", "99")]),
            ),
            (
                vec![(1, 102, "00", "01"), (2, 102, "01", "02")],
                vec![(1, 101, "00", "01")], // "01" is visited.
                vec![(1, 100, "00", "99")],
                Some(vec![(1, 100, "00", "99")]),
            ),
            (
                vec![(1, 100, "00", "01"), (2, 100, "01", "99")],
                vec![(2, 101, "01", "02"), (3, 101, "02", "03")],
                vec![(2, 101, "01", "02"), (3, 102, "02", "04")],
                Some(vec![(1, 100, "00", "01"), (2, 100, "01", "99")]),
            ),
            #[cfg_attr(rustfmt, rustfmt_skip)]
            (
                vec![(1, 100, "00", "01"), (2, 100, "01", "03"), (3, 100, "03", "04"), (4, 100, "04", "06"), (7, 102, "06", "07"), (8, 102, "07", "99")],
                vec![(1, 100, "00", "01"), (2, 101, "01", "02"), (3, 101, "02", "04"), (4, 100, "04", "05"), (5, 100, "05", "99")],
                vec![(1, 100, "00", "01"), (2, 100, "01", "03"), (3, 100, "03", "04"), (5, 101, "04", "06"), (6, 101, "06", "07")],
                Some(vec![
                    (1, 100, "00", "01"), (2, 101, "01", "02"), (3, 101, "02", "04"), (5, 101, "04", "06"), (7, 102, "06", "07"), (8, 102, "07", "99"),
                ]),
            ),
            (
                vec![(1, 100, "00", "01")],
                vec![(2, 100, "00", "02")],
                vec![(1, 100, "00", "01")],
                None,
            ),
            (
                vec![(1, 100, "01", "02")],
                vec![(1, 100, "01", "02")],
                vec![(1, 100, "01", "02")],
                None,
            ),
            (
                vec![(1, 100, "01", "99")],
                vec![(1, 100, "01", "99")],
                vec![(1, 100, "01", "99")],
                None,
            ),
        ];

        for (case_idx, (store0, store1, store2, expected)) in cases.into_iter().enumerate() {
            let mut leader_shards = HashMap::default();
            add_leader_shard(&mut leader_shards, 0, store0);
            add_leader_shard(&mut leader_shards, 1, store1);
            add_leader_shard(&mut leader_shards, 2, store2);

            let res = BackupCluster::handle_overlapping_shards("", leader_shards);
            if let Some(expected) = expected {
                let expected_sorted_shards: Vec<u64> = expected.iter().map(|x| x.0).collect();
                let expected_leader_shards: HashMap<u64, BackupShard> =
                    HashMap::from_iter(expected.into_iter().map(|x| (x.0, make_backup_shard(x))));

                let (leader_shards, sorted_shards) = res.unwrap_or_else(|err| {
                    panic!("case: {}: {:?}", case_idx, err);
                });
                assert_eq!(
                    BTreeMap::from_iter(leader_shards.into_iter()),
                    BTreeMap::from_iter(expected_leader_shards.into_iter()),
                    "case: {}",
                    case_idx
                );
                assert_eq!(sorted_shards, expected_sorted_shards, "case: {}", case_idx);
            } else {
                assert!(res.is_err(), "case: {}", case_idx);
            }
        }
    }

    #[test]
    fn test_online_snapshot_placement_plan_splits_source_boundaries_before_capture() {
        let source_keyspace_id = 1;
        let target_keyspace_id = 2;
        let (source_start, source_end) = ApiV2::get_keyspace_range_by_id(source_keyspace_id);
        let (target_start, target_end) = ApiV2::get_keyspace_range_by_id(target_keyspace_id);
        let source_m = [source_start.as_slice(), b"m"].concat();
        let source_t = [source_start.as_slice(), b"t"].concat();
        let target_h = [target_start.as_slice(), b"h"].concat();
        let target_m = [target_start.as_slice(), b"m"].concat();
        let target_t = [target_start.as_slice(), b"t"].concat();

        let make_source_region = |id, start: &[u8], end: &[u8], leader_store| {
            let mut region = metapb::Region::default();
            region.set_id(id);
            region.set_start_key(Key::from_raw(start).into_encoded());
            region.set_end_key(Key::from_raw(end).into_encoded());
            let peers = (1..=3)
                .map(|store_id| {
                    let mut peer = metapb::Peer::default();
                    peer.set_id(id * 10 + store_id);
                    peer.set_store_id(store_id);
                    peer
                })
                .collect::<Vec<_>>();
            let leader = peers
                .iter()
                .find(|peer| peer.get_store_id() == leader_store)
                .unwrap()
                .clone();
            region.set_peers(peers.into());
            (region, leader)
        };
        let sources = vec![
            make_source_region(1, &source_start, &source_m, 1),
            make_source_region(2, &source_m, &source_t, 2),
            make_source_region(3, &source_t, &source_end, 3),
        ];
        let plan = OnlineSnapshotPlacementPlan::new(&sources).unwrap();

        let unsplit_target = vec![RawRegion {
            id: 10,
            raw_start: target_start.clone(),
            raw_end: target_end.clone(),
            ..Default::default()
        }];
        let (coarse_split_keys, split_keys) = plan.split_keys(&unsplit_target, &target_start, 2);
        assert_eq!(coarse_split_keys, vec![target_t.clone()]);
        assert_eq!(split_keys, vec![target_m.clone()]);
        assert!(plan.placement_targets(unsplit_target).is_err());

        // An existing target-only boundary is harmless: both child Regions map
        // to the same source placement, while every source boundary is present.
        let target_regions = [
            (&target_start, &target_h),
            (&target_h, &target_m),
            (&target_m, &target_t),
            (&target_t, &target_end),
        ]
        .into_iter()
        .enumerate()
        .map(|(idx, (start, end))| RawRegion {
            id: 20 + idx as u64,
            raw_start: start.clone(),
            raw_end: end.clone(),
            ..Default::default()
        })
        .collect::<Vec<_>>();
        let targets = plan.placement_targets(target_regions).unwrap();
        let preferred_leader_store = |target: &OnlinePlacementTarget| {
            target
                .desired_peers
                .iter()
                .find_map(|(store_id, role)| {
                    matches!(role, &TransferRegionPeerRole::Leader).then_some(*store_id)
                })
                .unwrap()
        };
        assert_eq!(targets.len(), 4);
        assert_eq!(preferred_leader_store(&targets[&20]), 1);
        assert_eq!(preferred_leader_store(&targets[&21]), 1);
        assert_eq!(preferred_leader_store(&targets[&22]), 2);
        assert_eq!(preferred_leader_store(&targets[&23]), 3);
    }

    #[test]
    fn test_align_target_regions() {
        test_align_target_regions_helper(false);
        test_align_target_regions_helper(true);
    }

    fn test_align_target_regions_helper(inner_key_off_enabled: bool) {
        let inner_key_off = if inner_key_off_enabled {
            KEYSPACE_PREFIX_LEN
        } else {
            0
        };
        let key_prefix = b"x000";
        let cases: Vec<(Vec<&[u8]>, Vec<&[u8]>, Vec<Vec<u64>>)> = vec![
            (
                vec![b"0", b"1"], // backup_regions_boundary_keys, the index is the region id.
                vec![b"0", b"1"], // target_regions_boundary_keys
                vec![vec![0]],    // align_regions_id
            ),
            (
                vec![b"0", b"05", b"1"],
                vec![b"0", b"05", b"1"],
                vec![vec![0], vec![1]],
            ),
            (
                vec![b"1", b"13", b"16", b"2"],
                vec![b"1", b"15", b"2"],
                vec![vec![0, 1], vec![1, 2]],
            ),
            (
                vec![b"2", b"25", b"3"],
                vec![b"2", b"23", b"26", b"3"],
                vec![vec![0], vec![0, 1], vec![1]],
            ),
            #[cfg_attr(rustfmt, rustfmt_skip)]
            (
            vec![ b"3",    b"32",            b"34",  b"36", b"37", b"38", b"39", b"4"],
            vec![ b"3",    b"32",   b"33",   b"34",  b"36",                      b"4"],
            vec![vec![0], vec![1], vec![1], vec![2],         vec![3,4,5,6]],
            ),
            #[cfg_attr(rustfmt, rustfmt_skip)]
            (
                vec![b"2", b"25", b"41",               b"51", b"55"],
                vec!                    [b"42", b"5"],
                vec![vec![2]],
            ),
            #[cfg_attr(rustfmt, rustfmt_skip)]
            (
                vec![b"1", b"11", b"2", b"25",       b"41", b"5"],
                vec![                          b"4",        b"5"],
                vec![vec![3, 4]],
            ),
            #[cfg_attr(rustfmt, rustfmt_skip)]
            (
                vec![b"1", b"11", b"2", b"25", b"4", b"41", b"5"],
                vec![                          b"4",        b"5"],
                vec![vec![4, 5]],
            ),
        ];

        let make_regions = |keys: Vec<&[u8]>| -> Vec<RawRegion> {
            let mut regions = Vec::with_capacity(keys.len() - 1);
            for (idx, w) in keys.as_slice().windows(2).enumerate() {
                let (mut raw_start, mut raw_end) = (key_prefix.to_vec(), key_prefix.to_vec());
                raw_start.extend_from_slice(w[0]);
                raw_end.extend_from_slice(w[1]);
                regions.push(RawRegion {
                    id: idx as u64,
                    raw_start,
                    raw_end,
                    ..Default::default()
                });
            }

            regions
        };

        let make_backup_shards = |regions: Vec<RawRegion>| {
            let mut shards = HashMap::with_capacity(regions.len());
            let mut shards_id = Vec::with_capacity(regions.len());
            for r in regions {
                let mut shard = BackupShard::default();
                shard.region_id = r.id;
                shard.meta.range = ShardRange::new(r.get_start_key(), r.get_end_key());
                shard.meta.inner_key_off = inner_key_off;

                shards_id.push(shard.region_id);
                shards.insert(shard.region_id, shard);
            }

            (shards, shards_id)
        };

        for (case_idx, (backup_shards_key, target_regions_key, expected_regions_id)) in
            cases.into_iter().enumerate()
        {
            let backup_regions = make_regions(backup_shards_key);
            let (backup_shards, backup_shards_id) = make_backup_shards(backup_regions);

            let target_regions = make_regions(target_regions_key);

            let mut expected = Vec::with_capacity(target_regions.len());
            for (i, target_region) in target_regions.clone().into_iter().enumerate() {
                expected.push(AlignedRegion {
                    target_region,
                    backup_shards_id: expected_regions_id[i].clone(),
                });
            }

            let aligned_regions = BackupCluster::align_target_regions_impl(
                &backup_shards_id,
                &backup_shards,
                target_regions,
            );
            assert_eq!(aligned_regions, expected, "case: {}", case_idx);
        }
    }
}