// Copyright 2016 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2022 PingCAP, Inc.
//! Handles simple SQL query executors locally.
//!
//! Most TiDB read queries are processed by Coprocessor instead of KV interface.
//! By doing so, the CPU of TiKV nodes can be utilized for computing and the
//! amount of data to transfer can be reduced (i.e. filtered at TiKV side).
//!
//! Notice that Coprocessor handles more than simple SQL query executors (DAG
//! request). It also handles analyzing requests and checksum requests.
//!
//! The entry point of handling all coprocessor requests is `Endpoint`. Common
//! steps are:
//! - Parse the request into a DAG request, Checksum request or Analyze request.
//! - Retrieve a snapshot from the underlying engine according to the given
//!   timestamp.
//! - Build corresponding request handlers from the snapshot and request detail.
//! - Run request handlers once (for unary requests) or multiple times (for
//!   streaming requests) on a future thread pool.
//! - Return handling result as a response.
//!
//! Please refer to `Endpoint` for more details.

#![allow(clippy::diverging_sub_expression)]

mod cache;
mod checksum;
mod config_manager;
pub mod dag;
pub mod db9;
mod endpoint;
mod error;
mod interceptors;
mod keyspace_cpu;
pub use keyspace_cpu::register_keyspace_cpu_metric;
pub(crate) mod metrics;
pub mod readpool_impl;
pub mod remote_dispatcher;
mod runtime_budget;
mod scan_limiter;
mod statistics;
mod tracker;

use std::{ops::Deref, sync::Arc};

use async_trait::async_trait;
pub use checksum::checksum_crc64_xor;
pub use endpoint::{
    RemoteCopRequestKind, RemoteCopRequestType, RemoteCopResult, RemoteCopScanStats,
    make_error_response, parse_request_and_handle_remote_cop, prefetch_ia_remote_segments,
    remote_cop_request_type,
};
use kvproto::{coprocessor as coppb, kvrpcpb};
use lazy_static::lazy_static;
use protobuf::Message;
use rand::prelude::*;
pub use remote_dispatcher::RemoteRequest;
use resource_control::CpuTokenHandle;
use tidb_query_common::execute_stats::ExecSummary;
use tidb_query_executors::interface::CoprRuntimeBudgetRef;
use tikv_alloc::{Id, MemoryTrace, MemoryTraceGuard, mem_trace};
use tikv_util::{deadline::Deadline, memory::HeapSize, time::Duration};
use tipb::ResourceGroupTag;
use txn_types::TsSet;

pub use self::{
    endpoint::Endpoint,
    error::{Error, Result},
};
use crate::storage::{Statistics, mvcc::TimeStamp};

pub const REQ_TYPE_DAG: i64 = 103;
pub const REQ_TYPE_ANALYZE: i64 = 104;
pub const REQ_TYPE_CHECKSUM: i64 = 105;
pub const REQ_TYPE_DB9_DAG: i64 = 10001;

pub const REQ_FLAG_TIDB_SYSSESSION: u64 = 2048;

/// Decodes the SQL digest embedded in a request's resource group tag.
///
/// Returns an empty string when the tag is absent or malformed.
pub fn decode_query_digest(resource_group_tag: &[u8]) -> String {
    if resource_group_tag.is_empty() {
        return String::new();
    }

    let mut tag = ResourceGroupTag::new();
    if tag.merge_from_bytes(resource_group_tag).is_err() {
        return String::new();
    }
    hex::encode(tag.get_sql_digest())
}

type HandlerStreamStepResult = Result<(Option<coppb::Response>, bool)>;

/// An interface for all kind of Coprocessor request handlers.
#[async_trait]
pub trait RequestHandler: Send {
    /// Processes current request and produces a response.
    async fn handle_request(&mut self) -> Result<MemoryTraceGuard<coppb::Response>> {
        panic!("unary request is not supported for this handler");
    }

    /// Processes current request and produces streaming responses.
    async fn handle_streaming_request(&mut self) -> HandlerStreamStepResult {
        panic!("streaming request is not supported for this handler");
    }

