// Copyright 2025 PingCAP, Inc.

use std::{
    collections::{HashMap as StdHashMap, hash_map::Entry},
    fs, mem,
    net::SocketAddr,
    ops,
    ops::Deref,
    path::PathBuf,
    str::FromStr,
    sync::{Arc, atomic::AtomicBool},
    time::Duration,
};

use api_version::ApiV2;
use bytes::Bytes;
use cdc::{CdcEvent, Conn, ConnId, metrics::*};
use collections::{HashMap, HashMapExt, HashSet};
use grpcio::{ChannelBuilder, EnvBuilder, ServerBuilder};
use grpcio_health::{HealthService, ServingStatus, create_health};
use hyper::StatusCode;
use kvengine::{
    Engine, IdVer, LOCK_CF, ShardMeta, ShardTag, SnapAccess, UserMeta, WRITE_CF,
    dfs::Dfs,
    table::{InnerKey, SnapVersion, tiny_meta},
};
use kvproto::{
    cdcpb,
    cdcpb::{
        ChangeDataRequest, Event, EventLogType, EventRow, EventRowOpType, ResolvedTs,
        create_change_data,
    },
    metapb, pdpb,
    pdpb::StoreStats,
    raft_cmdpb::AdminRequest,
    tikvpb::create_tikv,
};
use log_wrappers::Value as LogValue;
use merged_engine::{
    ForceStop, KeyspaceShardMetaScope, KeyspaceShardMetas, LoadShardsMode, MergedEngine,
    MergedEngineContext, MergedEngineInitOpts,
};
use pd_client::{PdClient, RegionStat, util::check_resp_header};
use raftstore::{
    RegionInfoAccessor,
    coprocessor::{BoxRegionChangeObserver, RegionInfoProvider},
    store::util::is_region_initialized,
};
use rfengine::{RfEngine, TRUNCATE_ALL_INDEX};
use rfstore::store::{ApplyContext, GcRunner};
use security::{HttpClient, SecurityConfig};
use serde_json::{Value, json};
use tikv_util::{
    AsyncSafeRuntime, box_err, box_try, codec, debug, error,
    future::paired_future_callback,
    info,
    metrics::{ServerIsBusyType, new_server_is_busy},
    mpsc::{Receiver, SendError, Sender},
    thd_name,
    time::Instant,
    trace, warn, worker,
};
use tokio::sync::{OwnedSemaphorePermit, Semaphore};
use txn_types::{LockType, TimeStamp, parse_lock};

use crate::{
    CdcMsg, Deregister, Error, KeyspaceService, KeyspaceStates, ReplicationScheduler,
    ReplicationService, ReplicationWorkerConfig, Result,
    apply_observer::{CdcApplyObserver, RegionEvents, is_index_key},
    delegate::{
        InitAlive, InitId, RegionDelegate, RegionResolver, RequestId, RequestInfo, RequestKey,
    },
    error::error_label,
    kube::{KeyspaceKubeService, KubeApi},
    metrics::*,
    provisioned::KeyspaceProvisionedService,
    safepoint::{KeyspaceContext, ServiceSafepointManager},
    scheduler::get_cdc_status,
    ticdc_util::TiCdcError,
    util::{
        DISPATCH_CDC_TIMEOUT, ResolvedTsStats, build_request_range_for_keyspace,
        keyspace_prefix_len, post_to_ticdc, trim_keyspace_prefix_for_encoded_range,
    },
    wal_fetcher::{
        FetchedWalData, FetchedWalItem, FetchedWalsBatch, WalFetcherHandle, WalFetcherRequest,
        spawn_wal_fetcher,
    },
    wal_spool::cleanup_stale_wal_spool_dirs,
};

mod lazy_load;
pub use lazy_load::LazyLoadConfig;
#[path = "unsafe_recover.rs"]
mod unsafe_recover;

const MAX_INITIALIZE_SCAN_BATCH_BYTES: usize = 1024 * 1024;

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

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

pub struct ReplicationWorker {
    ctx: MergedEngineContext,
    pd_endpoints: Vec<String>,
    http_client: Arc<HttpClient>,
    config: ReplicationWorkerConfig,
    merged_engine: MergedEngine,
    grpc_server: Option<grpcio::Server>,
    health_service: Option<HealthService>,
    kube_api: Option<Arc<KubeApi>>,
    runtime: Arc<AsyncSafeRuntime>,
    gc_worker: Option<worker::Worker>,
    safepoint_mgr: Option<ServiceSafepointManager>,

    keyspaces: HashMap<u32, Keyspace>,
    cdc_addrs: Arc<dashmap::DashMap<u32, String>>,

    conns: HashMap<ConnId, Conn>,

    tx: Sender<CdcMsg>,
    rx: Receiver<CdcMsg>,

    apply_ctx: ApplyContext,

    region_delegates: HashMap<u64 /* region_id */, RegionDelegate>,
    conn_regions: HashMap<ConnId, HashSet<u64 /* region_id */>>,
    region_to_keyspace: HashMap<u64 /* region_id */, u32 /* keyspace_id */>,

    resolved_regions: HashMap<(TimeStamp, RequestKey), Vec<u64 /* region_id */>>,

    last_update_ts: TimeStamp,

    /// Synced target timestamp of last run, used for WalTargetFetcher
    /// initialization.
    init_target_ts: TimeStamp,
    wal_fetcher: Option<WalFetcherHandle>,
    pending_wals_batch: Option<PendingWalsBatch>,

    incr_scan_concurrency_limit: Arc<Semaphore>,

    shard_load_state: lazy_load::ShardLoadState,

    working_dir: PathBuf,
    stop: bool,
    force_stop: ForceStop,

    resolved_ts_stats: Option<ResolvedTsStats>,
}

struct PendingWalsBatch {
    batch: FetchedWalsBatch,
    appended: bool,
}

impl Drop for ReplicationWorker {
    fn drop(&mut self) {
        // Call shutdown again for safety (and handle the case of force stop).
        self.shutdown();
    }
}

impl ReplicationWorker {
    pub fn new(
        pd: Arc<dyn PdClient>,
        pd_endpoints: Vec<String>,
        fs: Arc<dyn Dfs>,
        data_dir: String,
        security: SecurityConfig,
        config: ReplicationWorkerConfig,
        runtime: Arc<AsyncSafeRuntime>,
        encryption_key_manager: Arc<cloud_encryption::EncryptionKeyManager>,
    ) -> Result<Self> {
        let data_dir = PathBuf::from(data_dir);
        let merged_engine_dir = data_dir.join("merged_engine");
        box_try!(fs::create_dir_all(&merged_engine_dir));
        let worker_dir = data_dir.join("rep_worker");
        box_try!(fs::create_dir_all(&worker_dir));
        cleanup_stale_wal_spool_dirs(&worker_dir);
        let master_key = box_try!(runtime.block_on(security.new_master_key()));
        let force_stop = ForceStop::default();
        let ctx = MergedEngineContext {
            pd,
            fs,
            local_dir: merged_engine_dir,
            security_config: Arc::new(security),
            config: config.merged_engine.clone(),
            master_key,
            force_stop: force_stop.clone(),
            encryption_key_manager,
        };
        let http_client = ctx
            .pd
            .get_security_mgr()
            .http_client(hyper::Client::builder())
            .unwrap();
        let load_shards_mode = if config.lazy_load.enabled {
            LoadShardsMode::Lazy {
                pinned_policy: Arc::new(lazy_load::TidbMetaPinnedShardPolicy::default()),
            }
        } else {
            LoadShardsMode::EagerAll
        };
        let mut merged_engine = box_try!(MergedEngine::new(
            ctx.clone(),
            None,
            MergedEngineInitOpts {
                load_shards_mode: load_shards_mode.clone(),
            },
        ));
        let init_target_ts = merged_engine.get_synced_target_ts();
        let keyspace_ids = merged_engine.get_keyspaces();
        let mut keyspaces = HashMap::with_capacity(keyspace_ids.len());
        let cdc_addrs = Arc::new(dashmap::DashMap::new());
        let kube_api = if config.is_kube_mode() {
            info!("init k8s api");
            let api = runtime.block_on(KubeApi::new(
                config.pd_sts_name.clone(),
                config.cdc_sts_name.clone(),
                config.namespace.clone(),
            ))?;
            Some(Arc::new(api))
        } else {
            None
        };
        for keyspace_id in keyspace_ids {
            let states_bin = merged_engine.get_keyspace_states(keyspace_id).unwrap();
            let states: KeyspaceStates = serde_json::from_slice(&states_bin).unwrap();
            let cdc_addr = states.cdc_addr.clone();
            let mut task_service: Box<dyn KeyspaceService> = if states.is_provisioned() {
                Box::new(KeyspaceProvisionedService::new(
                    keyspace_id,
                    &config,
                    &ctx.security_config,
                    states,
                ))
            } else {
                Box::new(KeyspaceKubeService::new(
                    keyspace_id,
                    kube_api.clone().unwrap(),
                    &config,
                    &ctx.security_config,
                    states,
                ))
            };
            if let Err(err) = runtime.block_on(task_service.start()) {
                error!("keyspace start service error {:?}", err; "keyspace" => keyspace_id);
                continue;
            }
            cdc_addrs.insert(keyspace_id, cdc_addr);
            keyspaces.insert(keyspace_id, Keyspace::from(task_service));
        }
        let gc_runner = GcRunner::new(merged_engine.get_kv(), None, config.local_file_gc_timeout.0);
        let gc_worker = worker::Builder::new("rep-gc-worker").create();
        gc_worker.start_with_timer("rep-gc-worker", gc_runner);
        let interval = config.report_region_interval.0;
        let incr_scan_concurrency_limit =
            Arc::new(Semaphore::new(config.incr_scan_concurrency_limit));
        let mut ks_ctxs = HashMap::with_capacity(keyspaces.len());
        let merged_store_id = ctx.config.merged_store_id;
        for (&keyspace_id, ks) in keyspaces.iter_mut() {
            let region_info_accessor = merged_engine.get_region_info_accessor();
            let rep_pd_cli = ks.get_pd_client();
            let ks_ctx = KeyspaceContext::new(keyspace_id, ks.get_states(), rep_pd_cli.clone());
            ks_ctxs.insert(keyspace_id, ks_ctx);
            ks.report_loop = Some(runtime.spawn(async move {
                Self::report_regions_loop(
                    keyspace_id,
                    merged_store_id,
                    region_info_accessor,
                    rep_pd_cli,
                    interval,
                )
                .await;
            }));
        }
        let safepoint_mgr = box_try!(ServiceSafepointManager::new(
            ctx.config.merged_store_id,
            ctx.pd.clone(),
            &config,
            runtime.handle().clone(),
            ks_ctxs,
        ));
        let mut apply_ctx = ApplyContext::new(
            merged_engine.get_kv(),
            Some(merged_engine.get_router()),
            None,
            None,
        );
        let (tx, rx) = tikv_util::mpsc::unbounded();
        let apply_observer =
            CdcApplyObserver::new(merged_engine.get_kv(), tx.clone(), runtime.handle().clone());
        apply_ctx.set_apply_observer(Box::new(apply_observer));
        merged_engine.register_region_change_observer(
            100,
            BoxRegionChangeObserver::new(lazy_load::DestroyedRegionObserver::new(tx.clone())),
        );

        let shard_load_state = match &load_shards_mode {
            LoadShardsMode::Lazy { pinned_policy } => {
                let enabled_keyspaces = Arc::new(dashmap::DashSet::new());
                for &keyspace_id in keyspaces.keys() {
                    enabled_keyspaces.insert(keyspace_id);
                }
                let shard_load_concurrency =
                    std::cmp::max(1, config.lazy_load.load_shard_concurrency_limit);
                let shard_load_concurrency_limit = Arc::new(Semaphore::new(shard_load_concurrency));
                merged_engine.register_region_change_observer(
                    100,
                    BoxRegionChangeObserver::new(lazy_load::PinnedRegionObserver::new(
                        tx.clone(),
                        pinned_policy.clone(),
                        enabled_keyspaces.clone(),
                    )),
                );
                lazy_load::ShardLoadState::Lazy(lazy_load::LazyLoadState::new(
                    pinned_policy.clone(),
                    enabled_keyspaces,
                    shard_load_concurrency_limit,
                ))
            }
            LoadShardsMode::EagerAll => lazy_load::ShardLoadState::Eager,
        };

        let mut worker = Self {
            config,
            ctx,
            pd_endpoints,
            http_client: Arc::new(http_client),
            merged_engine,
            runtime: runtime.clone(),
            grpc_server: None,
            health_service: None,
            kube_api,
            gc_worker: Some(gc_worker),
            safepoint_mgr: Some(safepoint_mgr),
            keyspaces,
            cdc_addrs,
            conns: Default::default(),
            tx: tx.clone(),
            rx,
            apply_ctx,
            region_delegates: Default::default(),
            conn_regions: Default::default(),
            region_to_keyspace: Default::default(),
            resolved_regions: Default::default(),
            last_update_ts: TimeStamp::zero(),
            init_target_ts,
            wal_fetcher: None,
            pending_wals_batch: None,
            incr_scan_concurrency_limit,
            shard_load_state,
            working_dir: worker_dir,
            stop: false,
            force_stop,
            resolved_ts_stats: Default::default(),
        };
        let env = Arc::new(
            EnvBuilder::new()
                .cq_count(2)
                .name_prefix(thd_name!("grpc-server"))
                .build(),
        );
        let channel_args = ChannelBuilder::new(env.clone())
            .stream_initial_window_size(2 * 1024 * 1024)
            .max_concurrent_stream(1024)
            .max_receive_message_len(-1)
            .max_send_message_len(-1)
            .http2_max_ping_strikes(i32::MAX) // For pings without data from clients.
            .keepalive_time(Duration::from_secs(10))
            .keepalive_timeout(Duration::from_secs(3))
            .build_args();
        let addr = SocketAddr::from_str(&worker.config.grpc_addr).unwrap();
        let security_mgr = worker.ctx.pd.get_security_mgr();
        let service = ReplicationService::new(
            worker.merged_engine.get_kv(),
            worker.merged_engine.get_raft(),
            tx.clone(),
            load_shards_mode,
            runtime.handle().clone(),
            worker.pd_endpoints.clone(),
            worker.ctx.security_config.as_ref().clone(),
            worker.config.enable_resolve_lock,
        )?;
        let health_service = HealthService::default();
        let sb = ServerBuilder::new(env)
            .channel_args(channel_args)
            .register_service(create_change_data(service.clone()))
            .register_service(create_tikv(service))
            .register_service(create_health(health_service.clone()));
        let sb = security_mgr.bind(sb, &addr.ip().to_string(), addr.port());
        let mut grpc_server = sb.build().unwrap();
        grpc_server.start();
        health_service.set_serving_status("", ServingStatus::Serving);
        worker.grpc_server = Some(grpc_server);
        worker.health_service = Some(health_service);
        Ok(worker)
    }

