// Copyright 20211 TiKV Project Authors. Licensed under Apache-2.0.
// Modifications Copyright 2023 PingCAP, Inc.

use api_version::KvFormat;
use collections::HashMap;
use kvproto::kvrpcpb::*;
use raftstore::store::is_load_based_split_enabled;
use tikv::{
    server::{
        metrics::{GrpcTypeKind, REQUEST_BATCH_SIZE_HISTOGRAM_VEC},
        service::{GrpcRequestDuration, MeasuredSingleResponse, batch_commands_response},
    },
    storage::{
        ResponseBatchConsumer, Result, Storage,
        errors::{extract_key_error, extract_region_error},
        kv::{Engine, Statistics},
        lock_manager::LockManager,
    },
};
use tikv_util::{
    future::poll_future_notify,
    mpsc::future::{Sender, WakePolicy},
    time::Instant,
};
use tracker::{GLOBAL_TRACKERS, RequestInfo, RequestType, Tracker, TrackerToken, with_tls_tracker};

pub const MAX_BATCH_GET_REQUEST_COUNT: usize = 10;
pub const MIN_BATCH_GET_REQUEST_COUNT: usize = 4;
pub const MAX_QUEUE_SIZE_PER_WORKER: usize = 16;

#[derive(Clone, Copy, Debug, Eq, Hash, PartialEq)]
struct GetBatchKey {
    keyspace_id: u32,
    region_id: Option<u64>,
}

impl GetBatchKey {
    fn new(req: &GetRequest, partition_by_region: bool) -> Self {
        let ctx = req.get_context();
        Self {
            keyspace_id: ctx.get_keyspace_id(),
            region_id: partition_by_region.then(|| ctx.get_region_id()),
        }
    }

    fn matches(self, req: &GetRequest) -> bool {
        let ctx = req.get_context();
        ctx.get_keyspace_id() == self.keyspace_id
            && self
                .region_id
                .is_none_or(|region_id| ctx.get_region_id() == region_id)
    }
}

struct PendingGetBatch {
    key: GetBatchKey,
    gets: Vec<GetRequest>,
    ids: Vec<u64>,
    trackers: Vec<TrackerToken>,
}

impl PendingGetBatch {
    fn new(key: GetBatchKey, capacity: usize) -> Self {
        Self {
            key,
            gets: Vec::with_capacity(capacity),
            ids: Vec::with_capacity(capacity),
            trackers: Vec::with_capacity(capacity),
        }
    }

    fn push(&mut self, req: GetRequest, id: u64, tracker: TrackerToken) {
        debug_assert!(self.key.matches(&req));
        self.gets.push(req);
        self.ids.push(id);
        self.trackers.push(tracker);
    }

    fn len(&self) -> usize {
        self.gets.len()
    }
}

pub struct ReqBatcher {
    get_batches: HashMap<GetBatchKey, PendingGetBatch>,
    ready_get_batches: Vec<PendingGetBatch>,
    begin_instant: Instant,
    batch_size: usize,
    partition_by_region: bool,
}

impl ReqBatcher {
    pub fn new(batch_size: usize, partition_by_region: bool) -> ReqBatcher {
        let begin_instant = Instant::now_coarse();
        ReqBatcher {
            get_batches: HashMap::default(),
            ready_get_batches: vec![],
            begin_instant,
            batch_size: std::cmp::min(batch_size, MAX_BATCH_GET_REQUEST_COUNT),
            partition_by_region,
        }
    }

    pub fn can_batch_get(&self, req: &GetRequest) -> bool {
        req.get_context().get_priority() == CommandPri::Normal
    }

    pub fn add_get_request(&mut self, req: GetRequest, id: u64) {
        let tracker = GLOBAL_TRACKERS.insert(Tracker::new(RequestInfo::new(
            req.get_context(),
            RequestType::KvBatchGetCommand,
            req.get_version(),
        )));
        let key = GetBatchKey::new(&req, self.partition_by_region);
        let batch_size = self.batch_size;
        let batch = self
            .get_batches
            .entry(key)
            .or_insert_with(|| PendingGetBatch::new(key, batch_size));
        batch.push(req, id, tracker);
        if batch.len() >= self.batch_size {
            self.ready_get_batches
                .push(self.get_batches.remove(&key).unwrap());
        }
    }

    pub fn maybe_commit<E: Engine, L: LockManager, F: KvFormat>(
        &mut self,
        storage: &Storage<E, L, F>,
        tx: &Sender<MeasuredSingleResponse>,
    ) {
        for batch in self.ready_get_batches.drain(..) {
            dispatch_get_batch(storage, tx, batch, self.begin_instant);
        }
    }

    pub fn commit<E: Engine, L: LockManager, F: KvFormat>(
        self,
        storage: &Storage<E, L, F>,
        tx: &Sender<MeasuredSingleResponse>,
    ) {
        for batch in self
            .ready_get_batches
            .into_iter()
            .chain(self.get_batches.into_values())
        {
            dispatch_get_batch(storage, tx, batch, self.begin_instant);
        }
    }
}

