// Copyright 2021 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.

use std::{
    sync::{
        Arc,
        atomic::{AtomicU64, Ordering},
    },
    time::{Duration, SystemTime, UNIX_EPOCH},
};

use api_version::{KvFormat, api_v2::DEFAULT_KEYSPACE_ID};
use futures::{
    compat::Future01CompatExt,
    future::{self, Future, FutureExt, TryFutureExt},
    sink::SinkExt,
    stream::{StreamExt, TryStreamExt},
};
use grpcio::{
    ClientStreamingSink, DuplexSink, Error as GrpcError, RequestStream, Result as GrpcResult,
    RpcContext, RpcStatus, RpcStatusCode, ServerStreamingSink, UnarySink, WriteFlags,
};
use health_controller::HealthController;
use kvproto::{
    coprocessor::*,
    kvrpcpb::*,
    mpp::{
        CancelTaskRequest, CancelTaskResponse, DispatchTaskRequest, DispatchTaskResponse,
        EstablishMppConnectionRequest, IsAliveRequest, IsAliveResponse, MppDataPacket,
    },
    raft_cmdpb::{CmdType, RaftCmdRequest, RaftRequestHeader, Request as RaftRequest},
    raft_serverpb::*,
    tikvpb::*,
};
use log_wrappers::hex_encode;
use protobuf::Message;
use rfstore::{
    Error as RaftStoreError,
    router::RaftStoreRouter,
    store::{Callback, CasualMessage, CheckLeaderTask, PeerTag, RegionIdVer, SplitSource},
};
use tikv::{
    coprocessor::Endpoint,
    forward_duplex, forward_unary,
    server::{Proxy, load_statistics::ThreadLoadPool, metrics::*},
    storage::{
        SecondaryLocksStatus, Storage, TxnStatus,
        errors::{
            extract_committed, extract_key_error, extract_key_errors, extract_region_error,
            map_kv_pairs,
        },
        kv::Engine,
        lock_manager::LockManager,
    },
};
use tikv_kv::StageLatencyStats;
use tikv_util::{
    future::{paired_future_callback, poll_future_notify},
    mpsc::future::{BatchReceiver, Sender, WakePolicy, unbounded},
    time::{Instant, nanos_to_secs},
    worker::Scheduler,
};
use trace_event::types::TraceContext;
use tracker::{
    GLOBAL_TRACKERS, RequestInfo, RequestType, Tracker, set_tls_trace_ctx, set_tls_tracker_token,
    with_tls_tracker,
};
use txn_types::{self, Key};

use super::batch::{BatcherBuilder, ReqBatcher};
use crate::server::{Error, Result as ServerResult};

const GRPC_MSG_MAX_BATCH_SIZE: usize = 128;
const GRPC_MSG_NOTIFY_SIZE: usize = 8;
const UNIMPLEMENTED_MSG: &str = "not implemented";

macro_rules! has_data_in_prefix {
    ($snap:expr, $prefix:expr) => {{
        if $snap.is_sync() {
            $snap.has_data_in_prefix($prefix)
        } else {
            $snap.has_data_in_prefix_async($prefix).await
        }
    }};
}

/// Service handles the RPC messages for the `Tikv` service.
pub struct Service<T: RaftStoreRouter, L: LockManager, F: KvFormat> {
    store_id: u64,
    // For handling KV requests.
    storage: Storage<RaftKv, L, F>,
    // For handling coprocessor requests.
    copr: Endpoint<RaftKv>,
    // For handling raft messages.
    ch: T,

    enable_req_batch: bool,

    grpc_thread_load: Arc<ThreadLoadPool>,

    proxy: Proxy,

    // For handling `CheckLeader` request.
    check_leader_scheduler: Option<Scheduler<CheckLeaderTask>>,

    // For returning health status to clients.
    health_controller: HealthController,
    health_feedback_interval: Option<Duration>,
    health_feedback_seq: Arc<AtomicU64>,
}

impl<T: RaftStoreRouter + Clone + 'static, L: LockManager + Clone, F: KvFormat + Clone> Clone
    for Service<T, L, F>
{
    fn clone(&self) -> Self {
        Service {
            store_id: self.store_id,
            storage: self.storage.clone(),
            copr: self.copr.clone(),
            ch: self.ch.clone(),
            enable_req_batch: self.enable_req_batch,
            grpc_thread_load: self.grpc_thread_load.clone(),
            proxy: self.proxy.clone(),
            check_leader_scheduler: self.check_leader_scheduler.clone(),
            health_controller: self.health_controller.clone(),
            health_feedback_seq: self.health_feedback_seq.clone(),
            health_feedback_interval: self.health_feedback_interval,
        }
    }
}

impl<T: RaftStoreRouter + 'static, L: LockManager, F: KvFormat> Service<T, L, F> {
    /// Constructs a new `Service` which provides the `Tikv` service.
    pub fn new(
        store_id: u64,
        storage: Storage<RaftKv, L, F>,
        copr: Endpoint<RaftKv>,
        ch: T,
        grpc_thread_load: Arc<ThreadLoadPool>,
        enable_req_batch: bool,
        proxy: Proxy,
        check_leader_scheduler: Option<Scheduler<CheckLeaderTask>>,
        health_controller: HealthController,
        health_feedback_interval: Option<Duration>,
    ) -> Self {
        let now_unix = std::time::SystemTime::now()
            .duration_since(std::time::UNIX_EPOCH)
            .unwrap()
            .as_millis() as u64;
        Service {
            store_id,
            storage,
            copr,
            ch,
            enable_req_batch,
            grpc_thread_load,
            proxy,
            check_leader_scheduler,
            health_controller,
            health_feedback_interval,
            health_feedback_seq: Arc::new(AtomicU64::new(now_unix)),
        }
    }
}

macro_rules! handle_request {
    ($fn_name: ident, $future_name: ident, $req_ty: ident, $resp_ty: ident) => {
        handle_request!($fn_name, $future_name, $req_ty, $resp_ty, no_time_detail);
    };
    ($fn_name: ident, $future_name: ident, $req_ty: ident, $resp_ty: ident, $time_detail: tt) => {
        fn $fn_name(&mut self, ctx: RpcContext<'_>, req: $req_ty, sink: UnarySink<$resp_ty>) {
            forward_unary!(self.proxy, $fn_name, ctx, req, sink);
            let begin_instant = Instant::now_coarse();
            // trace id should not be added here, it should be added to future_xxx rather than kv_xxx
            // That's because there are batch commands, and future_xxx are the common entry for both.
            let resp = $future_name(&self.storage, req);
            let task = async move {
                let resp = resp.await?;
                let elapsed = begin_instant.saturating_elapsed();
                set_total_time!(resp, elapsed, $time_detail);
                sink.success(resp).await?;
                GRPC_MSG_HISTOGRAM_STATIC
                    .$fn_name
                    .observe(elapsed.as_secs_f64());
                ServerResult::Ok(())
            }
            .map_err(|e| {
                debug!("kv rpc failed";
                    "request" => stringify!($fn_name),
                    "err" => ?e
                );
                GRPC_MSG_FAIL_COUNTER.$fn_name.inc();
            })
            .map(|_|());

            ctx.spawn(task);
        }
    }
}

macro_rules! fail_unimplemented {
    ($ctx:expr, $sink:expr, $name:expr) => {{
        let status = RpcStatus::new(RpcStatusCode::UNIMPLEMENTED);
        $ctx.spawn(
            $sink
                .fail(status)
                .unwrap_or_else(|e| error!("kv rpc failed"; "request" => $name, "err" => ?e)),
        );
    }};
}

macro_rules! set_total_time {
    ($resp:ident, $duration:expr,no_time_detail) => {};
    ($resp:ident, $duration:expr,has_time_detail) => {
        let mut $resp = $resp;
        $resp
            .mut_exec_details_v2()
            .mut_time_detail()
            .set_total_rpc_wall_time_ns($duration.as_nanos() as u64);
        $resp
            .mut_exec_details_v2()
            .mut_time_detail_v2()
            .set_total_rpc_wall_time_ns($duration.as_nanos() as u64);
    };
}