    // Note: `shutdown` will be called more than once.
    fn shutdown(&mut self) {
        // grpc_server should be dropped before merged_engine.
        let _ = self.grpc_server.take();
        if let Some(health_service) = self.health_service.take() {
            health_service.shutdown();
        }
        if let Some(gc_worker) = self.gc_worker.take() {
            gc_worker.stop()
        }
        if let Some(sp_mgr) = self.safepoint_mgr.take() {
            sp_mgr.shutdown();
        }
        self.shutdown_report_region_loops();
        if let Some(wal_fetcher) = self.wal_fetcher.take() {
            self.runtime.block_on(wal_fetcher.shutdown());
        }
        self.merged_engine.close();
    }

    fn shutdown_report_region_loops(&mut self) {
        let report_loop_handles = self
            .keyspaces
            .values_mut()
            .filter_map(|ks| ks.report_loop.take())
            .collect::<Vec<_>>();
        if !report_loop_handles.is_empty() {
            let mut join_set = tokio::task::JoinSet::new();
            for h in report_loop_handles {
                join_set.spawn_on(
                    async move {
                        h.abort();
                        if let Err(err) = h.await {
                            if err.is_panic() {
                                panic!("report region loop panic");
                            }
                        }
                    },
                    self.runtime.handle(),
                );
            }
            self.runtime
                .block_on(async { while (join_set.join_next().await).is_some() {} });
        }
    }

    #[inline]
    fn merged_store_id(&self) -> u64 {
        self.ctx.config.merged_store_id
    }

    #[inline]
    fn get_region_tag(&self, region_id: u64, region_version: u64) -> ShardTag {
        ShardTag::new(
            self.merged_store_id(),
            IdVer::new(region_id, region_version),
        )
    }

    pub fn scheduler(&self) -> ReplicationScheduler {
        ReplicationScheduler::new(
            self.tx.clone(),
            self.cdc_addrs.clone(),
            self.http_client.clone(),
            self.config.safepoint.sync_ticdc_timeout.0,
            self.force_stop.clone(),
        )
    }

    /// Returns downstream kvengine for test-only shard inspection.
    #[cfg(feature = "testexport")]
    pub fn test_kv_engine(&self) -> Engine {
        self.merged_engine.kv.clone()
    }

    /// Returns downstream raft engine for test-only dependency injection.
    #[cfg(feature = "testexport")]
    pub fn test_raft_engine(&self) -> rfengine::RfEngine {
        self.merged_engine.get_raft()
    }

    pub fn run(&mut self) {
        let _enter = self.runtime.enter();
        info!("replication worker started");

        if self.wal_fetcher.is_none() {
            self.wal_fetcher = Some(spawn_wal_fetcher(
                self.ctx.pd.clone(),
                self.ctx.fs.clone(),
                self.http_client.clone(),
                self.working_dir.clone(),
                self.init_target_ts,
                &self.config,
                self.runtime.handle().clone(),
            ));
        }
        self.send_fetch_request(None).expect("send fetch request");

        loop {
            let res = self.rx.recv_timeout(Duration::from_millis(100));
            let has_msg = res.is_ok();
            match res {
                Ok(msg) => {
                    try_force_stop!(self);
                    self.handle_msg(msg);
                    while let Ok(msg) = self.rx.try_recv() {
                        try_force_stop!(self);
                        self.handle_msg(msg);
                    }
                }
                Err(err) => {
                    if err.is_disconnected() {
                        return;
                    }
                }
            }
            if self.stop {
                info!("replication_worker stopped");
                return;
            }
            try_force_stop!(self);
            let should_sync = self.should_sync();
            if has_msg || should_sync {
                // When has message (e.g. `CdcMsg::Applied`), send resolved_ts in time.
                if let Err(err) = self.send_resolved_ts() {
                    error!("send resolved ts error"; "err" => ?err);
                }
            }
            try_force_stop!(self);
            if should_sync {
                match self.maybe_update_merged_engine() {
                    Ok(()) => {}
                    Err(Error::Stopped(cause)) => {
                        info!("replication_worker stopped"; "cause" => cause);
                        return;
                    }
                    Err(err) => {
                        error!("update merged engine error"; "err" => ?err);
                    }
                }
                self.report_metrics();
            }
        }
    }

    fn handle_msg(&mut self, msg: CdcMsg) {
        match msg {
            CdcMsg::AddKeyspace {
                keyspace_id,
                pd_url,
                cdc_addr,
                cb,
            } => {
                self.handle_add_keyspace(keyspace_id, pd_url, cdc_addr, cb);
            }
            CdcMsg::GetKeyspaces { cb } => {
                cb(self.keyspaces.keys().cloned().collect());
            }
            CdcMsg::LoadKeyspaceShards {
                keyspace_id,
                task_service,
                cb,
            } => {
                let res = self.handle_load_keyspace_shards(keyspace_id, task_service);
                cb(res);
            }
            CdcMsg::LoadKeyspaceShardMetas { keyspace_id, cb } => {
                let res = self.handle_load_keyspace_shard_metas(keyspace_id);
                cb(res);
            }
            CdcMsg::LoadRegionShardMeta { region_id, cb } => {
                let res = self.handle_load_region_shard_meta(region_id);
                cb(res);
            }
            CdcMsg::LoadRegionShard {
                region_id,
                attempt_token,
                prepare_res,
                permit,
                compact_guard,
            } => {
                tikv_util::set_current_region(region_id);
                let res = self.handle_load_region_shard(
                    region_id,
                    attempt_token,
                    prepare_res,
                    permit,
                    compact_guard,
                );
                self.handle_result(res, "load_region_shard");
            }
            CdcMsg::RemoveKeyspace { keyspace_id, cb } => {
                self.handle_remove_keyspace_service(keyspace_id, cb);
            }
            CdcMsg::NewTask {
                keyspace_id,
                changefeed_id,
                start_ts,
                body,
                cb,
            } => {
                self.handle_new_task(keyspace_id, changefeed_id, start_ts, body, cb);
            }
            CdcMsg::SyncSafepointKeyspace { keyspace_id } => {
                if let Some(sp_mgr) = self.safepoint_mgr.as_ref() {
                    sp_mgr.sync_keyspace(keyspace_id);
                }
            }
            CdcMsg::OpenConn(conn) => self.handle_open_conn(conn),
            CdcMsg::Register { request, conn_id } => {
                tikv_util::set_current_region(request.region_id);
                let res = self.handle_register(request, conn_id);
                self.handle_result(res, "register");
            }
            CdcMsg::SpawnRegisterHandler {
                request,
                conn_id,
                snap_access,
            } => {
                tikv_util::set_current_region(request.region_id);
                let res = self.handler_spawn_register_handler(request, conn_id, snap_access);
                self.handle_result(res, "spawn_register_handler");
            }
            CdcMsg::ResumeRegister {
                conn_id,
                request_id,
                snap_access,
                init_id,
            } => {
                tikv_util::set_current_region(snap_access.get_id());
                self.handle_resume_register(conn_id, request_id, snap_access, init_id);
            }
            CdcMsg::RegisterResult {
                region_id,
                conn_id,
                request_id,
                init_id,
                err_opt,
            } => {
                tikv_util::set_current_region(region_id);
                let res =
                    self.handle_register_result(region_id, conn_id, request_id, init_id, err_opt);
                if let Err(err) = &res {
                    self.deregister_region_on_error(conn_id, request_id, region_id, err);
                }
                self.handle_result(res.map_err(Into::into), "register_result");
            }
            CdcMsg::SpawnScanLocks { snap_access } => {
                tikv_util::set_current_region(snap_access.get_id());
                let res = self.handle_spawn_scan_locks(snap_access);
                self.handle_result(res, "spawn_scan_locks");
            }
            CdcMsg::ScanLocksResult {
                region_id,
                locks,
                snap_version,
            } => {
                tikv_util::set_current_region(region_id);
                let res = self.handle_scan_locks_result(region_id, locks, snap_version);
                self.handle_result(res, "scan_locks_result");
            }
            CdcMsg::Deregister(deregister) => {
                self.handle_deregister(deregister);
            }
            CdcMsg::Applied {
                region_id,
                region_events,
            } => {
                tikv_util::set_current_region(region_id);
                let sink_err_requests = self.handle_applied(region_id, region_events);
                for (req_key, err) in sink_err_requests {
                    self.deregister_region_on_error(
                        req_key.conn_id,
                        req_key.request_id,
                        region_id,
                        &err,
                    );
                }
            }
            CdcMsg::AppliedAdmin {
                region_id,
                region_version,
                admin,
            } => {
                tikv_util::set_current_region(region_id);
                self.handle_applied_admin(region_id, region_version, admin);
            }
            CdcMsg::RegionDestroyed {
                region_id,
                region_version,
            } => {
                tikv_util::set_current_region(region_id);
                self.handle_region_destroyed(region_id, region_version);
            }
            CdcMsg::EnsurePinnedShardLoaded { region_id, trigger } => {
                tikv_util::set_current_region(region_id);
                let res = self.handle_ensure_pinned_shard_loaded(region_id, trigger);
                self.handle_result(res, "ensure_pinned_shard_loaded");
            }
            CdcMsg::RemoveTask {
                keyspace_id,
                changefeed_id,
                cb,
            } => {
                let res = self.handle_remove_task(keyspace_id, changefeed_id);
                cb(res);
            }
            CdcMsg::UnsafeRecoverClear {
                target,
                dry_run,
                cb,
            } => {
                let res = self.handle_unsafe_recover_clear(target, dry_run);
                cb(res);
            }
            CdcMsg::Stop => {
                self.stop = true;
                self.shutdown();
            }
        }
    }

    fn handle_new_task(
        &mut self,
        keyspace_id: u32,
        changefeed_id: String,
        origin_start_ts: u64,
        origin_body: Bytes,
        cb: Box<dyn FnOnce(Result<(StatusCode, Bytes)>) + Send>,
    ) {
        let tag = format!("{keyspace_id}:new_task");

        let (start_ts, body) = match Self::handle_start_ts(
            &tag,
            origin_start_ts,
            origin_body,
            self.last_update_ts.into_inner(),
        ) {
            Ok(body) => body,
            Err(err) => {
                cb(Err(err));
                return;
            }
        };
        let body_string = String::from_utf8_lossy(&body).to_string();
        info!("handle_new_task"; "keyspace" => keyspace_id, "changefeed" => &changefeed_id, "req" => &body_string);
        #[allow(clippy::map_entry)]
        if !self.keyspaces.contains_key(&keyspace_id) {
            cb(Err(Error::OtherError("keyspace not found".into())));
            return;
        }

        let Some(sp_mgr) = self.safepoint_mgr.as_mut() else {
            cb(Err(Error::OtherError("replication worker stopped".into())));
            return;
        };
        if let Err(err) =
            sp_mgr.ensure_changefeed_start_ts_safety(keyspace_id, &changefeed_id, start_ts)
        {
            cb(Err(err));
            return;
        }

        let merged_store_id = self.merged_store_id();
        let ks = self.keyspaces.get_mut(&keyspace_id).unwrap();
        let pd_client = ks.get_pd_client();
        Self::report_store_to_pd(&pd_client, merged_store_id);
        let keyspace_region_ids = self.merged_engine.get_keyspace_regions(keyspace_id);
        let raft = self.merged_engine.get_raft();
        for region_id in keyspace_region_ids {
            Self::report_region_to_rep_pd_by_id(&raft, &pd_client, region_id);
        }

        // Still dispatch to TiCDC, as the previous request will failed.
        // TODO: Remove the changefeed in replication worker on error ?
        let sec_mgr = self.ctx.pd.get_security_mgr();
        let client = sec_mgr.http_client(hyper::Client::builder()).unwrap();
        let cdc_addr = ks.get_states().cdc_addr.clone();
        let tx = self.tx.clone();
        tokio::spawn(async move {
            fn is_err_retryable(err: &TiCdcError) -> bool {
                matches!(err, TiCdcError::ServerIsNotReady(_))
            }
            let new_cdc_task_uri = sec_mgr
                .build_uri(format!("{}/api/v2/changefeeds", &cdc_addr))
                .unwrap();
            let res = post_to_ticdc(&tag, &client, &new_cdc_task_uri, body, is_err_retryable).await;
            let _ = tx.send(CdcMsg::SyncSafepointKeyspace { keyspace_id });
            cb(res);
        });
    }

