// Copyright 2023 PingCAP, Inc.

use std::{
    collections::VecDeque,
    fmt::Write,
    future::Future,
    mem,
    ops::Deref,
    pin::Pin,
    sync::{
        Arc,
        atomic::{AtomicU8, Ordering},
    },
    time::Duration,
};

use async_trait::async_trait;
use bytes::{Buf, BufMut, Bytes};
use codec::number::NumberDecoder;
use http::{StatusCode, header};
use kvengine::{KvEngineConfig, LOCK_CF, SnapAccess, dfs::DFS_REMOTE_CACHE_ADDR_HEADER};
use kvproto::{
    coprocessor::{KeyRange, Response},
    db9_coprocessor::{Db9DagRequest, Db9ScanKind},
    kvrpcpb::ExecDetailsV2,
    metapb,
};
use protobuf::Message;
use resource_control::{CpuTokenHandle, TotalCpuMonitor};
use security::{HttpClientOptions, SecurityManager};
use tidb_query_common::execute_stats::ExecSummary;
use tidb_query_executors::interface::CoprRuntimeBudgetRef;
use tikv_alloc::MemoryTraceGuard;
use tikv_kv::Statistics;
use tikv_util::{
    backoff::ExponentialBackoff,
    config::VersionTrack,
    deadline::Deadline,
    http::{CONTENT_TYPE_PROTOBUF, HeaderExt},
    memory::HeapSize,
    retry::sleep_async,
};
use tipb::{DagRequest, ExecType, Executor, ExprType::ColumnRef};
use txn_types::{TimeStamp, TsSet};

use crate::{
    coprocessor::{
        Error, MEMTRACE_ROOT, REQ_TYPE_DAG, REQ_TYPE_DB9_DAG, ReqContext, RequestHandler, Result,
        decode_query_digest, metrics::*,
    },
    storage::{mvcc::NewerTsCheckState, txn::check_locks},
};

pub const REMOTE_REQUEST_TIMEOUT: Duration = Duration::from_secs(60 * 5);

pub const REMOTE_COP_FORMAT_V1: u32 = 1;

const RETRY_MAX_ATTEMPTS: usize = 10;
const REMOTE_MIN_PAGING_SIZE: u64 = 4096;
const RETRY_BASE_DELAY: Duration = Duration::from_secs(1);
const RETRY_MAX_DELAY: Duration = Duration::from_secs(30);
const REMOTE_NETWORK_RETRY_BASE_DELAY: Duration = Duration::from_millis(100);
const REMOTE_NETWORK_RETRY_MAX_DELAY: Duration = Duration::from_millis(500);
const COP_WORKER_FAILSAFE_MAX_ATTEMPTS: usize = 5;

const REMOTE_LOAD_STATE_UNDECIDED: u8 = 0;
const REMOTE_LOAD_STATE_LOCAL_PREFERRED: u8 = 1;
const REMOTE_LOAD_STATE_REMOTE_ALLOWED: u8 = 2;

type FallbackCpuTokenObserver = Box<dyn FnOnce(Arc<CpuTokenHandle>) + Send>;

#[derive(Clone)]
pub struct FallbackLocalAdmission {
    pub(crate) cpu_token_handle: Arc<CpuTokenHandle>,
    pub(crate) keyspace_id: u32,
    pub(crate) region_id: u64,
    pub(crate) peer: Option<metapb::Peer>,
    pub(crate) estimated_cpu_us: Option<u64>,
    pub(crate) is_low_priority_dag: bool,
    pub(crate) resource_publisher: Option<Arc<dyn resource_control::ResourcePublisher>>,
    pub(crate) resource_debug: bool,
}

pub(crate) type FallbackAdmissionFuture =
    Pin<Box<dyn Future<Output = Result<Option<FallbackLocalAdmission>>> + Send>>;
pub(crate) type FallbackAdmission = Box<dyn FnOnce() -> FallbackAdmissionFuture + Send>;

fn should_fallback_to_local(err: &Error) -> bool {
    // Local fallback is only for remote transport failures where TiKV cannot
    // get a complete response from the cop-worker, for example connection
    // refused, connection reset, or network unreachable.
    //
    // HTTP status codes are explicitly not fallback signals. 429/503 usually
    // mean cop-worker overload or autoscaling delay, and other 5xx responses
    // mean the request reached the worker but failed remotely.
    matches!(err, Error::RemoteNetwork(_))
}

fn remote_fallback_reason(err: &Error) -> &'static str {
    match err {
        Error::RemoteNetwork(_) => "remote_network",
        _ => "unknown",
    }
}

#[derive(Default)]
struct RemoteDispatchRetryState {
    retryable_errors: usize,
    consecutive_remote_network_errors: usize,
}

impl RemoteDispatchRetryState {
    fn record_retryable_error(&mut self, err: &Error) {
        self.retryable_errors += 1;
        if matches!(err, Error::RemoteNetwork(_)) {
            self.consecutive_remote_network_errors += 1;
        } else {
            self.consecutive_remote_network_errors = 0;
        }
    }

    fn should_fallback_to_local(&self) -> bool {
        self.consecutive_remote_network_errors >= COP_WORKER_FAILSAFE_MAX_ATTEMPTS
    }
}

pub(crate) fn remote_cop_request_tag(snap: &SnapAccess, req_ctx: &ReqContext) -> String {
    format!(
        "ks{}:{}:{}:{}",
        snap.get_keyspace_id(),
        snap.get_id(),
        snap.get_version(),
        req_ctx.txn_start_ts.into_inner()
    )
}

#[derive(Default, Debug)]
pub struct RemoteRequest {
    pub key: String,
    pub cop_req: Vec<u8>,
    pub req_body: Bytes,
}

pub async fn remote_request(
    remote_ctx: &RemoteContext,
    remote_addr: &str,
    tag: &str,
    req_body: Bytes,
    deadline: Deadline,
) -> Result<Vec<u8>> {
    let req = hyper::Request::builder()
        .method(hyper::Method::POST)
        .uri(remote_addr)
        .header(header::CONTENT_TYPE, "application/octet-stream")
        .header(header::ACCEPT, CONTENT_TYPE_PROTOBUF)
        .header(DFS_REMOTE_CACHE_ADDR_HEADER, remote_ctx.status_addr.clone())
        .body(hyper::Body::from(req_body))
        .map_err(|e| Error::Other(e.to_string()))?;
    let client = remote_ctx.client.clone();
    remote_ctx
        .runtime
        .spawn(async move {
            tokio::time::timeout_at(deadline.to_tokio_instant(), async move {
                let response = client
                    .request(req)
                    .await
                    .map_err(|e| Error::RemoteNetwork(e.to_string()))?;
                let is_pb_resp = response.headers().is_content_type_protobuf();
                let status = response.status();
                let body = hyper::body::to_bytes(response.into_body())
                    .await
                    .map_err(|e| Error::RemoteNetwork(e.to_string()))?;
                if !status.is_success() && !is_pb_resp {
                    let msg = String::from_utf8_lossy(body.chunk()).to_string();
                    let err = if status == StatusCode::SERVICE_UNAVAILABLE {
                        Error::RemoteServiceUnavailable(msg)
                    } else {
                        Error::Other(msg)
                    };
                    return Err(err);
                }
                Ok(body.to_vec())
            })
            .await
            .map_err(|_| Error::DeadlineExceeded)?
        })
        .await
        .map_err(|err| Error::Other(format!("{tag} join remote request task failed: {err:?}")))?
}

pub async fn remote_handle_request(
    req_type: &str,
    tag: &str,
    remote_ctx: &RemoteContext,
    remote_req: &RemoteRequest,
    deadline: Deadline,
) -> Result<Response> {
    info!("handle {} request, {}", req_type, tag);
    let req_body = remote_req.req_body.clone();
    let worker_url = remote_ctx
        .kvengine_config
        .value()
        .remote_worker_addr
        .clone();
    let resp_body = remote_request(remote_ctx, &worker_url, tag, req_body, deadline).await?;
    let mut resp = Response::default();
    resp.merge_from_bytes(&resp_body)
        .map_err(|err| Error::Other(format!("{tag}: decode response failed: {err:?}")))?;

    debug_assert!(
        !resp.has_region_error() && !resp.has_locked(),
        "unexpected error in remote cop resp: {:?}",
        resp
    );
    Ok(resp)
}

pub async fn remote_handle_request_with_retry(
    req_type: &str,
    tag: &str,
    remote_ctx: &RemoteContext,
    remote_req: &RemoteRequest,
) -> Result<Response> {
    let deadline = Deadline::from_now(REMOTE_REQUEST_TIMEOUT);
    let mut backoff =
        ExponentialBackoff::new(RETRY_BASE_DELAY, RETRY_MAX_DELAY, RETRY_MAX_ATTEMPTS);
    loop {
        match remote_handle_request(req_type, tag, remote_ctx, remote_req, deadline).await {
            Ok(resp) => return Ok(resp),
            Err(err @ Error::RemoteNetwork(_) | err @ Error::RemoteServiceUnavailable(_)) => {
                if deadline.check().is_err() {
                    warn!("{}:{}: deadline is exceeded", tag, req_type);
                    return Err(err);
                }

                if let Ok(delay) = backoff.next_delay() {
                    sleep_async(delay).await;
                } else {
                    warn!("{}:{}: retry limit exceeded", tag, req_type; "backoff" => ?backoff);
                    return Err(err);
                }
            }
            Err(err) => return Err(err),
        }
    }
}

#[derive(Clone)]
pub struct RemoteContext {
    pub core: Arc<RemoteContextCore>,
}

impl Deref for RemoteContext {
    type Target = RemoteContextCore;
    fn deref(&self) -> &Self::Target {
        &self.core
    }
}

pub struct RemoteContextCore {
    pub kvengine_config: Arc<VersionTrack<KvEngineConfig>>,
    pub cop_worker_provider: Arc<dyn CopWorkerProvider>,
    pub load_aware_config: RemoteLoadAwareConfig,
    pub runtime: tokio::runtime::Handle,
    pub client: security::HttpClient,
    lazy_remote_patterns: dashmap::DashMap<String, RemotePatternStats>,
    total_cpu_monitor: Option<Arc<TotalCpuMonitor>>,
    load_state: AtomicU8,
    #[cfg(test)]
    total_cpu_ratio_override: std::sync::RwLock<Option<Option<f64>>>,
    status_addr: String,
}

#[derive(Clone, Copy, Debug)]
pub struct RemoteLoadAwareConfig {
    pub enabled: bool,
    pub idle_threshold: f64,
    pub busy_threshold: f64,
    pub force_offload_blocks_size: usize,
    pub force_offload_num_ranges: usize,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum RemoteLoadState {
    Undecided,
    LocalPreferred,
    RemoteAllowed,
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum RemoteDecision {
    Local(RemoteDecisionReason),
    Remote(RemoteDecisionReason),
}

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
enum RemoteDecisionReason {
    BlocksSize,
    NumRanges,
    LongQueryBlocksSize,
    LazyPatternBlocksSize,
    LazyPatternDuration,
    CpuIdle,
    CpuBusy,
    ForceBlocksSize,
    ForceNumRanges,
    NoCpuSignal,
}

#[derive(Clone, Copy, Debug)]
enum RemoteBlocksSize {
    Estimated(usize),
    RangeCountFallback(usize),
}

impl RemoteBlocksSize {
    fn value(self) -> usize {
        match self {
            Self::Estimated(value) | Self::RangeCountFallback(value) => value,
        }
    }
}

#[derive(Clone, Copy, Debug)]
struct RemoteDagCandidate<'a> {
    lazy_remote_pattern: &'a str,
    eligibility_reason: Option<RemoteDecisionReason>,
}

pub trait CopWorkerProvider: Send + Sync {
    fn get(&self, keyspace_id: u32, start_ts: u64) -> Option<String>;
}

pub struct StaticCopWorkerProvider {
    config: Arc<VersionTrack<KvEngineConfig>>,
}

impl CopWorkerProvider for StaticCopWorkerProvider {
    fn get(&self, _keyspace_id: u32, _start_ts: u64) -> Option<String> {
        let addr = self.config.value().remote_coprocessor_addr.clone();
        if addr.is_empty() {
            return None;
        }
        Some(addr)
    }
}

impl RemoteContext {
    pub fn new(
        kvengine_config: Arc<VersionTrack<KvEngineConfig>>,
        security_mgr: Arc<SecurityManager>,
        runtime: tokio::runtime::Handle,
        total_cpu_monitor: Option<Arc<TotalCpuMonitor>>,
        status_addr: String,
    ) -> Self {
        let client = security_mgr
            .http_client_with_options(HttpClientOptions::default().pool_max_idle_per_host(0))
            .unwrap();
        let cop_worker_provider = Arc::new(StaticCopWorkerProvider {
            config: kvengine_config.clone(),
        });
        let lazy_remote_patterns = dashmap::DashMap::new();
        let load_aware_config = {
            let cfg = kvengine_config.value();
            RemoteLoadAwareConfig {
                enabled: cfg.remote_coprocessor_load_aware,
                idle_threshold: cfg.remote_coprocessor_load_aware_idle_threshold,
                busy_threshold: cfg.remote_coprocessor_load_aware_busy_threshold,
                force_offload_blocks_size: cfg.remote_coprocessor_load_aware_force_blocks_size,
                force_offload_num_ranges: cfg.remote_coprocessor_load_aware_force_num_ranges,
            }
        };
        Self {
            core: Arc::new(RemoteContextCore {
                kvengine_config,
                load_aware_config,
                cop_worker_provider,
                runtime,
                client,
                lazy_remote_patterns,
                total_cpu_monitor,
                load_state: AtomicU8::new(REMOTE_LOAD_STATE_UNDECIDED),
                #[cfg(test)]
                total_cpu_ratio_override: std::sync::RwLock::new(None),
                status_addr,
            }),
        }
    }