impl<T: RaftStoreRouter + 'static, L: LockManager, F: KvFormat> Tikv for Service<T, L, F> {
    handle_request!(kv_get, future_get, GetRequest, GetResponse, has_time_detail);
    handle_request!(kv_scan, future_scan, ScanRequest, ScanResponse);
    handle_request!(
        kv_prewrite,
        future_prewrite,
        PrewriteRequest,
        PrewriteResponse,
        has_time_detail
    );
    handle_request!(
        kv_pessimistic_lock,
        future_acquire_pessimistic_lock,
        PessimisticLockRequest,
        PessimisticLockResponse,
        has_time_detail
    );
    handle_request!(
        kv_pessimistic_rollback,
        future_pessimistic_rollback,
        PessimisticRollbackRequest,
        PessimisticRollbackResponse,
        has_time_detail
    );
    handle_request!(
        kv_commit,
        future_commit,
        CommitRequest,
        CommitResponse,
        has_time_detail
    );
    handle_request!(kv_cleanup, future_cleanup, CleanupRequest, CleanupResponse);
    handle_request!(
        kv_batch_get,
        future_batch_get,
        BatchGetRequest,
        BatchGetResponse
    );
    handle_request!(
        kv_batch_rollback,
        future_batch_rollback,
        BatchRollbackRequest,
        BatchRollbackResponse,
        has_time_detail
    );
    handle_request!(
        kv_txn_heart_beat,
        future_txn_heart_beat,
        TxnHeartBeatRequest,
        TxnHeartBeatResponse,
        has_time_detail
    );
    handle_request!(
        kv_check_txn_status,
        future_check_txn_status,
        CheckTxnStatusRequest,
        CheckTxnStatusResponse,
        has_time_detail
    );
    handle_request!(
        kv_check_secondary_locks,
        future_check_secondary_locks,
        CheckSecondaryLocksRequest,
        CheckSecondaryLocksResponse,
        has_time_detail
    );
    handle_request!(
        kv_scan_lock,
        future_scan_lock,
        ScanLockRequest,
        ScanLockResponse,
        has_time_detail
    );
    handle_request!(
        kv_resolve_lock,
        future_resolve_lock,
        ResolveLockRequest,
        ResolveLockResponse,
        has_time_detail
    );
    handle_request!(
        kv_delete_range,
        future_delete_range,
        DeleteRangeRequest,
        DeleteRangeResponse
    );
    handle_request!(
        mvcc_get_by_key,
        future_mvcc_get_by_key,
        MvccGetByKeyRequest,
        MvccGetByKeyResponse
    );
    handle_request!(
        mvcc_get_by_start_ts,
        future_mvcc_get_by_start_ts,
        MvccGetByStartTsRequest,
        MvccGetByStartTsResponse
    );

    fn kv_import(
        &mut self,
        ctx: RpcContext<'_>,
        _: ImportRequest,
        sink: UnarySink<ImportResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(kv_import));
    }

    fn kv_gc(&mut self, ctx: RpcContext<'_>, _: GcRequest, sink: UnarySink<GcResponse>) {
        fail_unimplemented!(ctx, sink, stringify!(kv_gc));
    }

    fn coprocessor(&mut self, ctx: RpcContext<'_>, req: Request, sink: UnarySink<Response>) {
        forward_unary!(self.proxy, coprocessor, ctx, req, sink);
        let begin_instant = Instant::now_coarse();
        let future = future_copr(&self.copr, Some(ctx.peer()), req);
        let task = async move {
            let resp = future.await?.consume();
            sink.success(resp).await?;
            GRPC_MSG_HISTOGRAM_STATIC
                .coprocessor
                .observe(begin_instant.saturating_elapsed().as_secs_f64());
            ServerResult::Ok(())
        }
        .map_err(|e| {
            debug!("kv rpc failed";
                "request" => "coprocessor",
                "err" => ?e
            );
            GRPC_MSG_FAIL_COUNTER.coprocessor.inc();
        })
        .map(|_| ());

        ctx.spawn(task);
    }

    fn delegate_coprocessor(
        &mut self,
        ctx: RpcContext<'_>,
        req: DelegateRequest,
        sink: UnarySink<DelegateResponse>,
    ) {
        let begin_instant = Instant::now_coarse();
        let future = self.copr.handle_delegate_request(req);
        let task = async move {
            let resp = future.await;
            sink.success(resp).await?;
            GRPC_MSG_HISTOGRAM_STATIC
                .delegate_coprocessor
                .observe(begin_instant.saturating_elapsed().as_secs_f64());
            ServerResult::Ok(())
        }
        .map_err(|e| {
            debug!("kv rpc failed";
                "request" => "delegate_coprocessor",
                "err" => ?e
            );
            GRPC_MSG_FAIL_COUNTER.delegate_coprocessor.inc();
        })
        .map(|_| ());

        ctx.spawn(task);
    }

    fn unsafe_destroy_range(
        &mut self,
        ctx: RpcContext<'_>,
        mut req: UnsafeDestroyRangeRequest,
        sink: UnarySink<UnsafeDestroyRangeResponse>,
    ) {
        let begin_instant = Instant::now_coarse();
        let start = req.take_start_key();
        let end = req.take_end_key();
        // DestroyRange is a very dangerous operation. We don't allow passing MIN_KEY as
        // start, or MAX_KEY as end here.
        if start.is_empty()
            || end.is_empty()
            // All unsafe destroy range issued by TiDB are prefix range.
            // To simplify the logic, we only handle prefix delete range.
            || !tidb_query_common::util::is_prefix_next(&start, &end)
        {
            ctx.spawn(
                sink.fail(RpcStatus::new(RpcStatusCode::INVALID_ARGUMENT))
                    .unwrap_or_else(|e| error!("kv rpc failed"; "err" => ?e)),
            );
            return;
        }
        let prefix = start.clone();
        let (callback, future) = paired_future_callback();
        self.ch.send_store_msg(StoreMsg::GetRegionsInRange {
            start,
            end,
            callback,
        });
        let ch = self.ch.clone();
        let kv = self.storage.get_engine().kv_engine().unwrap();
        let store_id = kv.get_engine_id();
        let task = async move {
            let mut resp = UnsafeDestroyRangeResponse::default();
            let regions = future.await?;
            info!(
                "{} unsafe_destroy_range: range prefix {:?} for regions {:?}",
                store_id, prefix, regions
            );
            let mut region_futures = vec![];
            let set_snapshot_error =
                |resp: &mut UnsafeDestroyRangeResponse, region: RegionIdVer, phase, err| {
                    warn!(
                        "{} unsafe_destroy_range: failed to check region snapshot {} delete prefix",
                        store_id, phase;
                        "region_id" => region.id(),
                        "region_ver" => region.ver(),
                        "err" => %err,
                    );
                    let mut errpb = kvproto::errorpb::Error::default();
                    errpb.set_message(format!(
                        "failed to check region {} snapshot {} delete range: {}",
                        region.id(),
                        phase,
                        err
                    ));
                    resp.set_region_error(errpb);
                };
            for region in &regions {
                tikv_util::set_current_region(region.id());
                let snap = match kv.try_get_snap_access(region.id(), region.ver()) {
                    Ok(Some(snap)) => snap,
                    Ok(None) => continue,
                    Err(err) => {
                        set_snapshot_error(&mut resp, *region, "before", err);
                        break;
                    }
                };
                if !has_data_in_prefix!(snap, &prefix) {
                    continue;
                }

                let region = *region;
                let ch = ch.clone();
                let prefix = prefix.clone();
                let region_task = async move {
                    let tag = PeerTag::new(store_id, region);
                    info!(
                        "{} unsafe_destroy_range: request delete prefix {:?}",
                        tag, prefix
                    );
                    let (cb, fu) = paired_future_callback();
                    let callback = Callback::write(Box::new(move |res| {
                        cb(res);
                    }));
                    ch.send_casual_msg(
                        region.id(),
                        CasualMessage::DeletePrefix {
                            region_version: region.ver(),
                            prefix,
                            callback,
                        },
                    );
                    match fu.await {
                        Ok(mut res) => {
                            if res.response.get_header().has_error() {
                                return Err(res.response.mut_header().take_error());
                            }
                        }
                        Err(e) => {
                            // Should happen only when TiKV is stopping.
                            // Can be ignored as the request is sent to all stores.
                            info!(
                                "{} unsafe_destroy_range: request delete prefix canceled: {:?}",
                                tag, e
                            );
                        }
                    }
                    Ok(())
                };
                region_futures.push(region_task);
            }

            tikv_util::set_current_region(0);
            if !resp.has_region_error() {
                if let Err(e) = futures::future::join_all(region_futures)
                    .await
                    .into_iter()
                    .collect::<Result<Vec<_>, _>>()
                {
                    resp.set_region_error(e);
                }
            }
            if !resp.has_region_error() {
                // Wait and check if all regions have applied delete prefix.
                let mut regions_applied = false;
                for i in 1..=5 {
                    let mut applied = true;
                    for region in &regions {
                        tikv_util::set_current_region(region.id());
                        let snap = match kv.try_get_snap_access(region.id(), region.ver()) {
                            Ok(Some(snap)) => snap,
                            Ok(None) => continue,
                            Err(err) => {
                                set_snapshot_error(&mut resp, *region, "after", err);
                                applied = false;
                                break;
                            }
                        };
                        if has_data_in_prefix!(snap, &prefix) {
                            applied = false;
                            break;
                        }
                    }
                    tikv_util::set_current_region(0);
                    if resp.has_region_error() {
                        break;
                    }
                    if !applied {
                        let deadline = std::time::Instant::now() + Duration::from_secs(i);
                        let _ = GLOBAL_TIMER_HANDLE.delay(deadline).compat().await.is_ok();
                        continue;
                    }
                    regions_applied = true;
                    break;
                }
                if regions_applied {
                    // At last, make sure all shards exists and not split.
                    for region in &regions {
                        let res = kv.get_shard_with_ver(region.id(), region.ver());
                        if res.is_err() {
                            let mut errpb = kvproto::errorpb::Error::default();
                            errpb.set_message(format!(
                                "region {} changed during delete range",
                                region.id()
                            ));
                            errpb.set_epoch_not_match(kvproto::errorpb::EpochNotMatch::default());
                            resp.set_region_error(errpb);
                            break;
                        }
                    }
                } else {
                    resp.set_error("some region failed to delete range".to_string());
                }
            }
            sink.success(resp).await?;
            GRPC_MSG_HISTOGRAM_STATIC
                .unsafe_destroy_range
                .observe(begin_instant.saturating_elapsed().as_secs_f64());
            ServerResult::Ok(())
        }
        .map_err(|e| {
            debug!("kv rpc failed";
                "request" => "unsafe_destroy_range",
                "err" => ?e
            );
            GRPC_MSG_FAIL_COUNTER.unsafe_destroy_range.inc();
        })
        .map(|_| ());
        self.storage.get_scheduler().force_spawn(
            task,
            // `unsafe_destroy_range` is an administrative internal task, not
            // tenant request work that should be charged to a specific
            // keyspace by read-pool fairness.
            DEFAULT_KEYSPACE_ID,
            0,
        );
    }

    fn coprocessor_stream(
        &mut self,
        ctx: RpcContext<'_>,
        req: Request,
        mut sink: ServerStreamingSink<Response>,
    ) {
        let begin_instant = Instant::now_coarse();

        let mut stream = self
            .copr
            .parse_and_handle_stream_request(req, Some(ctx.peer()))
            .map(|resp| {
                GrpcResult::<(Response, WriteFlags)>::Ok((
                    resp,
                    WriteFlags::default().buffer_hint(true),
                ))
            });
        let future = async move {
            match sink.send_all(&mut stream).await.map_err(Error::from) {
                Ok(_) => {
                    GRPC_MSG_HISTOGRAM_STATIC
                        .coprocessor_stream
                        .observe(begin_instant.saturating_elapsed().as_secs_f64());
                    let _ = sink.close().await;
                }
                Err(e) => {
                    debug!("kv rpc failed";
                        "request" => "coprocessor_stream",
                        "err" => ?e
                    );
                    GRPC_MSG_FAIL_COUNTER.coprocessor_stream.inc();
                }
            }
        };

        ctx.spawn(future);
    }

    fn raft(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<RaftMessage>,
        sink: ClientStreamingSink<Done>,
    ) {
        let store_id = self.store_id;
        let ch = self.ch.clone();
        ctx.spawn(async move {
            let res = stream.map_err(Error::from).try_for_each(move |msg| {
                RAFT_MESSAGE_RECV_COUNTER.inc();
                let to_store_id = msg.get_to_peer().get_store_id();
                if to_store_id != store_id {
                    future::err(Error::from(RaftStoreError::StoreNotMatch {
                        to_store_id,
                        my_store_id: store_id,
                    }))
                } else {
                    ch.send_raft_msg(msg);
                    future::ready(Ok(()))
                }
            });
            let status = match res.await {
                Err(e) => {
                    let msg = format!("{:?}", e);
                    error!("dispatch raft msg from gRPC to raftstore fail"; "err" => %msg);
                    RpcStatus::with_message(RpcStatusCode::UNKNOWN, msg)
                }
                Ok(_) => RpcStatus::new(RpcStatusCode::UNKNOWN),
            };
            let _ = sink
                .fail(status)
                .map_err(|e| error!("KvService::raft send response fail"; "err" => ?e))
                .await;
        });
    }

    fn batch_raft(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<BatchRaftMessage>,
        sink: ClientStreamingSink<Done>,
    ) {
        info!("batch_raft RPC is called, new gRPC stream established");
        let store_id = self.store_id;
        let ch = self.ch.clone();
        ctx.spawn(async move {
            let res = stream
                .map_err(Error::from)
                .try_for_each(move |mut batch_msg| {
                    let now = SystemTime::now()
                        .duration_since(UNIX_EPOCH)
                        .unwrap()
                        .as_nanos() as u64;
                    let elapsed = nanos_to_secs(now.saturating_sub(batch_msg.last_observed_time));
                    RAFT_MESSAGE_DURATION.receive_delay.observe(elapsed);

                    let len = batch_msg.get_msgs().len();
                    RAFT_MESSAGE_RECV_COUNTER.inc_by(len as u64);
                    RAFT_MESSAGE_BATCH_SIZE.observe(len as f64);
                    for msg in batch_msg.take_msgs().into_iter() {
                        let to_store_id = msg.get_to_peer().get_store_id();
                        if to_store_id != store_id {
                            return future::err(Error::from(RaftStoreError::StoreNotMatch {
                                to_store_id,
                                my_store_id: store_id,
                            }));
                        }
                        ch.send_raft_msg(msg);
                    }
                    future::ok(())
                });
            let status = match res.await {
                Err(e) => {
                    let msg = format!("{:?}", e);
                    error!("dispatch raft msg from gRPC to raftstore fail"; "err" => %msg);
                    RpcStatus::with_message(RpcStatusCode::UNKNOWN, msg)
                }
                Ok(_) => RpcStatus::new(RpcStatusCode::UNKNOWN),
            };
            let _ = sink
                .fail(status)
                .map_err(|e| error!("KvService::batch_raft send response fail"; "err" => ?e))
                .await;
        });
    }

    fn split_region(
        &mut self,
        ctx: RpcContext<'_>,
        mut req: SplitRegionRequest,
        sink: UnarySink<SplitRegionResponse>,
    ) {
        forward_unary!(self.proxy, split_region, ctx, req, sink);
        let begin_instant = Instant::now_coarse();

        let region_id = req.get_context().get_region_id();
        let (cb, f) = paired_future_callback();
        let mut split_keys = if !req.get_split_key().is_empty() {
            vec![Key::from_raw(req.get_split_key()).into_encoded()]
        } else {
            req.take_split_keys()
                .into_iter()
                .map(|x| Key::from_raw(&x).into_encoded())
                .collect()
        };
        split_keys.sort();
        let req = CasualMessage::SplitRegion {
            region_epoch: req.take_context().take_region_epoch(),
            split_keys,
            callback: Callback::write(cb),
            source: SplitSource::Peer(ctx.peer()),
        };

        self.ch.send_casual_msg(region_id, req);

        let task = async move {
            let mut res = f.await?;
            let mut resp = SplitRegionResponse::default();
            if let Some(key_errs) = res.key_errors {
                resp.set_errors(key_errs.into());
            } else if res.response.get_header().has_error() {
                resp.set_region_error(res.response.mut_header().take_error());
            } else {
                let admin_resp = res.response.mut_admin_response();
                let regions: Vec<_> = admin_resp.mut_splits().take_regions().into();
                if regions.len() < 2 {
                    error!(
                        "invalid split response";
                        "region_id" => region_id,
                        "resp" => ?admin_resp
                    );
                    resp.mut_region_error().set_message(format!(
                        "Internal Error: invalid response: {:?}",
                        admin_resp
                    ));
                } else {
                    if regions.len() == 2 {
                        resp.set_left(regions[0].clone());
                        resp.set_right(regions[1].clone());
                    }
                    resp.set_regions(regions.into());
                }
            }
            sink.success(resp).await?;
            GRPC_MSG_HISTOGRAM_STATIC
                .split_region
                .observe(begin_instant.saturating_elapsed().as_secs_f64());
            ServerResult::Ok(())
        }
        .map_err(|e| {
            debug!("kv rpc failed";
                "request" => "split_region",
                "err" => ?e
            );
            GRPC_MSG_FAIL_COUNTER.split_region.inc();
        })
        .map(|_| ());

        ctx.spawn(task);
    }

    fn read_index(
        &mut self,
        ctx: RpcContext<'_>,
        req: ReadIndexRequest,
        sink: UnarySink<ReadIndexResponse>,
    ) {
        forward_unary!(self.proxy, read_index, ctx, req, sink);
        let begin_instant = Instant::now_coarse();

        let region_id = req.get_context().get_region_id();
        let mut cmd = RaftCmdRequest::default();
        let mut header = RaftRequestHeader::default();
        let mut inner_req = RaftRequest::default();
        inner_req.set_cmd_type(CmdType::ReadIndex);
        inner_req.mut_read_index().set_start_ts(req.get_start_ts());
        for r in req.get_ranges() {
            let mut range = kvproto::kvrpcpb::KeyRange::default();
            range.set_start_key(Key::from_raw(r.get_start_key()).into_encoded());
            range.set_end_key(Key::from_raw(r.get_end_key()).into_encoded());
            inner_req.mut_read_index().mut_key_ranges().push(range);
        }
        header.set_region_id(req.get_context().get_region_id());
        header.set_peer(req.get_context().get_peer().clone());
        header.set_region_epoch(req.get_context().get_region_epoch().clone());
        if req.get_context().get_term() != 0 {
            header.set_term(req.get_context().get_term());
        }
        header.set_sync_log(req.get_context().get_sync_log());
        header.set_read_quorum(true);
        cmd.set_header(header);
        cmd.set_requests(vec![inner_req].into());

        let (cb, f) = paired_future_callback();
        let trace_ctx = TraceContext::from_proto(
            req.get_context().get_trace_id(),
            req.get_context().get_trace_control_flags(),
        );

        // We must deal with all requests which acquire read-quorum in raftstore-thread,
        // so just send it as an command.
        self.ch.send_command(cmd, trace_ctx, Callback::read(cb));

        let task = async move {
            let mut res = f.await?;
            let mut resp = ReadIndexResponse::default();
            if res.response.get_header().has_error() {
                resp.set_region_error(res.response.mut_header().take_error());
            } else {
                let mut raft_resps = res.response.take_responses();
                if raft_resps.len() != 1 {
                    error!(
                        "invalid read index response";
                        "region_id" => region_id,
                        "response" => ?raft_resps
                    );
                    resp.mut_region_error().set_message(format!(
                        "Internal Error: invalid response: {:?}",
                        raft_resps
                    ));
                } else {
                    let mut read_index_resp = raft_resps[0].take_read_index();
                    if read_index_resp.has_locked() {
                        resp.set_locked(read_index_resp.take_locked());
                    } else {
                        resp.set_read_index(read_index_resp.get_read_index());
                    }
                }
            }
            sink.success(resp).await?;
            GRPC_MSG_HISTOGRAM_STATIC
                .read_index
                .observe(begin_instant.saturating_elapsed_secs());
            ServerResult::Ok(())
        }
        .map_err(|e| {
            debug!("kv rpc failed";
                "request" => "read_index",
                "err" => ?e
            );
            GRPC_MSG_FAIL_COUNTER.read_index.inc();
        })
        .map(|_| ());

        ctx.spawn(task);
    }

    fn batch_commands(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<BatchCommandsRequest>,
        mut sink: DuplexSink<BatchCommandsResponse>,
    ) {
        forward_duplex!(self.proxy, batch_commands, ctx, stream, sink);
        let (tx, rx) = unbounded(WakePolicy::TillReach(GRPC_MSG_NOTIFY_SIZE));

        let ctx = Arc::new(ctx);
        let peer = ctx.peer();
        let storage = self.storage.clone();
        let copr = self.copr.clone();
        let pool_size = storage.get_normal_pool_size();
        let batch_builder = BatcherBuilder::new(self.enable_req_batch, pool_size);
        let mut health_feedback_attacher = HealthFeedbackAttacher::new(
            self.store_id,
            self.health_controller.clone(),
            self.health_feedback_seq.clone(),
            self.health_feedback_interval,
        );
        let request_handler = stream.try_for_each(move |mut req| {
            let request_ids = req.take_request_ids();
            let requests: Vec<_> = req.take_requests().into();
            let queue = storage.get_readpool_queue_per_worker();
            let mut batcher = batch_builder.build(queue, request_ids.len());
            GRPC_REQ_BATCH_COMMANDS_SIZE.observe(requests.len() as f64);
            for (id, req) in request_ids.into_iter().zip(requests) {
                handle_batch_commands_request(&mut batcher, &storage, &copr, &peer, id, req, &tx);
                if let Some(batch) = batcher.as_mut() {
                    batch.maybe_commit(&storage, &tx);
                }
            }
            if let Some(batch) = batcher {
                batch.commit(&storage, &tx);
            }
            future::ok(())
        });
        ctx.spawn(request_handler.unwrap_or_else(|e| error!("batch_commands error"; "err" => %e)));

        let grpc_thread_load = Arc::clone(&self.grpc_thread_load);
        let response_retriever = BatchReceiver::new(
            rx,
            GRPC_MSG_MAX_BATCH_SIZE,
            MeasuredBatchResponse::default,
            collect_batch_resp,
        );

        let mut response_retriever = response_retriever.map(move |mut item| {
            handle_measures_for_batch_commands(&mut item);
            let mut r = item.batch_resp;
            GRPC_RESP_BATCH_COMMANDS_SIZE.observe(r.request_ids.len() as f64);
            r.set_transport_layer_load(grpc_thread_load.total_load() as u64);
            health_feedback_attacher.attach_if_needed(&mut r);
            GrpcResult::<(BatchCommandsResponse, WriteFlags)>::Ok((
                r,
                WriteFlags::default().buffer_hint(false),
            ))
        });

        let send_task = async move {
            sink.send_all(&mut response_retriever).await?;
            sink.close().await?;
            Ok(())
        }
        .map_err(|e: grpcio::Error| {
            debug!("kv rpc failed";
                "request" => "batch_commands",
                "err" => ?e
            );
        })
        .map(|_| ());

        ctx.spawn(send_task);
    }

    fn batch_coprocessor(
        &mut self,
        ctx: RpcContext<'_>,
        _req: BatchRequest,
        sink: ServerStreamingSink<BatchResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(batch_coprocessor));
    }

    fn check_leader(
        &mut self,
        ctx: RpcContext<'_>,
        mut request: CheckLeaderRequest,
        sink: UnarySink<CheckLeaderResponse>,
    ) {
        let check_leader_scheduler = if let Some(s) = &self.check_leader_scheduler {
            s.clone()
        } else {
            fail_unimplemented!(ctx, sink, stringify!(check_leader));
            return;
        };
        let begin_instant = Instant::now();
        let addr = ctx.peer();
        let ts = request.get_ts();
        let leaders = request.take_regions().into();
        let (cb, resp) = paired_future_callback();
        let task = async move {
            check_leader_scheduler
                .schedule(CheckLeaderTask::CheckLeader { leaders, cb })
                .map_err(|e| Error::Other(format!("{}", e).into()))?;
            let regions = resp.await?;
            GRPC_MSG_HISTOGRAM_STATIC
                .check_leader
                .observe(begin_instant.saturating_elapsed().as_secs_f64());
            let mut resp = CheckLeaderResponse::default();
            resp.set_ts(ts);
            resp.set_regions(regions);
            if let Err(e) = sink.success(resp).await {
                // CheckLeader has a built-in fast-success mechanism, so `RemoteStopped`
                // can be treated as a general situation.
                if let GrpcError::RemoteStopped = e {
                    return ServerResult::Ok(());
                }
                return Err(Error::from(e));
            }
            ServerResult::Ok(())
        }
        .map_err(move |e| {
            // CheckLeader only needs quorum responses, remote may drops
            // requests early.
            info!("call CheckLeader failed"; "err" => ?e, "address" => addr);
        })
        .map(|_| ());

        ctx.spawn(task);
    }

    fn get_store_safe_ts(
        &mut self,
        ctx: RpcContext<'_>,
        _request: StoreSafeTsRequest,
        sink: UnarySink<StoreSafeTsResponse>,
    ) {
        // TODO(x) implement it.
        let task = async move {
            let response = StoreSafeTsResponse::default();
            sink.success(response).await?;
            ServerResult::Ok(())
        }
        .map_err(|e| {
            warn!("call get_store_safe_ts failed"; "err" => ?e);
        })
        .map(|_| ());
        ctx.spawn(task);
    }

    fn get_lock_wait_info(
        &mut self,
        ctx: RpcContext<'_>,
        request: GetLockWaitInfoRequest,
        sink: UnarySink<GetLockWaitInfoResponse>,
    ) {
        let (cb, f) = paired_future_callback();
        self.storage
            .dump_wait_for_entries(request.get_context(), cb);
        let task = async move {
            let res = f.await??;
            let mut response = GetLockWaitInfoResponse::default();
            response.set_entries(RepeatedField::from_vec(res));
            sink.success(response).await?;
            ServerResult::Ok(())
        }
        .map_err(|e| {
            warn!("call dump_wait_for_entries failed"; "err" => ?e);
        })
        .map(|_| ());
        ctx.spawn(task);
    }

    fn get_health_feedback(
        &mut self,
        ctx: RpcContext<'_>,
        _request: GetHealthFeedbackRequest,
        sink: UnarySink<GetHealthFeedbackResponse>,
    ) {
        let attacher = HealthFeedbackAttacher::new(
            self.store_id,
            self.health_controller.clone(),
            self.health_feedback_seq.clone(),
            self.health_feedback_interval,
        );

        let mut resp = GetHealthFeedbackResponse::default();
        resp.set_health_feedback(attacher.gen_health_feedback_pb());
        let task = sink
            .success(resp)
            .map_err(|e| {
                warn!("get_health_feedback failed"; "err" => ?e);
            })
            .map(|_| ());
        ctx.spawn(task);
    }

    fn raw_get(&mut self, ctx: RpcContext<'_>, _: RawGetRequest, sink: UnarySink<RawGetResponse>) {
        fail_unimplemented!(ctx, sink, stringify!(raw_get));
    }

    fn raw_batch_get(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawBatchGetRequest,
        sink: UnarySink<RawBatchGetResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_batch_get));
    }

    fn raw_put(&mut self, ctx: RpcContext<'_>, _: RawPutRequest, sink: UnarySink<RawPutResponse>) {
        fail_unimplemented!(ctx, sink, stringify!(raw_put));
    }

    fn raw_batch_put(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawBatchPutRequest,
        sink: UnarySink<RawBatchPutResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_batch_put));
    }

    fn raw_delete(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawDeleteRequest,
        sink: UnarySink<RawDeleteResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_delete));
    }

    fn raw_batch_delete(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawBatchDeleteRequest,
        sink: UnarySink<RawBatchDeleteResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_batch_delete));
    }

    fn raw_scan(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawScanRequest,
        sink: UnarySink<RawScanResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_scan));
    }

    fn raw_delete_range(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawDeleteRangeRequest,
        sink: UnarySink<RawDeleteRangeResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_delete_range));
    }

    fn raw_batch_scan(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawBatchScanRequest,
        sink: UnarySink<RawBatchScanResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_batch_scan));
    }

    fn raw_get_key_ttl(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawGetKeyTtlRequest,
        sink: UnarySink<RawGetKeyTtlResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_get_key_ttl));
    }

    fn raw_compare_and_swap(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawCasRequest,
        sink: UnarySink<RawCasResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_compare_and_swap));
    }

    fn raw_checksum(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawChecksumRequest,
        sink: UnarySink<RawChecksumResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_checksum));
    }

    fn register_lock_observer(
        &mut self,
        ctx: RpcContext<'_>,
        _: RegisterLockObserverRequest,
        sink: UnarySink<RegisterLockObserverResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(register_lock_observer));
    }

    fn check_lock_observer(
        &mut self,
        ctx: RpcContext<'_>,
        _: CheckLockObserverRequest,
        sink: UnarySink<CheckLockObserverResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(check_lock_observer));
    }

    fn remove_lock_observer(
        &mut self,
        ctx: RpcContext<'_>,
        _: RemoveLockObserverRequest,
        sink: UnarySink<RemoveLockObserverResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(remove_lock_observer));
    }

    fn physical_scan_lock(
        &mut self,
        ctx: RpcContext<'_>,
        _: PhysicalScanLockRequest,
        sink: UnarySink<PhysicalScanLockResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(physical_scan_lock));
    }

    fn raw_coprocessor(
        &mut self,
        ctx: RpcContext<'_>,
        _: RawCoprocessorRequest,
        sink: UnarySink<RawCoprocessorResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(raw_coprocessor));
    }

    fn snapshot(
        &mut self,
        ctx: RpcContext<'_>,
        _: RequestStream<SnapshotChunk>,
        sink: ClientStreamingSink<Done>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(snapshot));
    }

    fn dispatch_mpp_task(
        &mut self,
        ctx: RpcContext<'_>,
        _: DispatchTaskRequest,
        sink: UnarySink<DispatchTaskResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(dispatch_mpp_task));
    }

    fn cancel_mpp_task(
        &mut self,
        ctx: RpcContext<'_>,
        _: CancelTaskRequest,
        sink: UnarySink<CancelTaskResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(cancel_mpp_task));
    }

    fn establish_mpp_connection(
        &mut self,
        ctx: RpcContext<'_>,
        _: EstablishMppConnectionRequest,
        sink: ServerStreamingSink<MppDataPacket>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(establish_mpp_connection));
    }

    fn is_alive(
        &mut self,
        ctx: RpcContext<'_>,
        _: IsAliveRequest,
        sink: UnarySink<IsAliveResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(is_alive));
    }

    fn compact(
        &mut self,
        ctx: RpcContext<'_>,
        _: CompactRequest,
        sink: UnarySink<CompactResponse>,
    ) {
        fail_unimplemented!(ctx, sink, stringify!(compact));
    }
}