    fn handle_start_ts(
        tag: &str,
        origin_start_ts: u64,
        origin_body: Bytes,
        last_update_ts: u64,
    ) -> Result<(u64 /* new_start_ts */, Bytes /* new_body */)> {
        if last_update_ts == 0 {
            return Err(Error::OtherError("replication worker not ready".into()));
        }
        if origin_start_ts > 0 {
            if origin_start_ts <= last_update_ts {
                Ok((origin_start_ts, origin_body))
            } else {
                error!("{}: start_ts too large", tag; "start_ts" => origin_start_ts, "last_update_ts" => last_update_ts);
                Err(Error::OtherError(
                    format!("start_ts too large (> {})", last_update_ts).into(),
                ))
            }
        } else {
            let Ok(Value::Object(mut js_value)) = serde_json::from_slice(&origin_body) else {
                // The format has been verified in scheduler.
                unreachable!();
            };
            debug!("{}: set start_ts as {}", tag, last_update_ts);
            js_value.insert("start_ts".into(), json!(last_update_ts));
            Ok((
                last_update_ts,
                serde_json::to_vec(&json!(js_value)).unwrap().into(),
            ))
        }
    }

    fn handle_open_conn(&mut self, conn: Conn) {
        let conn_id = conn.get_id();
        self.conns.insert(conn_id, conn);
        self.conn_regions.insert(conn_id, HashSet::default());
    }

    // The process of register:
    // 1. Check region existence and epoch.
    // 2. Flush observer for region.
    // 3. Get snap access.
    // 4. Send message to spawn register handler.
    // 5. In message handler, set state to Initializing.
    // 6. Spawn register handler.
    fn handle_register(&mut self, request: ChangeDataRequest, conn_id: ConnId) -> Result<()> {
        let region_id = request.region_id;
        if let Some(shard) = self.merged_engine.get_kv().get_shard(region_id) {
            return self.handle_register_with_shard(request, conn_id, shard);
        }

        match &self.shard_load_state {
            lazy_load::ShardLoadState::Eager => {
                self.send_region_not_found(conn_id, &request);
                Ok(())
            }
            lazy_load::ShardLoadState::Lazy(_) => {
                self.handle_register_lazy_missing_shard(request, conn_id)
            }
        }
    }

    fn handle_register_with_shard(
        &mut self,
        request: ChangeDataRequest,
        conn_id: ConnId,
        shard: Arc<kvengine::Shard>,
    ) -> Result<()> {
        let region_id = request.region_id;
        let request_ver = request.get_region_epoch().get_version();
        if shard.ver != request_ver {
            self.send_epoch_not_match(conn_id, &request);
            return Ok(());
        }
        let tag = shard.tag();
        info!("{} cdc register", tag; "req" => ?request, "conn" => ?conn_id);

        let (need_scan_locks, need_incremental_scan) = if let Some(delegate) =
            self.region_delegates.get(&region_id)
        {
            if delegate.region_ver != request_ver {
                // Delegate is stale. Return error to let client retry.
                warn!("{} cdc register: delegate is stale", tag; "req" => ?request, "conn" => ?conn_id,
                    "delegate.ver" => delegate.region_ver, "request_ver" => request_ver);
                self.send_server_is_busy(conn_id, &request, "delegate is stale".to_string());
                return Ok(());
            }

            let need_scan_locks = delegate.resolver.is_none();
            let check_duplicated = delegate.requests.check_duplicated(&request, conn_id);
            if let Err(err) = &check_duplicated {
                info!("{} cdc register: ignore duplicated request", tag; "err" => ?err);
            }
            (need_scan_locks, check_duplicated.is_ok())
        } else {
            // Create delegate after `flush_observer_region`.
            // Otherwise, the new delegate would receive stale applied entries.
            (true, true)
        };

        if need_scan_locks || need_incremental_scan {
            // Flush observer for region.
            // Otherwise, the scan locks/incremental scan may get data duplicated with
            // pending locks/events.
            self.apply_ctx.flush_observer_region(region_id);
        }

        let snap_access = shard.new_snap_access();
        if need_scan_locks
            && self
                .tx
                .send(CdcMsg::SpawnScanLocks {
                    snap_access: snap_access.clone(),
                })
                .is_err()
        {
            self.send_server_is_busy(
                conn_id,
                &request,
                "handle_register: spawn scan locks failed".to_string(),
            );
            return Err(box_err!("{} spawn scan locks failed", tag));
        }
        if need_incremental_scan {
            if let Err(SendError(msg)) = self.tx.send(CdcMsg::SpawnRegisterHandler {
                request,
                conn_id,
                snap_access,
            }) {
                let CdcMsg::SpawnRegisterHandler {
                    request, conn_id, ..
                } = msg
                else {
                    unreachable!()
                };
                self.send_server_is_busy(
                    conn_id,
                    &request,
                    "handle_register: spawn handler failed".to_string(),
                );
                return Err(box_err!("{} spawn register handler failed", tag));
            }
        }
        Ok(())
    }

    fn handler_spawn_register_handler(
        &mut self,
        request: ChangeDataRequest,
        conn_id: ConnId,
        snap_access: SnapAccess,
    ) -> Result<()> {
        let tag = snap_access.get_tag();
        let merged_store_id = self.merged_store_id();
        let region_id = request.region_id;

        let Some(conn_regions) = self.conn_regions.get_mut(&conn_id) else {
            info!("{} cdc register: conn is closed, skip", tag;
                "conn" => ?conn_id, "request" => %request.request_id);
            return Ok(());
        };
        conn_regions.insert(region_id);
        let Some(conn) = self.conns.get(&conn_id) else {
            warn!("{} cdc register: conn not found", tag; "conn" => ?conn_id);
            debug_assert!(false);
            return Ok(());
        };
        let sink = conn.get_sink().clone();

        let delegate = self.region_delegates.entry(region_id).or_insert_with(|| {
            RegionDelegate::new(merged_store_id, region_id, snap_access.get_version())
        });
        if delegate.region_ver != request.get_region_epoch().get_version() {
            info!("{} cdc register: delegate version not match, skip", tag;
                "conn" => ?conn_id, "request" => %request.request_id, "delegate.ver" => delegate.region_ver);
            self.send_epoch_not_match(conn_id, &request);
            return Ok(());
        }

        // Check if all transactions with `start_ts <= checkpoint_ts` have been
        // seen (committed or prewritten). Prefer merged_engine `safe_target_ts`.
        // When it is unavailable (region is not synced), fall back to resolver
        // `resolved_ts`.
        let safe_target_ts = self
            .merged_engine
            .get_region_progress(region_id)
            .and_then(|progress| progress.get_safe_target_ts());
        let resolved_ts = delegate
            .resolver
            .as_ref()
            .and_then(|resolver| resolver.resolved_ts());
        if let (Some(safe_target_ts), Some(resolved_ts)) = (safe_target_ts, resolved_ts) {
            debug_assert!(resolved_ts <= safe_target_ts);
        }
        let seen_up_to_ts = safe_target_ts.or(resolved_ts);
        let can_initialize = seen_up_to_ts
            .is_some_and(|seen_up_to_ts| request.checkpoint_ts <= seen_up_to_ts.into_inner());

        let (init_id, init_alive_opt) = match delegate.requests.add(
            &request,
            conn_id,
            can_initialize,
        ) {
            Ok(init_id) => init_id,
            Err(current_state) => {
                info!("{} cdc register: ignore duplicate request", tag;
                    "conn" => ?conn_id, "request" => %request.request_id, "current" => current_state);
                return Ok(());
            }
        };

        // `init_alive_opt` is `Some` when state is `Initializing`.
        if let Some(init_alive) = init_alive_opt {
            debug_assert!(can_initialize);
            let mut register_handler = RegisterHandler::new(
                conn_id,
                &request,
                snap_access,
                init_id,
                init_alive,
                self.incr_scan_concurrency_limit.clone(),
                self.tx.clone(),
                sink,
            );
            debug!("{} cdc register: spawn handler", tag; "req" => ?request, "conn" => ?conn_id);
            self.runtime
                .spawn(async move { register_handler.handle_register().await });
        } else {
            info!("{} cdc register: blocked", tag;
                "req" => ?request,
                "conn" => ?conn_id,
                "safe_target_ts" => ?safe_target_ts,
                "resolved_ts" => ?resolved_ts);
        }
        Ok(())
    }

    fn handle_resume_register(
        &mut self,
        conn_id: ConnId,
        request_id: RequestId,
        snap_access: SnapAccess,
        init_id: InitId,
    ) {
        let tag = snap_access.get_tag();
        let region_id = snap_access.get_id();

        let Some(conn) = self.conns.get(&conn_id) else {
            info!("{} handle_resume_result: conn is closed, skip", tag;
                "conn" => ?conn_id, "request" => %request_id);
            return;
        };
        let sink = conn.get_sink().clone();

        let Some(request_info) =
            self.mut_blocked_request_info(region_id, conn_id, request_id, init_id)
        else {
            info!("{} handle_resume_register: ignore stale resume", tag;
                "conn" => ?conn_id, "request" => %request_id, "init_id" => ?init_id);
            return;
        };

        let (request, init_id_, init_alive) = request_info.state.must_resume();
        debug_assert_eq!(init_id, init_id_);
        debug_assert_eq!(request.region_id, region_id);
        // Version has been checked in `send_resolved_ts`.
        debug_assert_eq!(
            request.get_region_epoch().get_version(),
            snap_access.get_version()
        );

        let mut register_handler = RegisterHandler::new(
            conn_id,
            &request,
            snap_access,
            init_id,
            init_alive,
            self.incr_scan_concurrency_limit.clone(),
            self.tx.clone(),
            sink,
        );
        debug!("{} cdc register: resume register spawn handler", tag; "req" => ?request, "conn" => ?conn_id);
        self.runtime
            .spawn(async move { register_handler.handle_register().await });
    }

    fn handle_register_result(
        &mut self,
        region_id: u64,
        conn_id: ConnId,
        request_id: RequestId,
        init_id: InitId,
        err_opt: Option<cdc::Error>,
    ) -> cdc::Result<()> {
        let tag = self.get_region_tag(region_id, 0);
        let Some(delegate) = self.region_delegates.get_mut(&region_id) else {
            info!("{} handle_register_result: ignore stale result", tag; "conn" => ?conn_id, "request" => %request_id);
            return Ok(());
        };
        let Some(request_info) = delegate.mut_initializing_request(conn_id, request_id, init_id)
        else {
            info!("{} handle_register_result: ignore stale result", tag;
                "conn" => ?conn_id, "request" => %request_id, "init_id" => ?init_id);
            return Ok(());
        };

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

        let Some(conn) = self.conns.get(&conn_id) else {
            warn!("{} handle_register_result: conn not found", tag; "conn" => ?conn_id, "request" => %request_id);
            debug_assert!(false);
            // Unexpected conn not found. Return error to deregister for safe.
            return Err(box_err!("conn not found: {:?}", conn_id));
        };
        let sink = conn.get_sink();
        let (pending_events, events_bytes) = request_info.state.must_finish_initialize();
        let events_count = pending_events.len();
        for pending_event in pending_events {
            trace!(
                "{} handle_register_result: send pending event {:?}",
                tag, pending_event
            );
            sink.unbounded_send(CdcEvent::Event(pending_event), false)
                .map_err(|e| cdc::Error::from(e))?;
        }
        info!("{} handle_register_result: initialized", tag;
            "pending_count" => events_count, "pending_bytes" => events_bytes,
            "conn" => ?conn_id, "request" => %request_id);
        Ok(())
    }

    fn handle_spawn_scan_locks(&mut self, snap_access: SnapAccess) -> Result<()> {
        let tag = snap_access.get_tag();
        let merged_store_id = self.merged_store_id();
        let region_id = snap_access.get_id();

        let delegate = self.region_delegates.entry(region_id).or_insert_with(|| {
            RegionDelegate::new(merged_store_id, region_id, snap_access.get_version())
        });
        if delegate.region_ver != snap_access.get_version() {
            // There should be a coming `CdcMsg::AppliedAdmin` which will remove this stale
            // delegate.
            info!("{} handle_spawn_scan_locks: delegate is stale, skip", tag;
                "delegate.ver" => delegate.region_ver);
            return Ok(());
        }

        if delegate.resolver.is_none() {
            debug!("{} cdc register: spawn scan locks", tag);
            delegate.resolver = Some(RegionResolver::new_pending(
                snap_access.get_mem_table_snap_version(),
            ));
            let mut locks_handler = ScanLocksHandler::new(snap_access, self.tx.clone());
            self.runtime.spawn_blocking(move || {
                locks_handler.scan_locks();
            });
        } else {
            info!("{} handle_spawn_scan_locks: duplicated, skip", tag);
        }
        Ok(())
    }

    fn handle_scan_locks_result(
        &mut self,
        region_id: u64,
        locks: Result<Vec<(Vec<u8>, TimeStamp)>>,
        snap_version: SnapVersion,
    ) -> Result<()> {
        let Some(delegate) = self.region_delegates.get_mut(&region_id) else {
            warn!(
                "{} handle_scan_locks_result: region delegate not found",
                region_id
            );
            return Ok(());
        };

        match locks {
            Ok(locks) => {
                delegate.handle_scan_locks(locks, snap_version);
            }
            Err(err) => {
                delegate.handle_scan_locks_error(&err, snap_version, &self.conns);
            }
        }

        Ok(())
    }

