// Copyright 2025 PingCAP, Inc.

#![feature(box_patterns)]

mod bootstrap;
mod error;
mod manifest;
mod metrics;
mod preprocessor;
#[cfg(test)]
mod test_util;
mod unsafe_recover;
mod util;

use std::{
    cmp,
    collections::{
        HashMap as StdHashMap, HashSet as StdHashSet, VecDeque, hash_map::Entry as HashMapEntry,
    },
    fmt, fs, io, mem, ops,
    path::{Path, PathBuf},
    sync::Arc,
    thread,
    time::Duration,
};

use api_version::ApiV2;
use bytes::{Buf, BufMut, Bytes};
use cloud_encryption::MasterKey;
use collections::{HashMap, HashMapExt, HashSet};
pub use error::{Error, Result};
use file_system::{IoRateLimitMode, IoRateLimiter};
use kvengine::{
    IdVer, MetaIterator, Shard, ShardMeta, ShardTag, TERM_KEY,
    dfs::Dfs,
    ia::util::IaConfig,
    limiter::StoreLimiter,
    table::{
        InnerKey,
        tiny_meta::{CompactKeeper, MetaPackConfig, MetaPacker},
    },
};
use kvenginepb::ChangeSet;
use kvproto::{
    metapb,
    metapb::{Peer, PeerRole},
    raft_cmdpb::{AdminRequest, RaftCmdRequest},
    raft_serverpb::{MergeState, PeerState, RegionLocalState, StoreIdent},
};
use log_wrappers::Value as LogValue;
use native_br::{
    common::{
        ReplayWalLogsContext, collect_snapshot_meta_rlog_files, get_latest_backup_meta,
        replay_wal_logs_from_backup,
    },
    error::Error as BrError,
};
use pd_client::PdClient;
use protobuf::Message;
use raft::StateRole;
use raft_proto::{eraftpb, eraftpb::Entry};
use raftstore::coprocessor::{
    BoxRegionChangeObserver, CoprocessorHost, RegionChangeEvent, RegionChangeReason,
    RegionInfoAccessor,
};
use rfengine::{
    RaftLogOp, RfEngine, TRUNCATE_ALL_INDEX, WriteBatch, iterator::WalIterator, raft_state_key,
    region_state_key,
};
use rfenginepb::{ClusterBackupMeta, StoreBackupMeta};
use rfstore::{
    RaftRouter, RaftRouterReceiver, raft_router_peer_channel,
    store::{
        AffectMemtable, Applier, ApplyContext, ApplyMsgs, MetaChangeListener, PdIdAllocator,
        PeerMsg, PeerTag, PreprocessContext, PreprocessRef, PreprocessResult, PreprocessScene,
        RAFT_INIT_LOG_INDEX, RecoverHandler, RegionIdVer, StoreMsg, get_preprocess_cmd,
        is_region_initialized, load_engine_meta_from_wb, load_last_raft_state_from_wb, rlog,
        state::{RaftApplyState, RaftState},
        write_engine_meta,
    },
};
use security::SecurityConfig;
use serde_derive::{Deserialize, Serialize};
use tikv::config::TikvConfig;
use tikv_util::{
    box_err, box_try,
    config::{AbsoluteOrPercentSize, ReadableDuration, ReadableSize},
    debug, error, info, mpsc,
    time::Instant,
    trace, warn,
};
use txn_types::TimeStamp;

use crate::{
    manifest::{Manifest, UncommittedEntries},
    metrics::*,
    preprocessor::{Preprocessor, RecoverTxnFileLocksContext},
    util::RegionPersistProgress,
};

macro_rules! try_force_stop {
    ($self:ident, $expr:expr) => {{
        #[cfg(feature = "testexport")]
        if $self.ctx.force_stop.get() {
            info!("merged engine force stopped");
            return $expr;
        }
    }};
}

macro_rules! try_force_stop_err {
    ($self:ident) => {{
        try_force_stop!($self, Err(Error::ForceStopped));
    }};
}

// The quorum size when replicas number is 3.
// Used to check whether the Raft log is committed.
const QUORUM_SIZE: u8 = 2;

const MANIFEST_DIR_NAME: &str = "manifest";
const MERGED_DIR_NAME: &str = "merged";
const MERGED_STALE_DIR_NAME: &str = "merged.stale";

const ALL_KV_ENGINE_META_KEYS: &[&[u8]] = &[
    rfengine::KV_ENGINE_META_KEY,
    rfengine::KV_ENGINE_META_DIFF_KEY,
    rfengine::KV_ENGINE_META_SNAP_DIFF_KEY,
];

#[derive(Clone)]
pub struct MergedEngineContext {
    pub pd: Arc<dyn PdClient>,
    pub fs: Arc<dyn Dfs>,
    pub local_dir: PathBuf,
    pub master_key: MasterKey,
    pub config: MergedEngineConfig,
    pub security_config: Arc<SecurityConfig>,
    pub encryption_key_manager: Arc<cloud_encryption::EncryptionKeyManager>,

    #[allow(dead_code)]
    pub force_stop: ForceStop, // For test purpose.
}

impl MergedEngineContext {
    fn manifest_dir(&self) -> PathBuf {
        self.local_dir.join(MANIFEST_DIR_NAME)
    }

    fn merged_dir(&self) -> PathBuf {
        self.local_dir.join(MERGED_DIR_NAME)
    }

    fn merged_stale_dir(&self) -> PathBuf {
        self.local_dir.join(MERGED_STALE_DIR_NAME)
    }
}

/// Policy for selecting pinned shards in lazy shard loading mode.
///
/// A pinned shard is expected to stay available (loaded in kvengine) even when
/// lazy shard loading is enabled.
pub trait PinnedShardPolicy: Send + Sync {
    /// Returns whether the shard snapshot should be pinned.
    ///
    /// `snapshot.outer_start` / `snapshot.outer_end` are raw outer keys.
    fn should_pin_snapshot(&self, snapshot: &kvenginepb::Snapshot) -> bool;

    /// Returns whether the region should be pinned.
    ///
    /// `region.start_key` / `region.end_key` are outer encoded keys, matching
    /// what is stored in rfengine peer state.
    ///
    /// An empty end key represents an infinite upper bound.
    fn should_pin_region(&self, region: &metapb::Region) -> bool;

    /// Returns whether an inner key should be treated as pinned.
    ///
    /// `key` is an inner key (without keyspace prefix).
    fn should_pin_inner_key(&self, key: InnerKey<'_>) -> bool;

    /// Returns whether an inner key range should be treated as pinned.
    ///
    /// `start_key` / `end_key` are inner keys (without keyspace prefixes).
    ///
    /// An empty end key represents an infinite upper bound.
    fn should_pin_inner_range(&self, start_key: InnerKey<'_>, end_key: InnerKey<'_>) -> bool;
}

#[derive(Clone)]
pub enum LoadShardsMode {
    /// Load kvengine shards for all enabled keyspaces on startup.
    EagerAll,

    /// Do not load kvengine shards on startup.
    ///
    /// Shards will be loaded on-demand by replication_worker (e.g. on TiCDC
    /// register), while a bounded subset is still loaded eagerly by
    /// `pinned_policy`.
    Lazy {
        pinned_policy: Arc<dyn PinnedShardPolicy>,
    },
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum KeyspaceShardMetaScope {
    /// Prepare/load all shards for the keyspace.
    All,
    /// Prepare/load only pinned shards for the keyspace.
    Pinned,
}

impl KeyspaceShardMetaScope {
    pub fn label(&self) -> &'static str {
        match self {
            KeyspaceShardMetaScope::All => "all",
            KeyspaceShardMetaScope::Pinned => "pinned",
        }
    }
}

pub struct KeyspaceShardMetas {
    pub scope: KeyspaceShardMetaScope,
    pub metas: StdHashMap<u64 /* region_id */, ShardMeta>,
}

#[derive(Clone)]
pub struct MergedEngineInitOpts {
    pub load_shards_mode: LoadShardsMode,
}

impl Default for MergedEngineInitOpts {
    fn default() -> Self {
        Self {
            load_shards_mode: LoadShardsMode::EagerAll,
        }
    }
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct MergedEngineConfig {
    pub block_cache_size: AbsoluteOrPercentSize,
    pub timeout_fetch_wal: ReadableDuration,
    pub merged_store_id: u64,
    pub mem_table_size: ReadableSize,
    pub raft_write_batch_size: ReadableSize,
    pub force_ia: bool,
    pub get_latest_backup_timeout: ReadableDuration,
    pub meta_pack: MetaPackConfig,
}

impl Default for MergedEngineConfig {
    fn default() -> Self {
        Self {
            block_cache_size: AbsoluteOrPercentSize::Percent(10.0),
            timeout_fetch_wal: ReadableDuration::secs(30),
            merged_store_id: 1024,
            mem_table_size: ReadableSize::mb(128),
            raft_write_batch_size: ReadableSize::mb(4),
            force_ia: true,
            get_latest_backup_timeout: ReadableDuration::minutes(30),
            meta_pack: MetaPackConfig::default(),
        }
    }
}

const MAX_PENDING_TARGETS: usize = 8;

/// Target of region progress.
///
/// `ts`: The target timestamp of the corresponding WAL.
/// `log_idx`: The last log index of the region in the corresponding WAL.
///
/// When `commit_index` reaches `log_idx`, the region can safely advance
/// `safe_target_ts` (resolved_ts) to `ts`.
#[derive(Clone, Copy, Debug, PartialEq, PartialOrd)]
struct RegionProgressTarget {
    ts: TimeStamp,
    log_idx: u64,
}

#[derive(Clone)]
pub struct RegionProgress {
    pub keyspace_id: u32,
    pub region_id: u64,
    pub entries: HashMap<u64 /* log_index */, RaftLogOpWithCounter>,
    pub synced_index: u64,
    // Use `commit_index()`/`update_commit_index()` to read/write.
    commit_index: u64,
    // Last index of uncommitted raft logs (in `entries`).
    // Note that it's not necessary to be `commit_index <= last_index`.
    last_index: u64,
    persist_progress: RegionPersistProgress,

    safe_target_ts: TimeStamp,
    pending_targets: VecDeque<RegionProgressTarget>,
}

impl fmt::Debug for RegionProgress {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("RegionProgress")
            .field("keyspace", &self.keyspace_id)
            .field("region", &self.region_id)
            .field("entries", &self.entries.len())
            .field("synced", &self.synced_index)
            .field("commit", &self.commit_index)
            .field("last", &self.last_index)
            .field("persist", &self.persist_progress)
            .field("safe_target", &self.safe_target_ts)
            .field("pending_targets", &self.pending_targets)
            .finish()
    }
}

impl RegionProgress {
    pub fn new(keyspace_id: u32, region_id: u64) -> Self {
        Self {
            keyspace_id,
            region_id,
            entries: HashMap::default(),
            synced_index: 0,
            commit_index: 0,
            last_index: 0,
            persist_progress: RegionPersistProgress::default(),
            safe_target_ts: TimeStamp::zero(),
            pending_targets: VecDeque::default(),
        }
    }

    #[inline]
    pub fn commit_index(&self) -> u64 {
        self.commit_index
    }

    #[inline]
    pub fn update_commit_index(&mut self, new_commit_index: u64) -> bool {
        if self.commit_index < new_commit_index {
            self.commit_index = new_commit_index;
            true
        } else {
            false
        }
    }

    #[inline]
    fn update_last_index(&mut self, new_last_index: u64) {
        if self.last_index < new_last_index {
            self.last_index = new_last_index;
        }
    }

    // NOTE: this method is NOT reentrant.
    pub fn upsert_entry<F>(&mut self, tag: ShardTag, log_index: u64, term: u32, or_insert: F)
    where
        F: FnOnce() -> RaftLogOpWithCounter,
    {
        use std::cmp::Ordering::{Equal, Greater, Less};

        if let Some(existing_op) = self.entries.get_mut(&log_index) {
            debug_assert_eq!(log_index, existing_op.index);
            match existing_op.term.cmp(&term) {
                Greater => return,
                Equal => {
                    existing_op.inc_counter();
                    if existing_op.counter() >= QUORUM_SIZE && existing_op.index > self.commit_index
                    {
                        self.commit_index = existing_op.index;
                        debug!(
                            "{} upsert_entry: advance commit index: {}",
                            tag, self.commit_index
                        );
                    }
                    return;
                }
                Less => {}
            }
        }
        self.entries.insert(log_index, or_insert());
        self.update_last_index(log_index);
    }

    pub fn update_safe_target_ts(&mut self, tag: ShardTag, target_ts: TimeStamp) -> TimeStamp {
        debug_assert!(self.safe_target_ts <= target_ts);
        if self.safe_target_ts >= target_ts {
            return self.safe_target_ts;
        }

        // All logs are committed.
        if self.commit_index >= self.last_index {
            self.safe_target_ts = target_ts;
            debug!("{} update_safe_target_ts: {:?}", tag, self.safe_target_ts);
            if !self.pending_targets.is_empty() {
                self.pending_targets.clear();
                self.pending_targets.shrink_to_fit();
            }
            return self.safe_target_ts;
        }

        // Find the latest committed log and associated target ts.
        while let Some(pending) = self.pending_targets.front() {
            if self.commit_index >= pending.log_idx {
                self.safe_target_ts = pending.ts;
                debug!("{} update_safe_target_ts: {:?}", tag, self.safe_target_ts);
                let _ = self.pending_targets.pop_front();
            } else {
                break;
            }
        }

        // Append the new target.
        let new_pending = RegionProgressTarget {
            ts: target_ts,
            log_idx: self.last_index,
        };
        match self.pending_targets.back_mut() {
            Some(last) if *last >= new_pending => {}
            Some(last) if last.ts == target_ts => {
                last.log_idx = self.last_index;
            }
            Some(_) | None => {
                if self.pending_targets.len() >= MAX_PENDING_TARGETS {
                    self.pending_targets.pop_front();
                }
                self.pending_targets.push_back(new_pending);
            }
        }

        self.safe_target_ts
    }

    pub fn is_synced(&self) -> bool {
        debug_assert!(
            self.synced_index <= self.commit_index,
            "unexpected progress: {:?}",
            self
        );
        self.synced_index >= self.commit_index
    }

    /// Returns the region seen watermark for register gating.
    ///
    /// The returned timestamp means merged_engine has observed WAL up to this
    /// target for this region. It is only valid after `is_synced()` becomes
    /// true; otherwise `None` is returned.
    pub fn get_safe_target_ts(&self) -> Option<TimeStamp> {
        self.is_synced().then_some(self.safe_target_ts)
    }

    pub fn data_persisted_log_index(&self) -> u64 {
        self.persist_progress.persisted_idx()
    }
}

#[derive(Clone, Copy, Default, Debug, PartialEq, Eq)]
pub struct StoreProgress {
    pub store_id: u64,
    pub epoch: u32,
    pub offset: u64,
}

impl fmt::Display for StoreProgress {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(
            f,
            "{{ store: {}, epoch: {}, offset: {} }}",
            self.store_id, self.epoch, self.offset
        )
    }
}