    /// Collects scan statistics generated in this request handler so far.
    fn collect_scan_statistics(&mut self, _dest: &mut Statistics) {
        // Do nothing by default
    }

    /// Collects scan executor time in this request handler so far.
    fn collect_scan_summary(&mut self, _dest: &mut ExecSummary) {
        // Do nothing by default
    }

    /// Installs a runtime CPU budget for DAG batch-level enforcement.
    ///
    /// Returns true if the handler can enforce the budget.
    fn set_copr_runtime_budget(&mut self, _budget: Option<CoprRuntimeBudgetRef>) -> bool {
        false
    }

    /// Installs a one-shot hook for the CPU token allocated when a remote DAG
    /// request falls back to local execution.
    fn set_fallback_cpu_token_observer(
        &mut self,
        _observer: Box<dyn FnOnce(Arc<CpuTokenHandle>) + Send>,
    ) {
    }

    /// Collects executor inputs for RUV2 billing.
    fn collect_executor_inputs(&self, _dest: &mut kvproto::kvrpcpb::ExecutorInputs) {
        // Do nothing by default
    }

    /// Collects total executor iterations for RUV2 billing.
    fn collect_executor_iterations(&self) -> u64 {
        0
    }

    /// Populates RUV2 metrics for the current response.
    fn set_ru_v2_metrics(&mut self, ru_v2: &mut kvrpcpb::Ruv2, resp_bytes: u64) {
        let iterations = self.collect_executor_iterations();
        let mut inputs = kvrpcpb::ExecutorInputs::new();
        self.collect_executor_inputs(&mut inputs);
        ru_v2.set_coprocessor_response_bytes(resp_bytes);
        ru_v2.set_coprocessor_executor_iterations(iterations);
        ru_v2.set_executor_inputs(inputs);
    }

    fn into_boxed(self) -> Box<dyn RequestHandler>
    where
        Self: 'static + Sized,
    {
        Box::new(self)
    }
}

type RequestHandlerBuilder<Snap> =
    Box<dyn for<'a> FnOnce(Snap, &ReqContext) -> Result<Box<dyn RequestHandler>> + Send>;

/// The inner state of ReqContext.
#[derive(Debug, Clone, PartialEq)]
pub struct ReqContextInner {
    /// The rpc context carried in the request
    ///
    /// NOTE: Must be immutable. It's also used by remote cop.
    pub context: kvrpcpb::Context,

    /// Scan ranges of this request
    pub ranges: Vec<coppb::KeyRange>,

    /// The deadline of the request
    pub deadline: Deadline,

    /// The peer address of the request
    pub peer: Option<String>,

    /// Whether the request is a descending scan (only applicable to DAG)
    pub is_desc_scan: Option<bool>,

    /// The transaction start_ts of the request
    pub txn_start_ts: TimeStamp,

    /// The set of timestamps of locks that can be bypassed during the reading
    /// because either they will be rolled back or their commit_ts > read
    /// request's start_ts.
    pub bypass_locks: TsSet,

    /// The set of timestamps of locks that value in it can be accessed during
    /// the reading because they will be committed and their commit_ts <=
    /// read request's start_ts.
    pub access_locks: TsSet,

    /// The data version to match. If it matches the underlying data version,
    /// request will not be processed (i.e. cache hit).
    ///
    /// None means don't try to hit the cache.
    pub cache_match_version: Option<u64>,

    /// The lower bound key in ranges of the request
    pub lower_bound: Vec<u8>,

    /// The upper bound key in ranges of the request
    pub upper_bound: Vec<u8>,

    /// pattern for lazy remote execution based on actual execution cost.
    pub lazy_remote_pattern: Option<String>,
}

impl HeapSize for ReqContext {
    fn approximate_heap_size(&self) -> usize {
        self.context.approximate_heap_size()
            + self.ranges.approximate_heap_size()
            + self.peer.as_ref().map_or(0, |p| p.len())
            + self.lower_bound.approximate_heap_size()
            + self.upper_bound.approximate_heap_size()
    }
}

