// Copyright 2025 PingCAP, Inc.

mod apply_observer;
mod delegate;
mod error;
mod kube;
mod metrics;
mod provisioned;
mod safepoint;
mod scheduler;
mod ticdc_util;
mod util;
mod wal;
mod wal_fetcher;
mod wal_spool;
mod worker;

use std::sync::Arc;

use api_version::ApiV2;
use apply_observer::is_index_key;
pub use apply_observer::{CdcApplyObserver, RegionEvents};
use async_trait::async_trait;
use bytes::Bytes;
use cdc::{Conn, ConnId};
use collections::HashMap;
pub use error::{Error, Result};
use futures::{SinkExt, TryFutureExt, TryStreamExt, future};
use grpcio::{DuplexSink, RequestStream, RpcContext, RpcStatus, RpcStatusCode, UnarySink};
use http::StatusCode;
use kvengine::{
    LOCK_CF, Shard, ShardMeta, ShardTag, SnapAccess, WRITE_CF,
    table::{InnerKey, OwnedInnerKey, SnapVersion},
};
use kvproto::{
    cdcpb,
    cdcpb::{ChangeDataEvent, ChangeDataRequest},
    cdcpb_grpc::ChangeData,
    errorpb::EpochNotMatch,
    kvrpcpb::{
        GetRequest, GetResponse, ScanLockRequest, ScanLockResponse, ScanRequest, ScanResponse,
    },
    metapb,
    metapb::{NodeState, RegionEpoch},
    raft_cmdpb::AdminRequest,
    tikvpb_grpc::Tikv,
};
use merged_engine::{KeyspaceShardMetas, LoadShardsMode, MergedEngineConfig, PinnedShardPolicy};
#[cfg(feature = "testexport")]
pub use metrics::REP_SHARD_LOAD_TOTAL;
use pd_client::PdClient;
#[cfg(feature = "testexport")]
pub use provisioned::local_provider::LocalProvider;
use rfengine::RfEngine;
pub use scheduler::*;
use security::SecurityConfig;
use serde_derive::{Deserialize, Serialize};
use tikv::tikv_build_version;
use tikv_client::{
    Timestamp as TxnTimestamp, TimestampExt as _, TransactionClient as TxnClient,
    backoff::OPTIMISTIC_BACKOFF, transaction::ProtoLockInfo as TxnLockInfo,
};
use tikv_util::{
    Either, box_err, box_try,
    config::ReadableDuration,
    debug, error, info,
    memory::MemoryQuota,
    metrics::{ServerIsBusyType, new_server_is_busy},
    warn,
};
use txn_types::{LockType, TimeStamp};
pub use wal::WalConfig;
pub use worker::{LazyLoadConfig, ReplicationWorker};

use crate::{
    delegate::{InitId, RequestId},
    metrics::{REP_LIVE_LOCKS, REP_RESOLVED_LOCKS},
};

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct ReplicationWorkerConfig {
    pub enabled: bool,

    pub grpc_addr: String,
    pub advertise_addr: String,

    // used for k8s mode.
    pub pd_sts_name: String,
    pub cdc_sts_name: String,
    pub namespace: String,

    /// The maximum concurrency for incremental scan.
    ///
    /// For default value 1024, the maximum memory usage of incremental scan is
    /// about:
    ///
    /// 1152(MB) = 128 * 1MB (in channel) + 1024 * 1MB (waiting for channel).
    pub incr_scan_concurrency_limit: usize,

    /// The interval to sync changes from WAL.
    pub sync_interval: ReadableDuration,
    pub report_region_interval: ReadableDuration,
    pub local_file_gc_timeout: ReadableDuration,
    /// Overall timeout for the upstream-PD validation performed by the
    /// unsafe-recover clear admin API.
    ///
    /// The precheck runs on the worker main loop, so this bounds how long the
    /// admin request may stall ordinary CDC progress before failing.
    pub unsafe_recover_clear_pd_precheck_timeout: ReadableDuration,

    /// The address keywords of stores to be skipped during replication.
    pub skip_store_addr_keywords: Vec<String>,
    /// The address keywords of stores to be included during replication.
    /// When non-empty, only stores matching this list are kept, and this list
    /// has higher priority than `skip_store_addr_keywords`.
    pub include_only_store_addr_keywords: Vec<String>,

    /// Whether to resolve locks in `kv_scan_lock`.
    pub enable_resolve_lock: bool,

    pub wal: WalConfig,
    pub safepoint: SafepointConfig,
    pub merged_engine: MergedEngineConfig,
    pub lazy_load: LazyLoadConfig,
}

impl Default for ReplicationWorkerConfig {
    fn default() -> Self {
        Self {
            enabled: false,
            grpc_addr: "".to_string(),
            advertise_addr: "".to_string(),
            pd_sts_name: "".to_string(),
            cdc_sts_name: "".to_string(),
            namespace: "".to_string(),
            incr_scan_concurrency_limit: 1024,
            sync_interval: ReadableDuration::secs(3),
            report_region_interval: ReadableDuration::secs(60),
            local_file_gc_timeout: ReadableDuration::minutes(10),
            unsafe_recover_clear_pd_precheck_timeout: ReadableDuration::minutes(1),
            skip_store_addr_keywords: vec![],
            include_only_store_addr_keywords: vec![],
            enable_resolve_lock: false,
            wal: Default::default(),
            safepoint: Default::default(),
            merged_engine: Default::default(),
            lazy_load: Default::default(),
        }
    }
}

impl ReplicationWorkerConfig {
    pub fn override_from_env(&mut self) {
        Self::env_or_default("PD_STS_NAME", &mut self.pd_sts_name);
        Self::env_or_default("CDC_STS_NAME", &mut self.cdc_sts_name);
    }

    fn env_or_default(name: &str, val: &mut String) {
        if let Ok(v) = std::env::var(name) {
            *val = v;
        }
    }