impl PartialOrd for StoreProgress {
    fn partial_cmp(&self, other: &Self) -> Option<cmp::Ordering> {
        (self.store_id == other.store_id).then(|| {
            self.epoch
                .cmp(&other.epoch)
                .then_with(|| self.offset.cmp(&other.offset))
        })
    }
}

impl Ord for StoreProgress {
    fn cmp(&self, other: &Self) -> cmp::Ordering {
        self.partial_cmp(other).unwrap()
    }
}

impl StoreProgress {
    pub(crate) fn encode(&self, buf: &mut Vec<u8>) {
        buf.put_u64_le(self.store_id);
        buf.put_u32_le(self.epoch);
        buf.put_u64_le(self.offset);
    }

    pub(crate) fn decode(buf: &mut impl Buf) -> Result<Self> {
        if (buf.remaining()) < 20 {
            return Err(std::io::Error::new(
                std::io::ErrorKind::UnexpectedEof,
                "store progress buffer is too short",
            )
            .into());
        }
        let store_id = buf.get_u64_le();
        let epoch = buf.get_u32_le();
        let offset = buf.get_u64_le();
        Ok(Self {
            store_id,
            epoch,
            offset,
        })
    }
}

struct EmptyMetaIterator {}

impl MetaIterator for EmptyMetaIterator {
    fn iterate<F>(&mut self, _: F) -> kvengine::Result<()>
    where
        F: FnMut(ChangeSet),
    {
        Ok(())
    }

    fn engine_id(&self) -> u64 {
        0
    }
}

pub struct MergedEngine {
    pub ctx: MergedEngineContext,
    manifest: Manifest,
    region_progresses: HashMap<u64, RegionProgress>,
    updated_regions: HashSet<u64>,
    raft: RfEngine,
    pub kv: kvengine::Engine,
    pub recover_handler: RecoverHandler,
    load_shards_mode: LoadShardsMode,
    preprocessors: HashMap<u64, Preprocessor>,
    appliers: HashMap<u64, Applier>,
    delay_destroy_regions: HashMap<u64 /* region_id */, u32 /* keyspace_id */>,
    installable_region_bootstraps: HashMap<u64 /* region_id */, InstallableRegionBootstrap>,
    newly_bootstrapped_regions: Vec<u64>,
    pending_merge_states: HashMap<u64 /* source_region_id */, MergeState>,
    peer_receiver: RaftRouterReceiver,
    _store_receiver: mpsc::Receiver<StoreMsg>, // applier never send store message.
    router: RaftRouter,
    cop_host: CoprocessorHost,
    region_info_accessor: RegionInfoAccessor,
    meta_packer: Option<MetaPacker>,
    closed: bool,
}

impl MergedEngine {
    pub fn new(
        ctx: MergedEngineContext,
        backup_meta: Option<ClusterBackupMeta>,
        init_opts: MergedEngineInitOpts,
    ) -> Result<Self> {
        let manifest_dir = ctx.manifest_dir();
        let mut manifest = box_try!(Manifest::open(&manifest_dir));
        let is_bootstrap = manifest.store_progresses.is_empty();
        let merged_dir = ctx.merged_dir();
        if is_bootstrap && box_try!(merged_dir.try_exists()) {
            let merged_stale_dir = ctx.merged_stale_dir();
            box_try!(bootstrap::reset_merged_dir_on_bootstrap(
                &merged_dir,
                &merged_stale_dir,
            ));
        }

        let mut merged_cfg = rfengine::RfEngineConfig::default();
        merged_cfg.disable_rlog_offloading();
        let raft = box_try!(RfEngine::open(
            merged_dir.as_path(),
            &merged_cfg,
            None,
            None
        ));
        let merged_store_id = ctx.config.merged_store_id;
        if let Some(store_ident) = rfengine::load_store_ident(&raft) {
            if store_ident.get_store_id() != merged_store_id {
                panic!(
                    "store id mismatch, expect {}, got {}",
                    merged_store_id,
                    store_ident.get_store_id()
                );
            }
        } else {
            let mut store_ident = StoreIdent::new();
            store_ident.set_store_id(merged_store_id);
            rfengine::save_store_ident(&raft, &store_ident);
        }
        raft.set_engine_id(merged_store_id);
        let (mut region_progresses, tombstone_regions) = if is_bootstrap {
            let (region_progresses, store_progresses, backup_meta) =
                box_try!(Self::recover_from_backup(&ctx, backup_meta, &raft));
            manifest.store_progresses = store_progresses;
            // Note: If we persist the manifest and restart before sync to the `backup_ts`,
            // the time span of next target will exceed the
            // `wal.max_target_time_span`. The exceeded duration should be no more
            // that an interval of backup, i.e. 1 min in prod env, and should not
            // be a problem. So we do not handle this condition for easier.
            manifest.synced_target_ts = TimeStamp::new(backup_meta.backup_ts);

            // Note: manifest is not persisted here to avoid saving all entries. If
            // replication worker restart before next loop, we will recover from backup
            // again.
            (region_progresses, HashMap::default())
        } else {
            if manifest.synced_target_ts.is_zero() {
                // For backward compatibility. Then `WalProgressFetcher` will get latest
                // progress from rfengine.
                let now = TimeStamp::now();
                warn!("No synced_target_ts in manifest, use now: {}", now);
                manifest.synced_target_ts = now;
            }

            box_try!(Self::recover_from_merged_raft_engine(
                &raft,
                &manifest.uncommitted_entries
            ))
        };
        let mut pending_merge_states = HashMap::default();
        let mut preprocessors = HashMap::default();
        let mut cop_host = CoprocessorHost::default();
        let mut updated_regions = HashSet::default();
        let region_info_accessor = RegionInfoAccessor::new(&mut cop_host);
        let mut recover_txn_file_locks_ctx =
            RecoverTxnFileLocksContext::new(merged_store_id, &raft);
        for (region_id, _) in raft.get_region_peer_map() {
            if region_id == 0 || tombstone_regions.contains_key(&region_id) {
                continue;
            }
            tikv_util::set_current_region(region_id);
            let Some(mut processor) = Preprocessor::new(
                &raft,
                merged_store_id,
                region_id,
                &ctx.master_key,
                ctx.encryption_key_manager.clone(),
            ) else {
                error!(
                    "{}:{} failed to create preprocessor",
                    merged_store_id, region_id
                );
                debug_assert!(false);
                continue;
            };
            if let Some(merge_state) = processor.pending_merge_state() {
                pending_merge_states.insert(region_id, merge_state.clone());
            }

            let progress = region_progresses.get_mut(&region_id).unwrap_or_else(|| {
                panic!("{} no region progress", processor.tag());
            });
            let mut recover_pre_ctx = recover_txn_file_locks_ctx.as_ref();

            if let Some(shard_meta) = processor.shard_meta() {
                let persisted_idx = shard_meta.data_persisted_log_index();
                progress.persist_progress.reset(persisted_idx);

                // Recover parent first.
                if let Some(parent) = shard_meta.parent.as_deref() {
                    let mut parent_processor = Preprocessor::from_shard_meta(
                        &raft,
                        merged_store_id,
                        parent.clone(),
                        &ctx.master_key,
                        ctx.encryption_key_manager.clone(),
                    );
                    box_try!(parent_processor.recover(
                        &mut recover_pre_ctx,
                        progress,
                        shard_meta.data_sequence
                    ));
                }

                box_try!(processor.recover(
                    &mut recover_pre_ctx,
                    progress,
                    progress.synced_index.saturating_add(1)
                ));
            }

            if progress.data_persisted_log_index()
                > raft
                    .get_truncated_index(region_id)
                    .unwrap_or(RAFT_INIT_LOG_INDEX)
            {
                // Trigger truncate raft logs.
                updated_regions.insert(region_id);
            }

            cop_host.on_region_changed(
                processor.region(),
                RegionChangeEvent::Create,
                StateRole::Leader,
            );
            preprocessors.insert(region_id, processor);
        }
        drop(recover_txn_file_locks_ctx);
        let io_rate_limiter = Arc::new(IoRateLimiter::new(IoRateLimitMode::WriteOnly, true, true));
        let store_limiter = Arc::new(StoreLimiter::dummy());
        let mut meta_packer = if ctx.config.meta_pack.enabled {
            let path = ctx.local_dir.join("metas.pack");
            let meta_packer = MetaPacker::new(path, ctx.config.meta_pack.clone())?;
            Some(meta_packer)
        } else {
            None
        };
        let (meta_pack_scheduler, meta_pack_reader) = meta_packer
            .as_mut()
            .map(|x| (x.get_scheduler(), x.take_reader().unwrap()))
            .unzip();
        let mut recover_handler = RecoverHandler::new(raft.clone());
        recover_handler.set_merged_engine(true);
        recover_handler.set_meta_pack(meta_pack_scheduler, meta_pack_reader);
        let mut meta_iter = EmptyMetaIterator {};
        let kv = box_try!(Self::init_kv_engine(
            &ctx,
            io_rate_limiter,
            store_limiter,
            &mut meta_iter,
            recover_handler.clone(),
        ));
        box_try!(Self::load_shards_on_startup(
            &ctx,
            &kv,
            &recover_handler,
            &manifest.keyspace_states,
            &tombstone_regions,
            &init_opts.load_shards_mode,
        ));
        kv.set_loaded();

        if let Some(meta_pack) = meta_packer.as_mut() {
            box_try!(meta_pack.start_worker(kv.clone()));
        }

        // Should be invoked after dependency setup (`load_shards_on_startup`).
        let delay_destroy_regions =
            box_try!(Self::destroy_regions_on_startup(&raft, tombstone_regions));

        let (store_sender, store_receiver) = mpsc::unbounded();
        let (peer_sender, peer_receiver) = raft_router_peer_channel(0);
        let router = RaftRouter::new(peer_sender, store_sender);
        Ok(Self {
            ctx,
            manifest,
            region_progresses,
            updated_regions,
            raft,
            kv,
            recover_handler,
            load_shards_mode: init_opts.load_shards_mode,
            preprocessors,
            appliers: HashMap::default(),
            delay_destroy_regions,
            installable_region_bootstraps: HashMap::default(),
            newly_bootstrapped_regions: Vec::new(),
            pending_merge_states,
            _store_receiver: store_receiver,
            peer_receiver,
            router,
            cop_host,
            region_info_accessor,
            meta_packer,
            closed: false,
        })
    }

    fn load_shards_on_startup(
        ctx: &MergedEngineContext,
        kv: &kvengine::Engine,
        recover_handler: &RecoverHandler,
        keyspace_states: &HashMap<u32, Bytes>,
        tombstone_regions: &HashMap<u64 /* region_id */, u32 /* keyspace_id */>,
        load_shards_mode: &LoadShardsMode,
    ) -> Result<()> {
        match load_shards_mode {
            LoadShardsMode::EagerAll => {
                box_try!(tikv_util::init_task_local_sync(|| {
                    Self::load_all_shards_impl(
                        ctx,
                        kv,
                        recover_handler.clone(),
                        keyspace_states,
                        Some(tombstone_regions),
                    )
                }));
            }
            LoadShardsMode::Lazy {
                pinned_policy: policy,
            } => {
                box_try!(tikv_util::init_task_local_sync(|| {
                    Self::load_pinned_shards_impl(
                        ctx,
                        kv,
                        recover_handler.clone(),
                        keyspace_states,
                        Some(tombstone_regions),
                        policy.as_ref(),
                    )
                }));
            }
        }
        Ok(())
    }

    fn merged_store_id(&self) -> u64 {
        self.ctx.config.merged_store_id
    }

    fn get_region_tag(&self, region_id: u64, region_ver: u64) -> ShardTag {
        ShardTag::new(self.merged_store_id(), IdVer::new(region_id, region_ver))
    }

    fn notify_region_change(
        host: &CoprocessorHost,
        region: &metapb::Region,
        event: RegionChangeEvent,
    ) {
        host.on_region_changed(region, event, StateRole::Leader);
    }

    pub fn set_keyspace_states(&mut self, keyspace_id: u32, states: Bytes) {
        let old = self
            .manifest
            .set_keyspace_states(keyspace_id, states.clone());
        if old == Some(states.clone()) {
            return; // no change
        }
        self.manifest.persist().expect("persist manifest failed");
    }

    pub fn remove_keyspace(&mut self, keyspace_id: u32) {
        let shards = self.kv.remove_keyspace_shards(keyspace_id);
        for shard_id in shards.into_iter().flatten() {
            self.appliers.remove(&shard_id);
        }
        self.manifest.keyspace_states.remove(&keyspace_id);
        self.manifest.persist().expect("persist manifest failed");
    }

    pub fn get_keyspaces(&self) -> Vec<u32> {
        self.manifest.keyspace_states.keys().cloned().collect()
    }

    pub fn get_keyspace_states(&self, keyspace_id: u32) -> Option<Bytes> {
        self.manifest.keyspace_states.get(&keyspace_id).cloned()
    }

    pub fn get_router(&self) -> RaftRouter {
        self.router.clone()
    }

    pub fn get_meta_pack_compact_keeper(&self) -> Option<CompactKeeper> {
        self.meta_packer.as_ref().map(|x| x.compact_keeper())
    }

    fn load_all_shards_impl(
        ctx: &MergedEngineContext,
        kv: &kvengine::Engine,
        mut recoverer: RecoverHandler,
        keyspace_states: &HashMap<u32, Bytes>,
        tombstone_regions: Option<&HashMap<u64 /* region_id */, u32 /* keyspace_id */>>,
    ) -> Result<()> {
        let metas = Self::load_all_shard_metas_impl(
            ctx,
            &mut recoverer,
            keyspace_states,
            tombstone_regions,
        )?;
        kv.load_shards(metas, recoverer, None)?;
        Ok(())
    }

    fn load_pinned_shards_impl(
        ctx: &MergedEngineContext,
        kv: &kvengine::Engine,
        mut recoverer: RecoverHandler,
        keyspace_states: &HashMap<u32, Bytes>,
        tombstone_regions: Option<&HashMap<u64 /* region_id */, u32 /* keyspace_id */>>,
        pinned_shard_policy: &dyn PinnedShardPolicy,
    ) -> Result<()> {
        fail::fail_point!("merged_engine::skip_load_pinned_shards", |_| {
            // Keep kvengine empty (as if no shard matches the pinned policy), but still
            // populate the rfengine dependency graph and perform stale peer
            // cleanup.
            recoverer.iterate(|_| {})?;
            Ok(())
        });
        let metas = Self::load_pinned_shard_metas_impl(
            ctx,
            &mut recoverer,
            keyspace_states,
            tombstone_regions,
            pinned_shard_policy,
        )?;
        kv.load_shards(metas, recoverer, None)?;
        Ok(())
    }