fn unimplemented_region_error() -> kvproto::errorpb::Error {
    let mut err = kvproto::errorpb::Error::default();
    err.set_message(UNIMPLEMENTED_MSG.to_owned());
    err
}

fn respond_unimplemented_batch(
    id: u64,
    tx: &Sender<MeasuredSingleResponse>,
    label: GrpcTypeKind,
    cmd: batch_commands_response::response::Cmd,
) {
    let begin_instant = Instant::now();
    let resp = future::ok(batch_commands_response::Response {
        cmd: Some(cmd),
        ..Default::default()
    });
    response_batch_commands_request(id, resp, tx.clone(), begin_instant, label);
}

fn response_batch_commands_request<F, T>(
    id: u64,
    resp: F,
    tx: Sender<MeasuredSingleResponse>,
    begin: Instant,
    label: GrpcTypeKind,
) where
    MemoryTraceGuard<batch_commands_response::Response>: From<T>,
    F: Future<Output = Result<T, ()>> + Send + 'static,
{
    let task = async move {
        if let Ok(resp) = resp.await {
            let measure = GrpcRequestDuration::new(begin, label, String::default());
            let task = MeasuredSingleResponse::new(id, resp, measure);
            if let Err(e) = tx.send_with(task, WakePolicy::Immediately) {
                error!("KvService response batch commands fail"; "err" => ?e);
            }
        }
    };
    poll_future_notify(task);
}