    pub fn is_kube_mode(&self) -> bool {
        const K8S_SERVICE_HOST: &str = "KUBERNETES_SERVICE_HOST";

        !self.pd_sts_name.is_empty()
            && !self.cdc_sts_name.is_empty()
            && !self.namespace.is_empty()
            && std::env::var(K8S_SERVICE_HOST).is_ok()
    }
}

#[derive(Clone, Serialize, Deserialize, PartialEq, Debug)]
#[serde(default)]
#[serde(rename_all = "kebab-case")]
pub struct SafepointConfig {
    /// `gc_ttl` is the time-to-live when replication worker set service
    /// safepoint. It is defined the same as `gc-ttl` of TiCDC.
    pub gc_ttl: ReadableDuration,
    /// The time-to-live when replication worker set service safepoint during
    /// creating changefeed to ensure that the `start_ts` of the changefeed is
    /// available during changefeed initialization.
    ///
    /// The value should not be too small, or the safepoint will expire before
    /// changefeed initialized.
    ///
    /// And the value should not be too large, as we do not remove the safepoint
    /// for easier.
    pub create_changefeed_gc_ttl: ReadableDuration,
    /// The interval to sync GC safepoint of changefeeds from TiCDC.
    pub sync_safepoint_interval: ReadableDuration,
    pub sync_ticdc_timeout: ReadableDuration,
}

impl Default for SafepointConfig {
    fn default() -> Self {
        Self {
            gc_ttl: ReadableDuration::hours(24),
            create_changefeed_gc_ttl: ReadableDuration::minutes(10),
            sync_safepoint_interval: ReadableDuration::minutes(1),
            sync_ticdc_timeout: ReadableDuration::secs(30),
        }
    }
}

#[derive(Serialize, Deserialize, Clone, Debug, Default)]
pub struct KeyspaceStates {
    /// Compatibility-only field kept so rollback binaries can still deserialize
    /// manifest state written by newer binaries. Runtime changefeed state comes
    /// from TiCDC instead of this local manifest.
    #[serde(default)]
    pub(crate) feeds: HashMap<String, String>,
    pub(crate) pd_url: String,
    pub(crate) cdc_addr: String,

    // used in k8s.
    pub(crate) pd_sts_name: String,
    pub(crate) cdc_sts_name: String,
}

impl KeyspaceStates {
    #[cfg(feature = "testexport")]
    pub fn new_provisioned(pd_url: String, cdc_addr: String) -> Self {
        let mut states = Self::default();
        states.pd_url = pd_url;
        states.cdc_addr = cdc_addr;
        states
    }

    pub(crate) fn marshal(&self) -> Bytes {
        serde_json::to_vec(self).unwrap().into()
    }

    pub(crate) fn is_provisioned(&self) -> bool {
        self.pd_sts_name.is_empty()
    }
}

#[async_trait]
pub trait KeyspaceService: Send {
    fn keyspace_id(&self) -> u32;

    async fn start(&mut self) -> Result<()>;

    async fn destroy(&mut self) -> Result<()>;

    fn get_states(&self) -> &KeyspaceStates;

    fn get_states_mut(&mut self) -> &mut KeyspaceStates;

    fn get_pd_client(&self) -> Arc<dyn PdClient>;
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum PinnedLoadTrigger {
    /// Triggered by merged_engine region create/update events.
    RegionChange,
    /// Triggered by `kv_get` when a pinned shard is missing.
    KvGet,
    /// Triggered by `kv_scan` when a pinned shard is missing.
    KvScan,
}

pub enum CdcMsg {
    AddKeyspace {
        keyspace_id: u32,
        pd_url: String,
        cdc_addr: String,
        cb: Box<dyn FnOnce(Result<()>) + Send>,
    },
    LoadKeyspaceShards {
        keyspace_id: u32,
        task_service: Box<dyn KeyspaceService>,
        cb: Box<dyn FnOnce(Result<()>) + Send>,
    },
    LoadKeyspaceShardMetas {
        keyspace_id: u32,
        cb: Box<dyn FnOnce(Result<KeyspaceShardMetas>) + Send>,
    },
    LoadRegionShardMeta {
        region_id: u64,
        cb: Box<dyn FnOnce(Result<ShardMeta>) + Send>,
    },
    LoadRegionShard {
        region_id: u64,
        attempt_token: u64,
        prepare_res: Result<()>,
        permit: Option<tokio::sync::OwnedSemaphorePermit>,
        compact_guard: Option<kvengine::table::tiny_meta::CompactPauseGuard>,
    },
    NewTask {
        keyspace_id: u32,
        changefeed_id: String,
        start_ts: u64,
        body: Bytes,
        cb: Box<dyn FnOnce(Result<(StatusCode, Bytes)>) + Send>,
    },
    SyncSafepointKeyspace {
        keyspace_id: u32,
    },
    OpenConn(cdc::Conn),
    Register {
        request: ChangeDataRequest,
        conn_id: ConnId,
    },
    SpawnRegisterHandler {
        request: ChangeDataRequest,
        conn_id: ConnId,
        snap_access: SnapAccess,
    },
    ResumeRegister {
        conn_id: ConnId,
        request_id: RequestId,
        snap_access: SnapAccess,
        init_id: InitId,
    },
    RegisterResult {
        region_id: u64,
        conn_id: ConnId,
        request_id: RequestId,
        init_id: InitId,
        err_opt: Option<cdc::Error>,
    },
    SpawnScanLocks {
        snap_access: SnapAccess,
    },
    ScanLocksResult {
        region_id: u64,
        locks: Result<Vec<(Vec<u8>, TimeStamp)>>,
        snap_version: SnapVersion,
    },
    Applied {
        region_id: u64,
        region_events: RegionEvents,
    },
    AppliedAdmin {
        region_id: u64,
        region_version: u64,
        admin: AdminRequest,
    },
    /// Region is logically destroyed by a topology change in merged_engine.
    RegionDestroyed {
        region_id: u64,
        region_version: u64,
    },
    EnsurePinnedShardLoaded {
        region_id: u64,
        trigger: PinnedLoadTrigger,
    },
    Deregister(Deregister),
    RemoveTask {
        keyspace_id: u32,
        changefeed_id: String,
        cb: Box<dyn FnOnce(Result<()>) + Send>,
    },
    RemoveKeyspace {
        keyspace_id: u32,
        cb: Box<dyn FnOnce(Result<()>) + Send>,
    },
    GetKeyspaces {
        cb: Box<dyn FnOnce(Vec<u32>) + Send>,
    },
    UnsafeRecoverClear {
        target: UnsafeRecoverClearTarget,
        dry_run: bool,
        cb: Box<dyn FnOnce(Result<Vec<u64>>) + Send>,
    },
    Stop,
}

/// Explicit target selector for replication-worker unsafe recover clear.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum UnsafeRecoverClearTarget {
    /// Clear exactly one live local replay region.
    Region(u64),
    /// Clear every live local replay region in the keyspace.
    Keyspace(u32),
}