    pub fn active_worker_addr(&self) -> Option<String> {
        let addr = self.core.kvengine_config.value().remote_worker_addr.clone();
        (!addr.is_empty()).then_some(addr)
    }

    #[cfg(test)]
    pub(crate) fn set_total_cpu_ratio_for_test(&self, ratio: f64) {
        *self.total_cpu_ratio_override.write().unwrap() = Some(Some(ratio));
    }

    #[cfg(test)]
    pub(crate) fn set_total_cpu_ratio_unavailable_for_test(&self) {
        *self.total_cpu_ratio_override.write().unwrap() = Some(None);
    }

    pub fn report_lazy_remote_pattern(
        &self,
        pattern: String,
        total_size: usize,
        processed_duration: Duration,
    ) {
        let mut stats = self.lazy_remote_patterns.entry(pattern).or_default();
        stats.push(total_size, processed_duration);
    }

    pub fn should_offload(
        &self,
        pattern: &str,
        min_blocks_size: usize,
        min_process_duration: Duration,
    ) -> bool {
        self.lazy_pattern_offload_reason(pattern, min_blocks_size, min_process_duration)
            .is_some()
    }

    fn lazy_pattern_offload_reason(
        &self,
        pattern: &str,
        min_blocks_size: usize,
        min_process_duration: Duration,
    ) -> Option<RemoteDecisionReason> {
        if let Some(stats) = self.lazy_remote_patterns.get(pattern) {
            stats.offload_reason(min_blocks_size, min_process_duration)
        } else {
            None
        }
    }

    fn force_offload_reason(
        &self,
        blocks_size: RemoteBlocksSize,
        num_ranges: usize,
        eligibility_reason: RemoteDecisionReason,
    ) -> Option<RemoteDecisionReason> {
        if !self.load_aware_config.enabled {
            return None;
        }
        // Check ranges first because large-range requests use a synthetic block size.
        if num_ranges >= self.load_aware_config.force_offload_num_ranges {
            Some(RemoteDecisionReason::ForceNumRanges)
        } else {
            match blocks_size {
                RemoteBlocksSize::Estimated(size)
                    if size >= self.load_aware_config.force_offload_blocks_size =>
                {
                    Some(RemoteDecisionReason::ForceBlocksSize)
                }
                // Preserve the existing force-offload behavior for large-range
                // requests, but do not attribute it to a size that was not estimated.
                RemoteBlocksSize::RangeCountFallback(size)
                    if size >= self.load_aware_config.force_offload_blocks_size =>
                {
                    Some(eligibility_reason)
                }
                _ => None,
            }
        }
    }

    fn should_remote_based_on_load(
        &self,
        blocks_size: RemoteBlocksSize,
        num_ranges: usize,
        eligibility_reason: RemoteDecisionReason,
    ) -> bool {
        matches!(
            self.remote_decision_for_request(blocks_size, num_ranges, eligibility_reason),
            RemoteDecision::Remote(_)
        )
    }

    fn remote_decision_for_request(
        &self,
        blocks_size: RemoteBlocksSize,
        num_ranges: usize,
        eligibility_reason: RemoteDecisionReason,
    ) -> RemoteDecision {
        if !self.load_aware_config.enabled {
            record_remote_decision(eligibility_reason, true, None);
            return RemoteDecision::Remote(eligibility_reason);
        }
        if let Some(reason) = self.force_offload_reason(blocks_size, num_ranges, eligibility_reason)
        {
            record_remote_decision(reason, true, None);
            return RemoteDecision::Remote(reason);
        }
        match self.refresh_load_state() {
            RemoteLoadState::LocalPreferred => {
                record_remote_decision(RemoteDecisionReason::CpuIdle, false, self.cpu_ratio());
                RemoteDecision::Local(RemoteDecisionReason::CpuIdle)
            }
            RemoteLoadState::Undecided => {
                record_remote_decision(RemoteDecisionReason::NoCpuSignal, true, self.cpu_ratio());
                RemoteDecision::Remote(RemoteDecisionReason::NoCpuSignal)
            }
            RemoteLoadState::RemoteAllowed => {
                record_remote_decision(RemoteDecisionReason::CpuBusy, true, self.cpu_ratio());
                RemoteDecision::Remote(RemoteDecisionReason::CpuBusy)
            }
        }
    }

    fn refresh_load_state(&self) -> RemoteLoadState {
        let Some(total_cpu_ratio) = self.cpu_ratio() else {
            return RemoteLoadState::Undecided;
        };
        if total_cpu_ratio <= self.load_aware_config.idle_threshold {
            self.load_state
                .store(REMOTE_LOAD_STATE_LOCAL_PREFERRED, Ordering::Relaxed);
            return RemoteLoadState::LocalPreferred;
        }
        if total_cpu_ratio >= self.load_aware_config.busy_threshold {
            self.load_state
                .store(REMOTE_LOAD_STATE_REMOTE_ALLOWED, Ordering::Relaxed);
            return RemoteLoadState::RemoteAllowed;
        }
        self.load_state_from_atomic()
    }

    fn load_state_from_atomic(&self) -> RemoteLoadState {
        match self.load_state.load(Ordering::Relaxed) {
            REMOTE_LOAD_STATE_LOCAL_PREFERRED => RemoteLoadState::LocalPreferred,
            REMOTE_LOAD_STATE_REMOTE_ALLOWED => RemoteLoadState::RemoteAllowed,
            _ => RemoteLoadState::Undecided,
        }
    }

    fn cpu_ratio(&self) -> Option<f64> {
        #[cfg(test)]
        if let Some(ratio) = *self.total_cpu_ratio_override.read().unwrap() {
            return ratio;
        }

        self.total_cpu_monitor
            .as_ref()
            .and_then(|monitor| monitor.ratio())
    }
}

fn record_remote_decision(reason: RemoteDecisionReason, is_remote: bool, cpu_ratio: Option<f64>) {
    COPR_REMOTE_TOTAL_CPU_RATIO.set(cpu_ratio.unwrap_or(f64::NAN));
    COPR_REMOTE_DECISION_TOTAL
        .with_label_values(&[if is_remote { "remote" } else { "local" }, reason.as_str()])
        .inc();
}

impl RemoteDecisionReason {
    fn as_str(self) -> &'static str {
        match self {
            RemoteDecisionReason::BlocksSize => "blocks_size",
            RemoteDecisionReason::NumRanges => "num_ranges",
            RemoteDecisionReason::LongQueryBlocksSize => "long_query_blocks_size",
            RemoteDecisionReason::LazyPatternBlocksSize => "lazy_pattern_blocks_size",
            RemoteDecisionReason::LazyPatternDuration => "lazy_pattern_duration",
            RemoteDecisionReason::CpuIdle => "cpu_idle",
            RemoteDecisionReason::CpuBusy => "cpu_busy",
            RemoteDecisionReason::ForceBlocksSize => "load_aware_force_blocks_size",
            RemoteDecisionReason::ForceNumRanges => "load_aware_force_num_ranges",
            RemoteDecisionReason::NoCpuSignal => "no_cpu_signal",
        }
    }
}

fn remote_eligibility_reason(
    lazy_pattern_reason: Option<RemoteDecisionReason>,
    blocks_size: RemoteBlocksSize,
    configured_min_blocks_size: usize,
) -> RemoteDecisionReason {
    if let Some(reason) = lazy_pattern_reason {
        return reason;
    }
    match blocks_size {
        RemoteBlocksSize::RangeCountFallback(_) => RemoteDecisionReason::NumRanges,
        RemoteBlocksSize::Estimated(size) if size < configured_min_blocks_size => {
            // The caller has already checked the age-adjusted minimum, so an
            // estimate below the configured minimum is eligible only as a long query.
            RemoteDecisionReason::LongQueryBlocksSize
        }
        RemoteBlocksSize::Estimated(_) => RemoteDecisionReason::BlocksSize,
    }
}

pub(crate) fn try_remote_dag_handler(
    snap: Option<&kvengine::SnapAccess>,
    dag: &DagRequest,
    req_ctx: &ReqContext,
    remote_ctx: Option<RemoteContext>,
    paging_size: Option<u64>,
) -> Option<Box<dyn RequestHandler>> {
    let remote_ctx = remote_ctx?;
    let candidate = evaluate_remote_dag_static_candidate(dag, req_ctx, &remote_ctx, paging_size)?;
    let start_ts = req_ctx.txn_start_ts.into_inner();
    let snap = snap.cloned()?;
    let worker_addr = remote_ctx
        .cop_worker_provider
        .get(snap.get_keyspace_id(), start_ts)?;
    let kv_cfg = remote_ctx.kvengine_config.value();
    let mut ranges = Vec::with_capacity(req_ctx.ranges.len());
    for ran in &req_ctx.ranges {
        let start = Bytes::copy_from_slice(&ran.start);
        let end = Bytes::copy_from_slice(&ran.end);
        ranges.push((start, end))
    }
    let blocks_size = if ranges.len() < kv_cfg.remote_coprocessor_num_ranges {
        RemoteBlocksSize::Estimated(snap.estimated_range_blocks_size(&ranges))
    } else {
        // skip calculate blocks size if the number of ranges is too large.
        RemoteBlocksSize::RangeCountFallback(kv_cfg.remote_coprocessor_min_blocks_size)
    };
    let blocks_size_value = blocks_size.value();
    let configured_min_blocks_size = kv_cfg.remote_coprocessor_min_blocks_size;
    let min_blocks_size = calc_min_blocks_size(req_ctx.txn_start_ts, configured_min_blocks_size);
    let eligibility_reason = remote_eligibility_reason(
        candidate.eligibility_reason,
        blocks_size,
        configured_min_blocks_size,
    );
    let topn_max_pushdown_rows = kv_cfg.remote_coprocessor_topn_max_pushdown_rows;
    // Drop the read guard so remote_ctx can move into RemoteDagDispatcher::new.
    drop(kv_cfg);
    if blocks_size_value < min_blocks_size {
        return None;
    }
    if !remote_ctx.should_remote_based_on_load(blocks_size, ranges.len(), eligibility_reason) {
        return None;
    }
    COPR_REMOTE_DAG_ESTIMATE_BLOCKS_HISTOGRAM.observe(blocks_size_value as f64);
    let tag = remote_cop_request_tag(&snap, req_ctx);
    let query_digest = decode_query_digest(req_ctx.context.get_resource_group_tag());
    info!(
        "{} {} send remote coprocessor blocks_size:{} ranges:{}",
        tag,
        candidate.lazy_remote_pattern,
        blocks_size_value,
        ranges.len(),
        ;
        "query_digest" => &query_digest,
    );
    let cop_req = build_remote_dag_request(dag, req_ctx, topn_max_pushdown_rows);
    if cop_req.get_paging_size() != u32::MAX as u64 {
        COPR_REMOTE_TOPN_PAGING_REWRITE_TOTAL.inc();
    }
    Some(
        RemoteDagDispatcher::new(
            cop_req,
            snap,
            ranges,
            req_ctx.bypass_locks.clone(),
            remote_ctx,
            worker_addr,
            tag,
            req_ctx.deadline,
        )
        .into_boxed(),
    )
}