fn dispatch_get_batch<E: Engine, L: LockManager, F: KvFormat>(
    storage: &Storage<E, L, F>,
    tx: &Sender<MeasuredSingleResponse>,
    batch: PendingGetBatch,
    begin_instant: Instant,
) {
    debug_assert!(
        batch.gets.iter().all(|req| batch.key.matches(req)),
        "a scheduled Get batch must match its batching key"
    );
    future_batch_get_command(
        storage,
        batch.key.keyspace_id,
        batch.ids,
        batch.gets,
        batch.trackers,
        tx.clone(),
        begin_instant,
    );
}

pub struct BatcherBuilder {
    pool_size: usize,
    enable_batch: bool,
}

impl BatcherBuilder {
    pub fn new(enable_batch: bool, pool_size: usize) -> Self {
        BatcherBuilder {
            enable_batch,
            pool_size,
        }
    }
    pub fn build(&self, queue_per_worker: usize, req_batch_size: usize) -> Option<ReqBatcher> {
        let partition_by_region = is_load_based_split_enabled();
        fail::fail_point!("force_batch_get_command", |_| {
            self.enable_batch
                .then(|| ReqBatcher::new(req_batch_size, partition_by_region))
        });
        if !self.enable_batch {
            return None;
        }
        if req_batch_size > self.pool_size * MIN_BATCH_GET_REQUEST_COUNT
            && queue_per_worker >= MIN_BATCH_GET_REQUEST_COUNT
        {
            return Some(ReqBatcher::new(
                req_batch_size / self.pool_size,
                partition_by_region,
            ));
        }
        if req_batch_size >= MIN_BATCH_GET_REQUEST_COUNT
            && queue_per_worker >= MAX_QUEUE_SIZE_PER_WORKER
        {
            return Some(ReqBatcher::new(req_batch_size, partition_by_region));
        }
        None
    }
}

pub struct GetCommandResponseConsumer {
    tx: Sender<MeasuredSingleResponse>,
}

impl ResponseBatchConsumer<(Option<Vec<u8>>, Statistics)> for GetCommandResponseConsumer {
    fn consume(
        &self,
        id: u64,
        res: Result<(Option<Vec<u8>>, Statistics)>,
        begin: Instant,
        request_source: String,
    ) {
        let mut resp = GetResponse::default();
        if let Some(err) = extract_region_error(&res) {
            resp.set_region_error(err);
        } else {
            match res {
                Ok((val, statistics)) => {
                    let scan_detail_v2 = resp.mut_exec_details_v2().mut_scan_detail_v2();
                    statistics.write_scan_detail(scan_detail_v2);
                    with_tls_tracker(|tracker| tracker.write_scan_detail(scan_detail_v2));
                    match val {
                        Some(val) => resp.set_value(val),
                        None => resp.set_not_found(true),
                    }
                }
                Err(e) => resp.set_error(extract_key_error(&e)),
            }
        }
        with_tls_tracker(|tracker| {
            if let Some(details) = tracker.read_pool_task_details() {
                resp.mut_exec_details_v2()
                    .set_read_pool_task_details(details);
            }
        });

        let res = batch_commands_response::Response {
            cmd: Some(batch_commands_response::response::Cmd::Get(resp)),
            ..Default::default()
        };
        let mesure =
            GrpcRequestDuration::new(begin, GrpcTypeKind::kv_batch_get_command, request_source);
        let task = MeasuredSingleResponse::new(id, res, mesure);
        if self.tx.send_with(task, WakePolicy::Immediately).is_err() {
            error!("KvService response batch commands fail");
        }
    }
}

impl ResponseBatchConsumer<Option<Vec<u8>>> for GetCommandResponseConsumer {
    fn consume(
        &self,
        id: u64,
        res: Result<Option<Vec<u8>>>,
        begin: Instant,
        request_source: String,
    ) {
        let mut resp = RawGetResponse::default();
        if let Some(err) = extract_region_error(&res) {
            resp.set_region_error(err);
        } else {
            match res {
                Ok(Some(val)) => resp.set_value(val),
                Ok(None) => resp.set_not_found(true),
                Err(e) => resp.set_error(format!("{}", e)),
            }
        }
        let res = batch_commands_response::Response {
            cmd: Some(batch_commands_response::response::Cmd::RawGet(resp)),
            ..Default::default()
        };
        let mesure =
            GrpcRequestDuration::new(begin, GrpcTypeKind::raw_batch_get_command, request_source);
        let task = MeasuredSingleResponse::new(id, res, mesure);
        if self.tx.send_with(task, WakePolicy::Immediately).is_err() {
            error!("KvService response batch commands fail");
        }
    }
}