pub enum Deregister {
    Conn(ConnId),
    Request {
        conn_id: ConnId,
        request_id: RequestId,
    },
    Region {
        conn_id: ConnId,
        request_id: RequestId,
        region_id: u64,
    },
}

#[derive(Clone)]
struct ReplicationService {
    kv: kvengine::Engine,
    raft: RfEngine,
    scheduler: tikv_util::mpsc::Sender<CdcMsg>,
    load_shards_mode: LoadShardsMode,
    memory_quota: Arc<MemoryQuota>,
    runtime: tokio::runtime::Handle,
    txn_client: Arc<TxnClient>,
    enable_resolve_lock: bool,
}

enum PinnedShardCheck<'a> {
    InnerKey(InnerKey<'a>),
    InnerRange {
        start_key: InnerKey<'a>,
        end_key: InnerKey<'a>,
    },
}

struct BoundedScanRange {
    outer_start_key: Bytes,
    outer_end_key: Bytes,
}

impl ReplicationService {
    pub fn new(
        kv: kvengine::Engine,
        raft: RfEngine,
        scheduler: tikv_util::mpsc::Sender<CdcMsg>,
        load_shards_mode: LoadShardsMode,
        runtime: tokio::runtime::Handle,
        pd_endpoints: Vec<String>,
        security_config: SecurityConfig,
        enable_resolve_lock: bool,
    ) -> Result<Self> {
        const DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE: usize = 128 * 1024 * 1024; // 128MB

        let mut client_config = tikv_client::Config::default()
            .with_grpc_max_decoding_message_size(DEFAULT_GRPC_MAX_DECODING_MESSAGE_SIZE);
        if !security_config.ca_path.is_empty()
            || !security_config.cert_path.is_empty()
            || !security_config.key_path.is_empty()
        {
            let SecurityConfig {
                ca_path,
                cert_path,
                key_path,
                ..
            } = security_config;
            client_config = client_config.with_security(ca_path, cert_path, key_path);
        }
        let txn_client = box_try!(
            runtime.block_on(TxnClient::new_with_config_api_v2_no_prefix(
                pd_endpoints.clone(),
                client_config.clone(),
            ))
        );
        Ok(Self {
            kv,
            raft,
            scheduler,
            load_shards_mode,
            memory_quota: Arc::new(MemoryQuota::new(1024 * 1024 * 1024)),
            runtime,
            txn_client: Arc::new(txn_client),
            enable_resolve_lock,
        })
    }

    fn prepend_keyspace_prefix(shard: &Shard, req_key: &[u8]) -> Option<Bytes> {
        if req_key.is_empty() {
            None
        } else {
            let mut outer_req_key = ApiV2::get_keyspace_prefix_by_id(shard.keyspace_id);
            outer_req_key.extend_from_slice(req_key);
            Some(outer_req_key.into())
        }
    }

    fn get_region_inner_bounds(&self, region_id: u64) -> Option<(OwnedInnerKey, OwnedInnerKey)> {
        let region_state = rfstore::store::load_last_peer_state(&self.raft, region_id)?;
        if merged_engine::peer_is_skippable(&region_state) {
            return None;
        }
        crate::util::get_inner_bounds_from_region(region_state.get_region())
    }

    fn should_pin_inner_range_for_region(
        &self,
        pinned_policy: &dyn PinnedShardPolicy,
        region_id: u64,
        start_key: InnerKey<'_>,
        end_key: InnerKey<'_>,
    ) -> bool {
        let Some((region_inner_start, region_inner_end)) = self.get_region_inner_bounds(region_id)
        else {
            return false;
        };
        let region_inner_start = region_inner_start.as_ref();
        let region_inner_end = region_inner_end.as_ref();
        let (bounded_start_key, bounded_end_key) = crate::util::clamp_range_to_bounds(
            &start_key,
            &end_key,
            &region_inner_start,
            &region_inner_end,
        );
        pinned_policy.should_pin_inner_range(
            InnerKey::from_inner_buf(bounded_start_key),
            InnerKey::from_inner_buf(bounded_end_key),
        )
    }

    fn build_bounded_scan_range(
        shard: &Shard,
        start_key: &[u8],
        end_key: &[u8],
    ) -> BoundedScanRange {
        let shard_inner_start = shard.inner_start();
        let shard_inner_end = shard.inner_end();
        let (bounded_start_key, bounded_end_key) = crate::util::clamp_range_to_bounds(
            start_key,
            end_key,
            &shard_inner_start,
            &shard_inner_end,
        );
        let outer_start_key = Self::prepend_keyspace_prefix(shard, bounded_start_key)
            .unwrap_or_else(|| shard.outer_start.clone());
        let outer_end_key = Self::prepend_keyspace_prefix(shard, bounded_end_key)
            .unwrap_or_else(|| shard.outer_end.clone());
        BoundedScanRange {
            outer_start_key,
            outer_end_key,
        }
    }