pub(crate) fn try_remote_db9_handler(
    snap: Option<&kvengine::SnapAccess>,
    db9: &Db9DagRequest,
    req_ctx: &ReqContext,
    remote_ctx: Option<RemoteContext>,
) -> Option<Box<dyn RequestHandler>> {
    let remote_ctx = remote_ctx?;
    if !is_remote_db9_static_candidate(db9) {
        return None;
    }
    let start_ts = req_ctx.txn_start_ts.into_inner();
    let snap = snap.cloned()?;
    let worker_addr = remote_ctx
        .cop_worker_provider
        .get(snap.get_keyspace_id(), start_ts)?;
    let mut ranges = Vec::with_capacity(req_ctx.ranges.len());
    for ran in &req_ctx.ranges {
        let start = Bytes::copy_from_slice(&ran.start);
        let end = Bytes::copy_from_slice(&ran.end);
        ranges.push((start, end))
    }
    let kv_cfg = remote_ctx.kvengine_config.value();
    let blocks_size = if ranges.len() < kv_cfg.remote_coprocessor_num_ranges {
        RemoteBlocksSize::Estimated(snap.estimated_range_blocks_size(&ranges))
    } else {
        // skip calculate blocks size if the number of ranges is too large.
        RemoteBlocksSize::RangeCountFallback(kv_cfg.remote_coprocessor_min_blocks_size)
    };
    let blocks_size_value = blocks_size.value();
    let configured_min_blocks_size = kv_cfg.remote_coprocessor_min_blocks_size;
    let min_blocks_size = calc_min_blocks_size(req_ctx.txn_start_ts, configured_min_blocks_size);
    let eligibility_reason =
        remote_eligibility_reason(None, blocks_size, configured_min_blocks_size);
    // Drop the read guard so remote_ctx can move into RemoteDagDispatcher::new.
    drop(kv_cfg);
    if blocks_size_value < min_blocks_size {
        return None;
    }
    if !remote_ctx.should_remote_based_on_load(blocks_size, ranges.len(), eligibility_reason) {
        return None;
    }
    COPR_REMOTE_DAG_ESTIMATE_BLOCKS_HISTOGRAM.observe(blocks_size_value as f64);
    let tag = remote_cop_request_tag(&snap, req_ctx);
    let query_digest = decode_query_digest(req_ctx.context.get_resource_group_tag());
    info!(
        "{} send remote db9 coprocessor blocks_size:{} ranges:{}",
        tag,
        blocks_size_value,
        ranges.len(),
        ;
        "query_digest" => &query_digest,
    );
    let cop_req = build_remote_db9_request(db9, req_ctx);
    Some(
        RemoteDagDispatcher::new(
            cop_req,
            snap,
            ranges,
            req_ctx.bypass_locks.clone(),
            remote_ctx,
            worker_addr,
            tag,
            req_ctx.deadline,
        )
        .into_boxed(),
    )
}

pub(crate) fn unsupported_remote_db9_request_reason(db9: &Db9DagRequest) -> Option<&'static str> {
    // This guard is mirrored by the server-side dispatcher, which falls back
    // to local execution. The remote worker still keeps the same checks as a
    // hard fail-closed boundary for malformed or future-regressed callers.
    if db9.has_limit() {
        return Some("remote DB9 cop does not support pushed LIMIT requests");
    }
    // Index scans require base-table row fetches. The remote snapshot only covers
    // the index key ranges, and the worker uses StubAccessor::none(), so any
    // cross-region row fetch silently returns a miss. Keep the worker path
    // fail-closed even if a caller forgets to screen these requests first.
    if matches!(
        db9.get_scan().get_kind(),
        Db9ScanKind::Db9ScanKindIndex
            | Db9ScanKind::Db9ScanKindIndexRange
            | Db9ScanKind::Db9ScanKindIndexBoundedRange
            | Db9ScanKind::Db9ScanKindInList
    ) {
        return Some("remote DB9 cop does not support index scan requests");
    }
    None
}

fn build_remote_dag_request(
    dag: &DagRequest,
    req_ctx: &ReqContext,
    topn_max_pushdown_rows: u64,
) -> kvproto::coprocessor::Request {
    // Preserve the existing remote DAG paging behavior after removing the
    // shared request builder helper.
    let mut cop_req = kvproto::coprocessor::Request::default();
    cop_req.set_context(req_ctx.context.clone());
    cop_req.tp = REQ_TYPE_DAG;
    cop_req.start_ts = req_ctx.txn_start_ts.into_inner();
    cop_req.data = dag.write_to_bytes().unwrap();
    cop_req.ranges = req_ctx.ranges.clone().into();
    cop_req.paging_size =
        large_topn_paging_size(dag, topn_max_pushdown_rows).unwrap_or(u32::MAX as u64);
    // Preserve only cache admission for the remote execution. The local
    // endpoint has already handled cache-hit validation, so forwarding the
    // client's if-match version would let the worker return an empty response.
    cop_req.is_cache_enabled = req_ctx.cache_match_version.is_some();
    cop_req
}

/// Returns the paging size that bypasses every non-partition TopN above the
/// configured pushdown limit. Zero disables the policy. Limits already above
/// the remote `u32::MAX` paging sentinel need no rewrite.
pub fn large_topn_paging_size(dag: &DagRequest, max_pushdown_rows: u64) -> Option<u64> {
    if max_pushdown_rows == 0 {
        return None;
    }
    dag.get_executors()
        .iter()
        .filter(|executor| executor.get_tp() == ExecType::TypeTopN)
        .map(|executor| executor.get_top_n())
        .filter(|topn| topn.get_partition_by().is_empty())
        .map(|topn| topn.get_limit())
        .filter(|limit| *limit > max_pushdown_rows && *limit <= u32::MAX as u64)
        .map(|limit| limit - 1)
        .min()
}

fn build_remote_db9_request(
    db9: &Db9DagRequest,
    req_ctx: &ReqContext,
) -> kvproto::coprocessor::Request {
    // DB9 remote cop remains unary and must not participate in DAG paging.
    let mut cop_req = kvproto::coprocessor::Request::default();
    cop_req.set_context(req_ctx.context.clone());
    cop_req.tp = REQ_TYPE_DB9_DAG;
    cop_req.start_ts = req_ctx.txn_start_ts.into_inner();
    cop_req.data = db9.write_to_bytes().unwrap();
    cop_req.ranges = req_ctx.ranges.clone().into();
    cop_req
}

fn estimate_remote_ranges_copy_bytes(req_ctx: &ReqContext) -> usize {
    let copied_key_bytes = req_ctx
        .ranges
        .iter()
        .map(|range| range.get_start().len() + range.get_end().len())
        .sum::<usize>();
    let copied_bytes_ranges_vec = req_ctx.ranges.len() * mem::size_of::<(Bytes, Bytes)>();
    let copied_key_ranges_vec = req_ctx.ranges.len() * mem::size_of::<KeyRange>();

    // Remote dispatch keeps one `(Bytes, Bytes)` copy for estimation / lock checks
    // and clones the protobuf ranges again into the remote cop request.
    req_ctx.context.approximate_heap_size()
        + copied_key_bytes * 2
        + copied_bytes_ranges_vec
        + copied_key_ranges_vec
}

fn estimate_remote_dispatch_extra_bytes(req_ctx: &ReqContext, serialized_req_len: usize) -> usize {
    estimate_remote_ranges_copy_bytes(req_ctx) + serialized_req_len
}

fn evaluate_remote_dag_static_candidate<'a>(
    dag: &'a DagRequest,
    req_ctx: &'a ReqContext,
    remote_ctx: &'a RemoteContext,
    paging_size: Option<u64>,
) -> Option<RemoteDagCandidate<'a>> {
    if let Some(paging_size) = paging_size {
        if is_index_full_scan(dag) && paging_size <= REMOTE_MIN_PAGING_SIZE {
            // Do not offload index full scan with small paging size because the upper layer
            // operator may quickly finish the query.
            return None;
        }
    }
    if let Some(lazy_remote_pattern) = &req_ctx.lazy_remote_pattern {
        // Dag that extracts lazy remote pattern should be offloaded to remote
        // coprocessor based on actual execution cost.
        let kv_cfg = remote_ctx.kvengine_config.value();
        let reason = remote_ctx.lazy_pattern_offload_reason(
            lazy_remote_pattern,
            kv_cfg.remote_coprocessor_min_blocks_size,
            kv_cfg.remote_coprocessor_min_process_duration.0,
        )?;
        Some(RemoteDagCandidate {
            lazy_remote_pattern,
            eligibility_reason: Some(reason),
        })
    } else {
        let last_executor = dag.get_executors().last().unwrap();
        if last_executor.has_limit() {
            // Do not offload coprocessor with limit because the actual cost may be much
            // smaller.
            return None;
        }
        Some(RemoteDagCandidate {
            lazy_remote_pattern: "",
            eligibility_reason: None,
        })
    }
}

fn is_remote_db9_static_candidate(db9: &Db9DagRequest) -> bool {
    unsupported_remote_db9_request_reason(db9).is_none()
}

pub(crate) fn estimate_remote_dag_dispatch_extra_bytes(
    dag: &DagRequest,
    req_ctx: &ReqContext,
    remote_ctx: Option<&RemoteContext>,
    paging_size: Option<u64>,
) -> usize {
    let Some(remote_ctx) = remote_ctx else {
        return 0;
    };
    // Mirror the dispatch gate: an empty remote cop URL (operator's local
    // fallback signal) means the request runs local, so precharge nothing.
    if remote_ctx
        .cop_worker_provider
        .get(req_ctx.context.get_keyspace_id(), 0)
        .is_none()
    {
        return 0;
    }
    if evaluate_remote_dag_static_candidate(dag, req_ctx, remote_ctx, paging_size).is_none() {
        return 0;
    }
    estimate_remote_dispatch_extra_bytes(req_ctx, dag.compute_size() as usize)
}

pub(crate) fn estimate_remote_db9_dispatch_extra_bytes(
    db9: &Db9DagRequest,
    req_ctx: &ReqContext,
    remote_ctx: Option<&RemoteContext>,
) -> usize {
    let Some(remote_ctx) = remote_ctx else {
        return 0;
    };
    // Mirror the dispatch gate: an empty remote cop URL (operator's local
    // fallback signal) means the request runs local, so precharge nothing.
    if remote_ctx
        .cop_worker_provider
        .get(req_ctx.context.get_keyspace_id(), 0)
        .is_none()
    {
        return 0;
    }
    if !is_remote_db9_static_candidate(db9) {
        return 0;
    }
    estimate_remote_dispatch_extra_bytes(req_ctx, db9.compute_size() as usize)
}

/// Remote-then-local fail-safe wrapper for remote DAG offload.
pub(crate) struct RemoteThenLocalFallbackDagHandler {
    remote: Box<dyn RequestHandler>,
    local: Option<Box<dyn RequestHandler>>,
    local_builder: Option<Box<dyn FnOnce() -> Result<Box<dyn RequestHandler>> + Send>>,
    fallback_admission: Option<FallbackAdmission>,
    fallback_cpu_token_observer: Option<FallbackCpuTokenObserver>,
    copr_runtime_budget: Option<CoprRuntimeBudgetRef>,
    remote_tag: String,
}

impl RemoteThenLocalFallbackDagHandler {
    pub(crate) fn new(
        remote: Box<dyn RequestHandler>,
        local_builder: Box<dyn FnOnce() -> Result<Box<dyn RequestHandler>> + Send>,
        fallback_admission: Option<FallbackAdmission>,
        remote_tag: String,
    ) -> Self {
        Self {
            remote,
            local: None,
            local_builder: Some(local_builder),
            fallback_admission,
            fallback_cpu_token_observer: None,
            copr_runtime_budget: None,
            remote_tag,
        }
    }

    fn set_tracker_is_remote(is_remote: bool) {
        ::tracker::with_tls_tracker(|tracker| {
            tracker.req_info.is_remote = is_remote;
        });
    }

    fn build_local(&mut self) -> Result<&mut Box<dyn RequestHandler>> {
        if self.local.is_none() {
            let builder = self
                .local_builder
                .take()
                .ok_or_else(|| Error::Other("missing local builder".to_owned()))?;
            let mut local = builder()?;
            if self.copr_runtime_budget.is_some() {
                local.set_copr_runtime_budget(self.copr_runtime_budget.clone());
            }
            self.local = Some(local);
        }
        Ok(self.local.as_mut().unwrap())
    }

