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

use std::{
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
    time::{Duration, Instant},
};

use collections::{HashMap, HashMapEntry};
use crossbeam::atomic::AtomicCell;
use futures::{
    SinkExt,
    compat::Stream01CompatExt,
    stream::{StreamExt, TryStreamExt},
};
use grpcio::{DuplexSink, RequestStream, RpcContext, RpcStatus, RpcStatusCode};
use kvproto::cdcpb::{
    ChangeData, ChangeDataEvent, ChangeDataRequest, ChangeDataRequest_oneof_request,
};
use tikv_util::{error, info, memory::MemoryQuota, timer::GLOBAL_TIMER_HANDLE, warn, worker::*};

use crate::{
    channel::{CDC_CHANNEL_CAPACITY, Sink, channel},
    delegate::{Downstream, DownstreamId, DownstreamState, ObservedRange},
    endpoint::{Deregister, Task},
    metrics::CDC_ABORTED_CONNECTIONS,
};

static CONNECTION_ID_ALLOC: AtomicUsize = AtomicUsize::new(0);

// CDC connection monitoring constants in seconds
const CDC_WATCHDOG_CHECK_INTERVAL_SECS: u64 = 2;
const CDC_IDLE_DEREGISTER_THRESHOLD_SECS: u64 = 60 * 20; // 20 minutes
const CDC_MEMORY_QUOTA_ABORT_THRESHOLD: f64 = 0.999;

/// A unique identifier of a Connection.
#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct ConnId(usize);

#[derive(Clone, Copy, Debug, Eq, PartialEq, Hash)]
pub struct RequestId(pub u64);

impl ConnId {
    pub fn new() -> ConnId {
        ConnId(CONNECTION_ID_ALLOC.fetch_add(1, Ordering::SeqCst))
    }

    pub fn into_inner(self) -> usize {
        self.0
    }
}

impl Default for ConnId {
    fn default() -> Self {
        Self::new()
    }
}

pub struct Conn {
    id: ConnId,
    sink: Sink,
    downstreams: HashMap<DownstreamKey, DownstreamValue>,
    peer: String,
}

#[derive(PartialEq, Eq, Hash)]
struct DownstreamKey {
    request_id: RequestId,
    region_id: u64,
}

#[derive(Clone)]
struct DownstreamValue {
    id: DownstreamId,
    state: Arc<AtomicCell<DownstreamState>>,
}

impl Conn {
    pub fn new(conn_id: ConnId, sink: Sink, peer: String) -> Conn {
        Conn {
            id: conn_id,
            sink,
            downstreams: HashMap::default(),
            peer,
        }
    }

    pub fn get_peer(&self) -> &str {
        &self.peer
    }

    pub fn get_id(&self) -> ConnId {
        self.id
    }

    pub fn get_sink(&self) -> &Sink {
        &self.sink
    }

    pub fn get_downstream(&self, request_id: RequestId, region_id: u64) -> Option<DownstreamId> {
        let key = DownstreamKey {
            request_id,
            region_id,
        };
        self.downstreams.get(&key).map(|v| v.id)
    }

    pub fn subscribe(
        &mut self,
        request_id: RequestId,
        region_id: u64,
        downstream_id: DownstreamId,
        downstream_state: Arc<AtomicCell<DownstreamState>>,
    ) -> Option<DownstreamId> {
        let key = DownstreamKey {
            request_id,
            region_id,
        };
        match self.downstreams.entry(key) {
            HashMapEntry::Occupied(value) => Some(value.get().id),
            HashMapEntry::Vacant(v) => {
                v.insert(DownstreamValue {
                    id: downstream_id,
                    state: downstream_state,
                });
                None
            }
        }
    }

    pub fn unsubscribe(&mut self, request_id: RequestId, region_id: u64) -> Option<DownstreamId> {
        let key = DownstreamKey {
            request_id,
            region_id,
        };
        self.downstreams.remove(&key).map(|value| value.id)
    }