    fn handle_lazy_load_shard_not_found(
        &self,
        region_id: u64,
        trigger: PinnedLoadTrigger,
        pin_check: PinnedShardCheck<'_>,
    ) -> Option<kvproto::errorpb::Error> {
        let should_ensure_pinned_shard = match (&self.load_shards_mode, pin_check) {
            (LoadShardsMode::Lazy { pinned_policy }, PinnedShardCheck::InnerKey(key)) => {
                pinned_policy.should_pin_inner_key(key)
            }
            (
                LoadShardsMode::Lazy { pinned_policy },
                PinnedShardCheck::InnerRange { start_key, end_key },
            ) => self.should_pin_inner_range_for_region(
                pinned_policy.as_ref(),
                region_id,
                start_key,
                end_key,
            ),
            (LoadShardsMode::EagerAll, _) => false,
        };
        if !should_ensure_pinned_shard {
            return None;
        }

        let _ = self
            .scheduler
            .send(CdcMsg::EnsurePinnedShardLoaded { region_id, trigger });

        let mut region_err = kvproto::errorpb::Error::default();
        region_err.set_server_is_busy(new_server_is_busy(
            ServerIsBusyType::ReplicationWorker,
            format!("pinned shard not loaded: region {}", region_id),
        ));
        Some(region_err)
    }

    fn get_shard(
        kv: &kvengine::Engine,
        ctx: &kvproto::kvrpcpb::Context,
    ) -> std::result::Result<Arc<Shard>, kvproto::errorpb::Error> {
        let region_id = ctx.get_region_id();
        let shard = match kv.get_shard(region_id) {
            Some(shard) => shard,
            None => {
                let mut region_err = kvproto::errorpb::Error::default();
                region_err.mut_region_not_found().set_region_id(region_id);
                return Err(region_err);
            }
        };
        if shard.ver != ctx.get_region_epoch().get_version() {
            let mut region_err = kvproto::errorpb::Error::default();
            let mut epoch_not_match = EpochNotMatch::default();

            let mut current_region = metapb::Region::default();
            current_region.set_id(region_id);
            current_region.set_start_key(shard.inner_start().to_vec());
            current_region.set_end_key(shard.inner_end().to_vec());

            let mut peer = metapb::Peer::default();
            peer.set_id(region_id);
            peer.set_store_id(kv.get_engine_id());
            current_region.mut_peers().push(peer);

            let mut region_epoch = RegionEpoch::default();
            region_epoch.set_conf_ver(1);
            region_epoch.set_version(shard.ver);
            current_region.set_region_epoch(region_epoch);

            epoch_not_match.mut_current_regions().push(current_region);
            region_err.set_epoch_not_match(epoch_not_match);
            return Err(region_err);
        }
        Ok(shard)
    }

    fn get_shard_with_lazy_load(
        &self,
        rpc_ctx: &kvproto::kvrpcpb::Context,
        trigger: PinnedLoadTrigger,
        pin_check: PinnedShardCheck<'_>,
    ) -> std::result::Result<Arc<Shard>, kvproto::errorpb::Error> {
        match Self::get_shard(&self.kv, rpc_ctx) {
            Ok(shard) => Ok(shard),
            Err(region_err) if region_err.has_region_not_found() => {
                if let Some(err) = self.handle_lazy_load_shard_not_found(
                    rpc_ctx.get_region_id(),
                    trigger,
                    pin_check,
                ) {
                    return Err(err);
                }
                Err(region_err)
            }
            Err(region_err) => Err(region_err),
        }
    }

    async fn kv_scan_lock_impl(
        runtime: tokio::runtime::Handle,
        kv: kvengine::Engine,
        txn_client: Arc<TxnClient>,
        req: ScanLockRequest,
    ) -> Result<ScanLockResponse> {
        const SCAN_LOCKS_MAX_LIMIT: usize = 65536;

        let region_id = req.get_context().get_region_id();
        tikv_util::set_current_region(region_id);

        let mut resp = ScanLockResponse::default();
        let shard = match Self::get_shard(&kv, req.get_context()) {
            Ok(shard) => shard,
            Err(region_err) => {
                resp.set_region_error(region_err);
                return Ok(resp);
            }
        };
        let tag = shard.tag();
        let keyspace_id = shard.keyspace_id;
        let scan_range =
            Self::build_bounded_scan_range(&shard, req.get_start_key(), req.get_end_key());
        let max_version = req.max_version.into();
        let limit = (req.limit as usize).min(SCAN_LOCKS_MAX_LIMIT);

        let snap_access = shard.new_snap_access();
        let keyspace_prefix_len = ApiV2::get_keyspace_prefix_by_id(keyspace_id).len();

        let locks = match runtime
            .spawn_blocking(move || {
                tikv_util::init_task_local_sync(|| {
                    tikv_util::set_current_region(region_id);
                    Self::scan_locks(
                        snap_access,
                        keyspace_prefix_len,
                        scan_range.outer_start_key,
                        scan_range.outer_end_key,
                        max_version,
                        limit,
                    )
                })
            })
            .await
        {
            Ok(Ok(locks)) => locks,
            Ok(Err(err)) => return Err(box_err!(err)),
            Err(err) if err.is_panic() => std::panic::resume_unwind(err.into_panic()),
            Err(err) => return Err(box_err!(err)),
        };
        if locks.is_empty() {
            return Ok(ScanLockResponse::default());
        }
        let total_locks_count = locks.len();
        debug!("{} kv_scan_lock: resolving locks", tag;
            "count" => total_locks_count, "locks" => ?&locks[..locks.len().min(5)]);

        // Use a zero start_ts to avoid push `min_commit_ts`.
        let caller_start_ts = TxnTimestamp::from_version(0);
        let client_locks = locks.into_iter().map(convert_lock_info).collect::<Vec<_>>();
        let live_locks = box_try!(
            txn_client
                .resolve_locks(client_locks, caller_start_ts, OPTIMISTIC_BACKOFF)
                .await
        );

        let resolved_locks_count = total_locks_count.saturating_sub(live_locks.len());
        REP_RESOLVED_LOCKS.inc_by(resolved_locks_count as u64);

        if !live_locks.is_empty() {
            warn!("{} kv_scan_lock: still has live locks", tag;
                "live_locks" => live_locks.len(), "resolved_locks" => resolved_locks_count);
            debug!(
                "{} kv_scan_lock: live locks {:?}",
                tag,
                &live_locks[..live_locks.len().min(5)]
            );
            REP_LIVE_LOCKS.inc_by(live_locks.len() as u64);
        } else {
            info!("{} kv_scan_lock: locks resolved", tag; "resolved_locks" => resolved_locks_count);
        }
        Ok(ScanLockResponse::default())
    }