    async fn run_fallback_admission(&mut self) -> Result<Option<FallbackLocalAdmission>> {
        match self.fallback_admission.take() {
            Some(admission) => admission().await,
            None => Ok(None),
        }
    }

    async fn handle_local_request(
        admission: Option<FallbackLocalAdmission>,
        cpu_token_observer: Option<FallbackCpuTokenObserver>,
        local: &mut dyn RequestHandler,
    ) -> Result<MemoryTraceGuard<Response>> {
        let Some(admission) = admission else {
            return local.handle_request().await;
        };
        let record_actual_usage = if let Some(observer) = cpu_token_observer {
            observer(admission.cpu_token_handle.clone());
            false
        } else {
            true
        };

        let cpu_metrics_ctx = crate::coprocessor::endpoint::CpuThrottleMetricsContext {
            keyspace_id: admission.keyspace_id,
            region_id: admission.region_id,
            peer: admission.peer,
            request_type: ::tracker::RequestType::CoprocessorDag,
            is_remote: false,
            is_low_priority_dag: admission.is_low_priority_dag,
            estimated_cpu_us: admission.estimated_cpu_us,
            resource_publisher: admission.resource_publisher,
            resource_debug: admission.resource_debug,
        };
        let runtime_check_handle = crate::coprocessor::endpoint::prepare_cpu_token_runtime_check(
            local,
            Some(admission.cpu_token_handle.clone()),
            cpu_metrics_ctx.resource_debug,
            cpu_metrics_ctx.keyspace_id,
            cpu_metrics_ctx.region_id,
        );
        let (cpu_time, result) = crate::coprocessor::endpoint::run_with_cpu_token_check(
            local.handle_request(),
            None,
            runtime_check_handle,
            cpu_metrics_ctx.resource_debug,
            cpu_metrics_ctx.keyspace_id,
            cpu_metrics_ctx.region_id,
        )
        .await?;
        crate::coprocessor::endpoint::finalize_cpu_throttle_metrics(
            &cpu_metrics_ctx,
            &admission.cpu_token_handle,
            cpu_time,
            record_actual_usage,
        );
        result
    }
}

#[async_trait]
impl RequestHandler for RemoteThenLocalFallbackDagHandler {
    async fn handle_request(&mut self) -> Result<MemoryTraceGuard<Response>> {
        match self.remote.handle_request().await {
            Ok(resp) => Ok(resp),
            Err(err) if should_fallback_to_local(&err) => {
                Self::set_tracker_is_remote(false);
                COPR_REMOTE_FALLBACK_TOTAL
                    .with_label_values(&[remote_fallback_reason(&err)])
                    .inc();
                warn!(
                    "{} remote cop-worker is unreachable; fall back to local",
                    self.remote_tag;
                    "err" => ?err
                );
                let admission = self.run_fallback_admission().await?;
                let cpu_token_observer = self.fallback_cpu_token_observer.take();
                let local = self.build_local()?;
                Self::handle_local_request(admission, cpu_token_observer, local.as_mut()).await
            }
            Err(err) => Err(err),
        }
    }

    fn set_fallback_cpu_token_observer(&mut self, observer: FallbackCpuTokenObserver) {
        self.fallback_cpu_token_observer = Some(observer);
    }

    fn set_copr_runtime_budget(&mut self, budget: Option<CoprRuntimeBudgetRef>) -> bool {
        if let Some(local) = self.local.as_mut() {
            return local.set_copr_runtime_budget(budget);
        }
        self.copr_runtime_budget = budget;
        true
    }

    fn collect_scan_statistics(&mut self, dest: &mut Statistics) {
        if let Some(local) = self.local.as_mut() {
            local.collect_scan_statistics(dest);
        } else {
            self.remote.collect_scan_statistics(dest);
        }
    }

    fn collect_scan_summary(&mut self, dest: &mut ExecSummary) {
        if let Some(local) = self.local.as_mut() {
            local.collect_scan_summary(dest);
        } else {
            self.remote.collect_scan_summary(dest);
        }
    }

    fn collect_executor_inputs(&self, dest: &mut kvproto::kvrpcpb::ExecutorInputs) {
        if let Some(local) = self.local.as_ref() {
            local.collect_executor_inputs(dest);
        } else {
            self.remote.collect_executor_inputs(dest);
        }
    }

    fn collect_executor_iterations(&self) -> u64 {
        if let Some(local) = self.local.as_ref() {
            local.collect_executor_iterations()
        } else {
            self.remote.collect_executor_iterations()
        }
    }

    fn set_ru_v2_metrics(&mut self, ru_v2: &mut kvproto::kvrpcpb::Ruv2, resp_bytes: u64) {
        if let Some(local) = self.local.as_mut() {
            local.set_ru_v2_metrics(ru_v2, resp_bytes);
        } else {
            self.remote.set_ru_v2_metrics(ru_v2, resp_bytes);
        }
    }
}

fn is_index_full_scan(dag: &DagRequest) -> bool {
    let executors = dag.get_executors();
    executors.len() == 1 && executors[0].get_tp() == ExecType::TypeIndexScan
}

fn set_remote_cache_metadata(
    resp: &mut Response,
    is_cache_enabled: bool,
    snapshot_write_sequence: u64,
    origin_met_uncacheable_data: bool,
) {
    if is_cache_enabled {
        if origin_met_uncacheable_data {
            resp.set_can_be_cached(false);
        }
        resp.set_cache_last_version(snapshot_write_sequence);
    }
}

pub struct RemoteDagDispatcher {
    req: kvproto::coprocessor::Request,
    snap: SnapAccess,
    ranges: Vec<(Bytes, Bytes)>,
    bypass_locks: TsSet,
    remote_ctx: RemoteContext,
    worker_addr: String,
    tag: String,
    exec_details: Option<ExecDetailsV2>,
    deadline: Deadline,
}

impl RemoteDagDispatcher {
    fn new(
        req: kvproto::coprocessor::Request,
        snap: SnapAccess,
        ranges: Vec<(Bytes, Bytes)>,
        bypass_locks: TsSet,
        remote_ctx: RemoteContext,
        worker_addr: String,
        tag: String,
        deadline: Deadline,
    ) -> Self {
        Self {
            req,
            snap,
            ranges,
            bypass_locks,
            remote_ctx,
            worker_addr,
            tag,
            exec_details: None,
            deadline,
        }
    }

    fn check_locks(&self) -> Result<bool> {
        let mut stats = Statistics::default();
        let mut lock_iter = self.snap.new_iterator(LOCK_CF, false, false, None, true);
        let mut met_uncacheable_lock = if self.req.get_is_cache_enabled() {
            NewerTsCheckState::NotMetYet
        } else {
            NewerTsCheckState::Unknown
        };
        for (start, end) in &self.ranges {
            lock_iter.set_range(start.clone(), end.clone());
            check_locks(
                &mut lock_iter,
                &self.bypass_locks,
                self.req.start_ts,
                &mut stats,
                &mut met_uncacheable_lock,
            )?;
        }
        Ok(met_uncacheable_lock == NewerTsCheckState::Met)
    }

    async fn dispatch(&self, deadline: Deadline) -> Result<(Vec<u8>, bool)> {
        let cop_req: Vec<u8> = self.req.write_to_bytes().unwrap();
        let (_, req_body, met_newer_ts) = encode_remote_request_body(
            &cop_req,
            &self.ranges,
            self.req.start_ts,
            &self.snap,
            false,
        );
        let req_body = Bytes::from(req_body);

        let mut remote_service_unavailable_backoff =
            ExponentialBackoff::new(RETRY_BASE_DELAY, RETRY_MAX_DELAY, usize::MAX);
        // Remote network failures use a shorter backoff so TiKV can quickly
        // confirm consecutive transport failures and fall back to local.
        let mut remote_network_backoff = ExponentialBackoff::new(
            REMOTE_NETWORK_RETRY_BASE_DELAY,
            REMOTE_NETWORK_RETRY_MAX_DELAY,
            usize::MAX,
        );
        let mut retry_state = RemoteDispatchRetryState::default();
        loop {
            match remote_request(
                &self.remote_ctx,
                &self.worker_addr,
                &self.tag,
                req_body.clone(),
                deadline,
            )
            .await
            {
                Ok(data) => return Ok((data, met_newer_ts)),
                Err(err @ Error::RemoteNetwork(_) | err @ Error::RemoteServiceUnavailable(_)) => {
                    retry_state.record_retryable_error(&err);
                    let dur_left: Duration =
                        deadline.check().map_err(|_| Error::DeadlineExceeded)?;

                    warn!(
                        "{} dispatch remote coprocessor error {:?}, retry {}, consecutive remote network errors {}/{}, {:?} until deadline",
                        self.tag,
                        err,
                        retry_state.retryable_errors,
                        retry_state.consecutive_remote_network_errors,
                        COP_WORKER_FAILSAFE_MAX_ATTEMPTS,
                        dur_left
                    );
                    if retry_state.should_fallback_to_local() {
                        return Err(err);
                    }

                    let delay = match err {
                        Error::RemoteNetwork(_) => remote_network_backoff
                            .next_delay()
                            .unwrap_or(REMOTE_NETWORK_RETRY_MAX_DELAY),
                        Error::RemoteServiceUnavailable(_) => remote_service_unavailable_backoff
                            .next_delay()
                            .unwrap_or(RETRY_MAX_DELAY),
                        _ => unreachable!(),
                    };
                    if dur_left <= delay {
                        sleep_async(dur_left).await;
                        return Err(Error::DeadlineExceeded);
                    }
                    sleep_async(delay).await;
                }
                Err(err) => return Err(err),
            }
        }
    }
}

#[async_trait]
impl RequestHandler for RemoteDagDispatcher {
    async fn handle_request(&mut self) -> Result<MemoryTraceGuard<kvproto::coprocessor::Response>> {
        let met_uncacheable_lock = self.check_locks()?;
        let ret = self.dispatch(self.deadline).await;
        match ret {
            Ok((data, met_newer_ts)) => {
                let memory_size = data.capacity();
                let mut resp = kvproto::coprocessor::Response::default();
                resp.merge_from_bytes(&data).unwrap();
                // The worker snapshot message does not carry the origin
                // shard's write sequence. Veto worker-side cache admission if
                // the origin observed newer data omitted from the serialized
                // request, then stamp the cache token from the exact snapshot
                // used to build this request.
                set_remote_cache_metadata(
                    &mut resp,
                    self.req.get_is_cache_enabled(),
                    self.snap.get_write_sequence(),
                    met_uncacheable_lock || met_newer_ts,
                );
                let scan_detail_v2 = resp.get_exec_details_v2().get_scan_detail_v2();
                COPR_REMOTE_PROCESSED_SIZE.inc_by(scan_detail_v2.get_processed_versions_size());
                COPR_REMOTE_TOTAL_SIZE.inc_by(scan_detail_v2.get_total_versions_size());
                self.exec_details = Some(resp.get_exec_details_v2().clone());
                Ok(MEMTRACE_ROOT.trace_guard(resp, memory_size))
            }
            Err(Error::Other(e)) => {
                error!("{} remote coprocessor failed, error {}", self.tag, e);
                let mut resp = kvproto::coprocessor::Response::default();
                resp.set_other_error(e);
                Ok(resp.into())
            }
            Err(e) => {
                error!("{} remote coprocessor failed, error {:?}", self.tag, e);
                Err(e)
            }
        }
    }

    async fn handle_streaming_request(
        &mut self,
    ) -> Result<(Option<kvproto::coprocessor::Response>, bool)> {
        unimplemented!()
    }

    fn collect_scan_statistics(&mut self, dest: &mut Statistics) {
        if let Some(exec_details_v2) = self.exec_details.as_ref() {
            if let Some(scan_detail_v2) = exec_details_v2.scan_detail_v2.as_ref() {
                dest.processed_size = scan_detail_v2.processed_versions_size as usize;
                dest.write.processed_keys = scan_detail_v2.processed_versions as usize;
                // TODO: we treat arbitrary versions as write, this is not totally accurate,
                // but it should be enough here.
                dest.write.next = scan_detail_v2.total_versions as usize;
                dest.total_size = scan_detail_v2.total_versions_size as usize;
            }
        }
    }