    fn send_error_event(&self, conn_id: ConnId, request: &ChangeDataRequest, error: cdcpb::Error) {
        let Some(conn) = self.conns.get(&conn_id) else {
            warn!("send_error_event: conn not found"; "conn" => ?conn_id);
            return;
        };

        let mut event = Event::new();
        event.set_region_id(request.region_id);
        event.set_request_id(request.request_id);
        event.set_error(error);
        if let Err(err) = conn.get_sink().unbounded_send(CdcEvent::Event(event), true) {
            warn!("send_error_event: failed"; "conn" => ?conn_id, "err" => ?err);
        }
    }

    fn send_region_not_found(&self, conn_id: ConnId, request: &ChangeDataRequest) {
        let mut error = cdcpb::Error::new();
        error
            .mut_region_not_found()
            .set_region_id(request.region_id);
        self.send_error_event(conn_id, request, error);
    }

    fn send_epoch_not_match(&self, conn_id: ConnId, request: &ChangeDataRequest) {
        let mut error = cdcpb::Error::new();
        let epoch_not_match = error.mut_epoch_not_match();
        let rep_region =
            Self::get_region_for_rep(&self.merged_engine.get_raft(), request.region_id);
        if let Some(rep_region) = rep_region {
            epoch_not_match.mut_current_regions().push(rep_region);
        }
        self.send_error_event(conn_id, request, error);
    }

    fn send_server_is_busy(&self, conn_id: ConnId, request: &ChangeDataRequest, reason: String) {
        let mut error = cdcpb::Error::new();
        error.set_server_is_busy(new_server_is_busy(
            ServerIsBusyType::ReplicationWorker,
            reason,
        ));
        self.send_error_event(conn_id, request, error);
    }

    fn deregister_region_on_error(
        &mut self,
        conn_id: ConnId,
        request_id: RequestId,
        region_id: u64,
        err: &cdc::Error,
    ) {
        let mut err_event = cdcpb::Error::new();
        match err {
            cdc::Error::Sink(_) => {
                err_event.mut_congested().set_region_id(region_id);
            }
            _ => {
                err_event.set_server_is_busy(new_server_is_busy(
                    ServerIsBusyType::ReplicationWorker,
                    err.to_string(),
                ));
            }
        }
        self.handle_deregister_region(conn_id, request_id, region_id, Some(err_event));
    }

    fn handle_deregister(&mut self, deregister: Deregister) {
        match deregister {
            Deregister::Conn(conn_id) => self.handle_deregister_conn(conn_id),
            Deregister::Request {
                conn_id,
                request_id,
            } => self.handle_deregister_request(conn_id, request_id),
            Deregister::Region {
                conn_id,
                request_id,
                region_id,
            } => {
                tikv_util::set_current_region(region_id);
                self.handle_deregister_region(conn_id, request_id, region_id, None)
            }
        }
    }

    fn handle_deregister_conn(&mut self, conn_id: ConnId) {
        info!("deregister conn"; "conn" => ?conn_id);
        self.remove_pending_registers_for_conn(conn_id);
        self.conns.remove(&conn_id);
        if let Some(conn_regions) = self.conn_regions.remove(&conn_id) {
            for region_id in conn_regions {
                if let Some(delegate) = self.region_delegates.get_mut(&region_id) {
                    delegate.requests.retain(|k, _| k.conn_id != conn_id);
                    if delegate.requests.is_empty() {
                        self.remove_region(region_id);
                    }
                }
            }
        }
    }

    fn handle_deregister_request(&mut self, conn_id: ConnId, request_id: RequestId) {
        info!("deregister request"; "conn" => ?conn_id, "request" => %request_id);
        self.remove_pending_registers_for_request(conn_id, request_id);
        let mut remove_regions = vec![];
        if let Some(conn_regions) = self.conn_regions.get_mut(&conn_id) {
            let sink = self.conns.get(&conn_id).map(|c| c.get_sink());
            for &region_id in conn_regions.iter() {
                if let Some(delegate) = self.region_delegates.get_mut(&region_id) {
                    delegate.unsubscribe(conn_id, request_id, sink, None);
                    if delegate.requests.is_empty() {
                        // To work around mutable borrow limitation.
                        remove_regions.push(region_id);
                    }
                }
            }
        }
        for region_id in remove_regions {
            self.remove_region(region_id);
        }
    }

    fn handle_deregister_region(
        &mut self,
        conn_id: ConnId,
        request_id: RequestId,
        region_id: u64,
        err_event: Option<cdcpb::Error>,
    ) {
        info!("deregister request"; "conn" => ?conn_id, "request" => %request_id, "region" => region_id);
        self.remove_pending_register_for_region(region_id, conn_id, request_id);
        if let Some(delegate) = self.region_delegates.get_mut(&region_id) {
            let sink = self.conns.get(&conn_id).map(|c| c.get_sink());
            delegate.unsubscribe(conn_id, request_id, sink, err_event);
            if delegate.requests.is_empty() {
                self.remove_region(region_id);
            }
        }
    }

    fn handle_result(&mut self, res: Result<()>, tag: &str) {
        if let Err(err) = res {
            error!("handle {} error {:?}", tag, err);
        }
    }

    fn get_keyspace_id(&mut self, region_id: u64) -> Option<u32> {
        Some(match self.region_to_keyspace.entry(region_id) {
            Entry::Occupied(e) => *e.get(),
            Entry::Vacant(e) => {
                let keyspace_id = self
                    .merged_engine
                    .get_region_progress(region_id)?
                    .keyspace_id;
                *e.insert(keyspace_id)
            }
        })
    }

    fn send_resolved_ts(&mut self) -> Result<()> {
        if self.last_update_ts.is_zero() {
            debug!("send_resolved_ts: not initialized yet");
            return Ok(());
        }

        let merged_store_id = self.merged_store_id();
        info!("send_resolved_ts"; "last_update_ts" => self.last_update_ts, "store" => merged_store_id);

        if self.region_delegates.is_empty() {
            self.resolved_ts_stats = Some(ResolvedTsStats {
                min_ts: self.last_update_ts,
                ..Default::default()
            });
            return Ok(());
        }

        self.resolved_regions.clear();
        let mut stats = ResolvedTsStats::default();
        for (&region_id, delegate) in &mut self.region_delegates {
            try_force_stop_err!(self);
            tikv_util::set_current_region(region_id);
            let tag = ShardTag::new(merged_store_id, IdVer::new(region_id, 0));

            let Some(region_progress) = self.merged_engine.mut_region_progress(region_id) else {
                // The region is newly inserted but not start to sync yet.
                warn!(
                    "{} send_resolved_ts: region not in merged_engine, skip",
                    tag
                );
                stats.record_unresolved_region(region_id);
                continue;
            };
            if !region_progress.is_synced() {
                info!("{} send_resolved_ts: region is not synced, skip", tag;
                    "progress" => ?region_progress);
                if let Some(ts) = delegate.resolved_ts() {
                    stats.record_resolved_region(region_id, ts);
                } else {
                    stats.record_unresolved_region(region_id);
                }
                continue;
            }

            let safe_target_ts = region_progress.update_safe_target_ts(tag, self.last_update_ts);
            let Some(ts) = delegate
                .resolver
                .as_mut()
                .and_then(|r| r.resolve(self.last_update_ts.min(safe_target_ts)))
            else {
                stats.record_unresolved_region(region_id);
                continue;
            };
            stats.record_resolved_region(region_id, ts);

            debug!("{} send_resolved_ts: {}", tag, ts);
            let mut resume_reqs = vec![];
            for (req_key, req_info) in delegate.requests.iter_mut() {
                if req_info.resolved_ts != ts {
                    if req_info.state.is_initialized() {
                        debug!("{} send_resolved_ts: advance request resolved_ts: {} -> {}",
                            tag, req_info.resolved_ts, ts; "req" => ?req_key);
                        debug_assert!(req_info.resolved_ts < ts);
                        self.resolved_regions
                            .entry((ts, *req_key))
                            .or_default()
                            .push(region_id);
                        req_info.resolved_ts = ts;
                    } else if let Some(init_id) =
                        req_info.state.can_resume_by_safe_target_ts(safe_target_ts)
                    {
                        resume_reqs.push((*req_key, init_id));
                    }
                }
            }

            // Resume blocked requests.
            if !resume_reqs.is_empty() {
                let kv = self.merged_engine.get_kv();
                // If region not found / version not match, the delegate will be removed later.
                if let Ok(shard) = kv.get_shard_with_ver(region_id, delegate.region_ver) {
                    // Flush pending events in observer. Ref `handle_register`.
                    self.apply_ctx.flush_observer_region(region_id);

                    let snap_access = shard.new_snap_access();
                    for (req_key, init_id) in resume_reqs {
                        let conn_id = req_key.conn_id;
                        let request_id = req_key.request_id;
                        if self
                            .tx
                            .send(CdcMsg::ResumeRegister {
                                conn_id,
                                request_id,
                                snap_access: snap_access.clone(),
                                init_id,
                            })
                            .is_err()
                        {
                            // Just warning. The resume will be triggered again in next
                            // `send_resolved_ts`.
                            warn!("{} send_resolved_ts: send resume_register message failed", snap_access.get_tag();
                                "conn" => ?conn_id, "request" => %request_id);
                        }
                    }
                };
            }
        }

        // TODO: Send small ts with small batch first.
        // Ref: https://github.com/tikv/tikv/blob/release-7.5/components/cdc/src/endpoint.rs, on_min_ts.
        for ((ts, req_key), regions) in self.resolved_regions.drain() {
            try_force_stop_err!(self);
            let Some(conn) = self.conns.get(&req_key.conn_id) else {
                warn!("send_resolved_ts: conn not found"; "conn_id" => ?req_key.conn_id);
                continue;
            };

            let mut resolved_ts = ResolvedTs::default();
            resolved_ts.set_regions(regions);
            resolved_ts.set_request_id(req_key.request_id.into_inner());
            resolved_ts.set_ts(ts.into_inner());
            trace!("send_resolved_ts: msg: {:?}", resolved_ts; "conn" => ?req_key.conn_id);
            if let Err(err) = conn
                .get_sink()
                .unbounded_send(CdcEvent::ResolvedTs(resolved_ts), false)
            {
                // It's OK to drop resolved_ts event.
                warn!("send_resolved_ts: send failed: {:?}", err;
                    "conn" => ?req_key.conn_id, "request" => %req_key.request_id);
            }
        }

        self.resolved_ts_stats = Some(stats);
        Ok(())
    }

    fn maybe_update_merged_engine(&mut self) -> Result<()> {
        if self.wal_fetcher.is_none() {
            return Ok(());
        }

        if self.pending_wals_batch.is_none() {
            let batch = self.wal_fetcher.as_mut().unwrap().try_recv_batch()?;
            if let Some(batch) = batch {
                self.pending_wals_batch = Some(PendingWalsBatch {
                    batch,
                    appended: false,
                });
            } else {
                // No batch is ready yet: return so the main loop can keep handling CDC messages
                // and emitting resolved-ts.
                return Ok(());
            }
        }

        let mut pending = self.pending_wals_batch.take().unwrap();
        if !pending.appended {
            match self.append_fetched_wals_batch(&mut pending.batch) {
                Ok(()) => {
                    pending.appended = true;
                    self.send_fetch_request(pending.batch.advance_to)?;
                }
                Err(err) => {
                    error!("maybe_update_merged_engine: append WALs failed: {:?}", err);
                    // Append failures can leave a partially-appended in-memory state.
                    // Send a fetch request with `appended_target_ts=None` to let WalTargetFetcher
                    // re-build the WALs batch since latest progress on current
                    // target.
                    // And persist it by `sync_merged(advance_to=None)`.
                    pending.appended = true;
                    pending.batch.advance_to = None;
                    self.send_fetch_request(None)?;
                }
            }
        }

        let advance_to = pending.batch.advance_to;
        let sync_res = self
            .merged_engine
            .sync_merged(&mut self.apply_ctx, advance_to);
        match sync_res {
            Ok(()) => {
                self.apply_ctx.flush_observer();
                self.report_live_bootstrapped_regions_to_pd();
                if let Some(ts) = advance_to {
                    self.last_update_ts = ts;
                }
                Ok(())
            }
            Err(err) => {
                // Keep the pending batch for retry. WALs have already been appended, and must
                // not be re-append on retry.
                self.pending_wals_batch = Some(pending);
                Err(err.into())
            }
        }
    }

    // Return error only on stopped.
    fn send_fetch_request(&mut self, appended_target_ts: Option<TimeStamp>) -> Result<()> {
        let Some(wal_fetcher) = self.wal_fetcher.as_ref() else {
            return Ok(());
        };

        let mut store_progresses = Vec::new();
        self.merged_engine
            .collect_store_progresses(&mut store_progresses);
        if let Err(err) = wal_fetcher.request_tx().send(WalFetcherRequest::Fetch {
            store_progresses,
            appended_target_ts,
        }) {
            info!("send_fetch_request: WalFetcher stopped"; "err" => ?err);
            return Err(Error::Stopped("WalFetcher stopped".into()));
        }
        Ok(())
    }