    fn scan_locks(
        snap_access: SnapAccess,
        keyspace_prefix_len: usize,
        outer_start_key: Bytes,
        outer_end_key: Bytes,
        max_version: TimeStamp,
        limit: usize,
    ) -> Result<Vec<kvproto::kvrpcpb::LockInfo>> {
        let mut lock_iter = snap_access.new_iterator(LOCK_CF, false, false, None, false);
        lock_iter.set_range(outer_start_key, outer_end_key);
        let mut locks = Vec::new();
        while lock_iter.valid() && locks.len() < limit {
            let key = lock_iter.key().to_vec();
            let inner_lock_key = &key[keyspace_prefix_len..];
            if is_index_key(inner_lock_key) {
                lock_iter.next();
                continue;
            }
            let lock_or_shared = box_try!(txn_types::parse_lock(lock_iter.val()).map_err(|e| {
                debug_assert!(false, "failed to parse lock: {:?}", e);
                e
            }));
            if let Either::Left(lock) = lock_or_shared {
                // Only Put/Delete locks need to be resolved.
                if lock.ts <= max_version
                    && matches!(lock.lock_type, LockType::Put | LockType::Delete)
                {
                    locks.push(lock.into_lock_info(key));
                }
            }
            lock_iter.next();
        }
        Ok(locks)
    }
}

impl ReplicationService {
    fn handle_event_feed(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<ChangeDataRequest>,
        mut sink: DuplexSink<ChangeDataEvent>,
        _event_feed_v2: bool,
    ) {
        const CDC_CHANNEL_CAPACITY: usize = 128;

        // TODO: parse header for event_feed_v2.

        let conn_id = ConnId::new();
        let (event_sink, mut event_drain) =
            cdc::channel(conn_id, CDC_CHANNEL_CAPACITY, self.memory_quota.clone());
        let peer = ctx.peer();
        let conn = Conn::new(conn_id, event_sink, peer);

        if let Err(status) = self.scheduler.send(CdcMsg::OpenConn(conn)).map_err(|e| {
            RpcStatus::with_message(RpcStatusCode::INVALID_ARGUMENT, format!("{:?}", e))
        }) {
            error!("cdc connection initiate failed"; "error" => ?status);
            ctx.spawn(
                sink.fail(status)
                    .unwrap_or_else(|e| error!("cdc failed to send error"; "error" => ?e)),
            );
            return;
        }

        let scheduler = self.scheduler.clone();
        let recv_req = stream.try_for_each(move |request| {
            info!("got event feed request {:?}", request);
            let ret = Self::handle_request(&scheduler, request, conn_id)
                .map_err(|rpc_status| grpcio::Error::RpcFailure(rpc_status));
            future::ready(ret)
        });

        let peer = ctx.peer();
        let scheduler = self.scheduler.clone();
        ctx.spawn(async move {
            let res = recv_req.await;
            // Unregister this downstream only.
            if let Err(e) = scheduler.send(CdcMsg::Deregister(Deregister::Conn(conn_id))) {
                error!("cdc deregister failed"; "error" => ?e, "conn" => ?conn_id);
            }
            match res {
                Ok(()) => {
                    info!("cdc receive closed"; "downstream" => peer, "conn" => ?conn_id);
                }
                Err(e) => {
                    warn!("cdc receive failed"; "error" => ?e, "downstream" => peer, "conn" => ?conn_id);
                }
            }
        });

        let peer = ctx.peer();
        let scheduler = self.scheduler.clone();

        ctx.spawn(async move {
            let res = event_drain.forward(&mut sink, None).await;
            // Unregister this downstream only.
            if let Err(e) = scheduler.send(CdcMsg::Deregister(Deregister::Conn(conn_id))) {
                error!("cdc deregister failed"; "error" => ?e, "conn" => ?conn_id);
            }
            match res {
                Ok(_s) => {
                    info!("cdc send closed"; "downstream" => peer, "conn" => ?conn_id);
                    let _ = sink.close().await;
                }
                Err(e) => {
                    warn!("cdc send failed"; "error" => ?e, "downstream" => peer, "conn" => ?conn_id);
                }
            }
        });
        info!("cdc event feed started"; "conn" => ?conn_id);
    }

    fn handle_request(
        scheduler: &tikv_util::mpsc::Sender<CdcMsg>,
        request: ChangeDataRequest,
        conn_id: ConnId,
    ) -> std::result::Result<(), RpcStatus> {
        match request.request {
            None | Some(cdcpb::ChangeDataRequest_oneof_request::Register(_)) => scheduler
                .send(CdcMsg::Register { request, conn_id })
                .map_err(|e| {
                    RpcStatus::with_message(
                        RpcStatusCode::RESOURCE_EXHAUSTED,
                        format!("replication worker is busy: {:?}", e),
                    )
                }),
            Some(cdcpb::ChangeDataRequest_oneof_request::Deregister(_)) => {
                let deregister = if request.region_id == 0 {
                    Deregister::Request {
                        conn_id,
                        request_id: request.request_id.into(),
                    }
                } else {
                    Deregister::Region {
                        conn_id,
                        request_id: request.request_id.into(),
                        region_id: request.region_id,
                    }
                };
                scheduler.send(CdcMsg::Deregister(deregister)).map_err(|e| {
                    RpcStatus::with_message(
                        RpcStatusCode::RESOURCE_EXHAUSTED,
                        format!("replication worker is busy: {:?}", e),
                    )
                })
            }
            _ => Err(RpcStatus::with_message(
                RpcStatusCode::INVALID_ARGUMENT,
                format!("request not supported: {:?}", request),
            )),
        }
    }
}