    fn collect_scan_summary(&mut self, dest: &mut ExecSummary) {
        if let Some(exec_details_v2) = self.exec_details.as_ref() {
            if let Some(time_details) = exec_details_v2.time_detail.as_ref() {
                dest.time_processed_ns = time_details.process_wall_time_ms as usize * 1000000;
            }
        }
    }

    fn collect_executor_inputs(&self, dest: &mut kvproto::kvrpcpb::ExecutorInputs) {
        collect_remote_executor_inputs(self.exec_details.as_ref(), dest);
    }

    fn collect_executor_iterations(&self) -> u64 {
        collect_remote_executor_iterations(self.exec_details.as_ref())
    }
}

fn collect_remote_executor_inputs(
    exec_details: Option<&ExecDetailsV2>,
    dest: &mut kvproto::kvrpcpb::ExecutorInputs,
) {
    if let Some(remote_ru_v2) = exec_details.map(ExecDetailsV2::get_ru_v2) {
        *dest = remote_ru_v2.get_executor_inputs().clone();
    }
}

fn collect_remote_executor_iterations(exec_details: Option<&ExecDetailsV2>) -> u64 {
    exec_details.map(ExecDetailsV2::get_ru_v2).map_or(
        0,
        kvproto::kvrpcpb::Ruv2::get_coprocessor_executor_iterations,
    )
}

impl RemoteContext {
    pub fn extract(
        &self,
        keyspace_id: u32,
        ranges: &[KeyRange],
        dag: &DagRequest,
    ) -> Option<String> {
        let executors = dag.get_executors();
        if executors.is_empty() {
            return None;
        }
        let scan_executor = &executors[0];
        let cop_num_ranges = self.kvengine_config.value().remote_coprocessor_num_ranges;
        if ranges.len() >= cop_num_ranges {
            // on large range pattern.
            let mut str_buf = String::new();
            write!(str_buf, "k{}", keyspace_id).unwrap();
            Self::write_scan_digest(&mut str_buf, scan_executor);
            write!(str_buf, "r").unwrap();
            return Some(str_buf);
        }
        if executors.len() < 3 {
            return None;
        }

        let last_executor = &executors[executors.len() - 1];
        let second_last_executor = &executors[executors.len() - 2];
        if !last_executor.has_limit() || !second_last_executor.has_selection() {
            return None;
        }
        // on selection with limit pattern.
        let mut str_buf = String::new();
        write!(str_buf, "k{}", keyspace_id).unwrap();
        Self::write_scan_digest(&mut str_buf, scan_executor);
        let selection = second_last_executor.get_selection();
        Self::write_selection_digest(&mut str_buf, selection);
        Some(str_buf)
    }

    fn write_scan_digest(s: &mut String, executor: &Executor) {
        match executor.get_tp() {
            ExecType::TypeTableScan => {
                let tbl_scan = executor.get_tbl_scan();
                write!(s, "t{}", tbl_scan.get_table_id()).unwrap();
                for col in tbl_scan.get_columns() {
                    write!(s, "c{}", col.get_column_id()).unwrap();
                }
            }
            ExecType::TypeIndexScan => {
                let idx_scan = executor.get_idx_scan();
                write!(s, "t{}", idx_scan.get_table_id()).unwrap();
                write!(s, "i{}", idx_scan.get_index_id()).unwrap();
                for col in idx_scan.get_columns() {
                    write!(s, "c{}", col.get_column_id()).unwrap();
                }
            }
            _ => {}
        }
    }

    fn write_selection_digest(s: &mut String, selection: &tipb::Selection) {
        s.push('s');
        for expr in selection.get_conditions() {
            Self::extract_expr(s, expr);
        }
    }

    fn extract_expr(s: &mut String, expr: &tipb::Expr) {
        if expr.get_tp() == ColumnRef {
            write!(s, "c{}", expr.get_val().read_i64().unwrap_or_default()).unwrap();
        }
        if !expr.get_children().is_empty() {
            for child in expr.get_children() {
                Self::extract_expr(s, child);
            }
        }
    }
}

#[derive(Default)]
struct RemotePatternStats {
    process_stats: VecDeque<ProcessStat>,
}

#[derive(Default, Copy, Clone)]
struct ProcessStat {
    // The total size of data processed by the pattern.
    size: usize,
    duration: Duration,
}

impl RemotePatternStats {
    fn push(&mut self, total_size: usize, process_duration: Duration) {
        self.process_stats.push_back(ProcessStat {
            size: total_size,
            duration: process_duration,
        });
        if self.process_stats.len() > 5 {
            self.process_stats.pop_front();
        }
    }

    fn offload_reason(
        &self,
        threshold_size: usize,
        threshold_dur: Duration,
    ) -> Option<RemoteDecisionReason> {
        if self.process_stats.len() < 2 {
            return None;
        }
        let sum_size: usize = self.process_stats.iter().map(|x| x.size).sum();
        let avg_size = sum_size / self.process_stats.len();
        let sum_duration: Duration = self.process_stats.iter().map(|stat| stat.duration).sum();
        let avg_duration = sum_duration / self.process_stats.len() as u32;
        if avg_size > threshold_size {
            Some(RemoteDecisionReason::LazyPatternBlocksSize)
        } else if avg_duration > threshold_dur {
            Some(RemoteDecisionReason::LazyPatternDuration)
        } else {
            None
        }
    }
}

/// If the query has already run for a long time, the additional latency for
/// offloading to remote coprocessor is non-significant, we can decreases the
/// min blocks size to reduce the tikv-server resource consumption.
fn calc_min_blocks_size(start_ts: TimeStamp, config_min_blocks_size: usize) -> usize {
    const LONG_QUERY_MS: u64 = 15 * 1000;
    let elapsed_ms = TimeStamp::physical_now().saturating_sub(start_ts.physical());
    if elapsed_ms > LONG_QUERY_MS * 4 {
        // 60s
        config_min_blocks_size / 8
    } else if elapsed_ms > LONG_QUERY_MS * 2 {
        // 30s
        config_min_blocks_size / 4
    } else if elapsed_ms > LONG_QUERY_MS {
        // 15s
        config_min_blocks_size / 2
    } else {
        config_min_blocks_size
    }
}

pub fn encode_remote_request_body(
    cop_req: &[u8],
    ranges: &[(Bytes, Bytes)],
    start_ts: u64,
    snap: &SnapAccess,
    cache_key: bool,
) -> (String, Vec<u8>, bool) {
    let (mem_data, met_newer_ts) = snap.build_mem_data(ranges, start_ts);
    let (key, snap_data) = snap.marshal(ranges, true, cache_key);
    let req_body = encode_remote_cop_request(cop_req, &mem_data, &snap_data);
    (key, req_body, met_newer_ts)
}

pub fn encode_remote_cop_request(cop_req: &[u8], mem_data: &[u8], snap_data: &[u8]) -> Vec<u8> {
    let extra_len = 4 * 4;
    let mut req_body: Vec<u8> =
        Vec::with_capacity(extra_len + cop_req.len() + mem_data.len() + snap_data.len());
    req_body.put_u32_le(REMOTE_COP_FORMAT_V1);
    req_body.put_u32_le(cop_req.len() as u32);
    req_body.extend_from_slice(cop_req);
    req_body.put_u32_le(mem_data.len() as u32);
    req_body.extend_from_slice(mem_data);
    req_body.put_u32_le(snap_data.len() as u32);
    req_body.extend_from_slice(snap_data);
    req_body
}

pub fn decode_remote_cop_request(body: &[u8]) -> Result<(&[u8], &[u8], &[u8])> {
    let mut body_buf = body;
    let err = Error::Other("failed to decode cop request".to_string());
    if body_buf.len() < 4 {
        return Err(err);
    }
    let format = body_buf.get_u32_le();
    if format != REMOTE_COP_FORMAT_V1 {
        return Err(err);
    }
    if body_buf.len() < 4 {
        return Err(err);
    }
    let req_data_len = body_buf.get_u32_le() as usize;
    if body_buf.len() < req_data_len {
        return Err(err);
    }
    let req_data = &body_buf[..req_data_len];
    body_buf = &body_buf[req_data_len..];
    if body_buf.len() < 4 {
        return Err(Error::Other("invalid cop request".to_string()));
    }
    let mem_data_len = body_buf.get_u32_le() as usize;
    if body_buf.len() < mem_data_len {
        return Err(err);
    }
    let mem_data = &body_buf[..mem_data_len];
    body_buf = &body_buf[mem_data_len..];
    if body_buf.len() < 4 {
        return Err(Error::Other("invalid cop request".to_string()));
    }
    let snap_data_len = body_buf.get_u32_le() as usize;
    if body_buf.len() < snap_data_len {
        return Err(err);
    }
    let snap_data = &body_buf[..snap_data_len];
    Ok((req_data, mem_data, snap_data))
}

#[cfg(test)]
mod tests {
    use std::{
        sync::{
            Arc,
            atomic::{AtomicBool, AtomicUsize, Ordering},
        },
        thread::sleep,
    };

    use codec::number::NumberEncoder;
    use resource_control::{
        Config, Metric, ReadLimiter, Resource, ResourceController, ResourceEvent, ResourceType,
        Scope, Severity, TotalCpuSubscriber,
    };
    use tidb_query_datatype::{FieldTypeAccessor, FieldTypeTp};
    use tidb_query_executors::interface::CoprRuntimeBudget;
    use tikv_util::config::ReadableDuration;
    use tipb::{Expr, ResourceGroupTag, ScalarFuncSig};

    use super::*;

    struct OkRequestHandler {
        calls: Arc<AtomicUsize>,
    }

    #[async_trait]
    impl RequestHandler for OkRequestHandler {
        async fn handle_request(&mut self) -> Result<MemoryTraceGuard<Response>> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            Ok(MEMTRACE_ROOT.trace_guard(Response::default(), 0))
        }