    pub fn unsubscribe_request(&mut self, request_id: RequestId) -> Vec<(u64, DownstreamId)> {
        let mut downstreams = Vec::new();
        self.downstreams.retain(|key, value| -> bool {
            if key.request_id == request_id {
                downstreams.push((key.region_id, value.id));
                return false;
            }
            true
        });
        downstreams
    }

    pub fn iter_downstreams<F>(&self, mut f: F)
    where
        F: FnMut(RequestId, u64, DownstreamId, &Arc<AtomicCell<DownstreamState>>),
    {
        for (key, value) in &self.downstreams {
            f(key.request_id, key.region_id, value.id, &value.state);
        }
    }

    #[cfg(test)]
    pub fn downstreams_count(&self) -> usize {
        self.downstreams.len()
    }
}

/// Service implements the `ChangeData` service.
///
/// It's a front-end of the CDC service, schedules requests to the `Endpoint`.
#[derive(Clone)]
pub struct Service {
    scheduler: Scheduler<Task>,
    memory_quota: Arc<MemoryQuota>,
    pool: Arc<Worker>,
}

impl Service {
    /// Create a ChangeData service.
    ///
    /// It requires a scheduler of an `Endpoint` in order to schedule tasks.
    pub fn new(
        scheduler: Scheduler<Task>,
        memory_quota: Arc<MemoryQuota>,
        pool: Arc<Worker>,
    ) -> Service {
        Service {
            scheduler,
            memory_quota,
            pool,
        }
    }

    // ### Command types:
    // * Register registers a region. 1) both `request_id` and `region_id` must be
    //   specified; 2) `request_id` can be 0 but `region_id` can not.
    // * Deregister deregisters some regions in one same `request_id` or just one
    //   region. 1) if both `request_id` and `region_id` are specified, just
    //   deregister the region; 2) if only `request_id` is specified, all region
    //   subscriptions with the same `request_id` will be deregistered.
    fn handle_request(
        scheduler: &Scheduler<Task>,
        peer: &str,
        request: ChangeDataRequest,
        conn_id: ConnId,
    ) -> Result<(), String> {
        match request.request {
            None | Some(ChangeDataRequest_oneof_request::Register(_)) => {
                Self::handle_register(scheduler, peer, request, conn_id)
            }
            Some(ChangeDataRequest_oneof_request::Deregister(_)) => {
                Self::handle_deregister(scheduler, request, conn_id)
            }
            _ => unreachable!(),
        }
    }

    fn handle_register(
        scheduler: &Scheduler<Task>,
        peer: &str,
        request: ChangeDataRequest,
        conn_id: ConnId,
    ) -> Result<(), String> {
        let observed_range = ObservedRange::new(request.start_key.clone(), request.end_key.clone())
            .unwrap_or_else(|e| {
                warn!(
                    "cdc invalid observed start key or end key version";
                    "downstream" => ?peer,
                    "region_id" => request.region_id,
                    "request_id" => request.region_id,
                    "error" => ?e,
                    "start_key" => log_wrappers::Value::key(&request.start_key),
                    "end_key" => log_wrappers::Value::key(&request.end_key),
                );
                ObservedRange::default()
            });
        let downstream = Downstream::new(
            peer.to_owned(),
            request.get_region_epoch().clone(),
            RequestId(request.request_id),
            conn_id,
            request.filter_loop,
            observed_range,
        );
        let task = Task::Register {
            request,
            downstream,
        };
        scheduler.schedule(task).map_err(|e| format!("{:?}", e))
    }

    fn handle_deregister(
        scheduler: &Scheduler<Task>,
        request: ChangeDataRequest,
        conn_id: ConnId,
    ) -> Result<(), String> {
        let task = if request.region_id != 0 {
            Task::Deregister(Deregister::Region {
                conn_id,
                request_id: RequestId(request.request_id),
                region_id: request.region_id,
            })
        } else {
            Task::Deregister(Deregister::Request {
                conn_id,
                request_id: RequestId(request.request_id),
            })
        };
        scheduler.schedule(task).map_err(|e| format!("{:?}", e))
    }