    fn append_fetched_wals_batch(&mut self, batch: &mut FetchedWalsBatch) -> Result<()> {
        let start_time = Instant::now_coarse();

        // Append step 1: update_wal + optional rotate_wal (single-threaded), then
        // delete spools.
        // Append step 2: force advance store progress for tolerated
        // errors.
        //
        // Note: this consumes the batch items to ensure spool files are deleted even
        // when `sync_merged` later fails (the worker must not re-append on retry).
        let items = std::mem::take(&mut batch.items);
        let mut force_advances = Vec::new();
        for item in items {
            match item {
                FetchedWalItem::Wal {
                    store_id,
                    epoch,
                    start_off,
                    end_off,
                    rotate,
                    size,
                    source: _source,
                    data,
                    ..
                } => {
                    let store_start_time = Instant::now_coarse();
                    // New stores may appear in WAL targets before merged_engine has a persisted
                    // progress entry. Inserting a default progress (epoch=1, offset=0) on first
                    // encounter.
                    self.merged_engine.get_or_insert_store_progress(store_id);
                    let apply_res = (|| -> Result<()> {
                        match data {
                            FetchedWalData::Spool(mut spool) => {
                                let file = box_try!(spool.file(true));
                                let reader = std::io::BufReader::with_capacity(1024 * 1024, file);
                                self.merged_engine.update_wal(
                                    store_id,
                                    epoch,
                                    start_off,
                                    end_off,
                                    reader,
                                    &batch.target_stores,
                                )?;
                            }
                            FetchedWalData::SharedSpool(wal_data) => {
                                let reader = {
                                    let mut wal_data = wal_data.lock();
                                    box_try!(wal_data.range_reader(start_off, end_off))
                                };
                                self.merged_engine.update_wal(
                                    store_id,
                                    epoch,
                                    start_off,
                                    end_off,
                                    reader,
                                    &batch.target_stores,
                                )?;
                            }
                        }
                        if rotate {
                            self.merged_engine.rotate_wal(store_id, epoch, end_off)?;
                        }
                        Ok(())
                    })();

                    let store_id_str = store_id.to_string();
                    match &apply_res {
                        Ok(_) => {
                            REP_APPEND_STORE_WAL_BYTES
                                .with_label_values(&[&store_id_str])
                                .inc_by(size);
                            REP_APPEND_STORE_WAL_DURATION
                                .with_label_values(&[&store_id_str])
                                .observe(store_start_time.saturating_elapsed().as_secs_f64());
                        }
                        Err(err) => {
                            error!("maybe_update_merged_engine: update_wal failed";
                                "origin_store" => store_id, "epoch" => epoch,
                                "start_off" => start_off, "end_off" => end_off,
                                "err" => ?err);
                            REP_APPEND_STORE_WAL_COUNTER
                                .with_label_values(&[&store_id_str, error_label(err)])
                                .inc();
                        }
                    }

                    apply_res?;
                }
                FetchedWalItem::ForceAdvance {
                    store_id,
                    epoch,
                    offset,
                } => {
                    force_advances.push((store_id, epoch, offset));
                }
            }
        }
        for (store_id, epoch, offset) in force_advances {
            warn!("maybe_update_merged_engine: force advance";
                "origin_store" => store_id, "epoch" => epoch, "offset" => offset);
            self.merged_engine
                .update_store_progress(store_id, epoch, offset);

            let store_id_str = store_id.to_string();
            REP_APPEND_STORE_WAL_COUNTER
                .with_label_values(&[&store_id_str, "force_advance"])
                .inc();
        }

        REP_APPEND_STORE_WAL_DURATION
            .with_label_values(&["all"])
            .observe(start_time.saturating_elapsed().as_secs_f64());
        Ok(())
    }

    fn should_sync(&self) -> bool {
        TimeStamp::physical_now().saturating_sub(self.last_update_ts.physical())
            >= self.config.sync_interval.as_millis()
    }

    fn collect_report_regions(
        keyspace_id: u32,
        start_enc: &[u8],
        end_enc: &[u8],
        region_info_accessor: &RegionInfoAccessor,
        limit: usize,
    ) -> CollectRegions {
        if limit == 0 {
            return CollectRegions {
                regions: vec![],
                last_end_key: None,
            };
        }

        // `seek_region` requires a `'static` callback, so collect via a shared buffer.
        let range_end_key = end_enc.to_vec();
        let collected = Arc::new(std::sync::Mutex::new(CollectRegions {
            regions: Vec::with_capacity(limit),
            last_end_key: None,
        }));
        let collected_handle = Arc::clone(&collected);
        if let Err(err) = region_info_accessor.seek_region(
            start_enc,
            Box::new(move |regions| {
                const MAX_WARN_LOG: usize = 10;
                let mut collected = collected_handle.lock().unwrap();
                let mut err_count = 0;
                for info in regions {
                    let region = &info.region;
                    if !is_region_initialized(region) {
                        continue;
                    }
                    if region.get_start_key() >= range_end_key.as_slice() {
                        break;
                    }

                    let (start_key, end_key) = match trim_keyspace_prefix_for_encoded_range(
                        &region.start_key,
                        &region.end_key,
                    ) {
                        Ok((got_keyspace_id, start_key, end_key))
                            if got_keyspace_id == keyspace_id =>
                        {
                            (start_key, end_key)
                        }
                        Ok(_) => {
                            // Skip keyspace mismatch.
                            continue;
                        }
                        Err(err) => {
                            if err_count < MAX_WARN_LOG {
                                err_count += 1;
                                warn!(
                                    "report_regions_loop: invalid region keys";
                                    "keyspace" => keyspace_id,
                                    "region" => ?region,
                                    "err" => ?err,
                                );
                            }
                            continue;
                        }
                    };
                    let report_region = metapb::Region {
                        id: region.get_id(),
                        start_key,
                        end_key,
                        region_epoch: region.region_epoch.clone(),
                        peers: region.peers.clone(),
                        ..Default::default()
                    };
                    collected.regions.push(report_region);

                    if collected.regions.len() >= limit {
                        if !region.end_key.is_empty() {
                            collected.last_end_key = Some(region.end_key.clone());
                        }
                        break;
                    }
                }
            }),
        ) {
            warn!("report_regions_loop: seek_region failed"; "keyspace" => keyspace_id, "err" => ?err);
            return CollectRegions {
                regions: vec![],
                last_end_key: None,
            };
        }

        match Arc::try_unwrap(collected) {
            Ok(mutex) => mutex.into_inner().unwrap_or_else(|err| err.into_inner()),
            Err(arc) => arc.lock().unwrap().clone(), // Should not happen.
        }
    }

    async fn report_regions_loop(
        keyspace_id: u32,
        merged_store_id: u64,
        region_info_accessor: RegionInfoAccessor,
        rep_pd_cli: Arc<dyn PdClient>,
        interval: Duration,
    ) {
        const REPORT_REGIONS_BATCH: usize = 1024;

        let handle_heartbeat_resp =
            Self::handle_heartbeat_response(merged_store_id, keyspace_id, rep_pd_cli.clone());
        tikv_util::defer!({
            handle_heartbeat_resp.abort();
        });

        let (start_raw, end_raw) = ApiV2::get_keyspace_range_by_id(keyspace_id);
        let start_enc = codec::bytes::encode_bytes(&start_raw);
        let end_enc = codec::bytes::encode_bytes(&end_raw);

        // run a loop to report regions in case that the rep pd is restarted and lost
        // the region leader.
        loop {
            let start = Instant::now();
            let mut from = start_enc.clone();
            while from.as_slice() < end_enc.as_slice() {
                let batch = Self::collect_report_regions(
                    keyspace_id,
                    &from,
                    &end_enc,
                    &region_info_accessor,
                    REPORT_REGIONS_BATCH,
                );
                for region in batch.regions {
                    if region.get_peers().is_empty() {
                        // Redundant check for safety.
                        // collect_report_regions have filtered out uninitialized regions.
                        debug_assert!(false);
                        continue;
                    }
                    let tag = ShardTag::from_region(None, &region);
                    debug!("{} report_region", tag; "region" => ?region);

                    let leader = region.get_peers()[0].clone();
                    let mut stats = RegionStat::default();
                    stats.approximate_kv_size = 100 * 1024 * 1024;
                    stats.approximate_keys = 1000000;
                    stats.approximate_size = 100 * 1024 * 1024;
                    let res = rep_pd_cli
                        .region_heartbeat(1, region, leader, stats, None, None)
                        .await;
                    if let Err(err) = res {
                        warn!("{} report_region failed: {:?}", tag, err);
                    }
                }

                let Some(next_start_key) = batch.last_end_key else {
                    break;
                };
                if next_start_key.as_slice() <= from.as_slice() {
                    warn!("report_regions_loop: next_start_key did not advance";
                        "keyspace" => keyspace_id,
                        "from" => ?from,
                        "next" => ?next_start_key);
                    break;
                }
                from = next_start_key;
            }

            let elapsed = start.saturating_elapsed();
            tokio::time::sleep(interval.saturating_sub(elapsed)).await;
        }
    }

    fn handle_heartbeat_response(
        merged_store_id: u64,
        keyspace_id: u32,
        rep_pd_cli: Arc<dyn PdClient>,
    ) -> tokio::task::JoinHandle<()> {
        let fut = rep_pd_cli.handle_region_heartbeat_response(
            merged_store_id,
            Box::new(move |resp: pdpb::RegionHeartbeatResponse| {
                if let Err(err) = check_resp_header(resp.get_header()) {
                    warn!("region heartbeat response error: {:?}", err;
                        "store" => merged_store_id, "keyspace" => keyspace_id);
                } else {
                    debug!("region heartbeat response: {:?}", resp;
                        "store" => merged_store_id, "keyspace" => keyspace_id);
                }
            }),
        );
        let f = async move {
            match fut.await {
                Ok(_) => {
                    info!("region heartbeat response handler exit";
                        "store" => merged_store_id, "keyspace" => keyspace_id);
                }
                Err(e) => panic!("unexpected error: {:?}", e),
            }
        };
        tokio::spawn(f)
    }

    fn report_store_to_pd(pd_client: &Arc<dyn PdClient>, store_id: u64) {
        let mut store_stat = StoreStats::new();
        store_stat.set_store_id(store_id);
        store_stat.set_region_count(0);
        store_stat.set_capacity(100 * 1024 * 1024 * 1024);
        store_stat.set_used_size(50 * 1024 * 1024 * 1024);
        let resp = pd_client.store_heartbeat(store_stat, None, None);

        let new_pd_clinet = pd_client.clone();
        tokio::spawn(async move {
            if let Err(err) = resp.await {
                warn!("store heartbeat failed"; "err" => ?err);
                tokio::time::sleep(Duration::from_secs(10)).await;
                Self::report_store_to_pd(&new_pd_clinet, store_id);
            }
        });
    }

    fn report_region_to_rep_pd_by_id(
        raft: &RfEngine,
        rep_pd_cli: &Arc<dyn PdClient>,
        region_id: u64,
    ) {
        let Some(rep_region) = Self::get_region_for_rep(raft, region_id) else {
            return;
        };
        Self::report_region_to_rep_pd(rep_pd_cli, rep_region);
    }

    /// Drain regions that were live-bootstrapped during the last `sync_merged`
    /// call and immediately report them to rep-pd. Without this, newly
    /// installed regions would not be visible to TiCDC until the next
    /// `report_regions_loop` iteration (potentially minutes away), causing
    /// `GetRegion` to return nil for keys in those ranges.
    fn report_live_bootstrapped_regions_to_pd(&mut self) {
        let region_ids = self.merged_engine.take_newly_bootstrapped_regions();
        if region_ids.is_empty() {
            return;
        }
        let raft = self.merged_engine.get_raft();
        for region_id in region_ids {
            let Some(keyspace_id) = self.get_keyspace_id(region_id) else {
                warn!(
                    "report_live_bootstrapped_regions_to_pd: keyspace not found";
                    "region" => region_id
                );
                continue;
            };
            let Some(task_ctx) = self.keyspaces.get(&keyspace_id) else {
                // No active task for this keyspace yet; handle_new_task will report the region.
                continue;
            };
            let pd_client = task_ctx.get_pd_client();
            Self::report_region_to_rep_pd_by_id(&raft, &pd_client, region_id);
        }
    }

    fn report_region_to_rep_pd(rep_pd_cli: &Arc<dyn PdClient>, region: metapb::Region) {
        let tag = ShardTag::from_region(None, &region);
        info!("{} report_region_to_rep_pd: {:?}", tag, region);
        let leader = region.get_peers()[0].clone();
        let mut stats = RegionStat::default();
        stats.approximate_kv_size = 100 * 1024 * 1024;
        stats.approximate_keys = 1000000;
        stats.approximate_size = 100 * 1024 * 1024;
        let resp = rep_pd_cli.region_heartbeat(1, region, leader, stats, None, None);
        tokio::spawn(async move {
            if let Err(err) = resp.await {
                warn!("{} report_region_to_rep_pd failed", tag; "err" => ?err);
            }
        });
    }

    // The keyspace prefix will be trimmed for `rep-pd` & `rep-cdc`.
    fn get_region_for_rep(raft: &RfEngine, region_id: u64) -> Option<metapb::Region> {
        if raft.get_truncated_index(region_id) == Some(TRUNCATE_ALL_INDEX) {
            // region has been merged.
            return None;
        }

        let mut region_local_state = rfstore::store::load_last_peer_state(raft, region_id)?;
        let mut region = region_local_state.take_region();
        let (start_key, end_key) = match trim_keyspace_prefix_for_encoded_range(
            region.get_start_key(),
            region.get_end_key(),
        ) {
            Ok((_keyspace_id, start_key, end_key)) => (start_key, end_key),
            Err(err) => {
                warn!(
                    "get_region_for_rep: trim_keyspace_prefix_range failed";
                    "region" => ?region,
                    "err" => ?err,
                );
                return None;
            }
        };
        region.set_start_key(start_key);
        region.set_end_key(end_key);
        Some(region)
    }