impl ReqContextInner {
    #[inline]
    pub fn new(
        context: kvrpcpb::Context, // Must be immutable. It's also used by remote cop.
        ranges: Vec<coppb::KeyRange>,
        max_handle_duration: Duration,
        peer: Option<String>,
        is_desc_scan: Option<bool>,
        txn_start_ts: TimeStamp,
        cache_match_version: Option<u64>,
        lazy_remote_pattern: Option<String>,
    ) -> Self {
        let mut deadline_duration = max_handle_duration;
        if context.max_execution_duration_ms > 0 {
            deadline_duration = Duration::from_millis(context.max_execution_duration_ms);
        }
        let deadline = Deadline::from_now(deadline_duration);
        let bypass_locks = TsSet::from_u64s(context.get_resolved_locks().to_vec());
        let access_locks = TsSet::from_u64s(context.get_committed_locks().to_vec());
        let lower_bound = match ranges.first().as_ref() {
            Some(range) => range.start.clone(),
            None => vec![],
        };
        let upper_bound = match ranges.last().as_ref() {
            Some(range) => range.end.clone(),
            None => vec![],
        };
        Self::check_ranges(&context, txn_start_ts, &ranges);
        Self {
            context,
            deadline,
            peer,
            is_desc_scan,
            txn_start_ts,
            ranges,
            bypass_locks,
            access_locks,
            cache_match_version,
            lower_bound,
            upper_bound,
            lazy_remote_pattern,
        }
    }

    #[cfg(test)]
    pub fn default_for_test() -> Self {
        Self::new(
            Default::default(),
            Vec::new(),
            Duration::from_secs(100),
            None,
            None,
            TimeStamp::max(),
            None,
            None,
        )
    }

    fn check_ranges(ctx: &kvrpcpb::Context, txn_start_ts: TimeStamp, ranges: &[coppb::KeyRange]) {
        for ran in ranges.windows(2) {
            // Ranges should be in order. And this judge can also find disorder.
            if ran[0].end.is_empty() || ran[0].end > ran[1].start {
                warn!("coprocessor request range overlapped";
                    "region" => ctx.region_id,
                    "ts" => ?txn_start_ts,
                    "ran" => ?ran,
                );
                return;
            }
        }
    }
}

/// ReqContext provides the context for the coprocessor request.
/// It Encapsulate the `kvrpcpb::Context` to provide some extra properties.
/// It is immutable and can be shared across threads.
#[derive(Debug, Clone, PartialEq)]
pub struct ReqContext(Arc<ReqContextInner>);

impl Deref for ReqContext {
    type Target = Arc<ReqContextInner>;

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

impl From<ReqContextInner> for ReqContext {
    fn from(inner: ReqContextInner) -> Self {
        Self(Arc::new(inner))
    }
}

impl ReqContext {
    #[inline]
    pub fn new(
        context: kvrpcpb::Context,
        ranges: Vec<coppb::KeyRange>,
        max_handle_duration: Duration,
        peer: Option<String>,
        is_desc_scan: Option<bool>,
        txn_start_ts: TimeStamp,
        cache_match_version: Option<u64>,
        lazy_remote_pattern: Option<String>,
    ) -> Self {
        ReqContextInner::new(
            context,
            ranges,
            max_handle_duration,
            peer,
            is_desc_scan,
            txn_start_ts,
            cache_match_version,
            lazy_remote_pattern,
        )
        .into()
    }

    #[cfg(test)]
    pub fn default_for_test() -> Self {
        ReqContextInner::default_for_test().into()
    }