    fn load_all_shard_metas_impl(
        ctx: &MergedEngineContext,
        recoverer: &mut RecoverHandler,
        keyspace_states: &HashMap<u32, Bytes>,
        tombstone_regions: Option<&HashMap<u64 /* region_id */, u32 /* keyspace_id */>>,
    ) -> Result<StdHashMap<u64 /* region_id */, ShardMeta>> {
        let engine_id = ctx.config.merged_store_id;
        let mut metas = StdHashMap::default();
        recoverer.iterate(|cs| {
            debug_assert!(cs.has_snapshot());
            if tombstone_regions.is_some_and(|x| x.contains_key(&cs.shard_id)) {
                return;
            }
            let keyspace_id = get_keyspace_id_of_snapshot(cs.get_snapshot());
            if keyspace_states.contains_key(&keyspace_id) {
                let meta = ShardMeta::new(engine_id, &cs);
                info!("{} load shard meta: {:?}", meta.tag(), cs);
                metas.insert(meta.id, meta);
            }
        })?;
        Ok(metas)
    }

    fn load_pinned_shard_metas_impl(
        ctx: &MergedEngineContext,
        recoverer: &mut RecoverHandler,
        keyspace_states: &HashMap<u32, Bytes>,
        tombstone_regions: Option<&HashMap<u64 /* region_id */, u32 /* keyspace_id */>>,
        pinned_shard_policy: &dyn PinnedShardPolicy,
    ) -> Result<StdHashMap<u64 /* region_id */, ShardMeta>> {
        let engine_id = ctx.config.merged_store_id;
        let mut metas = StdHashMap::default();
        recoverer.iterate(|cs| {
            debug_assert!(cs.has_snapshot());
            if tombstone_regions.is_some_and(|x| x.contains_key(&cs.shard_id)) {
                return;
            }

            let snap = cs.get_snapshot();
            let keyspace_id = get_keyspace_id_of_snapshot(snap);
            if !keyspace_states.contains_key(&keyspace_id) {
                return;
            }
            if pinned_shard_policy.should_pin_snapshot(snap) {
                let meta = ShardMeta::new(engine_id, &cs);
                info!("{} load pinned shard meta: {:?}", meta.tag(), cs);
                metas.insert(meta.id, meta);
            }
        })?;

        Ok(metas)
    }

    pub fn load_all_shards(&self, keyspace_states: &HashMap<u32, Bytes>) -> Result<()> {
        Self::load_all_shards_impl(
            &self.ctx,
            &self.kv,
            self.recover_handler.clone(),
            keyspace_states,
            None,
        )
    }

    pub fn load_pinned_shards(&self, keyspace_states: &HashMap<u32, Bytes>) -> Result<()> {
        let LoadShardsMode::Lazy {
            pinned_policy: policy,
        } = &self.load_shards_mode
        else {
            return Ok(());
        };
        Self::load_pinned_shards_impl(
            &self.ctx,
            &self.kv,
            self.recover_handler.clone(),
            keyspace_states,
            None,
            policy.as_ref(),
        )
    }

    pub fn load_shards_for_load_mode(&self, keyspace_states: &HashMap<u32, Bytes>) -> Result<()> {
        match &self.load_shards_mode {
            LoadShardsMode::EagerAll => self.load_all_shards(keyspace_states),
            LoadShardsMode::Lazy { pinned_policy: _ } => self.load_pinned_shards(keyspace_states),
        }
    }

    /// Loads a single region shard into kvengine using the latest shard meta
    /// persisted in the merged rfengine.
    ///
    /// This is used by replication_worker lazy shard loading when TiCDC
    /// registers a region whose shard is not yet materialized in kvengine.
    ///
    /// Returns [`Error::ShardMetaNotReady`] if the region meta has not
    /// been persisted yet (e.g. the region is newly inserted and has not been
    /// preprocessed by `sync_merged`).
    pub fn load_region_shard(&self, region_id: u64) -> Result<()> {
        fail::fail_point!("merged_engine::load_region_shard_err", |_| {
            Err(Error::Other(box_err!(
                "failpoint merged_engine::load_region_shard_err: {}",
                region_id
            )))
        });
        let store_id = self.merged_store_id();
        let peer_id = region_id; // merged_engine uses `region_id` as `peer_id`.
        let Some(meta) =
            rfstore::store::peer_storage::load_engine_meta(&self.raft, store_id, peer_id)
        else {
            return Err(Error::ShardMetaNotReady(region_id));
        };
        if meta.id != region_id {
            return Err(Error::Other(box_err!(
                "load_region_shard: meta id mismatch, expect {}, got {}",
                region_id,
                meta.id
            )));
        }
        let mut metas = StdHashMap::with_capacity(1);
        metas.insert(meta.id, meta);
        self.kv
            .load_shards(metas, self.recover_handler.clone(), None)?;
        Ok(())
    }

    pub fn load_keyspace_all_shard_metas(
        &self,
        keyspace_id: u32,
    ) -> Result<StdHashMap<u64 /* region_id */, ShardMeta>> {
        let mut states = HashMap::with_capacity(1);
        states.insert(keyspace_id, Bytes::new());
        let mut recoverer = self.recover_handler.clone();
        Self::load_all_shard_metas_impl(&self.ctx, &mut recoverer, &states, None)
    }

    pub fn load_keyspace_pinned_shard_metas(
        &self,
        keyspace_id: u32,
        pinned_shard_policy: &dyn PinnedShardPolicy,
    ) -> Result<StdHashMap<u64 /* region_id */, ShardMeta>> {
        let mut states = HashMap::with_capacity(1);
        states.insert(keyspace_id, Bytes::new());
        let mut recoverer = self.recover_handler.clone();
        Self::load_pinned_shard_metas_impl(
            &self.ctx,
            &mut recoverer,
            &states,
            None,
            pinned_shard_policy,
        )
    }

    pub fn load_keyspace_shard_metas_for_load_mode(
        &self,
        keyspace_id: u32,
    ) -> Result<KeyspaceShardMetas> {
        match &self.load_shards_mode {
            LoadShardsMode::EagerAll => Ok(KeyspaceShardMetas {
                scope: KeyspaceShardMetaScope::All,
                metas: self.load_keyspace_all_shard_metas(keyspace_id)?,
            }),
            LoadShardsMode::Lazy {
                pinned_policy: policy,
            } => Ok(KeyspaceShardMetas {
                scope: KeyspaceShardMetaScope::Pinned,
                metas: self.load_keyspace_pinned_shard_metas(keyspace_id, policy.as_ref())?,
            }),
        }
    }

    pub fn close(&mut self) {
        if self.closed {
            return;
        }
        self.closed = true;

        self.region_info_accessor.stop();
        if let Some(meta_packer) = self.meta_packer.take() {
            meta_packer.stop();
        }
        self.raft.stop_worker(false);
        self.raft.close_writer();
        self.kv.close();
    }

    fn get_latest_backup(ctx: &MergedEngineContext) -> Result<ClusterBackupMeta> {
        let cluster_id = box_try!(ctx.pd.get_cluster_id());
        let runtime = ctx.fs.get_runtime();
        let start_time = Instant::now_coarse();
        while start_time.saturating_elapsed() < ctx.config.get_latest_backup_timeout.0 {
            match runtime.block_on(get_latest_backup_meta(ctx.fs.as_ref(), cluster_id)) {
                Ok(x) => return Ok(x),
                Err(BrError::MetaNotFound(_)) => {
                    warn!("get_latest_backup: not ready");
                    thread::sleep(Duration::from_secs(5));
                    continue;
                }
                Err(err) => return Err(box_err!(err)),
            }
        }
        Err(box_err!("get_latest_backup: timeout"))
    }

    fn recover_from_backup(
        ctx: &MergedEngineContext,
        backup_meta: Option<ClusterBackupMeta>,
        merged_raft: &RfEngine,
    ) -> Result<(
        HashMap<u64, RegionProgress>,
        HashMap<u64, StoreProgress>,
        ClusterBackupMeta,
    )> {
        let backup_meta = match backup_meta {
            Some(backup_meta) => backup_meta,
            None => box_try!(Self::get_latest_backup(ctx)),
        };

        let merged_store_id = ctx.config.merged_store_id;
        info!("recover_from_backup: {}", backup_meta; "store" => merged_store_id);

        let mut region_progresses = HashMap::default();
        let mut regions_raft_progress = RegionsRaftProgress::default();
        let mut store_progresses = HashMap::default();
        let mut raftdb_paths = Vec::new();
        let mut raft_wb = rfengine::WriteBatch::new();
        for store in backup_meta.get_stores() {
            let store_progress = StoreProgress {
                store_id: store.get_store_id(),
                epoch: store.get_epoch(),
                offset: store.get_offset(),
            };
            store_progresses.insert(store.get_store_id(), store_progress);
            let mut store_config = TikvConfig::default();
            let store_path = ctx.local_dir.join(store.store_id.to_string());
            store_config.storage.data_dir = store_path.to_str().unwrap().to_string();
            store_config.raft_store.raftdb_path =
                store_path.join("raft").to_str().unwrap().to_string();
            store_config.rfengine.lightweight_backup = false;
            let origin = box_try!(Self::setup_raft_engine(ctx, &backup_meta, store));
            raftdb_paths.push(store_config.raft_store.raftdb_path);
            let region_peers_map = origin.get_region_peer_map();
            for (&region_id, &peer_id) in &region_peers_map {
                if region_id == 0 {
                    continue;
                }
                let tag = ShardTag::new(merged_store_id, IdVer::new(region_id, 0));
                let origin_tag = ShardTag::new(store.store_id, IdVer::new(region_id, 0));

                // region state.
                let Some(region_state) = rfstore::store::load_last_peer_state(&origin, peer_id)
                else {
                    let states = origin.get_peer_all_states(peer_id, false);
                    warn!("{} recover_from_backup: no peer state", tag;
                        "origin" => %origin_tag, "peer" => peer_id, "states" => ?states);
                    debug_assert!(false);
                    continue;
                };
                let new_region = region_state.get_region();
                let region_version = new_region.get_region_epoch().get_version();
                let tag = tag.with_region_version(region_version);
                let origin_tag = origin_tag.with_region_version(region_version);

                if peer_is_skippable(&region_state) {
                    info!("{} recover_from_backup: skip peer", tag;
                        "origin" => %origin_tag, "region_state" => ?region_state);
                    continue;
                }
                let keyspace_id = ApiV2::get_u32_keyspace_id_by_key(new_region.get_start_key())
                    .unwrap_or_default();

                // raft state.
                let Some(raft_state) =
                    rfstore::store::load_peer_raft_state(&origin, peer_id, region_version)
                else {
                    let states = origin.get_peer_all_states(peer_id, false);
                    warn!("{} recover_from_backup: no raft state", tag;
                        "origin" => %origin_tag, "peer" => peer_id, "states" => ?states);
                    debug_assert!(false);
                    continue;
                };

                let Some(mut shard_meta) =
                    rfstore::store::load_engine_meta(&origin, store.store_id, peer_id)
                else {
                    let states = origin.get_peer_all_states(peer_id, false);
                    warn!("{} recover_from_backup: no engine meta", tag;
                        "origin" => %origin_tag, "peer" => peer_id, "states" => ?states);
                    debug_assert!(false);
                    continue;
                };
                debug_assert_eq!(shard_meta.range.keyspace_id, keyspace_id);
                if peer_is_stale(&shard_meta, &region_peers_map) {
                    info!("{} recover_from_backup: skip stale peer", tag; "origin" => %origin_tag);
                    continue;
                }
                shard_meta.engine_id = merged_store_id;

                let region_progress = region_progresses
                    .entry(region_id)
                    .or_insert(RegionProgress::new(keyspace_id, region_id));

                let preprocess_index = raft_state.get_last_preprocessed_index();
                let origin_commit = raft_state.get_commit();
                let origin_truncated_index = origin
                    .get_truncated_index(peer_id)
                    .unwrap_or_default()
                    .max(RAFT_INIT_LOG_INDEX);

                let merged_commit_index = region_progress.commit_index();

                // Fetch entries, and insert them into region progress, so that they
                // will be replayed when commit index advances (during sync_merged).
                // Entries earlier than truncated index of latest peer are also fetched, as they
                // would be required by dependents.
                {
                    let low_idx = origin_truncated_index + 1;
                    let high_idx = raft_state.get_last_index() + 1;
                    debug!(
                        "{} recover_from_backup: fetch entries: [{}, {})",
                        tag, low_idx, high_idx; "origin" => %origin_tag);
                    if low_idx < high_idx {
                        if let Err(err) = fetch_raft_entries_to_region_progress(
                            tag,
                            &origin,
                            peer_id,
                            low_idx,
                            high_idx,
                            region_progress,
                        ) {
                            panic!(
                                "{} recover_from_backup: fetch entries failed, low: {}, high: {}, err: {}, origin: {}",
                                tag, low_idx, high_idx, err, origin_tag
                            );
                        }
                    }
                }

                if merged_commit_index > origin_commit {
                    continue;
                }

                // merge states
                let peer_is_newer = regions_raft_progress.update(region_id, &raft_state);
                if peer_is_newer {
                    write_engine_meta(&mut raft_wb, region_id, &shard_meta);
                    origin.iterate_peer_states(peer_id, false, |k, v| {
                        trace!("{} recover_from_backup", tag;
                            "k" => LogValue::key(k), "v" => LogValue::value(v),
                            "origin" => %origin_tag, "peer" => peer_id,
                        );
                        update_peer_state_without_engine_meta(
                            &mut raft_wb,
                            k,
                            v,
                            merged_store_id,
                            region_id,
                            keyspace_id,
                        );
                        true
                    });

                    // Depends on other region.
                    if let Some(box parent) = &shard_meta.parent {
                        if parent.id != region_id {
                            let parent_tag = parent.tag();
                            let parent_peer_id =
                                *region_peers_map.get(&parent.id).unwrap_or_else(|| {
                                    panic!("{} recover_from_backup: parent region not found, origin: {}, parent: {}", tag, origin_tag, parent_tag);
                                });
                            // The state value of parent version should be the latest. It's safe to
                            // overwrite.
                            let state_keys =
                                [raft_state_key(parent.ver), region_state_key(parent.ver)];
                            for k in state_keys {
                                let state_val =
                                    origin.get_state(parent_peer_id, &k).unwrap_or_else(|| {
                                        panic!("{} recover_from_backup: parent state key not found, origin: {}, parent: {}", tag, origin_tag, parent_tag);
                                    });
                                update_peer_state_without_engine_meta(
                                    &mut raft_wb,
                                    &k,
                                    &state_val,
                                    merged_store_id,
                                    parent.id,
                                    keyspace_id,
                                );
                            }

                            debug!("{} recover_from_backup: set parent states", tag;
                                "origin" => %origin_tag, "parent" => %parent_tag, "parent_peer" => parent_peer_id);
                        }
                    }

                    region_progress.update_commit_index(origin_commit);
                    region_progress.synced_index = preprocess_index;
                }

                if raft_wb.estimated_size() >= ctx.config.raft_write_batch_size.0 as usize {
                    merged_raft
                        .write(mem::take(&mut raft_wb))
                        .expect("raft write");
                }

                debug!(
                    "{} recover_from_backup", tag;
                    "region" => ?region_state, "raft" => ?raft_state,
                    "progress" => ?region_progress,
                    "origin" => %origin_tag, "keyspace" => keyspace_id);
            }
        }

        // Append raft logs to merged rfengine.
        for progress in region_progresses.values_mut() {
            let region_id = progress.region_id;
            let keyspace_id = progress.keyspace_id;

            if !progress.entries.is_empty() {
                let commit_index = progress.commit_index();
                let mut committed_entries = vec![];
                progress.entries.retain(|_, v| {
                    if v.index <= commit_index {
                        committed_entries.push(v.to_entry());
                    }
                    // Retain not-synced, so that they will be replayed when commit index advances
                    // (during sync_merged).
                    v.index > progress.synced_index
                });

                committed_entries.sort_unstable_by_key(|e| e.index);
                for entry in committed_entries {
                    raft_wb.append_raft_log(region_id, region_id, keyspace_id, &entry);
                }
            }

            // Always truncate raft log to properly set truncated index.
            raft_wb.truncate_raft_log(region_id, region_id, keyspace_id, RAFT_INIT_LOG_INDEX);
            if raft_wb.estimated_size() >= ctx.config.raft_write_batch_size.0 as usize {
                merged_raft
                    .write(mem::take(&mut raft_wb))
                    .expect("raft write");
            }
        }

        if !raft_wb.is_empty() {
            merged_raft.write(raft_wb).expect("raft write");
        }

        // destroy original raft engines
        for raftdb_path in raftdb_paths {
            let raft_path = Path::new(&raftdb_path);
            // clean up dir
            if let Err(e) = std::fs::remove_dir_all(raft_path) {
                warn!("remove raft path failed: {:?}", e; "path" => ?raft_path);
            }
        }
        Ok((region_progresses, store_progresses, backup_meta))
    }