    fn handle_add_keyspace(
        &self,
        keyspace_id: u32,
        pd_url: String,
        cdc_addr: String,
        cb: Box<dyn FnOnce(Result<()>) + Send>,
    ) {
        if self.keyspaces.contains_key(&keyspace_id) {
            // Accept duplicated add_keyspace request.
            info!("handle_add_keyspace: keyspace already exists"; "keyspace" => keyspace_id);
            cb(Ok(()));
            return;
        }
        let mut task_service: Box<dyn KeyspaceService> = if self.config.is_kube_mode() {
            let kube_api = self.kube_api.clone().unwrap();
            Box::new(KeyspaceKubeService::new(
                keyspace_id,
                kube_api,
                &self.config,
                &self.ctx.security_config,
                KeyspaceStates::default(),
            ))
        } else if pd_url.is_empty() || cdc_addr.is_empty() {
            cb(Err(Error::OtherError("pd_url or cdc_addr is empty".into())));
            return;
        } else {
            let mut states = KeyspaceStates::default();
            states.pd_url = pd_url;
            states.cdc_addr = cdc_addr;
            Box::new(KeyspaceProvisionedService::new(
                keyspace_id,
                &self.config,
                &self.ctx.security_config,
                states,
            ))
        };
        let runtime = self.runtime.handle().clone();
        let scheduler = self.scheduler();
        let kv = self.merged_engine.kv.clone();
        let lazy_load_shards = matches!(self.shard_load_state, lazy_load::ShardLoadState::Lazy(_));
        // Pause meta pack compaction during prepare shards. Otherwise, the prepared
        // files may be compacted before load shards into kvengine.
        let meta_pack_compact_keeper = self.merged_engine.get_meta_pack_compact_keeper();
        let http_client = self.http_client.clone();
        tokio::spawn(tikv_util::init_task_local(async move {
            let start_time = Instant::now_coarse();
            // Start service.
            let res = task_service.start().await;
            if res.is_err() {
                cb(res);
                return;
            }
            let cdc_addr = task_service.get_states().cdc_addr.clone();
            let res = get_cdc_status(&http_client, &cdc_addr, DISPATCH_CDC_TIMEOUT).await;
            if let Err(e) = res {
                cb(Err(box_err!("cdc_status failed: {:?}", e)));
                return;
            }

            // Prepare shards.
            // Note that the rfengine will update at the same time. So during load shards,
            // we will prepare again.
            let prepare_time = Instant::now_coarse();
            let compact_guard: Option<tiny_meta::CompactPauseGuard> =
                meta_pack_compact_keeper.map(|x| x.pause());
            {
                let scheduler = scheduler.clone();
                let kv = kv.clone();
                let prepare_res = runtime
                    .spawn_blocking(move || -> Result<()> {
                        let mut metas = HashMap::default();
                        let scope = Self::prepare_keyspace_shard_metas(
                            keyspace_id,
                            &scheduler,
                            &kv,
                            &mut metas,
                        )?;
                        // Prepare again in case some shards are changed during last prepare.
                        let scope2 = Self::prepare_keyspace_shard_metas(
                            keyspace_id,
                            &scheduler,
                            &kv,
                            &mut metas,
                        )?;
                        debug_assert_eq!(scope, scope2);
                        Ok(())
                    })
                    .await;
                match prepare_res {
                    Ok(Ok(())) => {}
                    Ok(Err(e)) => {
                        cb(Err(box_err!("prepare keyspace metas failed: {:?}", e)));
                        return;
                    }
                    Err(e) => {
                        if e.is_panic() {
                            panic!("prepare_keyspace_shard_metas panicked");
                        }
                        cb(Err(box_err!("prepare keyspace metas failed: {:?}", e)));
                        return;
                    }
                }
            }

            // Load shards.
            let load_shards_time = Instant::now_coarse();
            let cb_with_log = move |res: Result<()>| {
                drop(compact_guard);

                let end_time = Instant::now_coarse();
                info!("add_keyspace: {:?}", &res;
                    "keyspace" => keyspace_id,
                    "lazy_load_shards" => lazy_load_shards,
                    "takes" => ?end_time.saturating_duration_since(start_time),
                    "start_svc" => ?prepare_time.saturating_duration_since(start_time),
                    "prepare" => ?load_shards_time.saturating_duration_since(prepare_time),
                    "load_shards" => ?end_time.saturating_duration_since(load_shards_time),
                );
                cb(res);
            };
            scheduler.schedule(CdcMsg::LoadKeyspaceShards {
                keyspace_id,
                task_service,
                cb: Box::new(cb_with_log),
            });
        }));
    }

    fn load_keyspace_shard_metas_via_scheduler(
        keyspace_id: u32,
        scheduler: &ReplicationScheduler,
    ) -> Result<KeyspaceShardMetas> {
        let (cb, fut) = tikv_util::mpsc::paired_callback();
        scheduler.schedule(CdcMsg::LoadKeyspaceShardMetas { keyspace_id, cb });
        fut.recv()
            .map_err(|_| -> Error { box_err!("load keyspace metas canceled") })?
            .map_err(|e| -> Error { box_err!("load keyspace metas failed: {:?}", e) })
    }

    fn load_region_shard_meta_via_scheduler(
        region_id: u64,
        scheduler: &ReplicationScheduler,
    ) -> Result<ShardMeta> {
        let (cb, fut) = tikv_util::mpsc::paired_callback();
        scheduler.schedule(CdcMsg::LoadRegionShardMeta { region_id, cb });
        fut.recv()
            .map_err(|_| -> Error { box_err!("load region {} shard meta canceled", region_id) })?
            .map_err(|e| -> Error {
                box_err!("load region {} shard meta failed: {:?}", region_id, e)
            })
    }

    fn retain_and_diff_shard_metas(
        metas: &mut HashMap<u64 /* region_id */, ShardMeta>,
        latest: StdHashMap<u64 /* region_id */, ShardMeta>,
    ) {
        if latest.is_empty() {
            metas.clear();
            return;
        }

        // Retain only ids that still exist in the latest scoped meta set.
        metas.retain(|region_id, _| latest.contains_key(region_id));

        // Diff: keep only new/changed metas; drop unchanged.
        for (region_id, meta) in latest {
            match metas.entry(region_id) {
                Entry::Vacant(e) => {
                    e.insert(meta);
                }
                Entry::Occupied(mut e) => {
                    let exist_meta = e.get_mut();
                    if exist_meta.seq < meta.seq {
                        *exist_meta = meta;
                    } else {
                        debug_assert_eq!(exist_meta.seq, meta.seq);
                        e.remove();
                    }
                }
            }
        }
    }

    fn prepare_keyspace_shard_metas(
        keyspace_id: u32,
        scheduler: &ReplicationScheduler,
        kv: &Engine,
        metas: &mut HashMap<u64 /* region_id */, ShardMeta>,
    ) -> Result<KeyspaceShardMetaScope> {
        let start_time = Instant::now_coarse();
        let KeyspaceShardMetas {
            scope,
            metas: latest,
        } = Self::load_keyspace_shard_metas_via_scheduler(keyspace_id, scheduler)?;
        let load_time = Instant::now_coarse();

        Self::retain_and_diff_shard_metas(metas, latest);

        let prepare_time = Instant::now_coarse();
        if !metas.is_empty() {
            kv.prepare_shards(metas).map_err(|e| -> Error {
                box_err!("prepare keyspace {} metas failed: {:?}", scope.label(), e)
            })?;
        }

        let end_time = Instant::now_coarse();
        info!("prepare_keyspace_shard_metas";
            "keyspace" => keyspace_id,
            "scope" => scope.label(),
            "count" => metas.len(),
            "load" => ?load_time.saturating_duration_since(start_time),
            "prepare" => ?end_time.saturating_duration_since(prepare_time),
        );
        Ok(scope)
    }

    fn handle_load_keyspace_shards(
        &mut self,
        keyspace_id: u32,
        task_service: Box<dyn KeyspaceService>,
    ) -> Result<()> {
        let mut states = HashMap::default();
        states.insert(keyspace_id, Bytes::new());
        self.merged_engine.load_shards_for_load_mode(&states)?;

        let states = task_service.get_states().marshal();
        let rep_pd_cli = task_service.get_pd_client();
        let merged_store_id = self.ctx.config.merged_store_id;
        let region_info_accessor = self.merged_engine.get_region_info_accessor();
        let interval = self.config.report_region_interval.0;
        let report_loop = self.runtime.spawn(async move {
            Self::report_regions_loop(
                keyspace_id,
                merged_store_id,
                region_info_accessor,
                rep_pd_cli,
                interval,
            )
            .await
        });
        let ks = Keyspace {
            service: task_service,
            report_loop: Some(report_loop),
        };
        if let Some(sp_mgr) = self.safepoint_mgr.as_ref() {
            sp_mgr.add_keyspace(
                keyspace_id,
                ks.get_states().cdc_addr.clone(),
                ks.get_pd_client(),
            );
        }
        self.cdc_addrs
            .insert(keyspace_id, ks.get_states().cdc_addr.clone());
        self.keyspaces.insert(keyspace_id, ks);
        if let lazy_load::ShardLoadState::Lazy(state) = &self.shard_load_state {
            state.enabled_keyspaces.insert(keyspace_id);
        }
        self.merged_engine.set_keyspace_states(keyspace_id, states);
        Ok(())
    }

    fn handle_load_keyspace_shard_metas(&mut self, keyspace_id: u32) -> Result<KeyspaceShardMetas> {
        self.merged_engine
            .load_keyspace_shard_metas_for_load_mode(keyspace_id)
            .map_err(Into::into)
    }

    fn handle_load_region_shard_meta(&mut self, region_id: u64) -> Result<ShardMeta> {
        let raft = self.merged_engine.get_raft();
        let merged_store_id = self.merged_store_id();
        let Some(meta) =
            rfstore::store::peer_storage::load_engine_meta(&raft, merged_store_id, region_id)
        else {
            return Err(Error::MergedEngineError(
                merged_engine::Error::ShardMetaNotReady(region_id),
            ));
        };
        if meta.id != region_id {
            return Err(Error::OtherError(box_err!(
                "load region shard meta: meta id mismatch, expect {}, got {}",
                region_id,
                meta.id
            )));
        }
        Ok(meta)
    }

    fn handle_remove_keyspace_service(
        &mut self,
        keyspace_id: u32,
        cb: Box<dyn FnOnce(Result<()>) + Send>,
    ) {
        info!("remove_keyspace"; "keyspace" => keyspace_id);
        if !self.keyspaces.contains_key(&keyspace_id) {
            info!("remove_keyspace: keyspace not found"; "keyspace" => keyspace_id);
            cb(Ok(()));
            return;
        }

        let ks = self.keyspaces.remove(&keyspace_id).unwrap();
        if let lazy_load::ShardLoadState::Lazy(state) = &self.shard_load_state {
            state.enabled_keyspaces.remove(&keyspace_id);
        }
        if let Some(sp_mgr) = self.safepoint_mgr.as_ref() {
            sp_mgr.remove_keyspace(keyspace_id);
        }
        self.cdc_addrs.remove(&keyspace_id);
        let keyspace_regions = self.merged_engine.get_keyspace_regions(keyspace_id);
        keyspace_regions.iter().for_each(|&region_id| {
            self.remove_region(region_id);
        });
        self.merged_engine.remove_keyspace(keyspace_id);
        self.runtime.spawn(async move {
            let res = ks.destroy().await;
            cb(res);
        });
    }

    fn handle_remove_task(&mut self, keyspace_id: u32, changefeed_id: String) -> Result<()> {
        info!("remove_task"; "keyspace" => keyspace_id, "changefeed" => &changefeed_id);
        if !self.keyspaces.contains_key(&keyspace_id) {
            return Err(Error::OtherError("keyspace service not found".into()));
        };
        if let Some(sp_mgr) = self.safepoint_mgr.as_ref() {
            sp_mgr.sync_keyspace(keyspace_id);
        }
        Ok(())
    }