    fn handle_event_feed(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<ChangeDataRequest>,
        mut sink: DuplexSink<ChangeDataEvent>,
    ) {
        sink.enhance_batch(true);
        let conn_id = ConnId::new();
        let (event_sink, mut event_drain) =
            channel(conn_id, CDC_CHANNEL_CAPACITY, self.memory_quota.clone());
        let conn = Conn::new(conn_id, event_sink, ctx.peer());

        if let Err(e) = self.scheduler.schedule(Task::OpenConn { conn }) {
            let peer = ctx.peer();
            error!("cdc connection initiate failed"; "error" => ?e, "downstream" => ?peer, "conn_id" => ?conn_id);
            ctx.spawn(async move {
                let status = RpcStatus::with_message(RpcStatusCode::UNKNOWN, format!("{:?}", e));
                if let Err(e) = sink.fail(status).await {
                    error!("cdc failed to send error"; "error" => ?e, "downstream" => ?peer, "conn_id" => ?conn_id);
                }
            });
            return;
        }

        info!("cdc connection created"; "downstream" => ctx.peer(), "conn_id" => ?conn_id);

        let peer = ctx.peer();
        let scheduler = self.scheduler.clone();
        let recv_req = async move {
            let mut stream = stream.map_err(|e| format!("{:?}", e));
            while let Some(request) = stream.try_next().await? {
                Self::handle_request(&scheduler, &peer, request, conn_id)?;
            }
            Ok::<(), String>(())
        };

        let scheduler = self.scheduler.clone();
        let peer = ctx.peer();
        ctx.spawn(async move {
            if let Err(e) = recv_req.await {
                warn!("cdc receive failed"; "error" => ?e, "downstream" => peer, "conn_id" => ?conn_id);
            } else {
                info!("cdc receive closed"; "downstream" => peer, "conn_id" => ?conn_id);
            }
            let deregister = Deregister::Conn(conn_id);
            if let Err(e) = scheduler.schedule(Task::Deregister(deregister)) {
                error!("cdc deregister failed"; "error" => ?e, "conn_id" => ?conn_id);
            }
        });

        let last_flush_time = Arc::new(AtomicCell::new(Instant::now()));
        let last_flush_time_for_forward = last_flush_time.clone();
        let peer_for_watchdog = ctx.peer();

        let peer = ctx.peer();
        let (cancel_tx, mut cancel_rx) = tokio::sync::oneshot::channel::<()>();

        // Create cancelCh for eventDrain.forward exit signal
        let (forward_exit_tx, forward_exit_rx) = tokio::sync::oneshot::channel::<()>();

        ctx.spawn(async move {
            #[cfg(feature = "failpoints")]
            sleep_before_drain_change_event().await;
            tokio::select! {
                _ = &mut cancel_rx => {
                    warn!("cdc send cancelled"; "downstream" => peer, "conn_id" => ?conn_id);
                    let status = RpcStatus::with_message(RpcStatusCode::UNKNOWN, "connection cancelled".to_string());
                    let _ = sink.fail(status).await;
                    CDC_ABORTED_CONNECTIONS.inc();
                }
                result = event_drain.forward(&mut sink, Some(&last_flush_time_for_forward)) => {
                    if let Err(e) = result {
                        warn!("cdc send failed"; "error" => ?e, "downstream" => peer, "conn_id" => ?conn_id);
                        let status = RpcStatus::with_message(RpcStatusCode::UNKNOWN,  format!("{:?}", e));
                        sink.fail(status).await.unwrap_or_else(|e| {
                            error!("cdc failed to send error"; "error" => ?e, "conn_id" => ?conn_id);
                        });
                    } else {
                        info!("cdc send closed"; "downstream" => peer, "conn_id" => ?conn_id);
                        let _ = sink.close().await;
                    }
                    // Send signal when eventDrain.forward exits
                    let _ = forward_exit_tx.send(());
                }
            }
        });

        // Start watchdog to monitor connection activity
        Self::start_connection_watchdog(
            self.pool.clone(),
            last_flush_time,
            peer_for_watchdog,
            conn_id,
            cancel_tx,
            forward_exit_rx,
            self.memory_quota.clone(),
        );
    }