    fn recover_from_merged_raft_engine(
        merged_raft: &RfEngine,
        uncommitted_entries: &UncommittedEntries,
    ) -> Result<(
        HashMap<u64, RegionProgress>,
        HashMap<u64, u32>, // tombstone_regions
    )> {
        let merged_store_id = merged_raft.get_engine_id();
        let mut region_progresses = HashMap::default();
        let mut tombstone_regions = HashMap::default();

        // Get region progresses from RfEngine.
        let region_peers_map = merged_raft.get_region_peer_map();
        for (region_id, peer_id) in region_peers_map {
            if region_id == 0 {
                continue;
            }
            let tag = ShardTag::new(merged_store_id, IdVer::new(region_id, 0));

            let Some(region_state) = rfstore::store::load_last_peer_state(merged_raft, peer_id)
            else {
                // The region is destroyed.
                let states = merged_raft.get_peer_all_states(peer_id, false);
                info!("{} recover_from_merged_raft_engine: no peer state, skip", tag;
                        "peer" => peer_id, "states" => ?states);
                continue;
            };
            let region_version = region_state.get_region().get_region_epoch().get_version();
            let tag = tag.with_region_version(region_version);
            let keyspace_id =
                ApiV2::get_u32_keyspace_id_by_key(region_state.get_region().get_start_key())
                    .unwrap_or_default();

            if region_state.state == PeerState::Tombstone {
                // The dependency is not setup yet. So tombstone regions should not be skipped
                // here.
                debug!("{} recover_from_merged_raft_engine: tombstone region", tag;
                    "region_state" => ?region_state);
                tombstone_regions.insert(region_id, keyspace_id);
            }
            if !is_region_initialized(region_state.get_region()) {
                // Should not happen, not initialized region is skipped during
                // `recover_from_backup` & `update_wal`.
                error!("{} recover_from_merged_raft_engine: region not initialized", tag;
                    "region_state" => ?region_state);
                debug_assert!(false);
                return Err(box_err!("region not initialized: {}", region_id));
            }

            let Some(raft_state) =
                rfstore::store::load_peer_raft_state(merged_raft, peer_id, region_version)
            else {
                let states = merged_raft.get_peer_all_states(peer_id, false);
                warn!("{} recover_from_merged_raft_engine: no raft state", tag;
                        "peer" => peer_id, "states" => ?states);
                debug_assert!(false);
                continue;
            };

            let region_progress = region_progresses
                .entry(region_id)
                .or_insert(RegionProgress::new(keyspace_id, region_id));
            region_progress.update_commit_index(raft_state.get_commit());
            region_progress.synced_index = raft_state.get_last_preprocessed_index();
            if let Some(entries) = uncommitted_entries.get_region_entries(region_id) {
                region_progress.entries = entries.clone();
            }
            if let Some(last_index) = region_progress
                .entries
                .iter()
                .map(|(&log_idx, _)| log_idx)
                .max()
            {
                region_progress.update_last_index(last_index);
            }

            if region_progress.commit_index() > region_progress.synced_index {
                // Fetch committed entries for `sync_merged`.
                let low_idx = region_progress.synced_index + 1;
                let high_idx = region_progress.commit_index() + 1;
                debug!(
                    "{} recover_from_merged_raft_engine: fetch committed entries: [{}, {})",
                    tag, low_idx, high_idx; "progress" => ?region_progress);
                if let Err(err) = fetch_raft_entries_to_region_progress(
                    tag,
                    merged_raft,
                    peer_id,
                    low_idx,
                    high_idx,
                    region_progress,
                ) {
                    panic!(
                        "{} recover_from_merged_raft_engine: fetch raft entries failed, low: {}, high: {}, err: {}",
                        tag, low_idx, high_idx, err
                    );
                }
            }

            debug!(
                "{} recover_from_merged_raft_engine", tag;
                "region" => ?region_state,
                "raft_state" => ?raft_state,
                "progress" => ?region_progress,
                "keyspace" => keyspace_id,
            );
        }

        Ok((region_progresses, tombstone_regions))
    }

    fn setup_raft_engine(
        ctx: &MergedEngineContext,
        backup_meta: &ClusterBackupMeta,
        store: &StoreBackupMeta,
    ) -> Result<RfEngine> {
        let mut store_config = TikvConfig::default();
        let store_path = ctx.local_dir.join(store.store_id.to_string());
        store_config.storage.data_dir = store_path.to_str().unwrap().to_string();
        store_config.raft_store.raftdb_path = store_path.join("raft").to_str().unwrap().to_string();
        store_config.rfengine.lightweight_backup = false;
        let store_id = store.get_store_id();
        let rlog_files = collect_snapshot_meta_rlog_files(
            ctx.fs.clone(),
            &ctx.fs.get_prefix(),
            backup_meta,
            store_id,
            None,
        )?;
        rfengine::lightweight_restore(
            store_id,
            None,
            Path::new(&store_config.raft_store.raftdb_path),
            rlog_files.snap_epoch,
            rlog_files.snap_meta,
            rlog_files.snap_rlog,
            store_config.rfengine.epoch_rotate_len,
        )?;
        let raft_db_path = Path::new(&store_config.raft_store.raftdb_path);
        let data_dir = Path::new(&store_config.storage.data_dir);
        let mut rfengine_cfg = store_config.rfengine.clone();
        rfengine_cfg.disable_rlog_offloading();
        let rf_engine = RfEngine::open(raft_db_path, &rfengine_cfg, Some(data_dir), None)?;
        let cache_dir = store_path.join("cache");
        box_try!(fs::create_dir_all(&cache_dir));
        let ctx = ReplayWalLogsContext {
            pd_client: ctx.pd.clone(),
            dfs: ctx.fs.clone(),
            store_id,
            cluster_backup: backup_meta,
            rf_engine: &rf_engine,
            complete_wal_chunks: false,
            full_restore: false,
            fetch_wal_timeout: ctx.config.timeout_fetch_wal.0,
            cache_dir: Some(cache_dir),
            wal_chunks_cache: None,
            from_archive: false,
        };
        let tag = &format!("merged_{}", store_id);
        replay_wal_logs_from_backup(tag, &ctx, rlog_files.snap_epoch)?;
        Ok(rf_engine)
    }

    fn init_kv_engine(
        ctx: &MergedEngineContext,
        rate_limiter: Arc<IoRateLimiter>,
        store_limiter: Arc<StoreLimiter>,
        meta_iter: &mut impl kvengine::MetaIterator,
        recoverer: impl kvengine::RecoverHandler + 'static,
    ) -> Result<kvengine::Engine> {
        let kv_engine_path = ctx.local_dir.join("db");
        if !kv_engine_path.exists() {
            fs::create_dir_all(&kv_engine_path)?;
        }
        let kv_opts = Self::replay_kv_engine_options(&ctx.config, kv_engine_path);
        let kv_conf = kvengine::KvEngineConfig::default();
        let opts = Arc::new(kv_opts);
        let id_allocator = Arc::new(PdIdAllocator::new(ctx.pd.clone()));
        let (sender, _) = mpsc::unbounded();
        let meta_change_listener = Box::new(MetaChangeListener { sender });
        let kv_engine = kvengine::Engine::open(
            ctx.fs.clone(),
            opts,
            kv_conf,
            meta_iter,
            recoverer,
            id_allocator,
            meta_change_listener,
            rate_limiter,
            store_limiter,
            None,
            ctx.master_key.clone(),
            ctx.pd.get_security_mgr(),
            false,
        )?;
        kv_engine.set_engine_id(ctx.config.merged_store_id);
        Ok(kv_engine)
    }

    fn replay_kv_engine_options(
        config: &MergedEngineConfig,
        kv_engine_path: PathBuf,
    ) -> kvengine::Options {
        let mut kv_opts = kvengine::Options::default();
        kv_opts.local_dirs = vec![kv_engine_path];
        kv_opts.max_mem_table_size = config.mem_table_size.0;
        // Replay must accept every entry that an online Crossbeam-enabled
        // cluster accepted. merged_engine has no serving-path rollout cost.
        kv_opts.enable_crossbeam_memtable = true;
        kv_opts.max_block_cache_size = config.block_cache_size.as_memory_size() as i64;
        kv_opts.for_restore = true;
        if config.force_ia {
            kv_opts.ia = IaConfig {
                mem_cap: AbsoluteOrPercentSize::Percent(10.0),
                disk_cap: AbsoluteOrPercentSize::Percent(60.0),
                dynamic_capacity: false,
                force_ia: true,
                ..Default::default()
            };
        }
        kv_opts
    }

    pub fn get_kv(&self) -> kvengine::Engine {
        self.kv.clone()
    }

    pub fn get_raft(&self) -> RfEngine {
        self.raft.clone()
    }

    /// Returns a shared `RegionInfoAccessor` for querying merged regions.
    pub fn get_region_info_accessor(&self) -> RegionInfoAccessor {
        self.region_info_accessor.clone()
    }

    /// Registers a `RegionChangeObserver` to receive merged-region
    /// create/update/destroy notifications.
    ///
    /// `RegionChangeEvent::Destroy` is emitted for merge-source logical
    /// destroy when commit-merge is accepted. It is not emitted for every
    /// physical region cleanup path.
    ///
    /// This is primarily used by replication_worker to keep pinned shards
    /// up-to-date without periodically scanning all persisted shard metas.
    pub fn register_region_change_observer(
        &mut self,
        priority: u32,
        observer: BoxRegionChangeObserver,
    ) {
        self.cop_host
            .registry
            .register_region_change_observer(priority, observer);
    }

    pub fn get_region_progress(&self, region_id: u64) -> Option<&RegionProgress> {
        self.region_progresses.get(&region_id)
    }

    pub fn mut_region_progress(&mut self, region_id: u64) -> Option<&mut RegionProgress> {
        self.region_progresses.get_mut(&region_id)
    }

    pub fn get_store_progress(&self, store_id: u64) -> Option<StoreProgress> {
        self.manifest.store_progresses.get(&store_id).cloned()
    }

    /// Collects all known store progresses into `out`.
    pub fn collect_store_progresses(&self, out: &mut Vec<StoreProgress>) {
        out.clear();
        out.extend(self.manifest.store_progresses.values().cloned());
    }

    pub fn get_or_insert_store_progress(&mut self, store_id: u64) -> StoreProgress {
        *self
            .manifest
            .store_progresses
            .entry(store_id)
            .or_insert_with(|| StoreProgress {
                store_id,
                epoch: 1,
                offset: 0,
            })
    }

    pub fn update_store_progress(&mut self, store_id: u64, epoch: u32, offset: u64) {
        self.manifest.update_store_progress(store_id, epoch, offset);
    }

    pub fn get_synced_target_ts(&self) -> TimeStamp {
        self.manifest.synced_target_ts
    }

    pub fn get_keyspace_regions(&self, keyspace_id: u32) -> Vec<u64> {
        let mut regions = Vec::new();
        for (&region_id, progress) in &self.region_progresses {
            if progress.keyspace_id == keyspace_id {
                regions.push(region_id);
            }
        }
        regions
    }