    fn handle_applied(
        &mut self,
        region_id: u64,
        region_events: RegionEvents,
    ) -> Vec<(RequestKey, cdc::Error)> /* sink_err_requests */ {
        let tag = self.get_region_tag(region_id, 0);
        let mut sink_err_requests: Vec<(RequestKey, cdc::Error)> = vec![];

        let Some(delegate) = self.region_delegates.get_mut(&region_id) else {
            debug!("{} handle_applied: region delegate not found, skip", tag; "events" => ?region_events);
            return sink_err_requests;
        };
        let tag = tag.with_region_version(delegate.region_ver);
        if region_events.region_version != delegate.region_ver {
            warn!("{} handle_applied: version not match", tag; "events.version" => region_events.region_version);
            debug_assert!(
                false,
                "{} version not match, events: {:?}",
                tag, region_events.events
            );
        }

        for (req_key, req_info) in delegate.requests.iter_mut() {
            let Some(conn) = self.conns.get(&req_key.conn_id) else {
                warn!("{} handle_applied: conn not found, skip", tag;
                    "conn" => ?req_key.conn_id, "request" => %req_key.request_id);
                continue;
            };
            'EVENTS_LOOP: for event in &region_events.events {
                // It's OK to not return error. Outer loop will exit before handle next message.
                try_force_stop!(self, vec![]);

                let mut event_to_send = Event::new();
                event_to_send.set_request_id(req_key.request_id.into_inner());
                event_to_send.set_region_id(region_id);
                event_to_send.set_index(event.get_index());
                if event.has_entries() {
                    let entries_to_send = event_to_send.mut_entries().mut_entries();
                    for entry in event.get_entries().get_entries() {
                        if req_info.in_range(entry.get_key()) {
                            entries_to_send.push(entry.clone());

                            trace!("send event";
                                "key" => LogValue::key(entry.get_key()),
                                "commit_ts" => entry.get_commit_ts(),
                                "region" => region_id,
                                "request" => %req_key.request_id,
                                "r_type" => ?entry.r_type,
                                "op_type" => ?entry.op_type);
                        }
                    }
                }
                if req_info.state.is_initialized() {
                    if let Err(err) = conn
                        .get_sink()
                        .unbounded_send(CdcEvent::Event(event_to_send), false)
                    {
                        error!("{} handle_applied: send event failed: {:?}", tag, err;
                            "conn" => ?req_key.conn_id, "request" => %req_key.request_id);
                        // When channel is full, simply send error to TiCDC may lead to cascade
                        // failure. Slow down would be better.
                        // TODO: find a better way to handle sink error.
                        sink_err_requests.push((*req_key, cdc::Error::from(err)));
                        break 'EVENTS_LOOP;
                    }
                } else if req_info.state.is_initializing() {
                    req_info.state.must_push_pending(event_to_send);
                }
            }
        }

        if let Some(resolver) = delegate.resolver.as_mut() {
            for (track_key, start_ts) in region_events.tracked_locks {
                if start_ts == 0 {
                    // TODO: handle error.
                    resolver.untrack_lock(&track_key).unwrap();
                } else {
                    // TODO: handle error.
                    resolver.track_lock(start_ts.into(), track_key).unwrap();
                }
            }
        }

        sink_err_requests
    }

    fn handle_applied_admin(&mut self, region_id: u64, region_version: u64, admin: AdminRequest) {
        let tag = self.get_region_tag(region_id, region_version);
        let raft = self.merged_engine.get_raft();
        let rep_region_opt = Self::get_region_for_rep(&raft, region_id);

        if let Some(rep_region) = &rep_region_opt {
            match self
                .get_keyspace_id(region_id)
                .map(|keyspace_id| (keyspace_id, self.keyspaces.get(&keyspace_id)))
            {
                None => {
                    let err_msg = format!("{} handle_applied_admin: region not found", tag);
                    warn!("{}", &err_msg);
                    debug_assert!(false, "{}", &err_msg);
                    // Still go on and sink error for safety.
                }
                Some((keyspace_id, None)) => {
                    // The keyspace is just removed.
                    warn!("{} handle_applied_admin: keyspace not found", tag; "keyspace" => keyspace_id);
                    debug_assert!(!self.region_delegates.contains_key(&region_id));
                    // Still go on and sink error for safety.
                    // TODO: return Ok(()).
                }
                Some((_, Some(task_ctx))) => {
                    let pd_client = task_ctx.get_pd_client();
                    Self::report_region_to_rep_pd(&pd_client, rep_region.clone());
                    if admin.has_splits() {
                        let split = admin.get_splits();
                        for req in split.get_requests() {
                            Self::report_region_to_rep_pd_by_id(
                                &raft,
                                &pd_client,
                                req.get_new_region_id(),
                            );
                        }
                    }
                }
            }
        } else {
            info!("{} handle_applied_admin: region is merged", tag);
        }

        let mut error = cdcpb::Error::new();
        if let Some(rep_region) = rep_region_opt {
            error
                .mut_epoch_not_match()
                .mut_current_regions()
                .push(rep_region);
        } else {
            error.mut_region_not_found().set_region_id(region_id);
        }

        self.remove_stale_delegate(tag, region_id, error);
    }

    fn handle_region_destroyed(&mut self, region_id: u64, region_version: u64) {
        let tag = self.get_region_tag(region_id, region_version);
        info!("{} handle_region_destroyed", tag);

        self.handle_region_destroyed_in_lazy_state(region_id);

        let mut error = cdcpb::Error::new();
        error.mut_region_not_found().set_region_id(region_id);
        if self.region_delegates.contains_key(&region_id) {
            self.remove_stale_delegate(tag, region_id, error);
        } else {
            self.remove_region(region_id);
        }
    }

    fn remove_stale_delegate(&mut self, tag: ShardTag, region_id: u64, error: cdcpb::Error) {
        if let Some(mut delegate) = self.region_delegates.remove(&region_id) {
            debug!("{} remove_stale_delegate: send error to requests", tag;
                "err" => ?error, "requests" => ?delegate.requests.keys());
            for (req_key, _) in delegate.requests.drain() {
                let Some(conn) = self.conns.get(&req_key.conn_id) else {
                    continue;
                };
                let mut event = Event::new();
                event.set_region_id(region_id);
                event.set_request_id(req_key.request_id.into_inner());
                event.set_error(error.clone());
                if let Err(err) = conn.get_sink().unbounded_send(CdcEvent::Event(event), true) {
                    warn!("{} remove_stale_delegate: send error failed: {:?}", tag, err;
                        "conn" => ?req_key.conn_id, "request" => %req_key.request_id);
                }
            }
            self.remove_region(region_id);
        }
    }

    fn remove_region(&mut self, region_id: u64) {
        self.region_to_keyspace.remove(&region_id);
        self.region_delegates.remove(&region_id);
        for conn_regions in self.conn_regions.values_mut() {
            conn_regions.remove(&region_id);
        }
        if let lazy_load::ShardLoadState::Lazy(state) = &mut self.shard_load_state {
            state.remove_region(region_id);
        }
    }

    fn report_metrics(&self) {
        CDC_ENDPOINT_PENDING_TASKS.set(self.rx.len() as i64);
        CDC_CAPTURED_REGION_COUNT.set(self.region_delegates.len() as i64);

        let last_update_physical = self.last_update_ts.physical();
        if let Some(stats) = &self.resolved_ts_stats {
            CDC_REGION_RESOLVE_STATUS_GAUGE_VEC
                .with_label_values(&["resolved"])
                .set(stats.resolved_regions as i64);
            CDC_REGION_RESOLVE_STATUS_GAUGE_VEC
                .with_label_values(&["unresolved"])
                .set(stats.unresolved_regions as i64);
            if !stats.min_ts.is_zero() {
                CDC_MIN_RESOLVED_TS_REGION.set(stats.min_ts_region_id as i64);
                CDC_MIN_RESOLVED_TS.set(stats.min_ts.physical() as i64);
                let lag = last_update_physical.saturating_sub(stats.min_ts.physical());
                CDC_MIN_RESOLVED_TS_LAG.set(lag as i64);
                let lag_in_secs = lag as f64 / 1000.0;
                CDC_RESOLVED_TS_GAP_HISTOGRAM.observe(lag_in_secs);
            }
        }

        REP_SYNC_WAL_TS.set(last_update_physical as i64);
        let lag = TimeStamp::physical_now().saturating_sub(last_update_physical);
        REP_SYNC_WAL_TS_LAG.set(lag as i64);
        let lag_in_secs = lag as f64 / 1000.0;
        REP_SYNC_WAL_TS_LAG_HISTOGRAM.observe(lag_in_secs);

        REP_LOADED_SHARDS.set(self.merged_engine.get_kv().shard_count() as i64);
        if let lazy_load::ShardLoadState::Lazy(state) = &self.shard_load_state {
            state.report_metrics();
        }
    }
}

// Misc helper functions.
impl ReplicationWorker {
    fn mut_blocked_request_info(
        &mut self,
        region_id: u64,
        conn_id: ConnId,
        request_id: RequestId,
        init_id: InitId,
    ) -> Option<&mut RequestInfo> {
        let delegate = self.region_delegates.get_mut(&region_id)?;
        delegate.mut_blocked_request(conn_id, request_id, init_id)
    }
}

#[derive(Clone)]
struct CollectRegions {
    regions: Vec<metapb::Region>,
    last_end_key: Option<Vec<u8>>,
}

struct RegisterHandler {
    conn_id: ConnId,
    request_id: RequestId,
    snap_access: SnapAccess,
    sender: Sender<CdcMsg>,
    sink: cdc::Sink,
    start_key: Bytes,
    end_key: Bytes,
    checkpoint_ts: u64,
    event_rows: Vec<EventRow>,
    initialized: bool,
    init_id: InitId,
    init_alive: InitAlive,
    concurrency_limit: Arc<Semaphore>,
}

impl RegisterHandler {
    fn new(
        conn_id: ConnId,
        request: &ChangeDataRequest,
        snap_access: SnapAccess,
        init_id: InitId,
        init_alive: InitAlive,
        concurrency_limit: Arc<Semaphore>,
        sender: Sender<CdcMsg>,
        sink: cdc::Sink,
    ) -> Self {
        let keyspace_id = snap_access.get_keyspace_id();
        let request_id = request.get_request_id().into();
        let (start_key, end_key) = build_request_range_for_keyspace(keyspace_id, request);
        let checkpoint_ts = request.get_checkpoint_ts();
        Self {
            conn_id,
            request_id,
            snap_access,
            sender,
            sink,
            start_key: start_key.into(),
            end_key: end_key.into(),
            checkpoint_ts,
            event_rows: vec![],
            initialized: false,
            init_id,
            init_alive,
            concurrency_limit,
        }
    }

    async fn handle_register(&mut self) {
        if let Err(err) = self.handle_register_impl().await {
            let tag = self.snap_access.get_tag();
            CDC_SCAN_TASKS.with_label_values(&["abort"]).inc();
            match err {
                Error::RegisterCancelled(msg) => {
                    info!("{} incremental scan cancelled: {}", tag, msg);
                }
                Error::CdcError(cdc_err) => {
                    error!("{} incremental scan failed: {:?}", tag, cdc_err;
                        "request" => %self.request_id, "conn" => ?self.conn_id);
                    if !self.send_register_result(Some(cdc_err)) {
                        warn!("{} incremental scan: send error result failed", tag);
                    }
                }
                err => {
                    // Should not reach here.
                    error!("{} incremental scan failed: {:?}", tag, err;
                        "request" => %self.request_id, "conn" => ?self.conn_id);
                    debug_assert!(false);
                    let cdc_err = cdc::Error::Other(box_err!(err));
                    if !self.send_register_result(Some(cdc_err)) {
                        warn!("{} incremental scan: send error result failed", tag);
                    }
                }
            }
        }
    }

    async fn handle_register_impl(&mut self) -> Result<()> {
        CDC_SCAN_TASKS.with_label_values(&["total"]).inc();

        let _permit: OwnedSemaphorePermit = self.acquire_permit().await?;
        self.check_alive()?;

        CDC_SCAN_TASKS.with_label_values(&["ongoing"]).inc();
        tikv_util::defer!({
            CDC_SCAN_TASKS.with_label_values(&["ongoing"]).dec();
        });

        let tag = self.snap_access.get_tag();
        let mut entries_bytes = 0;
        let keyspace_id = self.snap_access.get_keyspace_id();
        info!("{} start incremental scan", tag; "checkpoint_ts" => self.checkpoint_ts,
            "request" => %self.request_id, "conn" => ?self.conn_id);
        // scan incremental write after checkpoint ts;
        let start_time = Instant::now_coarse();
        let mut write_iter = self
            .snap_access
            .new_delta_write_iterator_async(self.checkpoint_ts)
            .await;
        write_iter
            .seek_async(InnerKey::from_outer_key(&self.start_key))
            .await;
        let end_key = self.end_key.clone();
        let inner_end_key = InnerKey::from_outer_end_key(&end_key);
        while write_iter.valid() {
            let key = write_iter.key();
            if key >= inner_end_key {
                break;
            }
            let val = write_iter.value();
            if is_index_key(key.deref()) || val.is_deleted() || val.version <= self.checkpoint_ts {
                write_iter.next_async().await;
                continue;
            }
            let um = UserMeta::from_slice(val.user_meta());
            let mut event_row = EventRow::new();
            event_row.set_start_ts(um.start_ts);
            event_row.set_commit_ts(um.commit_ts);
            event_row.set_key(key.deref().to_vec());
            if val.get_value().is_empty() {
                event_row.set_op_type(EventRowOpType::Delete);
            } else {
                event_row.set_op_type(EventRowOpType::Put);
            }
            event_row.set_value(val.get_value().to_vec());
            let mut outer_key = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
            outer_key.extend_from_slice(key.deref());
            let old_value = self
                .snap_access
                .get_async(WRITE_CF, &outer_key, um.commit_ts - 1)
                .await;
            if !old_value.get_value().is_empty() {
                event_row.set_old_value(old_value.get_value().to_vec());
            }
            event_row.set_type(EventLogType::Committed);
            entries_bytes += event_row.get_key().len() + event_row.get_value().len();
            self.event_rows.push(event_row);
            if entries_bytes > MAX_INITIALIZE_SCAN_BATCH_BYTES {
                self.send_rows().await?;
                CDC_SCAN_BYTES.inc_by(entries_bytes as u64);
                entries_bytes = 0;
            }
            write_iter.next_all_version_async().await;
        }
        let mut init_row = EventRow::new();
        init_row.set_type(EventLogType::Initialized);
        self.event_rows.push(init_row);
        self.initialized = true;
        self.send_rows().await?;
        CDC_SCAN_BYTES.inc_by(entries_bytes as u64);

        let takes = start_time.saturating_elapsed();
        CDC_SCAN_DURATION_HISTOGRAM.observe(takes.as_secs_f64());
        CDC_SCAN_TASKS.with_label_values(&["finish"]).inc();
        info!("{} incremental scan finished", tag; "checkpoint_ts" => self.checkpoint_ts,
            "request" => %self.request_id, "conn" => ?self.conn_id, "takes" => ?takes,
            "entries_bytes" => entries_bytes);
        Ok(())
    }