impl ChangeData for ReplicationService {
    fn event_feed(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<ChangeDataRequest>,
        sink: DuplexSink<ChangeDataEvent>,
    ) {
        self.handle_event_feed(ctx, stream, sink, false);
    }

    fn event_feed_v2(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<ChangeDataRequest>,
        sink: DuplexSink<ChangeDataEvent>,
    ) {
        self.handle_event_feed(ctx, stream, sink, true);
    }
}

impl Tikv for ReplicationService {
    fn kv_scan(&mut self, ctx: RpcContext<'_>, req: ScanRequest, sink: UnarySink<ScanResponse>) {
        if req.reverse {
            let status = RpcStatus::new(RpcStatusCode::INVALID_ARGUMENT);
            ctx.spawn(
                sink.fail(status.clone())
                    .unwrap_or_else(|e| error!("kv_scan failed"; "error" => ?e)),
            );
            return;
        }
        let mut resp = ScanResponse::default();
        let rpc_ctx = req.get_context();
        let region_id = rpc_ctx.get_region_id();
        let shard = match self.get_shard_with_lazy_load(
            rpc_ctx,
            PinnedLoadTrigger::KvScan,
            PinnedShardCheck::InnerRange {
                start_key: InnerKey::from_inner_buf(req.get_start_key()),
                end_key: InnerKey::from_inner_buf(req.get_end_key()),
            },
        ) {
            Ok(shard) => shard,
            Err(region_err) => {
                resp.set_region_error(region_err);
                ctx.spawn(
                    sink.success(resp)
                        .unwrap_or_else(|e| error!("kv_scan failed"; "error" => ?e)),
                );
                return;
            }
        };
        let scan_range =
            Self::build_bounded_scan_range(&shard, req.get_start_key(), req.get_end_key());
        let read_ts = Some(req.get_version());
        let snap_access = shard.new_snap_access();
        let task = async move {
            tikv_util::set_current_region(region_id);
            // TODO: handle locks.
            let mut iter = snap_access
                .new_iterator_async(WRITE_CF, false, false, read_ts, false)
                .await;
            iter.set_range_async(scan_range.outer_start_key, scan_range.outer_end_key)
                .await;
            let mut kv_pairs = vec![];
            let keyspace_prefix_len = ApiV2::get_keyspace_prefix_by_id(shard.keyspace_id).len();
            while iter.valid() {
                if kv_pairs.len() == req.get_limit() as usize {
                    break;
                }
                let mut kv_pair = kvproto::kvrpcpb::KvPair::default();
                let key = iter.key();
                // trim keyspace prefix.
                kv_pair.set_key(key[keyspace_prefix_len..].to_vec());
                let val = iter.val();
                if val.is_empty() {
                    iter.next_async().await;
                    continue;
                }
                kv_pair.set_value(val.to_vec());
                kv_pairs.push(kv_pair);
                iter.next_async().await;
            }
            resp.set_pairs(kv_pairs.into());
            resp
        };
        let resp = self.runtime.spawn(tikv_util::init_task_local(task));
        let fut = async move {
            match resp.await {
                Ok(resp) => sink.success(resp).await,
                Err(err) if err.is_panic() => {
                    panic!("kv_scan panic");
                }
                Err(err) => {
                    sink.fail(RpcStatus::with_message(
                        RpcStatusCode::CANCELLED,
                        format!("{}", err),
                    ))
                    .await
                }
            }
            .unwrap_or_else(|e| {
                error!("kv_scan rpc failed"; "error" => ?e);
            });
        };
        ctx.spawn(fut);
    }

    fn kv_get(&mut self, ctx: RpcContext<'_>, req: GetRequest, sink: UnarySink<GetResponse>) {
        let mut resp = GetResponse::default();
        let rpc_ctx = req.get_context();
        let region_id = rpc_ctx.get_region_id();
        let shard = match self.get_shard_with_lazy_load(
            rpc_ctx,
            PinnedLoadTrigger::KvGet,
            PinnedShardCheck::InnerKey(InnerKey::from_inner_buf(req.get_key())),
        ) {
            Ok(shard) => shard,
            Err(region_err) => {
                resp.set_region_error(region_err);
                ctx.spawn(
                    sink.success(resp)
                        .unwrap_or_else(|e| error!("kv_get failed"; "error" => ?e)),
                );
                return;
            }
        };
        let snap_access = shard.new_snap_access();
        let key = Self::prepend_keyspace_prefix(&shard, req.get_key()).unwrap();
        let task = async move {
            tikv_util::set_current_region(region_id);
            // TODO: handle locks.
            let item = snap_access
                .get_async(WRITE_CF, &key, req.get_version())
                .await;
            if !item.get_value().is_empty() {
                resp.set_value(item.get_value().to_vec());
            } else {
                resp.set_not_found(true);
            }
            resp
        };
        let resp = self.runtime.spawn(tikv_util::init_task_local(task));
        let fut = async move {
            match resp.await {
                Ok(resp) => sink.success(resp).await,
                Err(err) if err.is_panic() => {
                    panic!("kv_get panic");
                }
                Err(err) => {
                    sink.fail(RpcStatus::with_message(
                        RpcStatusCode::CANCELLED,
                        format!("{}", err),
                    ))
                    .await
                }
            }
            .unwrap_or_else(|e| {
                error!("kv_get rpc failed"; "error" => ?e);
            });
        };
        ctx.spawn(fut);
    }