    pub fn update_wal<R: io::Read>(
        &mut self,
        store_id: u64,
        epoch_id: u32,
        start_off: u64,
        end_off: u64,
        reader: R,
        target_stores: &HashSet<u64 /* store_id */>,
    ) -> Result<()> {
        let merged_store_id = self.merged_store_id();
        if let Some(store_progress) = self.manifest.store_progresses.get(&store_id) {
            if store_progress.epoch != epoch_id || store_progress.offset != start_off {
                let err_msg = format!(
                    "store {} expect ({}, {}), got ({}, {}), end_off: {}",
                    store_id,
                    epoch_id,
                    start_off,
                    store_progress.epoch,
                    store_progress.offset,
                    end_off,
                );
                error!("{}", &err_msg);
                debug_assert!(false);
                return Err(Error::StoreProgressMismatch(err_msg));
            }
        } else {
            return Err(Error::StoreProgressNotFound(store_id));
        };
        let mut wal_iterator = WalIterator::new_from_reader(reader, epoch_id, start_off);
        // Collect all write batches first, to avoid partial update if any error occurs.
        // RegionProgress::upsert_entry is NOT reentrant.
        let mut origin_batches = Vec::new();
        wal_iterator.iterate_write_batch(|origin_wb| {
            origin_batches.push(origin_wb);
        })?;
        fail::fail_point!("merged_engine::update_wal_err", |_| {
            Err(box_err!("injected error at merged_engine::update_wal_err"))
        });
        for origin_wb in origin_batches {
            let region_peer_map = origin_wb.get_region_peer_map();
            for (&region_id, &peer_id) in &region_peer_map {
                if region_id == 0 {
                    continue;
                }
                tikv_util::set_current_region(region_id);
                let tag = ShardTag::new(merged_store_id, IdVer::new(region_id, 0));
                let origin_tag = ShardTag::new(store_id, IdVer::new(region_id, 0));
                let origin_region_state = origin_wb.get_latest_peer_state(peer_id);

                let merged_bootstrap_exists = self.merged_bootstrap_exists(region_id);
                if !merged_bootstrap_exists {
                    if let Some(region_state) = origin_region_state.as_ref() {
                        if !peer_is_skippable(region_state) {
                            self.observe_origin_region_bootstrap_for_peer_moving_in(
                                merged_store_id,
                                store_id,
                                &origin_wb,
                                region_id,
                                peer_id,
                                region_state,
                                target_stores,
                            );
                        }
                    }
                }

                let progress = match self.region_progresses.entry(region_id) {
                    HashMapEntry::Occupied(e) => e.into_mut(),
                    HashMapEntry::Vacant(e) => {
                        let Some(region_state) = origin_region_state.as_ref() else {
                            info!("{} update_wal: no peer state in wb", tag; "origin" => %origin_tag);
                            continue;
                        };

                        if peer_is_skippable(region_state) {
                            // TODO: handle remove peer.
                            info!("{} update_wal: skip peer", tag;
                                "origin" => %origin_tag, "region_state" => ?region_state);
                            continue;
                        }
                        let new_region = region_state.get_region();
                        let keyspace_id =
                            ApiV2::get_u32_keyspace_id_by_key(new_region.get_start_key())
                                .unwrap_or_default();
                        debug!(
                            "{} update_wal: new region: {:?}", tag, region_state;
                            "origin" => %origin_tag, "keyspace" => keyspace_id);
                        e.insert(RegionProgress::new(keyspace_id, region_id))
                    }
                };
                if progress.data_persisted_log_index() == TRUNCATE_ALL_INDEX {
                    debug!("{} update_wal: truncate all", tag; "origin" => %origin_tag);
                    continue;
                }
                if let Some(raft_state) = load_last_raft_state_from_wb(&origin_wb, peer_id) {
                    let origin_commit = raft_state.get_commit();
                    if progress.update_commit_index(origin_commit) {
                        trace!(
                            "{} update_wal: advance commit index {}",
                            tag, progress.commit_index; "origin" => %origin_tag);
                    }
                }
                origin_wb.read_peer_logs(peer_id, |logs| {
                    for log_op in logs {
                        trace!("{} update_wal: insert log index {}", tag, log_op.index; "origin" => %origin_tag);
                        progress.upsert_entry(tag, log_op.index, log_op.term, || log_op.clone().into());
                    }
                });
                debug!("{} update_wal: progress: {:?}", tag, progress; "origin" => %origin_tag);
                self.updated_regions.insert(region_id);
            }
        }
        self.manifest
            .update_store_progress(store_id, epoch_id, end_off);
        Ok(())
    }

    pub fn rotate_wal(&mut self, store_id: u64, epoch_id: u32, offset: u64) -> Result<()> {
        if let Some(store_progress) = self.manifest.store_progresses.get_mut(&store_id) {
            if store_progress.epoch != epoch_id || store_progress.offset != offset {
                return Err(Error::StoreProgressMismatch(format!(
                    "store {} expect epoch {}, offset {}, got epoch {}, offset {}",
                    store_id, epoch_id, offset, store_progress.epoch, store_progress.offset
                )));
            }
            info!(
                "rotate store {} at epoch {}, offset {}",
                store_id, epoch_id, offset
            );
            store_progress.epoch += 1;
            store_progress.offset = 0;
        } else {
            return Err(Error::StoreProgressNotFound(store_id));
        }
        Ok(())
    }

    pub fn sync_merged(
        &mut self,
        apply_ctx: &mut ApplyContext,
        synced_target_ts: Option<TimeStamp>,
    ) -> Result<()> {
        // 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 = StdHashSet::default();
        let raft_engine = self.raft.clone();
        let router = self.router.clone();
        let pre_ctx = PreprocessContext {
            store_id: self.merged_store_id(),
            raft: &raft_engine,
            raft_wb: &mut raft_wb,
            remove_dependents: &mut remove_dependents,
            apply_msgs: &mut apply_msgs,
            cfg: &raft_cfg,
            destroying: &mut destroying,
            scene: PreprocessScene::MergedEngine {
                router: Some(&router),
            },
        };
        let mut ctx = SyncRegionsContext {
            pre_ctx,
            apply_ctx,
            prepared_msgs: HashMap::default(),
            destroyed_regions: HashMap::default(),
        };
        let updated_regions: HashSet<u64> = mem::take(&mut self.updated_regions);
        self.sync_merged_with_ctx(&mut ctx, updated_regions, synced_target_ts)
            .map_err(|e| {
                debug!("sync_merged: clear context on error: {:?}", ctx; "err" => ?e);
                ctx.clear();
                e
            })
    }

    fn sync_merged_with_ctx(
        &mut self,
        ctx: &mut SyncRegionsContext<'_>,
        mut updated_regions: HashSet<u64>,
        synced_target_ts: Option<TimeStamp>,
    ) -> Result<()> {
        self.sync_merged_for_regions(ctx, &updated_regions)?;
        self.handle_prepared_msgs(ctx);
        self.remove_dependents(ctx, &mut updated_regions);
        self.update_progress_and_truncate(&updated_regions, ctx.raft_wb)?;
        self.destroy_regions(ctx);
        if !ctx.raft_wb.is_empty() {
            try_force_stop_err!(self);
            self.raft.write(mem::take(ctx.raft_wb)).expect("raft write");
        }
        self.manifest
            .update_region_progresses(&self.region_progresses);
        if let Some(synced_target_ts) = synced_target_ts {
            self.manifest.update_synced_target_ts(synced_target_ts);
        }
        try_force_stop_err!(self);
        self.manifest.persist().expect("persist manifest");
        Ok(())
    }

    fn sync_merged_for_regions(
        &mut self,
        ctx: &mut SyncRegionsContext<'_>,
        updated_regions: &HashSet<u64>,
    ) -> Result<()> {
        let merged_store_id = self.merged_store_id();
        info!("sync merged for {} regions", updated_regions.len(); "store" => merged_store_id);
        SYNC_REGIONS_COUNTER.inc_by(updated_regions.len() as u64);
        let mut update_queue = VecDeque::from_iter(updated_regions.iter().cloned());
        let mut merged_wb = rfengine::WriteBatch::new();
        let mut merged_wb_estimated_size = 0;
        let mut continuous_postpone_count = 0;
        while let Some(updated_region) = update_queue.pop_front() {
            try_force_stop_err!(self);
            tikv_util::set_current_region(updated_region);
            let res = self.sync_region(
                ctx,
                updated_region,
                &mut merged_wb,
                &mut merged_wb_estimated_size,
            )?;
            match &res {
                SyncRegionResult::Finished => {}
                SyncRegionResult::Postponed => {
                    update_queue.push_back(updated_region);
                }
                SyncRegionResult::Resume => {
                    update_queue.push_front(updated_region);
                }
                SyncRegionResult::Dropped => {}
            }
            if res.is_postponed() {
                continuous_postpone_count += 1;
                if continuous_postpone_count >= update_queue.len() * 2 {
                    break;
                }
            } else {
                continuous_postpone_count = 0;
            }
        }
        for postponed_region in update_queue {
            debug!(
                "{}:{} sync_merged: region postponed to next round",
                merged_store_id, postponed_region
            );
            self.updated_regions.insert(postponed_region);
        }

        if !merged_wb.is_empty() {
            self.raft.persist(merged_wb).expect("raft persist");
        }
        Ok(())
    }

    fn sync_region(
        &mut self,
        ctx: &mut SyncRegionsContext<'_>,
        updated_region: u64,
        merged_wb: &mut rfengine::WriteBatch,
        merged_wb_estimated_size: &mut usize,
    ) -> Result<SyncRegionResult> {
        let merged_store_id = self.merged_store_id();
        let mut tag = PeerTag::new(merged_store_id, RegionIdVer::new(updated_region, 0));

        if self.merged_bootstrap_exists(updated_region) {
            self.clear_cached_bootstrap(updated_region);
        } else if let Some(bootstrap_shard_meta) =
            self.install_installable_region_bootstrap(ctx, updated_region)?
        {
            info!("{} sync_merged: installed pending live bootstrap", tag);
            self.flush_sync_region_wb(ctx, merged_wb, merged_wb_estimated_size);
            self.try_load_live_bootstrap_shard(tag, &bootstrap_shard_meta)?;
        } else {
            // region is newly inserted, should process parent first.
            info!("{} sync_merged: postpone newly inserted region", tag);
            return Ok(SyncRegionResult::Postponed);
        }

        let progress = self.region_progresses.get_mut(&updated_region).unwrap();
        let low = progress.synced_index.max(RAFT_INIT_LOG_INDEX) + 1;
        let high: u64 = progress.commit_index() + 1;
        debug!("{} sync_merged: [{}, {})", tag, low, high);
        if low >= high {
            return Ok(SyncRegionResult::Finished);
        }
        let preprocessor = if let Some(preprocessor) = self.preprocessors.get_mut(&updated_region) {
            preprocessor
        } else {
            let Some(preprocessor) = Preprocessor::new(
                &self.raft,
                ctx.store_id,
                updated_region,
                &self.ctx.master_key,
                self.ctx.encryption_key_manager.clone(),
            ) else {
                info!("{} sync_merged: region is merged or destroyed", tag);
                return Ok(SyncRegionResult::Dropped);
            };
            Self::notify_region_change(
                &self.cop_host,
                preprocessor.region(),
                RegionChangeEvent::Create,
            );
            match self.preprocessors.entry(updated_region) {
                HashMapEntry::Occupied(e) => e.into_mut(),
                HashMapEntry::Vacant(e) => e.insert(preprocessor),
            }
        };
        let mut preprocessor_ref = preprocessor.as_ref();
        tag = preprocessor_ref.tag();
        let mut entries = Vec::new();
        let mut wb_encoded_len = 0;
        let mut res = SyncRegionResult::Finished;
        // preprocess entries.
        for log_index in low..high {
            try_force_stop_err!(self);
            let Some(raft_log) = progress.entries.get(&log_index) else {
                // Ref: https://github.com/tidbcloud/cloud-storage-engine/issues/3654
                info!("{} sync_region: no raft log: {}", tag, log_index; "progress" => ?progress);
                break;
            };
            let mut entry = raft_log.to_entry();
            let mut preprocess_cmd = update_entry(&mut entry, merged_store_id);
            let is_preprocess_cmd = preprocess_cmd.is_some();
            let mut admin_req = preprocess_cmd
                .as_mut()
                .and_then(|req| req.has_admin_request().then(|| req.take_admin_request()));
            if let Some(admin) = admin_req.as_ref() {
                match Self::sync_region_admin_req(
                    tag,
                    updated_region,
                    admin,
                    log_index,
                    &mut self.pending_merge_states,
                ) {
                    SyncRegionResult::Finished => {}
                    SyncRegionResult::Postponed => {
                        res = SyncRegionResult::Postponed;
                        break;
                    }
                    _ => unreachable!(),
                }
            }
            let preprocessed_ok = match preprocessor_ref.preprocess_committed_entry(ctx, &entry) {
                PreprocessResult::Ok => {
                    if let Some(admin) = admin_req.as_ref() {
                        Self::sync_region_admin_req_post_preprocess(
                            tag,
                            updated_region,
                            admin,
                            log_index,
                            &preprocessor_ref,
                            &mut self.pending_merge_states,
                        );
                    }
                    true
                }
                PreprocessResult::Rejected { reason, report_err } => {
                    warn!("{} preprocess committed entry rejected", tag; "reason" => ?reason, "err" => ?report_err);
                    if is_preprocess_cmd {
                        // clear rejected command.
                        admin_req = None;
                        entry.set_data(Bytes::new());
                    }
                    false
                }
            };
            preprocessor_ref
                .raft_state
                .set_last_preprocessed_index(*preprocessor_ref.preprocessed_index);
            let mut hs = eraftpb::HardState::default();
            hs.set_term(1);
            hs.set_vote(updated_region);
            hs.set_commit(log_index);
            preprocessor_ref.raft_state.set_hard_state(&hs);
            preprocessor_ref.raft_state.set_last_index(log_index);
            let keyspace_id =
                ApiV2::get_u32_keyspace_id_by_key(preprocessor_ref.region.get_start_key())
                    .unwrap_or_default();
            wb_encoded_len +=
                ctx.raft_wb
                    .append_raft_log(updated_region, updated_region, keyspace_id, &entry);
            if preprocessed_ok {
                if let Some(admin_req) = admin_req {
                    if admin_req.has_commit_merge() {
                        let commit_merge = admin_req.get_commit_merge();
                        let source_region_id = commit_merge.get_source().get_id();
                        ctx.destroyed_regions.insert(source_region_id, keyspace_id);
                        // Emit logical topology change when commit-merge is accepted.
                        // Physical destroy may be delayed by dependents and retried later.
                        Self::notify_region_change(
                            &self.cop_host,
                            commit_merge.get_source(),
                            RegionChangeEvent::Destroy,
                        );
                        self.pending_merge_states.remove(&source_region_id);
                    }
                }
            }
            let entry_is_empty = entry.get_data().is_empty();
            entries.push(entry);

            progress.synced_index = log_index;

            let base_version = preprocessor_ref.shard_meta.as_ref().map(|m| m.base_version);
            Self::update_persist_progress(
                tag,
                progress,
                preprocessed_ok,
                entry_is_empty,
                preprocess_cmd.as_ref(),
                log_index,
                base_version,
            );

            if log_index + 1 < high
                && wb_encoded_len >= self.ctx.config.raft_write_batch_size.0 as i64
            {
                debug!("{} sync_region: wb exceed size limit, break at {}", tag, log_index;
                    "wb_size" => wb_encoded_len, "low" => low, "high" => high);
                res = SyncRegionResult::Resume;
                break;
            }
        }
        if ctx.raft_wb.is_empty() {
            return Ok(res);
        }
        preprocessor_ref.write_raft_state(ctx);
        let wb = mem::take(ctx.raft_wb);
        ctx.raft.apply(&wb);
        *merged_wb_estimated_size += wb.estimated_size();
        merged_wb.merge_write_batch(wb);
        if *merged_wb_estimated_size > self.ctx.config.raft_write_batch_size.0 as usize {
            ctx.raft.persist(mem::take(merged_wb)).expect("raft write");
            *merged_wb_estimated_size = 0;
        }
        let shard = self.kv.get_shard(updated_region);
        if shard.is_none() {
            // shard is not in the keyspace range, skip apply.
            debug!("{} sync_merged_for_regions: skip apply", tag);
            if preprocessor.sync_region() {
                Self::notify_region_change(
                    &self.cop_host,
                    preprocessor.region(),
                    RegionChangeEvent::Update(RegionChangeReason::Other),
                );
            }
            ctx.skip_apply();
            return Ok(res);
        }
        let shard = shard.unwrap();
        // apply committed entries.
        let applier = self
            .appliers
            .entry(updated_region)
            .or_insert_with(|| Self::new_applier(&shard, preprocessor_ref, low - 1));
        ctx.build_apply_msg_for_replication(entries);
        ctx.pre_ctx
            .handle_apply_msgs_for_replication(applier, ctx.apply_ctx);
        // We keep waiting for paused region because later region may depend on it.
        while applier.is_paused() {
            try_force_stop_err!(self);
            let msgs = if let Some(msgs) = ctx.prepared_msgs.remove(&updated_region) {
                // received by previous region, handle it now.
                msgs
            } else {
                let (id, peer_msg) = match self.peer_receiver.recv_timeout(Duration::from_secs(3)) {
                    Ok((id, msg)) => (id, msg),
                    Err(err) => {
                        if err.is_timeout() {
                            warn!("{} waiting for region to unpause", tag);
                            continue;
                        }
                        return Err(Error::Other(Box::new(err)));
                    }
                };
                if id != updated_region {
                    // For another region, handle it later.
                    ctx.prepared_msgs
                        .entry(id)
                        .or_default()
                        .push((id, peer_msg));
                    continue;
                }
                vec![(id, peer_msg)]
            };
            Self::apply_prepared_msgs(ctx, applier, msgs);
        }
        if preprocessor.sync_region() {
            Self::notify_region_change(
                &self.cop_host,
                preprocessor.region(),
                RegionChangeEvent::Update(RegionChangeReason::Other),
            );
        }

        Ok(res)
    }