    async fn send_rows(&mut self) -> Result<()> {
        self.check_alive()?;

        let event_rows = mem::take(&mut self.event_rows);
        let mut new_event = Event::new();
        new_event.set_region_id(self.snap_access.get_id());
        new_event.set_request_id(self.request_id.into_inner());
        new_event.mut_entries().set_entries(event_rows.into());

        let tag = self.snap_access.get_tag();
        trace!("{} incremental scan: send event {:?}", tag, new_event);
        let mut events = vec![CdcEvent::Event(new_event)];

        // CDC needs to make sure resolved ts events can only be sent after
        // incremental scan is finished.
        // Wait the barrier to ensure channel sends out all events.
        // Ref: cdc::Initializer::sink_scan_events.
        let barrier = if self.initialized {
            let (cb, fut) = paired_future_callback();
            events.push(CdcEvent::Barrier(Some(cb)));
            Some(fut)
        } else {
            None
        };

        // Use `send_all` for back-pressure.
        let truncated = Arc::new(AtomicBool::new(false));
        if let Err(e) = self.sink.send_all(events, truncated).await {
            return Err(Error::CdcError(cdc::Error::Sink(e)));
        }

        if let Some(barrier) = barrier {
            debug_assert!(self.initialized);
            if barrier.await.is_err() {
                // Should happen only when the sink is closed.
                return Err(Error::RegisterCancelled("barrier cancelled".into()));
            }

            // Send "initialized" event.
            if !self.send_register_result(None) {
                // Should happen only when the sender/receiver is closed.
                return Err(Error::RegisterCancelled(
                    "send register result failed".into(),
                ));
            }
        }

        Ok(())
    }

    fn send_register_result(&self, err_opt: Option<cdc::Error>) -> bool /* ok */ {
        debug!("{} incremental scan: send register result", self.snap_access.get_tag();
            "init_id" => ?self.init_id, "err" => ?err_opt,
            "request" => %self.request_id, "conn" => ?self.conn_id);
        self.sender
            .send(CdcMsg::RegisterResult {
                region_id: self.snap_access.get_id(),
                conn_id: self.conn_id,
                request_id: self.request_id,
                init_id: self.init_id,
                err_opt,
            })
            .is_ok()
    }

    async fn acquire_permit(&self) -> Result<OwnedSemaphorePermit> {
        self.concurrency_limit
            .clone()
            .acquire_owned()
            .await
            .map_err(|_| {
                // Semaphore closed, replication worker should be dropped.
                Error::RegisterCancelled("semaphore closed".into())
            })
    }

    fn check_alive(&self) -> Result<()> {
        if !self.init_alive.ok() {
            Err(Error::RegisterCancelled("init not alive".into()))
        } else {
            Ok(())
        }
    }
}

struct ScanLocksHandler {
    snap_access: SnapAccess,
    sender: Sender<CdcMsg>,
    locks: Vec<(Vec<u8>, TimeStamp)>,
}

impl ScanLocksHandler {
    fn new(snap_access: SnapAccess, sender: Sender<CdcMsg>) -> Self {
        Self {
            snap_access,
            sender,
            locks: vec![],
        }
    }

    fn scan_locks(&mut self) {
        REP_SCAN_LOCKS_TASKS.with_label_values(&["total"]).inc();
        REP_SCAN_LOCKS_TASKS.with_label_values(&["ongoing"]).inc();
        tikv_util::defer!({
            REP_SCAN_LOCKS_TASKS.with_label_values(&["ongoing"]).dec();
        });

        let region_id = self.snap_access.get_id();
        let res = self.scan_locks_impl();
        let label = if res.is_ok() { "finished" } else { "abort" };
        REP_SCAN_LOCKS_TASKS.with_label_values(&[label]).inc();
        let locks = res.map(|_| mem::take(&mut self.locks));
        let snap_version = self.snap_access.get_mem_table_snap_version();
        if let Err(e) = self.sender.send(CdcMsg::ScanLocksResult {
            region_id,
            locks,
            snap_version,
        }) {
            warn!("{} failed to send scan locks result", self.snap_access.get_tag(); "err" => ?e);
        }
    }

    fn scan_locks_impl(&mut self) -> Result<()> {
        let keyspace_id = self.snap_access.get_keyspace_id();
        let tag = self.snap_access.get_tag();
        info!("{} cdc scan locks", tag; "keyspace" => keyspace_id);

        let start_time = Instant::now_coarse();
        let mut bytes = 0;
        let keyspace_prefix_len = keyspace_prefix_len(keyspace_id);
        let mut lock_iter = self
            .snap_access
            .new_iterator(LOCK_CF, false, false, None, false);
        lock_iter.set_range(
            self.snap_access.clone_start_key(),
            self.snap_access.clone_end_key(),
        );
        while lock_iter.valid() {
            let inner_lock_key = &lock_iter.key()[keyspace_prefix_len..];
            if is_index_key(inner_lock_key) {
                lock_iter.next();
                continue;
            }
            let lock_val = lock_iter.val();
            let lock = match parse_lock(lock_val).unwrap() {
                tikv_util::Either::Left(lock) => lock,
                tikv_util::Either::Right(_) => {
                    lock_iter.next();
                    continue;
                }
            };
            if !matches!(lock.lock_type, LockType::Put | LockType::Delete) {
                lock_iter.next();
                continue;
            }
            let key = lock_iter.key()[keyspace_prefix_len..].to_vec();
            bytes += key.len() as u64 + mem::size_of::<TimeStamp>() as u64;
            self.locks.push((key, lock.ts));
            lock_iter.next();
        }

        let takes = start_time.saturating_elapsed();
        info!("{} cdc scan locks", tag; "keyspace" => keyspace_id,
            "locks" => self.locks.len(), "bytes" => bytes, "takes" => ?takes);
        REP_SCAN_LOCKS_BYTES.inc_by(bytes);
        REP_SCAN_LOCKS_DURATION_HISTOGRAM.observe(takes.as_secs_f64());
        Ok(())
    }
}

struct Keyspace {
    service: Box<dyn KeyspaceService>,
    report_loop: Option<tokio::task::JoinHandle<()>>,
}

impl From<Box<dyn KeyspaceService>> for Keyspace {
    fn from(service: Box<dyn KeyspaceService>) -> Self {
        Self {
            service,
            report_loop: None,
        }
    }
}

impl ops::Deref for Keyspace {
    type Target = Box<dyn KeyspaceService>;

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

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

impl Keyspace {
    async fn destroy(self) -> Result<()> {
        let Self {
            mut service,
            report_loop,
        } = self;
        let stop_report_loop = async {
            if let Some(report_loop) = report_loop {
                report_loop.abort();
                if let Err(e) = report_loop.await {
                    if e.is_panic() {
                        warn!("report loop panic: {:?}", e);
                        debug_assert!(false);
                    }
                }
            }
        };
        let (res_svc, _) = tokio::join!(service.destroy(), stop_report_loop);
        res_svc
    }
}

#[cfg(test)]
mod tests {
    use raft::StateRole;
    use raftstore::RegionInfo;

    use super::*;

    #[test]
    fn test_handle_start_ts() {
        {
            let req = json!({
                "changefeed_id": "cf1",
                "sink_uri": "sink",
                "a": 1,
                "b": "b",
            });
            let req_body = serde_json::to_vec(&req).unwrap();
            let (updated_start_ts, updated_body) =
                ReplicationWorker::handle_start_ts("tag", 0, req_body.into(), 1000).unwrap();
            assert_eq!(updated_start_ts, 1000);
            assert_eq!(
                String::from_utf8_lossy(&updated_body),
                r#"{"changefeed_id":"cf1","sink_uri":"sink","a":1,"b":"b","start_ts":1000}"#
            );
        }

        {
            let req = json!({
                "changefeed_id": "cf1",
                "sink_uri": "sink",
                "a": 1,
                "b": "b",
                "start_ts": 500,
            });
            let req_body = serde_json::to_vec(&req).unwrap();
            let err = ReplicationWorker::handle_start_ts("tag", 500, req_body.clone().into(), 499)
                .unwrap_err();
            assert!(err.to_string().contains("start_ts too large"));

            let (updated_start_ts, updated_body) =
                ReplicationWorker::handle_start_ts("tag", 500, req_body.into(), 1000).unwrap();
            assert_eq!(updated_start_ts, 500);
            assert_eq!(
                String::from_utf8_lossy(&updated_body),
                r#"{"changefeed_id":"cf1","sink_uri":"sink","a":1,"b":"b","start_ts":500}"#
            );
        }
    }

    fn encode_keyspace_key(keyspace_id: u32, suffix: &[u8]) -> Vec<u8> {
        let mut raw = ApiV2::get_keyspace_prefix_by_id(keyspace_id);
        raw.extend_from_slice(suffix);
        codec::bytes::encode_bytes(&raw)
    }

    fn decode_encoded_key(encoded: &[u8]) -> Vec<u8> {
        let mut slice = encoded;
        codec::bytes::decode_bytes(&mut slice, false).unwrap()
    }

    fn new_region(region_id: u64, start_key: Vec<u8>, end_key: Vec<u8>) -> metapb::Region {
        let mut region = metapb::Region::default();
        region.set_id(region_id);
        region.set_start_key(start_key);
        region.set_end_key(end_key);
        let mut peer = metapb::Peer::default();
        peer.set_id(region_id);
        peer.set_store_id(1);
        region.set_peers(vec![peer].into());
        region
    }

    #[test]
    fn test_collect_report_regions_trims_keyspace_prefix() {
        let keyspace_id = 10;
        let regions = vec![
            RegionInfo::new(
                new_region(
                    1,
                    encode_keyspace_key(keyspace_id, b"a"),
                    encode_keyspace_key(keyspace_id, b"m"),
                ),
                StateRole::Leader,
                0,
            ),
            RegionInfo::new(
                new_region(
                    2,
                    encode_keyspace_key(keyspace_id, b"m"),
                    encode_keyspace_key(keyspace_id, b"z"),
                ),
                StateRole::Leader,
                0,
            ),
            RegionInfo::new(
                new_region(
                    3,
                    encode_keyspace_key(keyspace_id + 1, b"a"),
                    encode_keyspace_key(keyspace_id + 1, b"b"),
                ),
                StateRole::Leader,
                0,
            ),
        ];
        let accessor = RegionInfoAccessor::new_with_regions_for_test(regions);
        let start_enc = codec::bytes::encode_bytes(&ApiV2::get_keyspace_prefix_by_id(keyspace_id));
        let end_enc =
            codec::bytes::encode_bytes(&ApiV2::get_keyspace_prefix_by_id(keyspace_id + 1));

        let batch = ReplicationWorker::collect_report_regions(
            keyspace_id,
            &start_enc,
            &end_enc,
            &accessor,
            16,
        );
        assert!(batch.last_end_key.is_none());
        let reported = batch.regions;
        assert_eq!(reported.len(), 2);

        let mut by_id = std::collections::HashMap::new();
        for region in reported {
            by_id.insert(region.get_id(), region);
        }
        let region1 = by_id.get(&1).unwrap();
        assert_eq!(decode_encoded_key(region1.get_start_key()), b"a");
        assert_eq!(decode_encoded_key(region1.get_end_key()), b"m");

        let region2 = by_id.get(&2).unwrap();
        assert_eq!(decode_encoded_key(region2.get_start_key()), b"m");
        assert_eq!(decode_encoded_key(region2.get_end_key()), b"z");
    }

    #[test]
    fn test_collect_report_regions_paginates_with_last_end_key() {
        let keyspace_id = 11;
        let regions = vec![
            RegionInfo::new(
                new_region(
                    1,
                    encode_keyspace_key(keyspace_id, b"a"),
                    encode_keyspace_key(keyspace_id, b"m"),
                ),
                StateRole::Leader,
                0,
            ),
            RegionInfo::new(
                new_region(
                    2,
                    encode_keyspace_key(keyspace_id, b"m"),
                    encode_keyspace_key(keyspace_id, b"z"),
                ),
                StateRole::Leader,
                0,
            ),
        ];
        let accessor = RegionInfoAccessor::new_with_regions_for_test(regions);
        let start_enc = codec::bytes::encode_bytes(&ApiV2::get_keyspace_prefix_by_id(keyspace_id));
        let end_enc =
            codec::bytes::encode_bytes(&ApiV2::get_keyspace_prefix_by_id(keyspace_id + 1));

        let batch = ReplicationWorker::collect_report_regions(
            keyspace_id,
            &start_enc,
            &end_enc,
            &accessor,
            1,
        );
        assert_eq!(batch.regions.len(), 1);
        assert_eq!(batch.regions[0].get_id(), 1);
        assert_eq!(
            batch.last_end_key,
            Some(encode_keyspace_key(keyspace_id, b"m"))
        );
    }
}