    /// Start a watchdog to monitor CDC connection activity.
    ///
    /// This function creates a background task that periodically checks if the
    /// connection has been idle for too long and takes appropriate action
    /// (warning or aborting).
    fn start_connection_watchdog(
        pool: Arc<Worker>,
        last_flush_time: Arc<AtomicCell<Instant>>,
        peer: String,
        conn_id: ConnId,
        cancel_tx: tokio::sync::oneshot::Sender<()>,
        mut forward_exit_rx: tokio::sync::oneshot::Receiver<()>,
        memory_quota: Arc<MemoryQuota>,
    ) {
        // Create a custom interval task that can be stopped
        let _ = pool.pool().spawn(async move {
            let mut interval = GLOBAL_TIMER_HANDLE
                .interval(
                    Instant::now(),
                    Duration::from_secs(CDC_WATCHDOG_CHECK_INTERVAL_SECS),
                )
                .compat();

            loop {
                tokio::select! {
                    _ = &mut forward_exit_rx => {
                        info!("cdc connection forward exit signal received, stopping watchdog");
                        break;
                    }
                    _ = interval.next() => {
                        let elapsed = last_flush_time.load().elapsed();
                        let _idle_threshold = CDC_IDLE_DEREGISTER_THRESHOLD_SECS;

                        #[cfg(feature = "failpoints")]
                        let _idle_threshold = {
                            let should_adjust = || {
                                fail::fail_point!("cdc_idle_deregister_threshold", |_| true);
                                false
                            };
                            if should_adjust() {
                                5
                            } else {
                                CDC_IDLE_DEREGISTER_THRESHOLD_SECS
                            }
                        };

                        // Check if last flush was more than the deregister threshold
                        // To prevent the case that the connection idle since there are a lot of
                        // incremental scan tasks queueing so won't send events, also check on the
                        // memory usage, if the memory quota is almost used up, we abort the connection.
                        if elapsed > Duration::from_secs(_idle_threshold)
                            && memory_quota.used_ratio() >= CDC_MEMORY_QUOTA_ABORT_THRESHOLD {
                            error!("cdc connection idle for too long, aborting connection";
                                   "downstream" => peer.clone(),
                                   "conn_id" => ?conn_id,
                                   "seconds_since_last_flush" => elapsed.as_secs());
                            // Cancel the gRPC connection
                            let _ = cancel_tx.send(());
                            break;
                        }
                    }
                }
            }
        });
    }
}

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

    // ### Why `v2`
    // `v2` is expected to resolve this problem: clients version is higher than
    // server, In which case, `v1` compatibility check mechanism doesn't work.
    //
    // ### How `v2`
    // In `v2`, clients tells requested features to connected servers. If a
    // server finds a client requires unavailable features, it can fail the
    // connection with an UNIMPLEMENTED status code.
    //
    // ### Details about `v2` features
    // * stream-multiplexing: a region can be subscribed multiple times in one
    //   `Conn` with different `request_id`.
    fn event_feed_v2(
        &mut self,
        ctx: RpcContext<'_>,
        stream: RequestStream<ChangeDataRequest>,
        sink: DuplexSink<ChangeDataEvent>,
    ) {
        self.handle_event_feed(ctx, stream, sink);
    }
}

#[cfg(feature = "failpoints")]
#[allow(unused_imports)]
async fn sleep_before_drain_change_event() {
    use std::time::{Duration, Instant};

    use tikv_util::timer::GLOBAL_TIMER_HANDLE;
    let should_sleep = || {
        fail::fail_point!("cdc_sleep_before_drain_change_event", |_| true);
        false
    };
    if should_sleep() {
        let dur = Duration::from_secs(5);
        let timer = GLOBAL_TIMER_HANDLE.delay(Instant::now() + dur);
        let _ = futures::compat::Compat01As03::new(timer).await;
    }
}