    fn sync_region_admin_req(
        tag: PeerTag,
        updated_region: u64,
        admin: &AdminRequest,
        log_index: u64,
        pending_merge_states: &mut HashMap<u64 /* source_region_id */, MergeState>,
    ) -> SyncRegionResult {
        if admin.has_commit_merge() {
            let commit_merge = admin.get_commit_merge();
            let source = commit_merge.get_source();
            let source_ready = if let HashMapEntry::Occupied(e) =
                pending_merge_states.entry(source.id)
            {
                let merge_state = e.get();
                let matched = merge_state.get_target().id == updated_region
                    && merge_state.commit == commit_merge.commit;
                if matched {
                    debug!("{} sync_region: commit merge at {}", tag, log_index;
                        "source" => source.id, "commit" => commit_merge.commit);
                } else {
                    debug!("{} sync_region: commit merge not match", tag;
                        "source" => source.id, "commit" => commit_merge.commit, "merge_state" => ?merge_state);
                }
                matched
            } else {
                false
            };
            if !source_ready {
                // need to process source region first.
                info!("{} sync_region: commit merge postponed at {}", tag, log_index;
                    "source" => source.id, "commit" => commit_merge.commit);
                return SyncRegionResult::Postponed;
            }
        }
        SyncRegionResult::Finished
    }

    fn sync_region_admin_req_post_preprocess(
        tag: PeerTag,
        updated_region: u64,
        admin: &AdminRequest,
        log_index: u64,
        preprocessor_ref: &PreprocessRef<'_>,
        pending_merge_states: &mut HashMap<u64 /* source_region_id */, MergeState>,
    ) {
        if admin.has_prepare_merge() {
            if let Some(merge_state) = preprocessor_ref.pending_merge_state.as_ref() {
                debug!("{} sync_region: prepare merge at {}", tag, log_index; "merge_state" => ?merge_state);
                pending_merge_states.insert(updated_region, merge_state.clone());
            } else {
                // Prepare merge is denied. Ref: `preprocess_prepare_merge`.
                debug!("{} sync_region: skip prepare merge at {}", tag, log_index);
            }
        } else if admin.has_rollback_merge() {
            if let Some(merge_state) = preprocessor_ref.pending_merge_state.as_ref() {
                // Should not happen. Rollback merge must succeed.
                // Ref: `preprocess_rollback_merge`.
                warn!("{} sync_region: skip rollback merge at {}", tag, log_index; "merge_state" => ?merge_state);
                debug_assert!(false);
            } else {
                let origin = pending_merge_states.remove(&updated_region);
                debug!("{} sync_region: rollback merge at {}", tag, log_index; "merge_state" => ?origin);
                debug_assert!(origin.is_some());
            }
        }
    }

    fn handle_prepared_msgs(&mut self, ctx: &mut SyncRegionsContext<'_>) {
        let msg_count = self.peer_receiver.len();
        for _ in 0..msg_count {
            let (id, peer_msg) = self.peer_receiver.recv().unwrap();
            ctx.prepared_msgs
                .entry(id)
                .or_default()
                .push((id, peer_msg));
        }
        let prepared_msgs = mem::take(&mut ctx.prepared_msgs);
        for (region_id, msgs) in prepared_msgs {
            tikv_util::set_current_region(region_id);
            let Some(applier) = self.appliers.get_mut(&region_id) else {
                continue;
            };
            Self::apply_prepared_msgs(ctx, applier, msgs);
        }
    }

    fn remove_dependents(
        &mut self,
        ctx: &mut SyncRegionsContext<'_>,
        updated_regions: &mut HashSet<u64>,
    ) {
        for (parent_id, dependent_id) in mem::take(ctx.remove_dependents) {
            let dependent_len = self.raft.remove_dependent(parent_id, dependent_id);
            if dependent_len == 0 {
                // Dependents became empty; the parent may now be eligible for raft-log
                // truncation even if it has no new WAL updates in this sync
                // round.
                updated_regions.insert(parent_id);
                if let Some(keyspace_id) = self.delay_destroy_regions.remove(&parent_id) {
                    ctx.destroyed_regions.insert(parent_id, keyspace_id);
                }
            }
        }
    }

    fn update_persist_progress(
        tag: PeerTag,
        progress: &mut RegionProgress,
        preprocessed_ok: bool,
        entry_is_empty: bool,
        cmd: Option<&RaftCmdRequest>,
        log_idx: u64,
        base_version: Option<u64>,
    ) {
        debug_assert_eq!(progress.synced_index, log_idx);
        let affect_memtable = if preprocessed_ok && !entry_is_empty {
            get_affect_memtable(tag, cmd, base_version)
        } else {
            // Conf changes or failed commands.
            AffectMemtable::None
        };
        progress.persist_progress.update(log_idx, affect_memtable);
    }

    fn update_progress_and_truncate(
        &mut self,
        regions: &HashSet<u64>,
        raft_wb: &mut WriteBatch,
    ) -> Result<()> {
        for &region_id in regions {
            tikv_util::set_current_region(region_id);
            let tag = ShardTag::new(self.merged_store_id(), IdVer::new(region_id, 0));
            let Some(progress) = self.region_progresses.get_mut(&region_id) else {
                warn!("{} update_progress_and_truncate: progress not found", tag);
                continue;
            };
            let synced_index = progress.synced_index;
            // We need to keep the not-synced logs for the next round.
            progress.entries.retain(|&index, _| index > synced_index);

            let truncate_index = progress.data_persisted_log_index();
            let truncate_raft_log =
                self.truncate_region_raft_log(tag, region_id, truncate_index, raft_wb);
            debug!(
                "{} update_progress_and_truncate: truncate entries <= {}, raft log <= {:?}",
                tag, synced_index, truncate_raft_log;
                "truncate_idx" => truncate_index,
            );

            if raft_wb.estimated_size() > self.ctx.config.raft_write_batch_size.0 as usize {
                try_force_stop_err!(self);
                self.raft.write(mem::take(raft_wb)).expect("raft write");
            }
        }
        Ok(())
    }

    fn truncate_region_raft_log(
        &mut self,
        tag: ShardTag,
        region_id: u64,
        mut truncate_index: u64,
        raft_wb: &mut WriteBatch,
    ) -> Option<u64> {
        // Skip truncate region with dependents. The parent region may need the old raft
        // logs on recover.
        if self.raft.has_dependents(region_id) {
            return None;
        }
        let rf_truncated_index = self.raft.get_truncated_index(region_id)?;
        if truncate_index <= rf_truncated_index {
            return None;
        }

        let shard_meta = self.preprocessors.get_mut(&region_id)?.mut_shard_meta()?;
        if shard_meta.data_sequence < truncate_index {
            debug!(
                "{} truncate_region_raft_log: advance data_sequence: {} -> {}",
                tag, shard_meta.data_sequence, truncate_index
            );
            shard_meta.data_sequence = truncate_index;
            write_engine_meta(raft_wb, region_id, shard_meta);
        }

        let persisted_index = shard_meta.data_persisted_log_index();
        truncate_index = cmp::min(truncate_index, persisted_index);
        if truncate_index <= rf_truncated_index {
            return None;
        }
        raft_wb.truncate_raft_log(
            region_id,
            region_id,
            shard_meta.range.keyspace_id,
            truncate_index,
        );
        Some(truncate_index)
    }

    fn destroy_regions(&mut self, ctx: &mut SyncRegionsContext<'_>) {
        for (region_id, keyspace_id) in ctx.destroyed_regions.drain() {
            tikv_util::set_current_region(region_id);

            if self.raft.has_dependents(region_id) {
                info!("{} delay destroy", self.get_region_tag(region_id, 0));
                self.delay_destroy_regions.insert(region_id, keyspace_id);
                continue;
            }

            info!("{} destroy", self.get_region_tag(region_id, 0));
            self.raft.iterate_peer_states(region_id, false, |k, _| {
                ctx.pre_ctx.raft_wb.set_state_bytes(
                    region_id,
                    region_id,
                    keyspace_id,
                    k.clone(),
                    Bytes::new(),
                );
                true
            });
            ctx.pre_ctx.raft_wb.truncate_raft_log(
                region_id,
                region_id,
                keyspace_id,
                TRUNCATE_ALL_INDEX,
            );
            self.remove_shard(region_id);
            self.preprocessors.remove(&region_id);
            self.clear_cached_bootstrap(region_id);
            self.pending_merge_states.remove(&region_id);
            let progress = self.region_progresses.get_mut(&region_id).unwrap();
            progress.persist_progress.truncate_all();
        }
    }

    pub fn remove_shard(&mut self, region_id: u64) {
        self.appliers.remove(&region_id);
        self.kv.remove_shard(region_id);
    }

    fn destroy_regions_on_startup(
        raft: &RfEngine,
        tombstone_regions: HashMap<u64, u32>,
    ) -> Result<HashMap<u64, u32> /* delay_destroy_regions */> {
        let mut delay_destroy_regions = HashMap::default();
        if tombstone_regions.is_empty() {
            return Ok(delay_destroy_regions);
        }

        let merged_store_id = raft.get_engine_id();
        let get_region_tag = |region_id| ShardTag::new(merged_store_id, IdVer::new(region_id, 0));

        let mut raft_wb = rfengine::WriteBatch::new();
        for (region_id, keyspace_id) in tombstone_regions {
            if raft.has_dependents(region_id) {
                info!("{} delay destroy", get_region_tag(region_id));
                delay_destroy_regions.insert(region_id, keyspace_id);
                continue;
            }

            info!("{} destroy", get_region_tag(region_id));
            raft.iterate_peer_states(region_id, false, |k, _| {
                raft_wb.set_state_bytes(region_id, region_id, keyspace_id, k.clone(), Bytes::new());
                true
            });
            raft_wb.truncate_raft_log(region_id, region_id, keyspace_id, TRUNCATE_ALL_INDEX);
        }

        if !raft_wb.is_empty() {
            raft.write(raft_wb).expect("raft write");
        }
        Ok(delay_destroy_regions)
    }

    fn new_applier(
        shard: &Shard,
        preprocess_ref: PreprocessRef<'_>,
        applied_index: u64,
    ) -> Applier {
        let encryption_state = shard.get_encryption_state();
        let term_val = shard.get_property(TERM_KEY).unwrap();
        let term = term_val.chunk().get_u64_le();
        Applier::new_for_replication(
            preprocess_ref.region.clone(),
            encryption_state,
            RaftApplyState::new(applied_index, term),
        )
    }

    fn apply_prepared_msgs(
        ctx: &mut SyncRegionsContext<'_>,
        applier: &mut Applier,
        msgs: Vec<(u64, Box<PeerMsg>)>,
    ) {
        for (_, msg) in msgs {
            ctx.build_prepared_msg_for_replication(msg);
        }
        ctx.pre_ctx
            .handle_apply_msgs_for_replication(applier, ctx.apply_ctx);
    }

    /// Returns true when merged rfengine already has reusable local bootstrap
    /// state for the region.
    ///
    /// `TRUNCATE_ALL_INDEX` means the local peer was cleared/destroyed. That is
    /// not a usable bootstrap for later replay re-entry.
    fn merged_bootstrap_exists(&self, region_id: u64) -> bool {
        self.raft
            .get_truncated_index(region_id)
            .is_some_and(|idx| idx != TRUNCATE_ALL_INDEX)
    }

    /// Cache an installable live bootstrap from origin WAL for a first-seen
    /// learner peer that is still moving into the merged-store target set.
    ///
    /// This only applies before merged rfengine has local bootstrap for the
    /// region. Ordinary live bootstraps can be installed directly from origin
    /// WAL. Any collected bootstrap whose `shard_meta.parent.is_some()` is
    /// cleared instead of cached, because replay does not install that origin
    /// payload directly.
    fn observe_origin_region_bootstrap_for_peer_moving_in(
        &mut self,
        merged_store_id: u64,
        store_id: u64,
        origin_wb: &WriteBatch,
        region_id: u64,
        peer_id: u64,
        region_state: &RegionLocalState,
        target_stores: &HashSet<u64 /* store_id */>,
    ) {
        let region = region_state.get_region();
        let tag = PeerTag::new(merged_store_id, RegionIdVer::from_region(region));
        if !is_peer_moving_in(region, store_id, target_stores) {
            debug!("{} observe_origin_region_bootstrap: not a moving in peer", tag;
                "origin_store" => store_id, "region" => ?region, "target_stores" => ?target_stores);
            return;
        }

        let bootstrap = match collect_installable_region_bootstrap(
            merged_store_id,
            store_id,
            origin_wb,
            region_id,
            peer_id,
            region_state,
        ) {
            Ok(bootstrap) => bootstrap,
            Err(err) => {
                debug!("{} observe_origin_region_bootstrap: not installable bootstrap: {:?}", tag, err;
                    "origin_store" => store_id);
                return;
            }
        };

        if bootstrap.shard_meta.parent.is_some() {
            debug!(
                "{} observe_origin_region_bootstrap: skip bootstrap for split/merge",
                tag
            );
            return;
        }

        insert_installable_region_bootstrap_into(
            &mut self.installable_region_bootstraps,
            region_id,
            bootstrap,
        );
    }