    pub fn build_task_id(&self) -> u64 {
        const ID_SHIFT: u32 = 16;
        const MASK: u64 = u64::MAX >> ID_SHIFT;
        const MAX_TS: u64 = u64::MAX;
        let base = match self.txn_start_ts.into_inner() {
            0 | MAX_TS => thread_rng().next_u64(),
            start_ts => start_ts,
        };
        let task_id: u64 = self.context.get_task_id();
        if task_id > 0 {
            // It is assumed that the lower bits of task IDs in a single transaction
            // tend to be different. So if task_id is provided, we concatenate the
            // low 16 bits of the task_id and the low 48 bits of the start_ts to build
            // the final task id.
            (task_id << (64 - ID_SHIFT)) | (base & MASK)
        } else {
            // Otherwise we use the start_ts as the task_id.
            base
        }
    }
}

lazy_static! {
    pub static ref MEMTRACE_ROOT: Arc<MemoryTrace> = mem_trace!(coprocessor, [analyze]);
    pub static ref MEMTRACE_ANALYZE: Arc<MemoryTrace> =
        MEMTRACE_ROOT.sub_trace(Id::Name("analyze"));
}

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

    fn default_req_ctx_with_ctx_duration(
        context: kvrpcpb::Context,
        max_handle_duration: Duration,
    ) -> ReqContext {
        ReqContext::new(
            context,
            Vec::new(),
            max_handle_duration,
            None,
            None,
            TimeStamp::max(),
            None,
            None,
        )
    }

    #[test]
    fn test_build_task_id() {
        let mut inner = ReqContextInner::default_for_test();
        let start_ts: u64 = 0x05C6_1BFA_2648_324A;
        inner.txn_start_ts = start_ts.into();
        inner.context.set_task_id(1);
        let ctx: ReqContext = inner.clone().into();
        assert_eq!(ctx.build_task_id(), 0x0001_1BFA_2648_324A);

        inner.context.set_task_id(0);
        let ctx: ReqContext = inner.clone().into();
        assert_eq!(ctx.build_task_id(), start_ts);
    }

    #[test]
    fn test_deadline_from_req_ctx() {
        let ctx = kvrpcpb::Context::default();
        let max_handle_duration = Duration::from_millis(100);
        let req_ctx = default_req_ctx_with_ctx_duration(ctx, max_handle_duration);
        // sleep at least 100ms
        std::thread::sleep(Duration::from_millis(200));
        req_ctx
            .deadline
            .check()
            .expect_err("deadline should exceed");

        let mut ctx = kvrpcpb::Context::default();
        ctx.max_execution_duration_ms = 100_000;
        let req_ctx = default_req_ctx_with_ctx_duration(ctx, max_handle_duration);
        // sleep at least 100ms
        std::thread::sleep(Duration::from_millis(200));
        req_ctx
            .deadline
            .check()
            .expect("deadline should not exceed");
    }

    #[test]
    fn test_req_ctx_new_match_inner() {
        let pb_ctx = kvrpcpb::Context {
            region_id: 12345,
            ..Default::default()
        };
        let ranges = vec![coppb::KeyRange {
            start: vec![1, 2, 3],
            end: vec![4, 5, 6],
            ..Default::default()
        }];
        let max_handle_duration = Duration::from_secs(100);
        let peer = Some(String::from("1223"));
        let is_desc_scan = Some(true);
        let txn_start_ts = TimeStamp::new(9898);
        let cache_match_version = Some(42);

        let mut inner = ReqContextInner::new(
            pb_ctx.clone(),
            ranges.clone(),
            max_handle_duration,
            peer.clone(),
            is_desc_scan,
            txn_start_ts,
            cache_match_version,
            None,
        );

        let ctx = ReqContext::new(
            pb_ctx,
            ranges,
            max_handle_duration,
            peer,
            is_desc_scan,
            txn_start_ts,
            cache_match_version,
            None,
        );

        // deadlines are not exactly equal, just compare the delta
        assert!(
            ctx.deadline.inner().duration_since(inner.deadline.inner()) < Duration::from_secs(1),
            "{:?} - {:?} > 1s",
            ctx.deadline,
            inner.deadline
        );
        inner.deadline = ctx.deadline;
        // test ReqContextInner::new and ReqContext::new should match
        assert_eq!(ctx.0.as_ref().clone(), inner);
    }
}