#[cfg(test)]
mod tests {
    use std::{sync::Arc, time::Duration};

    use futures::{SinkExt, executor::block_on};
    use grpcio::{self, ChannelBuilder, EnvBuilder, Server, ServerBuilder, WriteFlags};
    use kvproto::cdcpb::{ChangeDataClient, ResolvedTs, create_change_data};
    use tikv_util::future::block_on_timeout;

    use super::*;
    use crate::channel::{CdcEvent, recv_timeout};

    fn new_rpc_suite(capacity: usize) -> (Server, ChangeDataClient, ReceiverWrapper<Task>) {
        let memory_quota = Arc::new(MemoryQuota::new(capacity));
        let pool = Arc::new(Builder::new("cdc-watchdog-test").thread_count(1).create());
        let (scheduler, rx) = dummy_scheduler();
        let cdc_service = Service::new(scheduler, memory_quota, pool);
        let env = Arc::new(EnvBuilder::new().build());
        let builder =
            ServerBuilder::new(env.clone()).register_service(create_change_data(cdc_service));
        let mut server = builder.bind("127.0.0.1", 0).build().unwrap();
        server.start();
        let (_, port) = server.bind_addrs().next().unwrap();
        let addr = format!("127.0.0.1:{}", port);
        let channel = ChannelBuilder::new(env).connect(&addr);
        let client = ChangeDataClient::new(channel);
        (server, client, rx)
    }

    #[test]
    fn test_flow_control() {
        // Disable CDC sink memory quota.
        let capacity = usize::MAX;
        let (_server, client, mut task_rx) = new_rpc_suite(capacity);
        // Create a event feed stream.
        let (mut tx, mut rx) = client.event_feed_v2().unwrap();
        let mut req = ChangeDataRequest {
            region_id: 1,
            ..Default::default()
        };
        req.mut_header().set_ticdc_version("4.0.7".into());
        block_on(tx.send((req, WriteFlags::default()))).unwrap();
        let task = task_rx.recv_timeout(Duration::from_millis(100)).unwrap();
        let conn = if let Some(Task::OpenConn { conn }) = task {
            conn
        } else {
            panic!("expect to be Task::OpenConn");
        };
        let sink = conn.get_sink().clone();
        // Approximate 1 KB.
        let mut rts = ResolvedTs::default();
        rts.set_regions(vec![u64::MAX; 128]);

        let send = || {
            let rts_ = rts.clone();
            let mut sink_ = sink.clone();
            Box::pin(async move {
                sink_
                    .send_all(
                        vec![CdcEvent::ResolvedTs(rts_)],
                        Arc::new(Default::default()),
                    )
                    .await
            })
        };
        let must_fill_window = || {
            let mut window_size = 0;
            loop {
                if matches!(
                    block_on_timeout(send(), Duration::from_millis(100)),
                    Err(_) | Ok(Err(_))
                ) {
                    // Window is filled and flow control in sink is triggered.
                    break;
                }
                window_size += 1;
                // gRPC window size should not be larger than 1GB.
                assert!(window_size <= 1024 * 1024, "window_size: {}", window_size);
            }
            window_size
        };

        // Fill gRPC window.
        let window_size = must_fill_window();
        assert_ne!(window_size, 0);
        // After receiving a message, sink should be able to send again.
        recv_timeout(&mut rx, Duration::from_millis(100))
            .unwrap()
            .unwrap()
            .unwrap();
        block_on_timeout(send(), Duration::from_millis(100))
            .unwrap()
            .unwrap();
        // gRPC client may update window size after receiving a message,
        // though server should not be able to send messages infinitely.
        let window_size = must_fill_window();
        assert_ne!(window_size, 0);
    }
}