fn future_batch_get_command<E: Engine, L: LockManager, F: KvFormat>(
    storage: &Storage<E, L, F>,
    keyspace_id: u32,
    requests: Vec<u64>,
    gets: Vec<GetRequest>,
    trackers: Vec<TrackerToken>,
    tx: Sender<MeasuredSingleResponse>,
    begin_instant: tikv_util::time::Instant,
) {
    REQUEST_BATCH_SIZE_HISTOGRAM_VEC
        .kv_get
        .observe(gets.len() as f64);
    let id_sources: Vec<_> = requests
        .iter()
        .zip(gets.iter())
        .map(|(id, req)| (*id, req.get_context().get_request_source().to_string()))
        .collect();
    let res = storage.batch_get_command(
        keyspace_id,
        gets,
        requests,
        trackers.clone(),
        GetCommandResponseConsumer { tx: tx.clone() },
        begin_instant,
    );
    let f = async move {
        // This error can only cause by readpool busy.
        let res = res.await;
        for tracker in trackers {
            GLOBAL_TRACKERS.remove(tracker);
        }
        if let Some(e) = extract_region_error(&res) {
            let mut resp = GetResponse::default();
            resp.set_region_error(e);
            for (id, source) in id_sources {
                let res = batch_commands_response::Response {
                    cmd: Some(batch_commands_response::response::Cmd::Get(resp.clone())),
                    ..Default::default()
                };
                let measure = GrpcRequestDuration::new(
                    begin_instant,
                    GrpcTypeKind::kv_batch_get_command,
                    source,
                );
                let task = MeasuredSingleResponse::new(id, res, measure);
                if tx.send_with(task, WakePolicy::Immediately).is_err() {
                    error!("KvService response batch commands fail");
                }
            }
        }
    };
    poll_future_notify(f);
}

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

    fn new_get_request(keyspace_id: u32, region_id: u64) -> GetRequest {
        let mut req = GetRequest::default();
        req.mut_context().set_keyspace_id(keyspace_id);
        req.mut_context().set_region_id(region_id);
        req.mut_context().set_priority(CommandPri::Normal);
        req
    }

    fn remove_trackers(batcher: &ReqBatcher) {
        for batch in batcher
            .ready_get_batches
            .iter()
            .chain(batcher.get_batches.values())
        {
            for tracker in &batch.trackers {
                GLOBAL_TRACKERS.remove(*tracker);
            }
        }
    }

    #[test]
    fn test_req_batcher_partitions_gets_by_keyspace() {
        const KEYSPACE_A: u32 = 11;
        const KEYSPACE_B: u32 = 22;
        let mut batcher = ReqBatcher::new(2, false);

        batcher.add_get_request(new_get_request(KEYSPACE_A, 1), 1);
        batcher.add_get_request(new_get_request(KEYSPACE_B, 1), 2);
        batcher.add_get_request(new_get_request(KEYSPACE_A, 2), 3);

        assert_eq!(batcher.ready_get_batches.len(), 1);
        let ready_batch = &batcher.ready_get_batches[0];
        assert_eq!(ready_batch.key.keyspace_id, KEYSPACE_A);
        assert_eq!(ready_batch.key.region_id, None);
        assert_eq!(ready_batch.ids, vec![1, 3]);
        assert!(
            ready_batch
                .gets
                .iter()
                .all(|req| req.get_context().get_keyspace_id() == KEYSPACE_A)
        );

        let keyspace_a = GetBatchKey {
            keyspace_id: KEYSPACE_A,
            region_id: None,
        };
        let keyspace_b = GetBatchKey {
            keyspace_id: KEYSPACE_B,
            region_id: None,
        };
        assert!(!batcher.get_batches.contains_key(&keyspace_a));
        let pending_batch = batcher.get_batches.get(&keyspace_b).unwrap();
        assert_eq!(pending_batch.ids, vec![2]);
        assert!(
            pending_batch
                .gets
                .iter()
                .all(|req| req.get_context().get_keyspace_id() == KEYSPACE_B)
        );

        remove_trackers(&batcher);
    }

    #[test]
    fn test_req_batcher_partitions_gets_by_region_for_load_split() {
        const KEYSPACE_ID: u32 = 11;
        const REGION_A: u64 = 101;
        const REGION_B: u64 = 202;
        let mut batcher = ReqBatcher::new(2, true);

        batcher.add_get_request(new_get_request(KEYSPACE_ID, REGION_A), 1);
        batcher.add_get_request(new_get_request(KEYSPACE_ID, REGION_B), 2);
        batcher.add_get_request(new_get_request(KEYSPACE_ID, REGION_A), 3);

        assert_eq!(batcher.ready_get_batches.len(), 1);
        let ready_batch = &batcher.ready_get_batches[0];
        assert_eq!(ready_batch.key.keyspace_id, KEYSPACE_ID);
        assert_eq!(ready_batch.key.region_id, Some(REGION_A));
        assert_eq!(ready_batch.ids, vec![1, 3]);

        let region_a = GetBatchKey {
            keyspace_id: KEYSPACE_ID,
            region_id: Some(REGION_A),
        };
        let region_b = GetBatchKey {
            keyspace_id: KEYSPACE_ID,
            region_id: Some(REGION_B),
        };
        assert!(!batcher.get_batches.contains_key(&region_a));
        assert_eq!(batcher.get_batches.get(&region_b).unwrap().ids, vec![2]);

        remove_trackers(&batcher);
    }
}