    fn clear_cached_bootstrap(&mut self, region_id: u64) {
        self.installable_region_bootstraps.remove(&region_id);
    }

    fn try_load_live_bootstrap_shard(&self, tag: PeerTag, shard_meta: &ShardMeta) -> Result<()> {
        if self.kv.get_shard(shard_meta.id).is_some() {
            return Ok(());
        }

        let keyspace_id = shard_meta.range.keyspace_id;
        if !self.manifest.keyspace_states.contains_key(&keyspace_id) {
            debug!(
                "{} install live bootstrap: skip shard load for inactive keyspace",
                tag; "keyspace" => keyspace_id
            );
            return Ok(());
        }

        info!(
            "{} install live bootstrap: load shard into kvengine",
            tag; "keyspace" => keyspace_id
        );
        // `kv.load_shards` may prepare remote snapshot files on demand. Doing
        // that directly from `sync_region` blocks the merged-engine main
        // thread while replay is holding this sync round.
        //
        // TODO: Split this into an async prepare -> load pipeline so live
        // bootstrap shard reload does not block `sync_merged` on remote file
        // readiness.
        let mut load_metas = StdHashMap::with_capacity(1);
        load_metas.insert(shard_meta.id, shard_meta.clone());
        self.kv
            .load_shards(load_metas, self.recover_handler.clone(), None)?;
        Ok(())
    }

    fn install_installable_region_bootstrap(
        &mut self,
        ctx: &mut SyncRegionsContext<'_>,
        region_id: u64,
    ) -> Result<Option<ShardMeta>> {
        let Some(bootstrap) = self.installable_region_bootstraps.remove(&region_id) else {
            return Ok(None);
        };
        let keyspace_id = bootstrap.shard_meta.range.keyspace_id;
        write_engine_meta(ctx.raft_wb, region_id, &bootstrap.shard_meta);
        for (k, v) in &bootstrap.peer_states {
            update_peer_state_without_engine_meta(
                ctx.raft_wb,
                k,
                v,
                self.merged_store_id(),
                region_id,
                keyspace_id,
            );
        }
        ctx.raft_wb
            .truncate_raft_log(region_id, region_id, keyspace_id, bootstrap.truncated_index);

        let progress = self.region_progresses.get_mut(&region_id).unwrap();
        progress.synced_index = cmp::max(
            progress.synced_index,
            bootstrap.raft_state.get_last_preprocessed_index(),
        );
        progress.update_commit_index(bootstrap.raft_state.get_commit());
        if bootstrap.truncated_index > progress.data_persisted_log_index() {
            progress.persist_progress.reset(bootstrap.truncated_index);
        }
        self.newly_bootstrapped_regions.push(region_id);

        info!("{} install live bootstrap", bootstrap.shard_meta.tag();
            "bootstrap" => ?bootstrap, "progress" => ?progress);
        Ok(Some(bootstrap.shard_meta))
    }
}

#[cfg(test)]
mod large_memtable_tests {
    use super::*;

    #[test]
    fn test_replay_engine_enables_crossbeam_memtable() {
        let options = MergedEngine::replay_kv_engine_options(
            &MergedEngineConfig::default(),
            PathBuf::from("db"),
        );

        assert!(options.enable_crossbeam_memtable);
    }
}

// Helpers:
impl MergedEngine {
    fn flush_sync_region_wb(
        &self,
        ctx: &mut SyncRegionsContext<'_>,
        merged_wb: &mut rfengine::WriteBatch,
        merged_wb_estimated_size: &mut usize,
    ) {
        if ctx.raft_wb.is_empty() {
            return;
        }
        let wb = mem::take(ctx.raft_wb);
        ctx.raft.apply(&wb);
        *merged_wb_estimated_size += wb.estimated_size();
        merged_wb.merge_write_batch(wb);
        if *merged_wb_estimated_size > self.ctx.config.raft_write_batch_size.0 as usize {
            ctx.raft.persist(mem::take(merged_wb)).expect("raft write");
            *merged_wb_estimated_size = 0;
        }
    }

    /// Drains and returns region IDs that were installed via the live-bootstrap
    /// path since the last call. Called after each successful `sync_merged`
    /// so the replication worker can immediately report the regions to
    /// rep-pd without waiting for the next report loop.
    pub fn take_newly_bootstrapped_regions(&mut self) -> Vec<u64> {
        mem::take(&mut self.newly_bootstrapped_regions)
    }
}

enum SyncRegionResult {
    Finished,
    Postponed,
    Resume,
    Dropped,
}

impl SyncRegionResult {
    fn is_postponed(&self) -> bool {
        matches!(self, Self::Postponed)
    }
}

struct SyncRegionsContext<'a> {
    pre_ctx: PreprocessContext<'a>,
    apply_ctx: &'a mut ApplyContext,
    prepared_msgs: HashMap<u64 /* region_id */, Vec<(u64 /* region_id */, Box<PeerMsg>)>>,
    destroyed_regions: HashMap<u64, u32>,
}

impl<'a> ops::Deref for SyncRegionsContext<'a> {
    type Target = PreprocessContext<'a>;

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

impl ops::DerefMut for SyncRegionsContext<'_> {
    fn deref_mut(&mut self) -> &mut Self::Target {
        &mut self.pre_ctx
    }
}

impl Drop for SyncRegionsContext<'_> {
    fn drop(&mut self) {
        debug_assert!(self.apply_msgs.is_empty());
        debug_assert!(self.prepared_msgs.is_empty());
        debug_assert!(self.destroyed_regions.is_empty());
        debug_assert!(self.raft_wb.is_empty());
    }
}

impl SyncRegionsContext<'_> {
    fn skip_apply(&mut self) {
        self.apply_msgs.clear();
    }

    fn clear(&mut self) {
        self.apply_msgs.clear();
        self.prepared_msgs.clear();
        self.destroyed_regions.clear();
        self.raft_wb.reset();
    }
}

impl fmt::Debug for SyncRegionsContext<'_> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("SyncRegionsContext")
            .field("apply_msgs", &self.apply_msgs.len())
            .field("prepared_msgs", &self.prepared_msgs.len())
            .field("destroyed_regions", &self.destroyed_regions.len())
            .field("raft_wb", &self.raft_wb.len())
            .finish()
    }
}

fn update_peer_state_without_engine_meta(
    wb: &mut rfengine::WriteBatch,
    k: &Bytes,
    v: &Bytes,
    store_id: u64,
    region_id: u64,
    keyspace_id: u32,
) {
    if k.starts_with(rfengine::REGION_META_KEY_PREFIX) {
        let mut origin_state = kvproto::raft_serverpb::RegionLocalState::new();
        origin_state.merge_from_bytes(v).unwrap();
        let region_version = origin_state.get_region().get_region_epoch().get_version();
        let merged_region_state = merge_region_local_state(&origin_state, store_id);
        let data = merged_region_state.write_to_bytes().unwrap();
        wb.set_state_bytes(
            region_id,
            region_id,
            keyspace_id,
            rfengine::region_state_key(region_version),
            data.into(),
        );
    } else if !ALL_KV_ENGINE_META_KEYS.contains(&k.chunk()) {
        wb.set_state_bytes(region_id, region_id, keyspace_id, k.clone(), v.clone());
    }
}

// Update the entry if needed and return RaftCmdRequest for further
// pre-processing if the entry is normal command.
// Note: Change `get_affect_memtable` if the logic here is changed.
fn update_entry(entry: &mut Entry, merged_store_id: u64) -> Option<RaftCmdRequest> {
    if entry.get_entry_type() != raft_proto::eraftpb::EntryType::EntryNormal {
        // We don't need to handle conf change, set it to empty.
        entry.set_entry_type(raft_proto::eraftpb::EntryType::EntryNormal);
        entry.set_data(Bytes::new());
        return None;
    }
    if entry.get_data().is_empty() {
        return None;
    }
    let mut cmd = get_preprocess_cmd(entry)?;
    if !cmd.has_admin_request() {
        return Some(cmd);
    }
    let header = cmd.mut_header();
    let region_id = header.get_region_id();
    let peer = header.mut_peer();
    peer.set_store_id(merged_store_id);
    peer.set_id(region_id);
    let epoch = header.mut_region_epoch();
    epoch.set_conf_ver(1);
    let admin_cmd = cmd.mut_admin_request();
    if admin_cmd.has_splits() {
        let splits = admin_cmd.mut_splits().mut_requests();
        for req in splits.iter_mut() {
            req.set_new_peer_ids(vec![req.new_region_id]);
        }
    } else if admin_cmd.has_prepare_merge() {
        let prepare_merge = admin_cmd.mut_prepare_merge();
        let new_target = merged_region_meta(prepare_merge.get_target(), merged_store_id);
        prepare_merge.set_target(new_target);
    } else if admin_cmd.has_commit_merge() {
        let commit_merge = admin_cmd.mut_commit_merge();
        let new_source = merged_region_meta(commit_merge.get_source(), merged_store_id);
        commit_merge.set_source(new_source);
    }
    let new_cmd = cmd.write_to_bytes().unwrap();
    entry.set_data(new_cmd.into());
    Some(cmd)
}

// Get affect memtable type of the raft command.
// Depends on the logic of `update_entry`.
fn get_affect_memtable(
    tag: PeerTag,
    cmd: Option<&RaftCmdRequest>,
    base_version: Option<u64>,
) -> AffectMemtable {
    let Some(cmd) = cmd else {
        // Normal writes (no preprocess).
        return AffectMemtable::Write;
    };
    let Some(custom) = rlog::get_custom_log(cmd) else {
        // Admin requests.
        return AffectMemtable::None;
    };
    custom.is_affect_memtable(tag, base_version, None)
}

// merged region meta has a single peer with id same as region id.
fn merged_region_meta(origin_region: &metapb::Region, store_id: u64) -> metapb::Region {
    let mut merged_peer = Peer::new();
    merged_peer.set_id(origin_region.get_id());
    merged_peer.set_store_id(store_id);
    let mut merged_region = origin_region.clone();
    merged_region.set_peers(vec![merged_peer].into());
    merged_region.mut_region_epoch().set_conf_ver(1);
    merged_region
}

fn merge_region_local_state(
    origin: &kvproto::raft_serverpb::RegionLocalState,
    store_id: u64,
) -> kvproto::raft_serverpb::RegionLocalState {
    let mut merged = origin.clone();
    merged.set_region(merged_region_meta(origin.get_region(), store_id));
    if merged.has_merge_state() {
        let merge_state = merged.get_merge_state();
        if merge_state.has_target() {
            let new_target = merged_region_meta(merge_state.get_target(), store_id);
            merged.mut_merge_state().set_target(new_target);
        }
    }
    merged
}

#[derive(Clone, Debug, PartialEq)]
pub struct RaftLogOpWithCounter {
    pub op: RaftLogOp,
    counter: u8,
}

impl ops::Deref for RaftLogOpWithCounter {
    type Target = RaftLogOp;

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

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

impl From<RaftLogOp> for RaftLogOpWithCounter {
    fn from(op: RaftLogOp) -> Self {
        RaftLogOpWithCounter { op, counter: 1 }
    }
}

impl RaftLogOpWithCounter {
    pub fn inc_counter(&mut self) {
        self.counter = self.counter.saturating_add(1);
    }

    pub fn counter(&self) -> u8 {
        self.counter
    }
}

fn get_keyspace_id_of_snapshot(snap: &kvenginepb::Snapshot) -> u32 {
    ApiV2::get_u32_keyspace_id_by_key(snap.get_outer_start()).unwrap_or_default()
}

#[cfg(feature = "testexport")]
#[derive(Default, Clone)]
pub struct ForceStop(Arc<std::sync::atomic::AtomicBool>);

#[cfg(feature = "testexport")]
impl ForceStop {
    pub fn set(&self) {
        use std::sync::atomic::Ordering;
        self.0.store(true, Ordering::Relaxed);
    }