    fn kv_scan_lock(
        &mut self,
        ctx: RpcContext<'_>,
        req: ScanLockRequest,
        sink: UnarySink<ScanLockResponse>,
    ) {
        let tag = ShardTag::from_pb_ctx(req.get_context());
        info!("{} kv_scan_lock", tag; "req" => ?req);
        if !self.enable_resolve_lock {
            let fut = async move {
                sink.success(ScanLockResponse::default())
                    .await
                    .unwrap_or_else(|e| {
                        error!("{} kv_scan_lock failed", tag; "error" => ?e);
                    });
            };
            ctx.spawn(fut);
            return;
        }
        let task = Self::kv_scan_lock_impl(
            self.runtime.clone(),
            self.kv.clone(),
            self.txn_client.clone(),
            req,
        );
        let resp = self.runtime.spawn(tikv_util::init_task_local(task));
        let fut = async move {
            match resp.await {
                Ok(Ok(resp)) => sink.success(resp).await,
                Ok(Err(err)) => {
                    error!("{} kv_scan_lock failed", tag; "error" => ?err);
                    let mut resp = ScanLockResponse::default();
                    let mut region_err = kvproto::errorpb::Error::default();
                    region_err.set_message(format!("{}", err));
                    resp.set_region_error(region_err);
                    sink.success(resp).await
                }
                Err(err) if err.is_panic() => {
                    panic!("kv_scan_lock panic");
                }
                Err(err) => {
                    sink.fail(RpcStatus::with_message(
                        RpcStatusCode::CANCELLED,
                        format!("{}", err),
                    ))
                    .await
                }
            }
            .unwrap_or_else(|e| {
                error!("{} kv_scan_lock failed", tag; "error" => ?e);
            });
        };
        ctx.spawn(fut);
    }
}

pub async fn bootstrap(
    pd_client: Arc<dyn PdClient>,
    store_id: u64,
    advertise_addr: String,
) -> Result<()> {
    let bootstrapped = pd_client.is_cluster_bootstrapped()?;
    if !bootstrapped {
        let mut store = metapb::Store::new();
        store.set_id(store_id);
        store.set_node_state(NodeState::Serving);
        store.set_address(advertise_addr);
        store.set_version(tikv_build_version().to_string());
        let mut initial_region = metapb::Region::new();
        initial_region.set_id(1);
        let mut initial_peer = metapb::Peer::new();
        initial_peer.set_id(1);
        initial_peer.set_store_id(store_id);
        initial_region.set_peers(vec![initial_peer].into());
        let mut initial_epoch = RegionEpoch::new();
        initial_epoch.set_version(1);
        initial_epoch.set_conf_ver(1);
        initial_region.set_region_epoch(initial_epoch);
        pd_client.bootstrap_cluster(store, initial_region)?;
    }
    Ok(())
}

fn convert_lock_info(mut lock: kvproto::kvrpcpb::LockInfo) -> TxnLockInfo {
    TxnLockInfo {
        primary_lock: lock.take_primary_lock(),
        lock_version: lock.get_lock_version(),
        key: lock.take_key(),
        lock_ttl: lock.get_lock_ttl(),
        txn_size: lock.get_txn_size(),
        lock_type: lock.get_lock_type() as i32,
        lock_for_update_ts: lock.get_lock_for_update_ts(),
        use_async_commit: lock.get_use_async_commit(),
        min_commit_ts: lock.get_min_commit_ts(),
        secondaries: lock.take_secondaries().into_iter().collect(),
        duration_to_last_update_ms: lock.get_duration_to_last_update_ms(),
        shared_lock_infos: lock
            .take_shared_lock_infos()
            .into_iter()
            .map(convert_lock_info)
            .collect(),
        is_txn_file: lock.get_is_txn_file(),
    }
}

#[cfg(test)]
mod tests {
    use api_version::ApiV2;
    use kvengine::{Engine, WriteBatch, test_engine};
    use kvproto::kvrpcpb::Op;
    use txn_types::{Lock, SharedLocks};

    use super::*;

    const SHARD_ID: u64 = 1;
    const KEYSPACE_ID: u32 = test_engine::KEYSPACE_ID;
    const TABLE_ID: i64 = 42;

    #[test]
    fn test_convert_lock_info_preserves_shared_lock_holders() {
        let mut holder = kvproto::kvrpcpb::LockInfo::default();
        holder.set_primary_lock(b"primary".to_vec());
        holder.set_key(b"key".to_vec());
        holder.set_lock_version(42);

        let mut wrapper = kvproto::kvrpcpb::LockInfo::default();
        wrapper.mut_shared_lock_infos().push(holder);

        let converted = convert_lock_info(wrapper);
        assert_eq!(converted.shared_lock_infos.len(), 1);
        assert_eq!(converted.shared_lock_infos[0].primary_lock, b"primary");
        assert_eq!(converted.shared_lock_infos[0].key, b"key");
        assert_eq!(converted.shared_lock_infos[0].lock_version, 42);
    }

    fn lock_value(lock_type: LockType, start_ts: u64) -> Vec<u8> {
        Lock::new(
            lock_type,
            b"primary".to_vec(),
            start_ts.into(),
            100,
            None,
            TimeStamp::zero(),
            1,
            TimeStamp::zero(),
        )
        .to_bytes()
    }

    fn shared_locks_value(start_tss: &[u64]) -> Vec<u8> {
        let mut shared_locks = SharedLocks::new();
        for start_ts in start_tss {
            shared_locks
                .insert_lock(Lock::new(
                    LockType::Lock,
                    b"primary".to_vec(),
                    (*start_ts).into(),
                    100,
                    None,
                    TimeStamp::zero(),
                    1,
                    TimeStamp::zero(),
                ))
                .unwrap();
        }
        shared_locks.to_bytes()
    }

    fn write_locks(engine: &Engine, entries: &[(Vec<u8>, Vec<u8>)]) {
        let shard = engine.get_shard(SHARD_ID).unwrap();
        let mut wb = WriteBatch::default();
        wb.reset(SHARD_ID);
        wb.set_sequence(shard.get_write_sequence() + 1);
        for (key, value) in entries {
            wb.put(LOCK_CF, key, value, 0, &[], 0);
        }
        engine.write(&mut wb, &[]);
    }