        async fn handle_streaming_request(&mut self) -> Result<(Option<Response>, bool)> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            Ok((Some(Response::default()), false))
        }
    }

    struct RuntimeBudgetAwareRequestHandler {
        calls: Arc<AtomicUsize>,
        budget_installed: Arc<AtomicBool>,
    }

    #[async_trait]
    impl RequestHandler for RuntimeBudgetAwareRequestHandler {
        async fn handle_request(&mut self) -> Result<MemoryTraceGuard<Response>> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            Ok(MEMTRACE_ROOT.trace_guard(Response::default(), 0))
        }

        fn set_copr_runtime_budget(&mut self, budget: Option<CoprRuntimeBudgetRef>) -> bool {
            self.budget_installed
                .store(budget.is_some(), Ordering::SeqCst);
            true
        }
    }

    struct ErrRequestHandler {
        calls: Arc<AtomicUsize>,
        err: Error,
    }

    #[async_trait]
    impl RequestHandler for ErrRequestHandler {
        async fn handle_request(&mut self) -> Result<MemoryTraceGuard<Response>> {
            self.calls.fetch_add(1, Ordering::SeqCst);
            Err(self.err.clone())
        }
    }

    struct NoopRuntimeBudget;

    #[async_trait]
    impl CoprRuntimeBudget for NoopRuntimeBudget {
        async fn on_batch(&self, _cpu_delta_us: u64) -> tidb_query_common::Result<()> {
            Ok(())
        }
    }

    #[tokio::test]
    async fn test_remote_then_local_fallback_no_local_build_on_success() {
        let remote_calls = Arc::new(AtomicUsize::new(0));
        let local_builds = Arc::new(AtomicUsize::new(0));
        let local_calls = Arc::new(AtomicUsize::new(0));

        let remote = Box::new(OkRequestHandler {
            calls: remote_calls.clone(),
        });
        let local_builder = {
            let local_builds = local_builds.clone();
            let local_calls = local_calls.clone();
            Box::new(move || {
                local_builds.fetch_add(1, Ordering::SeqCst);
                Ok(Box::new(OkRequestHandler { calls: local_calls }) as _)
            })
        };

        let mut handler = RemoteThenLocalFallbackDagHandler::new(
            remote,
            local_builder,
            None,
            "ks1:1:1:1".to_owned(),
        );
        handler.handle_request().await.unwrap();

        assert_eq!(remote_calls.load(Ordering::SeqCst), 1);
        assert_eq!(local_builds.load(Ordering::SeqCst), 0);
        assert_eq!(local_calls.load(Ordering::SeqCst), 0);
    }

    #[tokio::test]
    async fn test_remote_then_local_fallback_builds_local_on_network_error() {
        let remote_calls = Arc::new(AtomicUsize::new(0));
        let local_builds = Arc::new(AtomicUsize::new(0));
        let local_calls = Arc::new(AtomicUsize::new(0));

        let remote = Box::new(ErrRequestHandler {
            calls: remote_calls.clone(),
            err: Error::RemoteNetwork("mock network error".to_string()),
        });
        let local_builder = {
            let local_builds = local_builds.clone();
            let local_calls = local_calls.clone();
            Box::new(move || {
                local_builds.fetch_add(1, Ordering::SeqCst);
                Ok(Box::new(OkRequestHandler { calls: local_calls }) as _)
            })
        };

        let mut handler = RemoteThenLocalFallbackDagHandler::new(
            remote,
            local_builder,
            None,
            "ks1:1:1:1".to_owned(),
        );
        handler.handle_request().await.unwrap();

        assert_eq!(remote_calls.load(Ordering::SeqCst), 1);
        assert_eq!(local_builds.load(Ordering::SeqCst), 1);
        assert_eq!(local_calls.load(Ordering::SeqCst), 1);
    }

    #[tokio::test]
    async fn test_remote_then_local_fallback_forwards_runtime_budget_to_local() {
        let remote_calls = Arc::new(AtomicUsize::new(0));
        let local_calls = Arc::new(AtomicUsize::new(0));
        let local_budget_installed = Arc::new(AtomicBool::new(false));

        let remote = Box::new(ErrRequestHandler {
            calls: remote_calls.clone(),
            err: Error::RemoteNetwork("mock network error".to_string()),
        });
        let local_builder = {
            let local_calls = local_calls.clone();
            let local_budget_installed = local_budget_installed.clone();
            Box::new(move || {
                Ok(Box::new(RuntimeBudgetAwareRequestHandler {
                    calls: local_calls,
                    budget_installed: local_budget_installed,
                }) as _)
            })
        };

        let mut handler = RemoteThenLocalFallbackDagHandler::new(
            remote,
            local_builder,
            None,
            "ks1:1:1:1".to_owned(),
        );
        assert!(handler.set_copr_runtime_budget(Some(Arc::new(NoopRuntimeBudget))));

        handler.handle_request().await.unwrap();

        assert_eq!(remote_calls.load(Ordering::SeqCst), 1);
        assert_eq!(local_calls.load(Ordering::SeqCst), 1);
        assert!(local_budget_installed.load(Ordering::SeqCst));
    }

    #[tokio::test]
    async fn test_remote_then_local_fallback_observes_fallback_cpu_token() {
        let remote_calls = Arc::new(AtomicUsize::new(0));
        let local_calls = Arc::new(AtomicUsize::new(0));
        let observed_token = Arc::new(AtomicUsize::new(0));
        let read_limiter = ReadLimiter::new(Config::default());
        let cpu_token_handle = Arc::new(
            read_limiter
                .try_allocate_cpu(
                    1,
                    1000,
                    std::time::Instant::now() + std::time::Duration::from_secs(5),
                )
                .unwrap(),
        );

        let remote = Box::new(ErrRequestHandler {
            calls: remote_calls.clone(),
            err: Error::RemoteNetwork("mock network error".to_string()),
        });
        let local_builder = {
            let local_calls = local_calls.clone();
            Box::new(move || Ok(Box::new(OkRequestHandler { calls: local_calls }) as _))
        };
        let fallback_admission = {
            let cpu_token_handle = cpu_token_handle.clone();
            Box::new(move || {
                let cpu_token_handle = cpu_token_handle.clone();
                Box::pin(async move {
                    Ok(Some(FallbackLocalAdmission {
                        cpu_token_handle,
                        keyspace_id: 1,
                        region_id: 2,
                        peer: None,
                        estimated_cpu_us: Some(1000),
                        is_low_priority_dag: false,
                        resource_publisher: None,
                        resource_debug: false,
                    }))
                }) as FallbackAdmissionFuture
            }) as FallbackAdmission
        };

        let mut handler = RemoteThenLocalFallbackDagHandler::new(
            remote,
            local_builder,
            Some(fallback_admission),
            "ks1:2:3:4".to_owned(),
        );
        let expected_token = cpu_token_handle.clone();
        let observed_token_for_observer = observed_token.clone();
        handler.set_fallback_cpu_token_observer(Box::new(move |token| {
            if Arc::ptr_eq(&token, &expected_token) {
                observed_token_for_observer.store(1, Ordering::SeqCst);
            }
        }));

        handler.handle_request().await.unwrap();

        assert_eq!(remote_calls.load(Ordering::SeqCst), 1);
        assert_eq!(local_calls.load(Ordering::SeqCst), 1);
        assert_eq!(observed_token.load(Ordering::SeqCst), 1);
        assert_eq!(cpu_token_handle.actual_used(), 0);
    }

    #[tokio::test]
    async fn test_remote_then_local_fallback_passthrough_service_unavailable() {
        let remote_calls = Arc::new(AtomicUsize::new(0));
        let local_builds = Arc::new(AtomicUsize::new(0));

        let remote = Box::new(ErrRequestHandler {
            calls: remote_calls.clone(),
            err: Error::RemoteServiceUnavailable("mock overload".to_string()),
        });
        let local_builder = {
            let local_builds = local_builds.clone();
            Box::new(move || {
                local_builds.fetch_add(1, Ordering::SeqCst);
                Ok(Box::new(OkRequestHandler {
                    calls: Arc::new(AtomicUsize::new(0)),
                }) as _)
            })
        };

        let mut handler = RemoteThenLocalFallbackDagHandler::new(
            remote,
            local_builder,
            None,
            "ks1:1:1:1".to_owned(),
        );
        let err = handler.handle_request().await.unwrap_err();
        assert!(matches!(err, Error::RemoteServiceUnavailable(_)));

        assert_eq!(remote_calls.load(Ordering::SeqCst), 1);
        assert_eq!(local_builds.load(Ordering::SeqCst), 0);
    }

    #[test]
    fn test_should_fallback_to_local_only_on_network_error() {
        assert!(should_fallback_to_local(&Error::RemoteNetwork(
            "connection reset".to_owned()
        )));
        assert!(!should_fallback_to_local(&Error::RemoteServiceUnavailable(
            "queue full".to_owned()
        )));
        assert!(!should_fallback_to_local(&Error::Other(
            "remote execution failed".to_owned()
        )));
    }

    #[test]
    fn test_remote_dispatch_retry_state_keeps_pure_service_unavailable_out_of_fallback_window() {
        let mut retry_state = RemoteDispatchRetryState::default();
        for _ in 0..COP_WORKER_FAILSAFE_MAX_ATTEMPTS * 2 {
            retry_state
                .record_retryable_error(&Error::RemoteServiceUnavailable("queue full".to_owned()));
        }

        assert_eq!(
            retry_state.retryable_errors,
            COP_WORKER_FAILSAFE_MAX_ATTEMPTS * 2
        );
        assert_eq!(retry_state.consecutive_remote_network_errors, 0);
        assert!(!retry_state.should_fallback_to_local());
    }

    #[test]
    fn test_remote_dispatch_retry_state_requires_consecutive_network_errors() {
        let mut retry_state = RemoteDispatchRetryState::default();
        for _ in 1..COP_WORKER_FAILSAFE_MAX_ATTEMPTS {
            retry_state
                .record_retryable_error(&Error::RemoteNetwork("connection reset".to_owned()));
        }
        assert_eq!(
            retry_state.consecutive_remote_network_errors,
            COP_WORKER_FAILSAFE_MAX_ATTEMPTS - 1
        );

        retry_state
            .record_retryable_error(&Error::RemoteServiceUnavailable("queue full".to_owned()));
        assert_eq!(retry_state.consecutive_remote_network_errors, 0);
        assert!(!retry_state.should_fallback_to_local());
    }

    #[test]
    fn test_remote_dispatch_retry_state_falls_back_after_consecutive_network_errors() {
        let mut retry_state = RemoteDispatchRetryState::default();
        for _ in 0..COP_WORKER_FAILSAFE_MAX_ATTEMPTS {
            retry_state
                .record_retryable_error(&Error::RemoteNetwork("connection reset".to_owned()));
        }
        assert_eq!(
            retry_state.consecutive_remote_network_errors,
            COP_WORKER_FAILSAFE_MAX_ATTEMPTS
        );
        assert!(retry_state.should_fallback_to_local());
    }

    #[allow(clippy::too_many_arguments)]
    fn test_kvengine_cfg(
        remote_worker_addr: &str,
        remote_cop_addr: &str,
        min_blocks_size: usize,
        num_ranges: usize,
        min_process_duration: Duration,
        load_aware: bool,
        idle_threshold: f64,
        busy_threshold: f64,
        force_blocks_size: usize,
        force_num_ranges: usize,
    ) -> Arc<VersionTrack<KvEngineConfig>> {
        Arc::new(VersionTrack::new(KvEngineConfig {
            remote_worker_addr: remote_worker_addr.to_string(),
            remote_coprocessor_addr: remote_cop_addr.to_string(),
            remote_coprocessor_min_blocks_size: min_blocks_size,
            remote_coprocessor_num_ranges: num_ranges,
            remote_coprocessor_min_process_duration: ReadableDuration(min_process_duration),
            remote_coprocessor_load_aware: load_aware,
            remote_coprocessor_load_aware_idle_threshold: idle_threshold,
            remote_coprocessor_load_aware_busy_threshold: busy_threshold,
            remote_coprocessor_load_aware_force_blocks_size: force_blocks_size,
            remote_coprocessor_load_aware_force_num_ranges: force_num_ranges,
            ..KvEngineConfig::default()
        }))
    }

    #[test]
    fn test_remote_cop_coded() {
        let cop_req = b"cop_req".to_vec();
        let mem_data = b"mem_data".to_vec();
        let snap_data = b"snap_data".to_vec();
        let req_body = encode_remote_cop_request(&cop_req, &mem_data, &snap_data);
        let (cop_req, mem_data, snap_data) = decode_remote_cop_request(&req_body).unwrap();
        assert_eq!(cop_req, "cop_req".as_bytes());
        assert_eq!(mem_data, "mem_data".as_bytes());
        assert_eq!(snap_data, "snap_data".as_bytes());
    }

    #[test]
    fn test_remote_cop_coded_with_empty_mem() {
        let cop_req = b"cop_req".to_vec();
        let mem_data = b"".to_vec();
        let snap_data = b"snap_data".to_vec();
        let req_body = encode_remote_cop_request(&cop_req, &mem_data, &snap_data);
        let (cop_req, mem_data, snap_data) = decode_remote_cop_request(&req_body).unwrap();
        assert_eq!(cop_req, "cop_req".as_bytes());
        assert_eq!(mem_data, "".as_bytes());
        assert_eq!(snap_data, "snap_data".as_bytes());
    }

    #[test]
    fn test_decode_query_digest() {
        let mut resource_group_tag = ResourceGroupTag::new();
        resource_group_tag.set_sql_digest(vec![0x12, 0xab, 0x00]);

        assert_eq!(
            decode_query_digest(&resource_group_tag.write_to_bytes().unwrap()),
            "12ab00"
        );
        assert_eq!(decode_query_digest(&[]), "");
        assert_eq!(decode_query_digest(&[0x0a, 0x02, 0x01]), "");
    }

    #[test]
    fn test_calc_min_blocks() {
        let conf_min_blocks_size = 512usize;
        let ts = TimeStamp::default();
        assert_eq!(calc_min_blocks_size(ts, conf_min_blocks_size), 64);
        let ts = TimeStamp::max();
        assert_eq!(calc_min_blocks_size(ts, conf_min_blocks_size), 512);

        let phys_now = TimeStamp::physical_now();
        let ts = TimeStamp::compose(phys_now - 70 * 1000, 1);
        assert_eq!(calc_min_blocks_size(ts, conf_min_blocks_size), 64);
        let ts = TimeStamp::compose(phys_now - 40 * 1000, 1);
        assert_eq!(calc_min_blocks_size(ts, conf_min_blocks_size), 128);
        let ts = TimeStamp::compose(phys_now - 20 * 1000, 1);
        assert_eq!(calc_min_blocks_size(ts, conf_min_blocks_size), 256);
        let ts = TimeStamp::compose(phys_now - 10 * 1000, 1);
        assert_eq!(calc_min_blocks_size(ts, conf_min_blocks_size), 512);
        let ts = TimeStamp::compose(phys_now + 20 * 1000, 1);
        assert_eq!(calc_min_blocks_size(ts, conf_min_blocks_size), 512);
    }

    #[test]
    fn test_remote_eligibility_reason() {
        assert_eq!(
            remote_eligibility_reason(None, RemoteBlocksSize::Estimated(100), 100),
            RemoteDecisionReason::BlocksSize
        );
        assert_eq!(
            remote_eligibility_reason(None, RemoteBlocksSize::RangeCountFallback(100), 100),
            RemoteDecisionReason::NumRanges
        );
        assert_eq!(
            remote_eligibility_reason(None, RemoteBlocksSize::Estimated(50), 100),
            RemoteDecisionReason::LongQueryBlocksSize
        );
        assert_eq!(
            remote_eligibility_reason(
                Some(RemoteDecisionReason::LazyPatternDuration),
                RemoteBlocksSize::RangeCountFallback(100),
                100,
            ),
            RemoteDecisionReason::LazyPatternDuration
        );
    }

    #[test]
    fn test_remote_decision_reason_labels() {
        let cases = [
            (RemoteDecisionReason::BlocksSize, "blocks_size"),
            (RemoteDecisionReason::NumRanges, "num_ranges"),
            (
                RemoteDecisionReason::LongQueryBlocksSize,
                "long_query_blocks_size",
            ),
            (
                RemoteDecisionReason::LazyPatternBlocksSize,
                "lazy_pattern_blocks_size",
            ),
            (
                RemoteDecisionReason::LazyPatternDuration,
                "lazy_pattern_duration",
            ),
            (RemoteDecisionReason::CpuIdle, "cpu_idle"),
            (RemoteDecisionReason::CpuBusy, "cpu_busy"),
            (
                RemoteDecisionReason::ForceBlocksSize,
                "load_aware_force_blocks_size",
            ),
            (
                RemoteDecisionReason::ForceNumRanges,
                "load_aware_force_num_ranges",
            ),
            (RemoteDecisionReason::NoCpuSignal, "no_cpu_signal"),
        ];

        for (reason, label) in cases {
            assert_eq!(reason.as_str(), label);
        }
    }

    #[test]
    fn test_remote_patten() {
        let mut dag = DagRequest::default();
        dag.mut_executors()
            .push(new_table_scan_executor(1, vec![1, 2, 3]));
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let cfg = KvEngineConfig {
            remote_worker_addr: "a".to_string(),
            remote_coprocessor_addr: "b".to_string(),
            remote_coprocessor_min_blocks_size: 100,
            remote_coprocessor_num_ranges: 2,
            remote_coprocessor_min_process_duration: tikv_util::config::ReadableDuration::millis(
                10,
            ),
            ..KvEngineConfig::default()
        };
        let remote_ctx = RemoteContext::new(
            Arc::new(VersionTrack::new(cfg)),
            Arc::new(security::SecurityManager::default()),
            runtime.handle().clone(),
            None,
            String::new(),
        );
        let make_key_range = |start: u8, end: u8| KeyRange {
            start: vec![start],
            end: vec![end],
            ..Default::default()
        };
        let small_key_ranges = vec![make_key_range(1, 2)];
        assert!(remote_ctx.extract(1, &small_key_ranges, &dag).is_none());
        let large_key_ranges = vec![
            make_key_range(1, 2),
            make_key_range(3, 4),
            make_key_range(5, 6),
        ];
        let large_key_range_pattern = remote_ctx.extract(1, &large_key_ranges, &dag).unwrap();
        assert_eq!(large_key_range_pattern, "k1t1c1c2c3r");

        assert!(remote_ctx.extract(1, &small_key_ranges, &dag).is_none());
        dag.mut_executors().push(new_selection(vec![1, 2]));
        assert!(remote_ctx.extract(1, &small_key_ranges, &dag).is_none());
        dag.mut_executors().push(new_limit(100));
        let tbl_pattern = remote_ctx.extract(1, &small_key_ranges, &dag).unwrap();
        assert_eq!(tbl_pattern, "k1t1c1c2c3sc1c2");

        let mut idx_dag = dag.clone();
        let executors = idx_dag.mut_executors();
        executors[0] = new_index_scan_executor(3, 2, vec![1, 2, 3]);
        let idx_pattern = remote_ctx.extract(1, &small_key_ranges, &idx_dag).unwrap();
        assert_eq!(idx_pattern, "k1t3i2c1c2c3sc1c2");

        let short_process_dur = Duration::from_millis(5);
        let long_process_dur = Duration::from_millis(100);
        let cfg_min_blocks = 100;
        let cfg_min_dur = Duration::from_millis(10);
        remote_ctx.report_lazy_remote_pattern(tbl_pattern.clone(), 101, long_process_dur);
        assert!(!remote_ctx.should_offload(&tbl_pattern, cfg_min_blocks, cfg_min_dur));
        assert_eq!(
            remote_ctx.lazy_pattern_offload_reason(&tbl_pattern, cfg_min_blocks, cfg_min_dur),
            None
        );
        remote_ctx.report_lazy_remote_pattern(tbl_pattern.clone(), 101, long_process_dur);
        assert!(remote_ctx.should_offload(&tbl_pattern, cfg_min_blocks, cfg_min_dur));
        assert_eq!(
            remote_ctx.lazy_pattern_offload_reason(&tbl_pattern, cfg_min_blocks, cfg_min_dur),
            Some(RemoteDecisionReason::LazyPatternBlocksSize)
        );
        #[derive(Debug)]
        struct Case {
            processed_size: usize,
            process_duration: Duration,
            expected_reason: Option<RemoteDecisionReason>,
        }
        let cases = vec![
            Case {
                processed_size: 50,
                process_duration: short_process_dur,
                expected_reason: None,
            },
            Case {
                processed_size: 50,
                process_duration: long_process_dur,
                expected_reason: Some(RemoteDecisionReason::LazyPatternDuration),
            },
            Case {
                processed_size: 101,
                process_duration: short_process_dur,
                expected_reason: Some(RemoteDecisionReason::LazyPatternBlocksSize),
            },
            Case {
                processed_size: 101,
                process_duration: long_process_dur,
                expected_reason: Some(RemoteDecisionReason::LazyPatternBlocksSize),
            },
        ];
        for case in cases {
            for _ in 0..10 {
                remote_ctx.report_lazy_remote_pattern(
                    tbl_pattern.clone(),
                    case.processed_size,
                    case.process_duration,
                );
            }
            assert_eq!(
                remote_ctx.lazy_pattern_offload_reason(&tbl_pattern, cfg_min_blocks, cfg_min_dur),
                case.expected_reason,
                "case: {:?}",
                case
            );
        }
    }

    #[test]
    fn test_collect_remote_ru_v2_metrics() {
        let mut exec_details = ExecDetailsV2::default();
        exec_details
            .mut_ru_v2()
            .set_coprocessor_executor_iterations(7);
        exec_details
            .mut_ru_v2()
            .mut_executor_inputs()
            .set_tikv_coprocessor_executor_work_total_batch_table_scan(11);

        let mut inputs = kvproto::kvrpcpb::ExecutorInputs::default();
        collect_remote_executor_inputs(Some(&exec_details), &mut inputs);

        assert_eq!(collect_remote_executor_iterations(Some(&exec_details)), 7);
        assert_eq!(
            inputs.get_tikv_coprocessor_executor_work_total_batch_table_scan(),
            11
        );
    }

    #[test]
    fn test_remote_offload_without_load_aware_reports_eligibility_reason() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let remote_ctx = RemoteContext::new(
            test_kvengine_cfg(
                "a",
                "b",
                100,
                2,
                Duration::from_millis(10),
                false,
                0.3,
                0.5,
                400,
                4,
            ),
            Arc::new(security::SecurityManager::default()),
            runtime.handle().clone(),
            None,
            String::new(),
        );

        for reason in [
            RemoteDecisionReason::BlocksSize,
            RemoteDecisionReason::NumRanges,
            RemoteDecisionReason::LongQueryBlocksSize,
            RemoteDecisionReason::LazyPatternBlocksSize,
            RemoteDecisionReason::LazyPatternDuration,
        ] {
            assert_eq!(
                remote_ctx.remote_decision_for_request(RemoteBlocksSize::Estimated(200), 1, reason),
                RemoteDecision::Remote(reason)
            );
        }
    }

    #[test]
    fn test_load_aware_remote_offload_uses_cpu_hysteresis() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let remote_ctx = RemoteContext::new(
            test_kvengine_cfg(
                "a",
                "b",
                100,
                2,
                Duration::from_millis(10),
                true,
                0.3,
                0.5,
                400,
                4,
            ),
            Arc::new(security::SecurityManager::default()),
            runtime.handle().clone(),
            None,
            String::new(),
        );

        remote_ctx.set_total_cpu_ratio_for_test(0.20);
        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(200),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Local(RemoteDecisionReason::CpuIdle)
        );
        assert_eq!(
            remote_ctx.refresh_load_state(),
            RemoteLoadState::LocalPreferred
        );

        remote_ctx.set_total_cpu_ratio_for_test(0.40);
        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(200),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Local(RemoteDecisionReason::CpuIdle)
        );
        assert_eq!(
            remote_ctx.refresh_load_state(),
            RemoteLoadState::LocalPreferred
        );

        remote_ctx.set_total_cpu_ratio_for_test(0.80);
        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(200),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Remote(RemoteDecisionReason::CpuBusy)
        );
        assert_eq!(
            remote_ctx.refresh_load_state(),
            RemoteLoadState::RemoteAllowed
        );

        remote_ctx.set_total_cpu_ratio_for_test(0.40);
        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(200),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Remote(RemoteDecisionReason::CpuBusy)
        );
        assert_eq!(
            remote_ctx.refresh_load_state(),
            RemoteLoadState::RemoteAllowed
        );
    }

    #[test]
    fn test_load_aware_remote_offload_force_offload_and_fallback() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let remote_ctx = RemoteContext::new(
            test_kvengine_cfg(
                "a",
                "b",
                100,
                2,
                Duration::from_millis(10),
                true,
                0.3,
                0.5,
                400,
                4,
            ),
            Arc::new(security::SecurityManager::default()),
            runtime.handle().clone(),
            None,
            String::new(),
        );

        remote_ctx.set_total_cpu_ratio_for_test(0.20);
        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(200),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Local(RemoteDecisionReason::CpuIdle)
        );
        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(400),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Remote(RemoteDecisionReason::ForceBlocksSize)
        );
        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(800),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Remote(RemoteDecisionReason::ForceBlocksSize)
        );

        remote_ctx.set_total_cpu_ratio_unavailable_for_test();
        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(200),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Remote(RemoteDecisionReason::NoCpuSignal)
        );
        assert_eq!(remote_ctx.refresh_load_state(), RemoteLoadState::Undecided);

        let no_signal_ctx = RemoteContext::new(
            test_kvengine_cfg(
                "a",
                "b",
                100,
                2,
                Duration::from_millis(10),
                true,
                0.3,
                0.5,
                400,
                4,
            ),
            Arc::new(security::SecurityManager::default()),
            runtime.handle().clone(),
            None,
            String::new(),
        );
        no_signal_ctx.set_total_cpu_ratio_unavailable_for_test();
        assert_eq!(
            no_signal_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(200),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Remote(RemoteDecisionReason::NoCpuSignal)
        );
        assert_eq!(
            no_signal_ctx.refresh_load_state(),
            RemoteLoadState::Undecided
        );
    }

    #[test]
    fn test_load_aware_remote_offload_force_offload_by_num_ranges() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let remote_ctx = RemoteContext::new(
            test_kvengine_cfg(
                "a",
                "b",
                100,
                2,
                Duration::from_millis(10),
                true,
                0.3,
                0.5,
                400,
                4,
            ),
            Arc::new(security::SecurityManager::default()),
            runtime.handle().clone(),
            None,
            String::new(),
        );

        remote_ctx.set_total_cpu_ratio_for_test(0.20);
        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::RangeCountFallback(200),
                4,
                RemoteDecisionReason::NumRanges,
            ),
            RemoteDecision::Remote(RemoteDecisionReason::ForceNumRanges)
        );
    }

    #[test]
    fn test_load_aware_remote_offload_does_not_report_unmeasured_blocks_size() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let remote_ctx = RemoteContext::new(
            test_kvengine_cfg(
                "a",
                "b",
                100,
                2,
                Duration::from_millis(10),
                true,
                0.3,
                0.5,
                100,
                4,
            ),
            Arc::new(security::SecurityManager::default()),
            runtime.handle().clone(),
            None,
            String::new(),
        );

        remote_ctx.set_total_cpu_ratio_for_test(0.20);
        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::RangeCountFallback(100),
                2,
                RemoteDecisionReason::NumRanges,
            ),
            RemoteDecision::Remote(RemoteDecisionReason::NumRanges)
        );
    }

    #[test]
    fn test_load_aware_remote_offload_uses_resource_control_total_cpu_signal() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let mut resource_control_config = resource_control::Config::default();
        resource_control_config.enabled = true;
        resource_control_config.stats_interval = ReadableDuration(Duration::from_secs(1));
        let resource_controller = ResourceController::new(resource_control_config.clone());
        let remote_ctx = RemoteContext::new(
            test_kvengine_cfg(
                "a",
                "b",
                100,
                2,
                Duration::from_millis(10),
                true,
                0.3,
                0.5,
                400,
                4,
            ),
            Arc::new(security::SecurityManager::default()),
            runtime.handle().clone(),
            Some(resource_controller.total_cpu_monitor()),
            String::new(),
        );
        let total_cpu_subscriber = TotalCpuSubscriber::new(
            resource_controller.total_cpu_monitor(),
            resource_control_config.clone(),
        );
        resource_controller.register_subscriber(
            ResourceType::Cpu {
                cpu_type: resource_control::CpuType::Total,
            },
            Box::new(total_cpu_subscriber),
        );

        let publisher = resource_controller.publisher();
        let publish_total_cpu = |cpu_time: f64| {
            publisher.publish(ResourceEvent::CollectMetric(Metric {
                resource: Resource::Cpu {
                    cpu_type: resource_control::CpuType::Total,
                    cpu_time,
                },
                scope: Scope::Global,
                severity: Severity::Normal,
            }));
        };

        // Keep the cumulative CPU time flat across the window to produce an idle
        // signal.
        publish_total_cpu(10.0);
        sleep(Duration::from_millis(100));
        publish_total_cpu(10.0);
        publisher.publish(ResourceEvent::Tick);
        sleep(Duration::from_millis(100));

        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(200),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Local(RemoteDecisionReason::CpuIdle)
        );
        assert_eq!(
            remote_ctx.refresh_load_state(),
            RemoteLoadState::LocalPreferred
        );

        // Increase cumulative CPU time sharply to push the ratio well above the busy
        // threshold.
        publish_total_cpu(10.0);
        sleep(Duration::from_millis(100));
        publish_total_cpu(1000.0);
        publisher.publish(ResourceEvent::Tick);
        sleep(Duration::from_millis(100));

        assert_eq!(
            remote_ctx.remote_decision_for_request(
                RemoteBlocksSize::Estimated(200),
                1,
                RemoteDecisionReason::BlocksSize,
            ),
            RemoteDecision::Remote(RemoteDecisionReason::CpuBusy)
        );
        assert_eq!(
            remote_ctx.refresh_load_state(),
            RemoteLoadState::RemoteAllowed
        );
    }

    fn topn_dag(limit: u64, partitioned: bool) -> DagRequest {
        let mut topn = tipb::TopN::default();
        topn.set_limit(limit);
        if partitioned {
            topn.mut_partition_by().push(Default::default());
        }
        let mut executor = Executor::default();
        executor.set_tp(ExecType::TypeTopN);
        executor.set_top_n(topn);
        let mut dag = DagRequest::default();
        dag.mut_executors().push(executor);
        dag
    }

    #[test]
    fn test_build_remote_dag_request_keeps_sentinel_without_large_topn() {
        let req_ctx = ReqContext::default_for_test();
        let dag = DagRequest::default();
        let req = build_remote_dag_request(&dag, &req_ctx, 50_000);

        assert_eq!(req.get_tp(), REQ_TYPE_DAG);
        assert_eq!(req.get_start_ts(), req_ctx.txn_start_ts.into_inner());
        assert_eq!(req.get_data(), dag.write_to_bytes().unwrap().as_slice());
        assert_eq!(req.get_paging_size(), u32::MAX as u64);
        assert_eq!(req.get_paging_size_bytes(), 0);
    }

    #[test]
    fn test_build_remote_dag_request_preserves_cache_enablement() {
        let req_ctx = ReqContext::new(
            Default::default(),
            Vec::new(),
            Duration::from_secs(100),
            None,
            None,
            TimeStamp::max(),
            Some(42),
            None,
        );
        let dag = DagRequest::default();
        let req = build_remote_dag_request(&dag, &req_ctx, 50_000);

        assert!(req.get_is_cache_enabled());
        assert_eq!(req.get_cache_if_match_version(), 0);
    }

    #[test]
    fn test_set_remote_cache_metadata_uses_origin_evidence() {
        let mut resp = Response::default();
        resp.set_can_be_cached(true);
        resp.set_cache_last_version(7);

        set_remote_cache_metadata(&mut resp, true, 89, false);
        assert!(resp.get_can_be_cached());
        assert_eq!(resp.get_cache_last_version(), 89);

        set_remote_cache_metadata(&mut resp, true, 144, true);
        assert!(!resp.get_can_be_cached());
        assert_eq!(resp.get_cache_last_version(), 144);

        set_remote_cache_metadata(&mut resp, false, 233, false);
        assert_eq!(resp.get_cache_last_version(), 144);
    }

    #[test]
    fn test_build_remote_dag_request_bypasses_large_topn() {
        let req_ctx = ReqContext::default_for_test();
        let dag = topn_dag(50_001, false);
        let req = build_remote_dag_request(&dag, &req_ctx, 50_000);

        assert_eq!(req.get_paging_size(), 50_000);
        assert_eq!(req.get_data(), dag.write_to_bytes().unwrap().as_slice());
    }

    #[test]
    fn test_large_topn_paging_size_boundaries() {
        assert_eq!(
            large_topn_paging_size(&topn_dag(50_001, false), 50_000),
            Some(50_000)
        );
        assert_eq!(
            large_topn_paging_size(&topn_dag(50_000, false), 50_000),
            None
        );
        assert_eq!(
            large_topn_paging_size(&topn_dag(50_001, true), 50_000),
            None
        );
        assert_eq!(large_topn_paging_size(&topn_dag(50_001, false), 0), None);
        assert_eq!(
            large_topn_paging_size(&topn_dag(u32::MAX as u64 + 1, false), 50_000),
            None
        );
    }

    #[test]
    fn test_build_remote_db9_request_keeps_paging_unset() {
        let req_ctx = ReqContext::default_for_test();
        let db9 = Db9DagRequest::default();
        let req = build_remote_db9_request(&db9, &req_ctx);

        assert_eq!(req.get_tp(), REQ_TYPE_DB9_DAG);
        assert_eq!(req.get_start_ts(), req_ctx.txn_start_ts.into_inner());
        assert_eq!(req.get_data(), db9.write_to_bytes().unwrap().as_slice());
        assert_eq!(req.get_paging_size(), 0);
        assert_eq!(req.get_paging_size_bytes(), 0);
    }

    #[test]
    fn test_estimate_remote_dispatch_extra_bytes_skips_static_local_requests() {
        let runtime = tokio::runtime::Runtime::new().unwrap();
        let remote_ctx = RemoteContext::new(
            test_kvengine_cfg(
                "a",
                "b",
                100,
                2,
                Duration::from_millis(10),
                true,
                0.3,
                0.5,
                400,
                4,
            ),
            Arc::new(security::SecurityManager::default()),
            runtime.handle().clone(),
            None,
            String::new(),
        );

        let req_ctx = ReqContext::default_for_test();

        let mut limit_dag = DagRequest::default();
        limit_dag.mut_executors().push(tipb::Executor::default());
        limit_dag.mut_executors().last_mut().unwrap().mut_limit();
        assert_eq!(
            estimate_remote_dag_dispatch_extra_bytes(&limit_dag, &req_ctx, Some(&remote_ctx), None,),
            0
        );

        let mut paging_dag = DagRequest::default();
        let mut index_scan = tipb::Executor::default();
        index_scan.set_tp(ExecType::TypeIndexScan);
        paging_dag.mut_executors().push(index_scan);
        assert_eq!(
            estimate_remote_dag_dispatch_extra_bytes(
                &paging_dag,
                &req_ctx,
                Some(&remote_ctx),
                Some(REMOTE_MIN_PAGING_SIZE),
            ),
            0
        );

        let mut db9 = Db9DagRequest::default();
        db9.set_codec_version(1);
        db9.mut_scan().set_kind(Db9ScanKind::Db9ScanKindIndex);
        assert_eq!(
            estimate_remote_db9_dispatch_extra_bytes(&db9, &req_ctx, Some(&remote_ctx)),
            0
        );

        // Empty remote cop URL means the request runs local, so a non-limit
        // DAG and a non-index db9 candidate must precharge nothing.
        let local_ctx = RemoteContext::new(
            test_kvengine_cfg(
                "a",
                "",
                100,
                2,
                Duration::from_millis(10),
                true,
                0.3,
                0.5,
                400,
                4,
            ),
            Arc::new(security::SecurityManager::default()),
            runtime.handle().clone(),
            None,
            String::new(),
        );
        let mut select_dag = DagRequest::default();
        select_dag
            .mut_executors()
            .push(new_table_scan_executor(1, vec![1, 2, 3]));
        assert_eq!(
            estimate_remote_dag_dispatch_extra_bytes(&select_dag, &req_ctx, Some(&local_ctx), None,),
            0
        );
        let mut full_scan_db9 = Db9DagRequest::default();
        full_scan_db9.set_codec_version(1);
        assert_eq!(
            estimate_remote_db9_dispatch_extra_bytes(&full_scan_db9, &req_ctx, Some(&local_ctx)),
            0
        );
    }

    fn new_table_scan_executor(tbl_id: i64, col_ids: Vec<i64>) -> tipb::Executor {
        let mut executor = tipb::Executor::default();
        executor.set_tp(tipb::ExecType::TypeTableScan);
        let tbl_scan = executor.mut_tbl_scan();
        tbl_scan.set_table_id(tbl_id);
        for col_id in col_ids {
            let mut col = tipb::ColumnInfo::new();
            col.set_column_id(col_id);
            tbl_scan.mut_columns().push(col);
        }
        executor
    }

    fn new_index_scan_executor(tbl_id: i64, index_id: i64, col_ids: Vec<i64>) -> tipb::Executor {
        let mut executor = tipb::Executor::default();
        executor.set_tp(tipb::ExecType::TypeIndexScan);
        let idx_scan = executor.mut_idx_scan();
        idx_scan.set_table_id(tbl_id);
        idx_scan.set_index_id(index_id);
        for col_id in col_ids {
            let mut col = tipb::ColumnInfo::new();
            col.set_column_id(col_id);
            idx_scan.mut_columns().push(col);
        }
        executor
    }

    fn new_selection(col_offs: Vec<i64>) -> tipb::Executor {
        let mut executor = tipb::Executor::default();
        executor.set_tp(tipb::ExecType::TypeSelection);
        let conditions = executor.mut_selection().mut_conditions();
        for col_off in col_offs {
            let mut condition = tipb::Expr::default();
            condition.set_tp(tipb::ExprType::Int64);
            condition.set_sig(ScalarFuncSig::EqInt);
            condition
                .mut_field_type()
                .as_mut_accessor()
                .set_tp(FieldTypeTp::LongLong);

            let mut col = Expr::default();
            col.set_tp(tipb::ExprType::ColumnRef);
            col.mut_val().write_i64(col_off).unwrap();
            col.mut_field_type()
                .as_mut_accessor()
                .set_tp(FieldTypeTp::LongLong);

            let mut value = Expr::default();
            value.set_tp(tipb::ExprType::Int64);
            value.mut_val().write_i64(1).unwrap();
            value
                .mut_field_type()
                .as_mut_accessor()
                .set_tp(FieldTypeTp::LongLong);
            condition.mut_children().push(col);
            condition.mut_children().push(value);
            conditions.push(condition);
        }
        executor
    }

    fn new_limit(limit: u64) -> tipb::Executor {
        let mut executor = tipb::Executor::default();
        executor.set_tp(tipb::ExecType::TypeLimit);
        let limit_executor = executor.mut_limit();
        limit_executor.set_limit(limit);
        executor
    }
}