    pub fn get(&self) -> bool {
        use std::sync::atomic::Ordering;
        self.0.load(Ordering::Relaxed)
    }
}

#[cfg(not(feature = "testexport"))]
pub type ForceStop = ();

struct RaftProgress {
    last_preprocess_index: u64,
    last_index: u64,
    term: u64,
}

impl From<&RaftState> for RaftProgress {
    fn from(rs: &RaftState) -> Self {
        RaftProgress {
            last_preprocess_index: rs.get_last_preprocessed_index(),
            last_index: rs.get_last_index(),
            term: rs.get_term(),
        }
    }
}

fn raft_progress_cmp_key(rs: &RaftState) -> (u64, u64, u64) {
    (
        rs.get_last_preprocessed_index(),
        rs.get_last_index(),
        rs.get_term(),
    )
}

fn raft_progress_cmp_raw(
    last_preprocessed_index: u64,
    last_index: u64,
    term: u64,
) -> (u64, u64, u64) {
    (last_preprocessed_index, last_index, term)
}

fn raft_progress_is_newer(lhs: &RaftState, rhs: &RaftState) -> bool {
    raft_progress_cmp_key(lhs) > raft_progress_cmp_key(rhs)
}

#[derive(Default)]
struct RegionsRaftProgress {
    inner: HashMap<u64, RaftProgress>,
}

impl RegionsRaftProgress {
    fn update(&mut self, region_id: u64, rs: &RaftState) -> bool /* updated (is_newer) */ {
        match self.inner.entry(region_id) {
            HashMapEntry::Vacant(e) => {
                e.insert(RaftProgress::from(rs));
                true
            }
            HashMapEntry::Occupied(mut e) => {
                let curr = e.get();
                let is_newer = raft_progress_cmp_key(rs)
                    > raft_progress_cmp_raw(curr.last_preprocess_index, curr.last_index, curr.term);
                if is_newer {
                    e.insert(RaftProgress::from(rs));
                }
                is_newer
            }
        }
    }
}

fn fetch_raft_entries_to_region_progress(
    tag: ShardTag,
    raft: &RfEngine,
    peer_id: u64,
    low_idx: u64,
    high_idx: u64,
    region_progress: &mut RegionProgress,
) -> Result<()> {
    let mut entry_buf = Vec::with_capacity((high_idx - low_idx) as usize);
    box_try!(raft.fetch_raft_entries_to(peer_id, low_idx, high_idx, None, &mut entry_buf, None)
        .map_err(|err| {
            let stats = raft.get_peer_stats(peer_id);
            let truncated_state = rfstore::store::load_raft_truncated_state(raft, peer_id);
            error!("{} fetch_raft_entries failed: {:?}", tag, err;
                "low" => low_idx, "high" => high_idx, "stats" => ?stats, "truncated_state" => ?truncated_state);
            err
        }));
    for entry in entry_buf {
        debug!(
            "{} fetch_raft_entries: insert log index {}",
            tag, entry.index
        );
        region_progress.upsert_entry(tag, entry.index, entry.term as u32, || {
            RaftLogOp::new(&entry).into()
        });
    }
    Ok(())
}

#[inline]
pub fn peer_is_skippable(region_local_state: &RegionLocalState) -> bool {
    // `start_key` is empty if region is not initialized, and we will get incorrect
    // keyspace.
    region_local_state.state == PeerState::Tombstone
        || !is_region_initialized(region_local_state.get_region())
}

// Ref: https://github.com/tidbcloud/cloud-storage-engine/issues/4319.
fn peer_is_stale(
    shard_meta: &ShardMeta,
    region_peers_map: &StdHashMap<u64 /* region_id */, u64 /* peer_id */>,
) -> bool {
    if let Some(box parent) = &shard_meta.parent {
        if parent.id != shard_meta.id && !region_peers_map.contains_key(&parent.id) {
            warn!("{} peer is stale", shard_meta.tag(); "parent" => %parent.tag());
            return true;
        }
    }
    false
}

#[derive(Clone, Debug)]
struct InstallableRegionBootstrap {
    raft_state: RaftState,
    truncated_index: u64,
    shard_meta: ShardMeta,
    peer_states: Vec<(Bytes, Bytes)>,
}

fn live_bootstrap_error(
    merged_store_id: u64,
    store_id: u64,
    region_id: u64,
    region_state: &RegionLocalState,
    field: &str,
) -> Error {
    let region_version = region_state.get_region().get_region_epoch().get_version();
    let tag = ShardTag::new(merged_store_id, IdVer::new(region_id, region_version));
    let origin_tag = ShardTag::new(store_id, IdVer::new(region_id, region_version));
    Error::Other(box_err!(
        "{} update_wal: missing {} for first-seen live bootstrap, origin: {}",
        tag,
        field,
        origin_tag
    ))
}

fn collect_installable_region_bootstrap(
    merged_store_id: u64,
    store_id: u64,
    origin_wb: &WriteBatch,
    region_id: u64,
    peer_id: u64,
    region_state: &RegionLocalState,
) -> Result<InstallableRegionBootstrap> {
    let raft_state = load_last_raft_state_from_wb(origin_wb, peer_id).ok_or_else(|| {
        live_bootstrap_error(
            merged_store_id,
            store_id,
            region_id,
            region_state,
            "raft_state",
        )
    })?;
    let mut shard_meta =
        load_engine_meta_from_wb(origin_wb, store_id, peer_id).ok_or_else(|| {
            live_bootstrap_error(
                merged_store_id,
                store_id,
                region_id,
                region_state,
                "engine_meta",
            )
        })?;
    let truncated_index = origin_wb
        .get_truncated_idx(peer_id)
        .unwrap_or_default()
        .max(RAFT_INIT_LOG_INDEX);

    let keyspace_id = ApiV2::get_u32_keyspace_id_by_key(region_state.get_region().get_start_key())
        .unwrap_or_default();
    debug_assert_eq!(shard_meta.range.keyspace_id, keyspace_id);
    shard_meta.engine_id = merged_store_id;

    let mut peer_states = Vec::with_capacity(4);
    origin_wb.iterate_peer_states_bytes(peer_id, |key, value| {
        peer_states.push((key.clone(), value.clone()));
    });

    Ok(InstallableRegionBootstrap {
        raft_state,
        truncated_index,
        shard_meta,
        peer_states,
    })
}

fn insert_installable_region_bootstrap_into(
    bootstraps: &mut HashMap<u64, InstallableRegionBootstrap>,
    region_id: u64,
    bootstrap: InstallableRegionBootstrap,
) {
    match bootstraps.entry(region_id) {
        HashMapEntry::Vacant(entry) => {
            info!(
                "{} insert_installable_region_bootstrap_into",
                bootstrap.shard_meta.tag()
            );
            entry.insert(bootstrap);
        }
        HashMapEntry::Occupied(mut entry) => {
            if raft_progress_is_newer(&bootstrap.raft_state, &entry.get().raft_state) {
                info!(
                    "{} insert_installable_region_bootstrap_into",
                    bootstrap.shard_meta.tag()
                );
                entry.insert(bootstrap);
            }
        }
    }
}

/// Returns whether the learner peer on `store_id` still looks like a peer move
/// into `target_stores`.
fn is_peer_moving_in(
    region: &metapb::Region,
    store_id: u64,
    target_stores: &HashSet<u64 /* store_id */>,
) -> bool {
    let Some(peer) = tikv_util::store::find_peer(region, store_id) else {
        return false;
    };
    if peer.role != PeerRole::Learner {
        return false;
    }
    let Some(max_overlapped) = region.get_peers().len().checked_sub(1) else {
        return false;
    };
    let max_overlapped = cmp::min(max_overlapped, 2);
    let overlapped = region
        .get_peers()
        .iter()
        .filter(|p| target_stores.contains(&p.store_id))
        .count();
    overlapped <= max_overlapped
}

#[cfg(test)]
mod regions_raft_progress_tests {
    use kvenginepb::{ChangeSet, Snapshot};
    use kvproto::{metapb, raft_serverpb::PeerState};
    use rfengine::raft_state_key;
    use rfstore::store::write_peer_state;
    use tikv_util::{
        codec::bytes::encode_bytes,
        store::{new_learner_peer, new_peer},
    };

    use super::*;

    fn new_region_state(region_id: u64, region_ver: u64, keyspace_id: u32) -> RegionLocalState {
        let mut region = metapb::Region::default();
        region.set_id(region_id);
        region.set_start_key(encode_bytes(&ApiV2::get_txn_keyspace_prefix(keyspace_id)));
        region.set_end_key(encode_bytes(&ApiV2::get_txn_keyspace_prefix(
            keyspace_id + 1,
        )));
        region.mut_region_epoch().set_version(region_ver);
        region.mut_region_epoch().set_conf_ver(1);
        let mut peer = metapb::Peer::default();
        peer.set_id(region_id);
        peer.set_store_id(1);
        region.set_peers(vec![peer].into());

        let mut state = RegionLocalState::default();
        state.set_region(region);
        state.set_state(PeerState::Normal);
        state
    }

    fn new_shard_meta(
        store_id: u64,
        region_id: u64,
        region_ver: u64,
        keyspace_id: u32,
    ) -> ShardMeta {
        let mut snap = Snapshot::default();
        snap.set_outer_start(ApiV2::get_txn_keyspace_prefix(keyspace_id));
        snap.set_outer_end(ApiV2::get_txn_keyspace_prefix(keyspace_id + 1));

        let mut cs = ChangeSet::default();
        cs.set_shard_id(region_id);
        cs.set_shard_ver(region_ver);
        cs.set_sequence(10);
        cs.set_snapshot(snap);
        ShardMeta::new(store_id, &cs)
    }

    fn new_raft_state(peer_id: u64) -> RaftState {
        let mut raft_state = RaftState::default();
        let mut hs = eraftpb::HardState::default();
        hs.set_term(7);
        hs.set_vote(peer_id);
        hs.set_commit(RAFT_INIT_LOG_INDEX);
        raft_state.set_hard_state(&hs);
        raft_state.set_last_index(RAFT_INIT_LOG_INDEX);
        raft_state.set_last_preprocessed_index(RAFT_INIT_LOG_INDEX);
        raft_state
    }

    fn encode_raft_state(raft_state: &RaftState) -> Vec<u8> {
        let mut buf = Vec::with_capacity(40);
        let hs = raft_state.get_hard_state();
        buf.extend_from_slice(&raft_state.get_term().to_le_bytes());
        buf.extend_from_slice(&hs.get_vote().to_le_bytes());
        buf.extend_from_slice(&raft_state.get_commit().to_le_bytes());
        buf.extend_from_slice(&raft_state.get_last_index().to_le_bytes());
        buf.extend_from_slice(&raft_state.get_last_preprocessed_index().to_le_bytes());
        buf
    }

    fn new_region_with_peers(peers: Vec<metapb::Peer>) -> metapb::Region {
        let mut region = metapb::Region::default();
        region.set_peers(peers.into());
        region
    }

    fn new_target_stores(store_ids: &[u64]) -> HashSet<u64> {
        store_ids.iter().copied().collect()
    }

    #[test]
    fn test_regions_raft_progress_update_prefers_last_preprocessed_index() {
        let mut progress = RegionsRaftProgress::default();

        let mut a = RaftState::default();
        let mut a_hs = eraftpb::HardState::default();
        a_hs.set_term(10);
        a.set_hard_state(&a_hs);
        a.set_last_index(100);
        a.set_last_preprocessed_index(7);

        let mut b = RaftState::default();
        let mut b_hs = eraftpb::HardState::default();
        b_hs.set_term(1);
        b.set_hard_state(&b_hs);
        b.set_last_index(100);
        b.set_last_preprocessed_index(11);

        assert!(progress.update(42, &a));
        assert!(progress.update(42, &b));
        assert!(!progress.update(42, &a));
    }

    #[test]
    fn test_collect_installable_region_bootstrap_uses_init_log_index_when_truncated_state_missing()
    {
        let merged_store_id = 1024;
        let store_id = 1;
        let region_id = 11;
        let peer_id = 11;
        let region_ver = 9;
        let keyspace_id = 6;

        let region_state = new_region_state(region_id, region_ver, keyspace_id);
        let shard_meta = new_shard_meta(store_id, region_id, region_ver, keyspace_id);
        let raft_state = new_raft_state(peer_id);

        let mut wb = WriteBatch::new();
        write_peer_state(
            &mut wb,
            peer_id,
            region_state.get_region(),
            PeerState::Normal,
            None,
        );
        write_engine_meta(&mut wb, peer_id, &shard_meta);
        wb.set_state(
            peer_id,
            region_id,
            keyspace_id,
            &raft_state_key(region_ver),
            &encode_raft_state(&raft_state),
        );

        let bootstrap = collect_installable_region_bootstrap(
            merged_store_id,
            store_id,
            &wb,
            region_id,
            peer_id,
            &region_state,
        )
        .unwrap();
        assert_eq!(bootstrap.truncated_index, RAFT_INIT_LOG_INDEX);
        assert_eq!(bootstrap.shard_meta.engine_id, merged_store_id);
    }

    #[test]
    fn test_collect_installable_region_bootstrap_requires_engine_meta() {
        let merged_store_id = 1024;
        let store_id = 1;
        let region_id = 12;
        let peer_id = 12;
        let region_ver = 3;
        let keyspace_id = 8;

        let region_state = new_region_state(region_id, region_ver, keyspace_id);
        let raft_state = new_raft_state(peer_id);

        let mut wb = WriteBatch::new();
        write_peer_state(
            &mut wb,
            peer_id,
            region_state.get_region(),
            PeerState::Normal,
            None,
        );
        wb.set_state(
            peer_id,
            region_id,
            keyspace_id,
            &raft_state_key(region_ver),
            &encode_raft_state(&raft_state),
        );

        let err = collect_installable_region_bootstrap(
            merged_store_id,
            store_id,
            &wb,
            region_id,
            peer_id,
            &region_state,
        )
        .unwrap_err();
        assert!(err.to_string().contains("missing engine_meta"));
    }

    #[test]
    fn test_is_peer_moving_in_keeps_target_overlap_below_settled_placement() {
        let two_peer_region = new_region_with_peers(vec![new_peer(1, 11), new_learner_peer(2, 12)]);
        let three_peer_region = new_region_with_peers(vec![
            new_peer(1, 21),
            new_peer(2, 22),
            new_learner_peer(3, 23),
        ]);
        let four_peer_region = new_region_with_peers(vec![
            new_peer(1, 31),
            new_peer(2, 32),
            new_peer(3, 33),
            new_learner_peer(4, 34),
        ]);
        let cases = [
            (
                "empty region has no moving peer",
                metapb::Region::default(),
                1,
                vec![1],
                false,
            ),
            (
                "peer not found is not moving in",
                two_peer_region.clone(),
                9,
                vec![2],
                false,
            ),
            (
                "peer not learner is not moving in",
                three_peer_region.clone(),
                1,
                vec![2, 3],
                false,
            ),
            (
                "two-peer region is moving in while target stores cover only one peer",
                two_peer_region.clone(),
                2,
                vec![2],
                true,
            ),
            (
                "two-peer region is settled once target stores cover both peers",
                two_peer_region,
                2,
                vec![1, 2],
                false,
            ),
            (
                "three-peer region is moving in while target stores cover only two peers",
                three_peer_region.clone(),
                3,
                vec![2, 3],
                true,
            ),
            (
                "three-peer region is settled once target stores cover all peers",
                three_peer_region,
                3,
                vec![1, 2, 3],
                false,
            ),
            (
                "four-peer region still counts as moving in when target stores overlap with only two peers",
                four_peer_region.clone(),
                4,
                vec![3, 4],
                true,
            ),
            (
                "four-peer region is settled once target stores overlap with three peers",
                four_peer_region,
                4,
                vec![2, 3, 4],
                false,
            ),
        ];

        for (name, region, store_id, target_stores, expected) in cases {
            let actual = is_peer_moving_in(&region, store_id, &new_target_stores(&target_stores));
            assert_eq!(
                actual, expected,
                "{name}: store {store_id} in {region:?} on {target_stores:?}"
            );
        }
    }

    #[test]
    fn test_installable_region_bootstrap_prefers_newer_raft_progress() {
        let region_id = 42;
        let keyspace_id = 1;
        let shard_meta = new_shard_meta(1, region_id, 1, keyspace_id);

        let mut older = new_raft_state(region_id);
        older.set_last_preprocessed_index(7);
        older.set_last_index(100);

        let mut newer = new_raft_state(region_id);
        let mut newer_hs = eraftpb::HardState::default();
        newer_hs.set_term(1);
        newer_hs.set_vote(region_id);
        newer_hs.set_commit(RAFT_INIT_LOG_INDEX);
        newer.set_hard_state(&newer_hs);
        newer.set_last_preprocessed_index(11);
        newer.set_last_index(100);

        let mut bootstraps = HashMap::default();
        insert_installable_region_bootstrap_into(
            &mut bootstraps,
            region_id,
            InstallableRegionBootstrap {
                raft_state: older,
                truncated_index: RAFT_INIT_LOG_INDEX,
                shard_meta: shard_meta.clone(),
                peer_states: Vec::new(),
            },
        );
        insert_installable_region_bootstrap_into(
            &mut bootstraps,
            region_id,
            InstallableRegionBootstrap {
                raft_state: newer,
                truncated_index: RAFT_INIT_LOG_INDEX,
                shard_meta,
                peer_states: Vec::new(),
            },
        );

        assert_eq!(
            bootstraps
                .get(&region_id)
                .unwrap()
                .raft_state
                .get_last_preprocessed_index(),
            11
        );
    }
}