    fn scan_locks(
        engine: &Engine,
        max_version: u64,
        limit: usize,
    ) -> Vec<kvproto::kvrpcpb::LockInfo> {
        let shard = engine.get_shard(SHARD_ID).unwrap();
        ReplicationService::scan_locks(
            shard.new_snap_access(),
            ApiV2::get_keyspace_prefix_by_id(KEYSPACE_ID).len(),
            shard.outer_start.clone(),
            shard.outer_end.clone(),
            max_version.into(),
            limit,
        )
        .unwrap()
    }

    #[test]
    fn test_keyspace_states_deserialize_legacy_feeds_for_compat() {
        let states: KeyspaceStates = serde_json::from_str(
            r#"{
                "feeds": {"cf-1": "{\"changefeed_id\":\"cf-1\"}"},
                "pd_url": "http://127.0.0.1:2379",
                "cdc_addr": "127.0.0.1:8300",
                "pd_sts_name": "pd-0",
                "cdc_sts_name": "cdc-0"
            }"#,
        )
        .unwrap();

        assert_eq!(states.feeds.len(), 1);
        assert_eq!(
            states.feeds.get("cf-1").unwrap(),
            "{\"changefeed_id\":\"cf-1\"}"
        );
        assert_eq!(states.pd_url, "http://127.0.0.1:2379");
        assert_eq!(states.cdc_addr, "127.0.0.1:8300");
        assert_eq!(states.pd_sts_name, "pd-0");
        assert_eq!(states.cdc_sts_name, "cdc-0");
    }

    #[test]
    fn test_keyspace_states_deserialize_missing_feeds_defaults_empty() {
        let states: KeyspaceStates = serde_json::from_str(
            r#"{
                "pd_url": "http://127.0.0.1:2379",
                "cdc_addr": "127.0.0.1:8300",
                "pd_sts_name": "pd-0",
                "cdc_sts_name": "cdc-0"
            }"#,
        )
        .unwrap();

        assert!(states.feeds.is_empty());
        assert_eq!(states.pd_url, "http://127.0.0.1:2379");
        assert_eq!(states.cdc_addr, "127.0.0.1:8300");
        assert_eq!(states.pd_sts_name, "pd-0");
        assert_eq!(states.cdc_sts_name, "cdc-0");
    }

    #[test]
    fn test_keyspace_states_serialize_includes_legacy_feeds_for_rollback() {
        let states = KeyspaceStates {
            pd_url: "http://127.0.0.1:2379".to_owned(),
            cdc_addr: "127.0.0.1:8300".to_owned(),
            pd_sts_name: "pd-0".to_owned(),
            cdc_sts_name: "cdc-0".to_owned(),
            ..Default::default()
        };

        let encoded = states.marshal();
        let value: serde_json::Value = serde_json::from_slice(&encoded).unwrap();

        assert_eq!(value.get("feeds"), Some(&serde_json::json!({})));
        assert_eq!(
            value.get("pd_url"),
            Some(&serde_json::json!("http://127.0.0.1:2379"))
        );
        assert_eq!(
            value.get("cdc_addr"),
            Some(&serde_json::json!("127.0.0.1:8300"))
        );
        assert_eq!(value.get("pd_sts_name"), Some(&serde_json::json!("pd-0")));
        assert_eq!(value.get("cdc_sts_name"), Some(&serde_json::json!("cdc-0")));
    }

    #[test]
    fn test_scan_locks_returns_only_stale_row_put_and_delete_locks() {
        let (engine, _apply_tx) = test_engine::new_test_engine_api_v2();
        let key = |i| engine.key_builder.gen_row_outer_key(TABLE_ID, i);
        let put_key = key(10);
        let delete_key = key(11);

        write_locks(
            &engine,
            &[
                (put_key.clone(), lock_value(LockType::Put, 10)),
                (delete_key.clone(), lock_value(LockType::Delete, 20)),
                (key(12), lock_value(LockType::Lock, 30)),
                (key(13), lock_value(LockType::Pessimistic, 40)),
                (key(14), shared_locks_value(&[10, 20])),
                (key(15), lock_value(LockType::Put, 60)),
                (key(16), lock_value(LockType::Delete, 70)),
            ],
        );

        let locks = scan_locks(&engine, 50, 100);
        let returned = locks
            .iter()
            .map(|lock| (lock.get_key().to_vec(), lock.get_lock_type()))
            .collect::<Vec<_>>();

        assert_eq!(returned, vec![(put_key, Op::Put), (delete_key, Op::Del)]);
        assert!(
            locks
                .iter()
                .all(|lock| lock.get_lock_type() != Op::SharedLock),
            "replication_worker scan_locks must not return shared lock wrappers: {:?}",
            locks
        );
    }

    #[test]
    fn test_scan_locks_limit_counts_only_returned_locks() {
        let (engine, _apply_tx) = test_engine::new_test_engine_api_v2();
        let key = |i| engine.key_builder.gen_row_outer_key(TABLE_ID, i);
        let first_put_key = key(13);
        let second_put_key = key(15);

        write_locks(
            &engine,
            &[
                (key(10), shared_locks_value(&[10, 11])),
                (key(11), lock_value(LockType::Pessimistic, 20)),
                (key(12), lock_value(LockType::Put, 200)),
                (first_put_key.clone(), lock_value(LockType::Put, 30)),
                (key(14), lock_value(LockType::Lock, 40)),
                (second_put_key.clone(), lock_value(LockType::Put, 50)),
            ],
        );

        let locks = scan_locks(&engine, 100, 1);
        assert_eq!(locks.len(), 1);
        assert_eq!(locks[0].get_key(), first_put_key.as_slice());
        assert_eq!(locks[0].get_lock_type(), Op::Put);

        let locks = scan_locks(&engine, 100, 2);
        let returned = locks
            .iter()
            .map(|lock| (lock.get_key().to_vec(), lock.get_lock_type()))
            .collect::<Vec<_>>();
        assert_eq!(
            returned,
            vec![(first_put_key, Op::Put), (second_put_key, Op::Put)]
        );
    }
}