fn handle_batch_commands_request<L: LockManager, F: KvFormat>(
    batcher: &mut Option<ReqBatcher>,
    storage: &Storage<RaftKv, L, F>,
    copr: &Endpoint<RaftKv>,
    peer: &str,
    id: u64,
    req: batch_commands_request::Request,
    tx: &Sender<MeasuredSingleResponse>,
) {
    match req.cmd {
        None => {
            let begin_instant = Instant::now();
            let resp = future::ok(batch_commands_response::Response::default());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::invalid,
            );
        }
        Some(batch_commands_request::request::Cmd::Get(req)) => {
            if batcher
                .as_mut()
                .is_some_and(|req_batch| req_batch.can_batch_get(&req))
            {
                batcher.as_mut().unwrap().add_get_request(req, id);
            } else {
                let begin_instant = Instant::now();
                let resp = future_get(storage, req)
                    .map_ok(|resp| batch_commands_response::Response {
                        cmd: Some(batch_commands_response::response::Cmd::Get(resp)),
                        ..Default::default()
                    })
                    .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_get.inc());
                response_batch_commands_request(
                    id,
                    resp,
                    tx.clone(),
                    begin_instant,
                    GrpcTypeKind::kv_get,
                );
            }
        }
        Some(batch_commands_request::request::Cmd::Coprocessor(req)) => {
            let begin_instant = Instant::now();
            let resp = future_copr(copr, Some(peer.to_string()), req)
                .map_ok(|resp| {
                    resp.map(|resp| batch_commands_response::Response {
                        cmd: Some(batch_commands_response::response::Cmd::Coprocessor(resp)),
                        ..Default::default()
                    })
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.coprocessor.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::coprocessor,
            );
        }
        Some(batch_commands_request::request::Cmd::Scan(req)) => {
            let begin_instant = Instant::now();
            let resp = future_scan(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::Scan(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_scan.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_scan,
            );
        }
        Some(batch_commands_request::request::Cmd::Prewrite(req)) => {
            let begin_instant = Instant::now();
            let resp = future_prewrite(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::Prewrite(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_prewrite.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_prewrite,
            );
        }
        Some(batch_commands_request::request::Cmd::Commit(req)) => {
            let begin_instant = Instant::now();
            let resp = future_commit(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::Commit(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_commit.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_commit,
            );
        }
        Some(batch_commands_request::request::Cmd::Cleanup(req)) => {
            let begin_instant = Instant::now();
            let resp = future_cleanup(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::Cleanup(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_cleanup.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_cleanup,
            );
        }
        Some(batch_commands_request::request::Cmd::BatchGet(req)) => {
            let begin_instant = Instant::now();
            let resp = future_batch_get(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::BatchGet(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_batch_get.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_batch_get,
            );
        }
        Some(batch_commands_request::request::Cmd::BatchRollback(req)) => {
            let begin_instant = Instant::now();
            let resp = future_batch_rollback(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::BatchRollback(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_batch_rollback.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_batch_rollback,
            );
        }
        Some(batch_commands_request::request::Cmd::TxnHeartBeat(req)) => {
            let begin_instant = Instant::now();
            let resp = future_txn_heart_beat(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::TxnHeartBeat(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_txn_heart_beat.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_txn_heart_beat,
            );
        }
        Some(batch_commands_request::request::Cmd::CheckTxnStatus(req)) => {
            let begin_instant = Instant::now();
            let resp = future_check_txn_status(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::CheckTxnStatus(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_check_txn_status.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_check_txn_status,
            );
        }
        Some(batch_commands_request::request::Cmd::CheckSecondaryLocks(req)) => {
            let begin_instant = Instant::now();
            let resp = future_check_secondary_locks(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::CheckSecondaryLocks(
                        resp,
                    )),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_check_secondary_locks.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_check_secondary_locks,
            );
        }
        Some(batch_commands_request::request::Cmd::ScanLock(req)) => {
            let begin_instant = Instant::now();
            let resp = future_scan_lock(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::ScanLock(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_scan_lock.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_scan_lock,
            );
        }
        Some(batch_commands_request::request::Cmd::ResolveLock(req)) => {
            let begin_instant = Instant::now();
            let resp = future_resolve_lock(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::ResolveLock(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_resolve_lock.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_resolve_lock,
            );
        }
        Some(batch_commands_request::request::Cmd::Gc(req)) => {
            let begin_instant = Instant::now();
            let resp = future_gc(req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::Gc(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_gc.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_gc,
            );
        }
        Some(batch_commands_request::request::Cmd::DeleteRange(req)) => {
            let begin_instant = Instant::now();
            let resp = future_delete_range(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::DeleteRange(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_delete_range.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_delete_range,
            );
        }
        Some(batch_commands_request::request::Cmd::Import(_)) => {
            let mut resp = ImportResponse::default();
            resp.set_error(UNIMPLEMENTED_MSG.to_owned());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::invalid,
                batch_commands_response::response::Cmd::Import(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::RawGet(_)) => {
            let mut resp = RawGetResponse::default();
            resp.set_error(UNIMPLEMENTED_MSG.to_owned());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::raw_get,
                batch_commands_response::response::Cmd::RawGet(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::RawBatchGet(_)) => {
            let mut resp = RawBatchGetResponse::default();
            resp.set_region_error(unimplemented_region_error());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::raw_batch_get,
                batch_commands_response::response::Cmd::RawBatchGet(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::RawPut(_)) => {
            let mut resp = RawPutResponse::default();
            resp.set_error(UNIMPLEMENTED_MSG.to_owned());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::raw_put,
                batch_commands_response::response::Cmd::RawPut(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::RawBatchPut(_)) => {
            let mut resp = RawBatchPutResponse::default();
            resp.set_error(UNIMPLEMENTED_MSG.to_owned());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::raw_batch_put,
                batch_commands_response::response::Cmd::RawBatchPut(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::RawDelete(_)) => {
            let mut resp = RawDeleteResponse::default();
            resp.set_error(UNIMPLEMENTED_MSG.to_owned());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::raw_delete,
                batch_commands_response::response::Cmd::RawDelete(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::RawBatchDelete(_)) => {
            let mut resp = RawBatchDeleteResponse::default();
            resp.set_error(UNIMPLEMENTED_MSG.to_owned());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::raw_batch_delete,
                batch_commands_response::response::Cmd::RawBatchDelete(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::RawScan(_)) => {
            let mut resp = RawScanResponse::default();
            resp.set_region_error(unimplemented_region_error());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::raw_scan,
                batch_commands_response::response::Cmd::RawScan(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::RawDeleteRange(_)) => {
            let mut resp = RawDeleteRangeResponse::default();
            resp.set_error(UNIMPLEMENTED_MSG.to_owned());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::raw_delete_range,
                batch_commands_response::response::Cmd::RawDeleteRange(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::RawBatchScan(_)) => {
            let mut resp = RawBatchScanResponse::default();
            resp.set_region_error(unimplemented_region_error());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::raw_batch_scan,
                batch_commands_response::response::Cmd::RawBatchScan(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::RawCoprocessor(_)) => {
            let mut resp = RawCoprocessorResponse::default();
            resp.set_error(UNIMPLEMENTED_MSG.to_owned());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::raw_coprocessor,
                batch_commands_response::response::Cmd::RawCoprocessor(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::PrepareFlashbackToVersion(_)) => {
            let mut resp = PrepareFlashbackToVersionResponse::default();
            resp.set_error(UNIMPLEMENTED_MSG.to_owned());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::kv_prepare_flashback_to_version,
                batch_commands_response::response::Cmd::PrepareFlashbackToVersion(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::FlashbackToVersion(_)) => {
            let mut resp = FlashbackToVersionResponse::default();
            resp.set_error(UNIMPLEMENTED_MSG.to_owned());
            respond_unimplemented_batch(
                id,
                tx,
                GrpcTypeKind::kv_flashback_to_version,
                batch_commands_response::response::Cmd::FlashbackToVersion(resp),
            );
        }
        Some(batch_commands_request::request::Cmd::PessimisticLock(req)) => {
            let begin_instant = Instant::now();
            let resp = future_acquire_pessimistic_lock(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::PessimisticLock(
                        resp,
                    )),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_pessimistic_lock.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_pessimistic_lock,
            );
        }
        Some(batch_commands_request::request::Cmd::PessimisticRollback(req)) => {
            let begin_instant = Instant::now();
            let resp = future_pessimistic_rollback(storage, req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::PessimisticRollback(
                        resp,
                    )),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.kv_pessimistic_rollback.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::kv_pessimistic_rollback,
            );
        }
        Some(batch_commands_request::request::Cmd::Empty(req)) => {
            let begin_instant = Instant::now();
            let resp = future_handle_empty(req)
                .map_ok(|resp| batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::Empty(resp)),
                    ..Default::default()
                })
                .map_err(|_| GRPC_MSG_FAIL_COUNTER.invalid.inc());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::invalid,
            );
        }
        Some(batch_commands_request::request::Cmd::GetHealthFeedback(_req)) => {
            let begin_instant = Instant::now();
            // The response is empty at this place, and will be filled when collected to
            // the batch response. See `HealthFeedbackAttacher::attach_if_needed`.
            let resp =
                std::future::ready(Ok(GetHealthFeedbackResponse::default())).map_ok(|resp| {
                    batch_commands_response::Response {
                        cmd: Some(batch_commands_response::response::Cmd::GetHealthFeedback(
                            resp,
                        )),
                        ..Default::default()
                    }
                });
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::get_health_feedback,
            );
        }
        Some(_) => {
            let begin_instant = Instant::now();
            let resp = future::ok(batch_commands_response::Response::default());
            response_batch_commands_request(
                id,
                resp,
                tx.clone(),
                begin_instant,
                GrpcTypeKind::invalid,
            );
        }
    }
}

fn handle_measures_for_batch_commands(measures: &mut MeasuredBatchResponse) {
    use BatchCommandsResponse_Response_oneof_cmd::*;
    for (resp, measure) in measures
        .batch_resp
        .mut_responses()
        .iter_mut()
        .zip(std::mem::take(&mut measures.measures))
    {
        let GrpcRequestDuration {
            label,
            begin,
            source,
            sent,
        } = measure;
        let elapsed = begin.saturating_elapsed();
        let wait = sent.saturating_elapsed();
        GRPC_MSG_HISTOGRAM_STATIC
            .get(label)
            .observe(elapsed.as_secs_f64());
        GRPC_BATCH_COMMANDS_WAIT_HISTOGRAM.observe(wait.as_secs_f64());
        record_request_source_metrics(source, elapsed);
        let exec_details = resp.cmd.as_mut().and_then(|cmd| match cmd {
            Get(resp) => Some(resp.mut_exec_details_v2()),
            Prewrite(resp) => Some(resp.mut_exec_details_v2()),
            Commit(resp) => Some(resp.mut_exec_details_v2()),
            BatchGet(resp) => Some(resp.mut_exec_details_v2()),
            ResolveLock(resp) => Some(resp.mut_exec_details_v2()),
            Coprocessor(resp) => Some(resp.mut_exec_details_v2()),
            PessimisticLock(resp) => Some(resp.mut_exec_details_v2()),
            CheckTxnStatus(resp) => Some(resp.mut_exec_details_v2()),
            TxnHeartBeat(resp) => Some(resp.mut_exec_details_v2()),
            CheckSecondaryLocks(resp) => Some(resp.mut_exec_details_v2()),
            _ => None,
        });
        if let Some(exec_details) = exec_details {
            exec_details
                .mut_time_detail()
                .set_total_rpc_wall_time_ns(elapsed.as_nanos() as u64);
            exec_details
                .mut_time_detail_v2()
                .set_total_rpc_wall_time_ns(elapsed.as_nanos() as u64);
            exec_details
                .mut_time_detail_v2()
                .set_kv_grpc_wait_time_ns(wait.as_nanos() as u64);
        }
    }
}

async fn future_handle_empty(
    req: BatchCommandsEmptyRequest,
) -> ServerResult<BatchCommandsEmptyResponse> {
    let mut res = BatchCommandsEmptyResponse::default();
    res.set_test_id(req.get_test_id());
    // `BatchCommandsWaker` processes futures in notify. If delay_time is too small,
    // notify can be called immediately, so the future is polled recursively and
    // lead to deadlock.
    if req.get_delay_time() >= 10 {
        let _ = tikv_util::timer::GLOBAL_TIMER_HANDLE
            .delay(
                std::time::Instant::now() + std::time::Duration::from_millis(req.get_delay_time()),
            )
            .compat()
            .await;
    }
    Ok(res)
}

fn future_get<L: LockManager, F: KvFormat>(
    storage: &Storage<RaftKv, L, F>,
    mut req: GetRequest,
) -> impl Future<Output = ServerResult<GetResponse>> {
    if recovery::check_request_rejected(req.get_context().get_keyspace_id()) {
        return future::ready(Err(box_err!("rejected in recovery mode"))).boxed();
    }
    let ctx = req.get_context();
    set_tls_trace_ctx(TraceContext::from_proto(
        ctx.get_trace_id(),
        ctx.get_trace_control_flags(),
    ));
    let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::new(
        req.get_context(),
        RequestType::KvGet,
        req.get_version(),
    )));
    set_tls_tracker_token(tracker);
    with_tls_tracker(|tracker| {
        tracker.metrics.grpc_req_size = req.compute_size() as u64;
    });
    let start = Instant::now();
    let v = storage.get(
        req.take_context(),
        Key::from_raw(req.get_key()),
        req.get_version().into(),
    );

    async move {
        let v = v.await;
        let duration = start.saturating_elapsed();
        let mut resp = GetResponse::default();
        if let Some(err) = extract_region_error(&v) {
            resp.set_region_error(err);
        } else {
            match v {
                Ok((val, stats)) => {
                    let exec_detail_v2 = resp.mut_exec_details_v2();
                    let scan_detail_v2 = exec_detail_v2.mut_scan_detail_v2();
                    stats.stats.write_scan_detail(scan_detail_v2);
                    GLOBAL_TRACKERS.with_tracker(tracker, |tracker| {
                        tracker.write_scan_detail(exec_detail_v2.mut_scan_detail_v2());
                        tracker.write_ru_v2(exec_detail_v2.mut_ru_v2());
                    });
                    exec_detail_v2.mut_ru_v2().set_storage_processed_keys_get(1);
                    set_time_detail(exec_detail_v2, duration, &stats.latency_stats);
                    match val {
                        Some(val) => resp.set_value(val),
                        None => resp.set_not_found(true),
                    }
                }
                Err(e) => resp.set_error(extract_key_error(&e)),
            }
        }
        GLOBAL_TRACKERS.with_tracker(tracker, |tracker| {
            if let Some(details) = tracker.read_pool_task_details() {
                resp.mut_exec_details_v2()
                    .set_read_pool_task_details(details);
            }
        });
        GLOBAL_TRACKERS.remove(tracker);
        Ok(resp)
    }
    .boxed()
}

fn set_time_detail(
    exec_detail_v2: &mut ExecDetailsV2,
    total_dur: Duration,
    stats: &StageLatencyStats,
) {
    let duration_ns = total_dur.as_nanos() as u64;
    // deprecated. we will remove the `time_detail` field in future version.
    {
        let time_detail = exec_detail_v2.mut_time_detail();
        time_detail.set_kv_read_wall_time_ms(duration_ns / 1_000_000);
        time_detail.set_wait_wall_time_ms(stats.wait_wall_time_ns / 1_000_000);
        time_detail.set_process_wall_time_ms(stats.process_wall_time_ns / 1_000_000);
    }

    let time_detail_v2 = exec_detail_v2.mut_time_detail_v2();
    time_detail_v2.set_kv_read_wall_time_ns(duration_ns);
    time_detail_v2.set_wait_wall_time_ns(stats.wait_wall_time_ns);
    time_detail_v2.set_process_wall_time_ns(stats.process_wall_time_ns);
}

fn future_scan<L: LockManager, F: KvFormat>(
    storage: &Storage<RaftKv, L, F>,
    mut req: ScanRequest,
) -> impl Future<Output = ServerResult<ScanResponse>> {
    if recovery::check_request_rejected(req.get_context().get_keyspace_id()) {
        return future::ready(Err(box_err!("rejected in recovery mode"))).boxed();
    }
    let ctx = req.get_context();
    set_tls_trace_ctx(TraceContext::from_proto(
        ctx.get_trace_id(),
        ctx.get_trace_control_flags(),
    ));
    let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::new(
        req.get_context(),
        RequestType::KvScan,
        req.get_version(),
    )));
    set_tls_tracker_token(tracker);
    with_tls_tracker(|tracker| {
        tracker.metrics.grpc_req_size = req.compute_size() as u64;
    });
    let end_key = Key::from_raw_maybe_unbounded(req.get_end_key());
    let rev_range = if req.reverse && req.version == u64::MAX {
        Some((req.start_key.clone(), req.end_key.clone()))
    } else {
        None
    };
    let v = storage.scan(
        req.take_context(),
        Key::from_raw(req.get_start_key()),
        end_key,
        req.get_limit() as usize,
        req.get_sample_step() as usize,
        req.get_version().into(),
        req.get_key_only(),
        req.get_reverse(),
    );

    Box::pin(async move {
        let v = v.await;
        let mut resp = ScanResponse::default();
        if let Some(err) = extract_region_error(&v) {
            resp.set_region_error(err);
        } else {
            match v {
                Ok(kv_res) => {
                    if let Some((start_key, end_key)) = rev_range {
                        info!("reverse scan";
                            "start_key" => hex_encode(start_key),
                            "end_key" => hex_encode(end_key),
                            "count" => kv_res.len(),
                        );
                    }
                    resp.set_pairs(map_kv_pairs(kv_res).into());
                }
                Err(e) => {
                    if let Some((start_key, end_key)) = rev_range {
                        warn!("reverse scan error";
                            "start_key" => hex_encode(start_key),
                            "end_key" => hex_encode(end_key),
                            "error" => ?e,
                        );
                    }
                    let key_error = extract_key_error(&e);
                    resp.set_error(key_error.clone());
                    // Set key_error in the first kv_pair for backward compatibility.
                    let mut pair = KvPair::default();
                    pair.set_error(key_error);
                    resp.mut_pairs().push(pair);
                }
            }
        }
        GLOBAL_TRACKERS.remove(tracker);
        Ok(resp)
    })
}

fn future_batch_get<L: LockManager, F: KvFormat>(
    storage: &Storage<RaftKv, L, F>,
    mut req: BatchGetRequest,
) -> impl Future<Output = ServerResult<BatchGetResponse>> {
    if recovery::check_request_rejected(req.get_context().get_keyspace_id()) {
        return future::ready(Err(box_err!("rejected in recovery mode"))).boxed();
    }
    let ctx = req.get_context();
    set_tls_trace_ctx(TraceContext::from_proto(
        ctx.get_trace_id(),
        ctx.get_trace_control_flags(),
    ));
    let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::new(
        req.get_context(),
        RequestType::KvBatchGet,
        req.get_version(),
    )));
    set_tls_tracker_token(tracker);
    with_tls_tracker(|tracker| {
        tracker.metrics.grpc_req_size = req.compute_size() as u64;
    });
    let start = Instant::now();
    let keys: Vec<_> = req.get_keys().iter().map(|x| Key::from_raw(x)).collect();
    let num_keys = keys.len() as u64;
    let v = storage.batch_get(req.take_context(), keys, req.get_version().into());

    Box::pin(async move {
        let v = v.await;
        let duration = start.saturating_elapsed();
        let mut resp = BatchGetResponse::default();
        if let Some(err) = extract_region_error(&v) {
            resp.set_region_error(err);
        } else {
            match v {
                Ok((kv_res, stats)) => {
                    let pairs = map_kv_pairs(kv_res);
                    let exec_detail_v2 = resp.mut_exec_details_v2();
                    let scan_detail_v2 = exec_detail_v2.mut_scan_detail_v2();
                    stats.stats.write_scan_detail(scan_detail_v2);
                    GLOBAL_TRACKERS.with_tracker(tracker, |tracker| {
                        tracker.write_scan_detail(exec_detail_v2.mut_scan_detail_v2());
                        tracker.write_ru_v2(exec_detail_v2.mut_ru_v2());
                    });
                    exec_detail_v2
                        .mut_ru_v2()
                        .set_storage_processed_keys_batch_get(num_keys);
                    set_time_detail(exec_detail_v2, duration, &stats.latency_stats);
                    resp.set_pairs(pairs.into());
                }
                Err(e) => {
                    let key_error = extract_key_error(&e);
                    resp.set_error(key_error.clone());
                    // Set key_error in the first kv_pair for backward compatibility.
                    let mut pair = KvPair::default();
                    pair.set_error(key_error);
                    resp.mut_pairs().push(pair);
                }
            }
        }
        GLOBAL_TRACKERS.with_tracker(tracker, |tracker| {
            if let Some(details) = tracker.read_pool_task_details() {
                resp.mut_exec_details_v2()
                    .set_read_pool_task_details(details);
            }
        });
        GLOBAL_TRACKERS.remove(tracker);
        Ok(resp)
    })
}

fn future_scan_lock<L: LockManager, F: KvFormat>(
    storage: &Storage<RaftKv, L, F>,
    mut req: ScanLockRequest,
) -> impl Future<Output = ServerResult<ScanLockResponse>> {
    if recovery::check_request_rejected(req.get_context().get_keyspace_id()) {
        return future::ready(Err(box_err!("rejected in recovery mode"))).boxed();
    }
    let ctx = req.get_context();
    set_tls_trace_ctx(TraceContext::from_proto(
        ctx.get_trace_id(),
        ctx.get_trace_control_flags(),
    ));
    let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::new(
        req.get_context(),
        RequestType::KvScanLock,
        req.get_max_version(),
    )));
    set_tls_tracker_token(tracker);
    with_tls_tracker(|tracker| {
        tracker.metrics.grpc_req_size = req.compute_size() as u64;
    });
    let start_key = Key::from_raw_maybe_unbounded(req.get_start_key());
    let end_key = Key::from_raw_maybe_unbounded(req.get_end_key());

    let v = storage.scan_lock(
        req.take_context(),
        req.get_max_version().into(),
        start_key,
        end_key,
        req.get_limit() as usize,
    );

    Box::pin(async move {
        let v = v.await;
        let mut resp = ScanLockResponse::default();
        if let Some(err) = extract_region_error(&v) {
            resp.set_region_error(err);
        } else {
            match v {
                Ok(locks) => resp.set_locks(locks.into()),
                Err(e) => resp.set_error(extract_key_error(&e)),
            }
        }
        GLOBAL_TRACKERS.with_tracker(tracker, |tracker| {
            if let Some(details) = tracker.read_pool_task_details() {
                resp.mut_exec_details_v2()
                    .set_read_pool_task_details(details);
            }
        });
        GLOBAL_TRACKERS.remove(tracker);
        Ok(resp)
    })
}

async fn future_gc(_: GcRequest) -> ServerResult<GcResponse> {
    let mut resp = GcResponse::default();
    resp.set_region_error(unimplemented_region_error());
    Ok(resp)
}

fn future_delete_range<L: LockManager, F: KvFormat>(
    storage: &Storage<RaftKv, L, F>,
    mut req: DeleteRangeRequest,
) -> impl Future<Output = ServerResult<DeleteRangeResponse>> {
    if recovery::check_request_rejected(req.get_context().get_keyspace_id()) {
        return future::ready(Err(box_err!("rejected in recovery mode"))).boxed();
    }
    let ctx = req.get_context();
    set_tls_trace_ctx(TraceContext::from_proto(
        ctx.get_trace_id(),
        ctx.get_trace_control_flags(),
    ));
    let (cb, f) = paired_future_callback();
    let res = storage.delete_range(
        req.take_context(),
        Key::from_raw(req.get_start_key()),
        Key::from_raw(req.get_end_key()),
        req.get_notify_only(),
        cb,
    );

    Box::pin(async move {
        let v = match res {
            Err(e) => Err(e),
            Ok(_) => f.await?,
        };
        let mut resp = DeleteRangeResponse::default();
        if let Some(err) = extract_region_error(&v) {
            resp.set_region_error(err);
        } else if let Err(e) = v {
            resp.set_error(format!("{}", e));
        }
        Ok(resp)
    })
    .boxed()
}

fn future_copr<E: Engine>(
    copr: &Endpoint<E>,
    peer: Option<String>,
    req: Request,
) -> impl Future<Output = ServerResult<MemoryTraceGuard<Response>>> {
    if recovery::check_request_rejected(req.get_context().get_keyspace_id()) {
        return future::ready(Err(box_err!("rejected in recovery mode"))).boxed();
    }
    let ctx = req.get_context();
    set_tls_trace_ctx(TraceContext::from_proto(
        ctx.get_trace_id(),
        ctx.get_trace_control_flags(),
    ));
    let ret = copr.parse_and_handle_unary_request(req, peer);
    async move { Ok(ret.await) }.boxed()
}

macro_rules! txn_command_future {
    ($fn_name: ident, $req_ty: ident, $resp_ty: ident, ($req: ident) {$($prelude: stmt)*}; ($v: ident, $resp: ident, $tracker: ident) $else_branch: block) => {
        txn_command_future!(inner $fn_name, $req_ty, $resp_ty, ($req) {$($prelude)*}; ($v, $resp, $tracker) {
            $else_branch
            GLOBAL_TRACKERS.with_tracker($tracker, |tracker| {
                tracker.write_scan_detail($resp.mut_exec_details_v2().mut_scan_detail_v2());
                tracker.write_write_detail($resp.mut_exec_details_v2().mut_write_detail());
                tracker.write_ru_v2($resp.mut_exec_details_v2().mut_ru_v2());
            });
        });
    };

    ($fn_name: ident, $req_ty: ident, $resp_ty: ident, ($v: ident, $resp: ident, $tracker: ident) $else_branch: block ) => {
        txn_command_future!(inner $fn_name, $req_ty, $resp_ty, (req) {}; ($v, $resp, $tracker) {
            $else_branch
            GLOBAL_TRACKERS.with_tracker($tracker, |tracker| {
                tracker.write_scan_detail($resp.mut_exec_details_v2().mut_scan_detail_v2());
                tracker.write_write_detail($resp.mut_exec_details_v2().mut_write_detail());
                tracker.write_ru_v2($resp.mut_exec_details_v2().mut_ru_v2());
            });
        });
    };

    ($fn_name: ident, $req_ty: ident, $resp_ty: ident, ($v: ident, $resp: ident) $else_branch: block ) => {
        txn_command_future!(inner $fn_name, $req_ty, $resp_ty, (req) {}; ($v, $resp, tracker) { $else_branch });
    };

    (inner $fn_name: ident, $req_ty: ident, $resp_ty: ident, ($req: ident) {$($prelude: stmt)*}; ($v: ident, $resp: ident, $tracker: ident) $else_branch: block) => {
        fn $fn_name<E: Engine, L: LockManager, F: KvFormat>(
            storage: &Storage<E, L, F>,
            $req: $req_ty,
        ) -> impl Future<Output = ServerResult<$resp_ty>> {
            $($prelude)*
            let ctx = $req.get_context();
            set_tls_trace_ctx(TraceContext::from_proto(ctx.get_trace_id(), ctx.get_trace_control_flags()));
            let $tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::new(
                $req.get_context(),
                RequestType::Unknown,
                0,
            )));
            set_tls_tracker_token($tracker);
            with_tls_tracker(|tracker| {
                tracker.metrics.grpc_req_size = $req.compute_size() as u64;
            });
            let (cb, f) = paired_future_callback();
            let res = storage.sched_txn_command($req.into(), cb);

            Box::pin(async move {
                defer!{{
                    GLOBAL_TRACKERS.remove($tracker);
                }};
                let $v = match res {
                    Err(e) => Err(e),
                    Ok(_) => f.await?,
                };
                let mut $resp = $resp_ty::default();
                if let Some(err) = extract_region_error(&$v) {
                    $resp.set_region_error(err);
                } else {
                    $else_branch;
                }
                Ok($resp)
            })
        }
    };
}

txn_command_future!(future_prewrite, PrewriteRequest, PrewriteResponse, (v, resp, tracker) {{
    if let Ok(v) = &v {
        resp.set_min_commit_ts(v.min_commit_ts.into_inner());
        resp.set_one_pc_commit_ts(v.one_pc_commit_ts.into_inner());
    }
    resp.set_errors(extract_key_errors(v.map(|v| v.locks)).into());
}});
txn_command_future!(future_acquire_pessimistic_lock, PessimisticLockRequest, PessimisticLockResponse, (v, resp, tracker) {
    match v {
        Ok(Ok(res)) => match res.try_into_legacy_values_and_not_founds() {
            Ok(legacy) => {
                resp.set_values(legacy.values.into());
                resp.set_not_founds(legacy.not_founds);
            }
            Err(e) => resp.set_errors(vec![e.extract_key_error()].into()),
        },
        Err(e) | Ok(Err(e)) => resp.set_errors(vec![extract_key_error(&e)].into()),
    }
});
txn_command_future!(future_pessimistic_rollback, PessimisticRollbackRequest, PessimisticRollbackResponse, (v, resp, tracker) {
    resp.set_errors(extract_key_errors(v).into())
});
txn_command_future!(future_batch_rollback, BatchRollbackRequest, BatchRollbackResponse, (v, resp, tracker) {
    if let Err(e) = v {
        resp.set_error(extract_key_error(&e));
    }
});
txn_command_future!(future_resolve_lock, ResolveLockRequest, ResolveLockResponse, (v, resp, tracker) {
    if let Err(e) = v {
        resp.set_error(extract_key_error(&e));
    }
});
txn_command_future!(future_commit, CommitRequest, CommitResponse, (v, resp, tracker) {
    match v {
        Ok(TxnStatus::Committed { commit_ts }) => {
            resp.set_commit_version(commit_ts.into_inner())
        }
        Ok(_) => unreachable!(),
        Err(e) => resp.set_error(extract_key_error(&e)),
    }
});
txn_command_future!(future_cleanup, CleanupRequest, CleanupResponse, (v, resp) {
    if let Err(e) = v {
        if let Some(ts) = extract_committed(&e) {
            resp.set_commit_version(ts.into_inner());
        } else {
            resp.set_error(extract_key_error(&e));
        }
    }
});
txn_command_future!(future_txn_heart_beat, TxnHeartBeatRequest, TxnHeartBeatResponse, (v, resp, tracker) {
    match v {
        Ok(txn_status) => {
            if let TxnStatus::Uncommitted { lock, .. } = txn_status {
                resp.set_lock_ttl(lock.ttl);
            } else {
                unreachable!();
            }
        }
        Err(e) => resp.set_error(extract_key_error(&e)),
    }
});
txn_command_future!(future_check_txn_status, CheckTxnStatusRequest, CheckTxnStatusResponse,
    (v, resp, tracker) {
        match v {
            Ok(txn_status) => match txn_status {
                TxnStatus::RolledBack => resp.set_action(Action::NoAction),
                TxnStatus::TtlExpire => resp.set_action(Action::TtlExpireRollback),
                TxnStatus::LockNotExist => resp.set_action(Action::LockNotExistRollback),
                TxnStatus::Committed { commit_ts } => {
                    resp.set_commit_version(commit_ts.into_inner())
                }
                TxnStatus::Uncommitted { lock, min_commit_ts_pushed } => {
                    if min_commit_ts_pushed {
                        resp.set_action(Action::MinCommitTsPushed);
                    }
                    resp.set_lock_ttl(lock.ttl);
                    let primary = lock.primary.clone();
                    resp.set_lock_info(lock.into_lock_info(primary));
                }
                TxnStatus::PessimisticRollBack => resp.set_action(Action::TtlExpirePessimisticRollback),
                TxnStatus::LockNotExistDoNothing => resp.set_action(Action::LockNotExistDoNothing),
            },
            Err(e) => resp.set_error(extract_key_error(&e)),
        }
});
txn_command_future!(future_check_secondary_locks, CheckSecondaryLocksRequest, CheckSecondaryLocksResponse, (status, resp, tracker) {
    match status {
        Ok(SecondaryLocksStatus::Locked(locks)) => {
            resp.set_locks(locks.into());
        },
        Ok(SecondaryLocksStatus::Committed(ts)) => {
            resp.set_commit_ts(ts.into_inner());
        },
        Ok(SecondaryLocksStatus::RolledBack) => {},
        Err(e) => resp.set_error(extract_key_error(&e)),
    }
});
txn_command_future!(future_mvcc_get_by_key, MvccGetByKeyRequest, MvccGetByKeyResponse, (v, resp) {
    match v {
        Ok(mvcc) => resp.set_info(mvcc.into_proto()),
        Err(e) => resp.set_error(format!("{}", e)),
    }
});
txn_command_future!(future_mvcc_get_by_start_ts, MvccGetByStartTsRequest, MvccGetByStartTsResponse, (v, resp) {
    match v {
        Ok(Some((k, vv))) => {
            resp.set_key(k.into_raw().unwrap());
            resp.set_info(vv.into_proto());
        }
        Ok(None) => {
            resp.set_info(Default::default());
        }
        Err(e) => resp.set_error(format!("{}", e)),
    }
});

pub mod batch_commands_response {
    pub type Response = kvproto::tikvpb::BatchCommandsResponseResponse;

    pub mod response {
        pub type Cmd = kvproto::tikvpb::BatchCommandsResponse_Response_oneof_cmd;
    }
}

pub mod batch_commands_request {
    pub type Request = kvproto::tikvpb::BatchCommandsRequestRequest;

    pub mod request {
        pub type Cmd = kvproto::tikvpb::BatchCommandsRequest_Request_oneof_cmd;
    }
}

use protobuf::RepeatedField;
use rfstore::store::StoreMsg;
use tikv::server::service::{GrpcRequestDuration, MeasuredBatchResponse, MeasuredSingleResponse};
use tikv_alloc::MemoryTraceGuard;
use tikv_util::timer::GLOBAL_TIMER_HANDLE;

use crate::RaftKv;

fn collect_batch_resp(v: &mut MeasuredBatchResponse, mut e: MeasuredSingleResponse) {
    v.batch_resp.mut_request_ids().push(e.id);
    v.batch_resp.mut_responses().push(e.resp.consume());
    v.measures.push(e.measure);
}

struct HealthFeedbackAttacher {
    store_id: u64,
    health_controller: HealthController,
    last_feedback_time: Option<Instant>,
    seq: Arc<AtomicU64>,
    feedback_interval: Option<Duration>,
}

impl HealthFeedbackAttacher {
    fn new(
        store_id: u64,
        health_controller: HealthController,
        seq: Arc<AtomicU64>,
        feedback_interval: Option<Duration>,
    ) -> Self {
        Self {
            store_id,
            health_controller,
            last_feedback_time: None,
            seq,
            feedback_interval,
        }
    }

    fn attach_if_needed(&mut self, resp: &mut BatchCommandsResponse) {
        let mut requested_feedback = None;

        for r in resp.responses.iter_mut().flat_map(|r| &mut r.cmd) {
            if let BatchCommandsResponse_Response_oneof_cmd::GetHealthFeedback(r) = r {
                if requested_feedback.is_none() {
                    requested_feedback = Some(self.gen_health_feedback_pb());
                }
                r.set_health_feedback(requested_feedback.clone().unwrap());
            }
        }

        // If there's any explicit request for health feedback information, attach it
        // to the BatchCommandsResponse no matter how it's configured.
        if let Some(feedback) = requested_feedback {
            self.attach_with(resp, Instant::now_coarse(), feedback);
            return;
        }

        let feedback_interval = match self.feedback_interval {
            Some(i) => i,
            None => return,
        };

        let now = Instant::now_coarse();

        if let Some(last_feedback_time) = self.last_feedback_time {
            if now - last_feedback_time < feedback_interval {
                return;
            }
        }

        self.attach(resp, now);
    }

    fn attach(&mut self, resp: &mut BatchCommandsResponse, now: Instant) {
        self.attach_with(resp, now, self.gen_health_feedback_pb())
    }

    fn attach_with(
        &mut self,
        resp: &mut BatchCommandsResponse,
        now: Instant,
        health_feedback_pb: HealthFeedback,
    ) {
        self.last_feedback_time = Some(now);
        resp.set_health_feedback(health_feedback_pb);
    }

    fn gen_health_feedback_pb(&self) -> HealthFeedback {
        let mut feedback = HealthFeedback::default();
        feedback.set_store_id(self.store_id);
        feedback.set_feedback_seq_no(self.seq.fetch_add(1, Ordering::Relaxed));
        feedback.set_slow_score(self.health_controller.get_rfstore_slow_score() as i32);
        // TODO: set network slow score?
        feedback
    }
}

#[cfg(test)]
mod tests {
    use std::thread;

    use futures::{channel::oneshot, executor::block_on};
    use tikv_util::sys::thread::StdThreadBuildWrapper;

    use super::*;

    #[test]
    fn test_poll_future_notify_with_slow_source() {
        let (tx, rx) = oneshot::channel::<usize>();
        let (signal_tx, signal_rx) = oneshot::channel();

        thread::Builder::new()
            .name("source".to_owned())
            .spawn_wrapper(move || {
                block_on(signal_rx).unwrap();
                tx.send(100).unwrap();
            })
            .unwrap();

        let (tx1, rx1) = oneshot::channel::<usize>();
        let task = async move {
            let i = rx.await.unwrap();
            assert_eq!(thread::current().name(), Some("source"));
            tx1.send(i + 100).unwrap();
        };
        poll_future_notify(task);
        signal_tx.send(()).unwrap();
        assert_eq!(block_on(rx1).unwrap(), 200);
    }

    #[test]
    fn test_poll_future_notify_with_slow_poller() {
        let (tx, rx) = oneshot::channel::<usize>();
        let (signal_tx, signal_rx) = oneshot::channel();
        thread::Builder::new()
            .name("source".to_owned())
            .spawn_wrapper(move || {
                tx.send(100).unwrap();
                signal_tx.send(()).unwrap();
            })
            .unwrap();

        let (tx1, rx1) = oneshot::channel::<usize>();
        block_on(signal_rx).unwrap();
        let task = async move {
            let i = rx.await.unwrap();
            assert_ne!(thread::current().name(), Some("source"));
            tx1.send(i + 100).unwrap();
        };
        poll_future_notify(task);
        assert_eq!(block_on(rx1).unwrap(), 200);
    }

    #[test]
    fn test_health_feedback_attacher() {
        let health_controller = HealthController::new();
        let test_reporter = health_controller::reporters::TestReporter::new(&health_controller);
        let seq = Arc::new(AtomicU64::new(1));

        let mut a = HealthFeedbackAttacher::new(1, health_controller.clone(), seq.clone(), None);
        let mut resp = BatchCommandsResponse::default();
        a.attach_if_needed(&mut resp);
        assert!(!resp.has_health_feedback());

        let mut a =
            HealthFeedbackAttacher::new(1, health_controller, seq, Some(Duration::from_secs(1)));
        resp = BatchCommandsResponse::default();
        a.attach_if_needed(&mut resp);
        assert!(resp.has_health_feedback());
        assert_eq!(resp.get_health_feedback().get_store_id(), 1);
        assert_eq!(resp.get_health_feedback().get_feedback_seq_no(), 1);
        assert_eq!(resp.get_health_feedback().get_slow_score(), 1);

        // Skips attaching feedback because last attaching was just done.
        test_reporter.set_rfstore_slow_score(50.);
        resp = BatchCommandsResponse::default();
        a.attach_if_needed(&mut resp);
        assert!(!resp.has_health_feedback());

        // Simulate elapsing enough time by changing the recorded last update time.
        *a.last_feedback_time.as_mut().unwrap() -= Duration::from_millis(1001);
        a.attach_if_needed(&mut resp);
        assert!(resp.has_health_feedback());
        assert_eq!(resp.get_health_feedback().get_store_id(), 1);
        // seq_no increased.
        assert_eq!(resp.get_health_feedback().get_feedback_seq_no(), 2);
        assert_eq!(resp.get_health_feedback().get_slow_score(), 50);

        // Do not skip if feedback info is explicitly requested, i.e.
        // `GetHealthFeedback` request is received from the client.
        let resp_with_get_health_feedback = || {
            let mut resp = BatchCommandsResponse::default();
            let mut resp_item = BatchCommandsResponseResponse::default();
            resp_item.cmd = Some(BatchCommandsResponse_Response_oneof_cmd::GetHealthFeedback(
                GetHealthFeedbackResponse::default(),
            ));
            resp.responses.push(resp_item);
            resp.request_ids.push(1);
            resp
        };
        resp = resp_with_get_health_feedback();
        a.attach_if_needed(&mut resp);
        assert!(resp.has_health_feedback());
        assert_eq!(resp.get_health_feedback().get_store_id(), 1);
        assert_eq!(resp.get_health_feedback().get_feedback_seq_no(), 3);
        assert_eq!(resp.get_health_feedback().get_slow_score(), 50);
        // And the `GetHealthFeedbackResponse` will also be filled with the feedback
        // info.
        fn get_feedback_in_response(resp: &BatchCommandsResponseResponse) -> &HealthFeedback {
            match resp.cmd {
                Some(BatchCommandsResponse_Response_oneof_cmd::GetHealthFeedback(ref f)) => {
                    assert!(!f.has_region_error());
                    f.get_health_feedback()
                }
                _ => panic!("unexpected response body: {:?}", resp),
            }
        }
        assert_eq!(
            get_feedback_in_response(&resp.get_responses()[0]),
            resp.get_health_feedback()
        );

        // If requested, it attaches the feedback info even `feedback_interval` is not
        // given, which means never attaching.
        a.feedback_interval = None;
        a.last_feedback_time = None;
        resp = BatchCommandsResponse::default();
        a.attach_if_needed(&mut resp);
        assert!(!resp.has_health_feedback());
        resp = resp_with_get_health_feedback();
        a.attach_if_needed(&mut resp);
        assert!(resp.has_health_feedback());
        assert_eq!(resp.get_health_feedback().get_store_id(), 1);
        assert_eq!(resp.get_health_feedback().get_feedback_seq_no(), 4);
        assert_eq!(resp.get_health_feedback().get_slow_score(), 50);
        assert_eq!(
            get_feedback_in_response(&resp.get_responses()[0]),
            resp.get_health_feedback()
        );

        // It also works when there are multiple `GetHealthFeedbackResponse` in the
        // batch, and each single `GetHealthFeedback` call gets the same result.
        test_reporter.set_rfstore_slow_score(80.);
        resp = resp_with_get_health_feedback();
        resp.responses.push(resp.responses[0].clone());
        resp.request_ids.push(2);
        a.attach_if_needed(&mut resp);
        assert!(resp.has_health_feedback());
        assert_eq!(resp.get_health_feedback().get_store_id(), 1);
        assert_eq!(resp.get_health_feedback().get_feedback_seq_no(), 5);
        assert_eq!(resp.get_health_feedback().get_slow_score(), 80);
        for i in 0..2 {
            assert_eq!(
                get_feedback_in_response(&resp.get_responses()[i]),
                resp.get_health_feedback(),
                "{}-th response in BatchCommandsResponse mismatches",
                i
            );
        }
    }
}
