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

use std::{
    collections::VecDeque,
    ffi::CString,
    marker::Unpin,
    pin::Pin,
    result,
    sync::{
        Arc, Mutex,
        atomic::{AtomicI32, AtomicU8, Ordering},
    },
    time::{Duration, Instant, SystemTime, UNIX_EPOCH},
};

use collections::{HashMap, HashSet};
use crossbeam::queue::ArrayQueue;
use futures::{
    Future, Sink,
    channel::oneshot,
    compat::Future01CompatExt,
    ready,
    task::{Context, Poll, Waker},
};
use futures_timer::Delay;
use grpcio::{
    Channel, ChannelBuilder, ClientCStreamReceiver, ClientCStreamSender, Environment,
    RpcStatusCode, WriteFlags,
};
use kvproto::{
    raft_serverpb::{Done, RaftMessage},
    tikvpb::{BatchRaftMessage, TikvClient},
};
use rfstore::{errors::DiscardReason, router::RaftStoreRouter};
use security::SecurityManager;
use tikv::server::{self, Config, StoreAddrResolver, load_statistics::ThreadLoadPool, metrics::*};
use tikv_util::{
    config::{Tracker, VersionTrack},
    lru::LruCache,
    time::InstantExt,
    timer::GLOBAL_TIMER_HANDLE,
};
use yatp::{ThreadPool, task::future::TaskCell};

static CONN_ID: AtomicI32 = AtomicI32::new(0);

const _ON_RESOLVE_FP: &str = "transport_snapshot_on_resolve";

#[repr(u8)]
enum ConnState {
    Established = 0,
    /// The connection is paused and may be resumed later.
    Paused = 1,
    /// The connection is closed and removed from the connection pool.
    Disconnected = 2,
}

impl From<u8> for ConnState {
    fn from(state: u8) -> ConnState {
        match state {
            0 => ConnState::Established,
            1 => ConnState::Paused,
            2 => ConnState::Disconnected,
            _ => unreachable!(),
        }
    }
}

/// A quick queue for sending raft messages.
struct Queue {
    buf: ArrayQueue<(RaftMessage, Instant)>,
    conn_state: AtomicU8,
    waker: Mutex<Option<Waker>>,
}

impl Queue {
    /// Creates a Queue that can store at lease `cap` messages.
    fn with_capacity(cap: usize) -> Queue {
        Queue {
            buf: ArrayQueue::new(cap),
            conn_state: AtomicU8::new(ConnState::Established as u8),
            waker: Mutex::new(None),
        }
    }

    /// Pushes message into the tail of the Queue.
    ///
    /// You are supposed to call `notify` to make sure the message will be sent
    /// finally.
    ///
    /// True when the message is pushed into queue otherwise false.
    fn push(&self, msg_with_time: (RaftMessage, Instant)) -> Result<(), DiscardReason> {
        match self.conn_state.load(Ordering::SeqCst).into() {
            ConnState::Established => match self.buf.push(msg_with_time) {
                Ok(()) => Ok(()),
                Err(_) => Err(DiscardReason::Full),
            },
            ConnState::Paused => Err(DiscardReason::Paused),
            ConnState::Disconnected => Err(DiscardReason::Disconnected),
        }
    }

    fn set_conn_state(&self, s: ConnState) {
        self.conn_state.store(s as u8, Ordering::SeqCst);
    }

    /// Wakes up consumer to retrive message.
    fn notify(&self) {
        if !self.buf.is_empty() {
            let t = self.waker.lock().unwrap().take();
            if let Some(t) = t {
                t.wake();
            }
        }
    }

    /// Gets the buffer len.
    #[inline]
    fn len(&self) -> usize {
        self.buf.len()
    }

    /// Gets message from the head of the queue.
    fn try_pop(&self) -> Option<(RaftMessage, Instant)> {
        self.buf.pop()
    }

    /// Same as `try_pop` but register interest on readiness when `None` is
    /// returned.
    ///
    /// The method should be called in polling context. If the queue is empty,
    /// it will register current polling task for notifications.
    #[inline]
    fn pop(&self, ctx: &Context<'_>) -> Option<(RaftMessage, Instant)> {
        self.buf.pop().or_else(|| {
            {
                let mut waker = self.waker.lock().unwrap();
                *waker = Some(ctx.waker().clone());
            }
            self.buf.pop()
        })
    }
}

trait Buffer {
    type OutputMessage;

    /// Tests if it is full.
    ///
    /// A full buffer should be flushed successfully before calling `push`.
    fn full(&self) -> bool;
    /// Pushes the message into buffer.
    fn push(&mut self, msg: (RaftMessage, Instant));
    /// Checks if the batch is empty.
    fn empty(&self) -> bool;
    /// Flushes the message to grpc.
    ///
    /// `sender` should be able to accept messages.
    fn flush(
        &mut self,
        sender: &mut ClientCStreamSender<Self::OutputMessage>,
    ) -> grpcio::Result<()>;

    /// If the buffer is not full, suggest whether sender should wait
    /// for next message.
    fn wait_hint(&mut self) -> Option<Duration> {
        None
    }
}

/// A buffer for BatchRaftMessage.
struct BatchMessageBuffer {
    batch: Vec<(RaftMessage, Instant)>,
    overflowing: Option<(RaftMessage, Instant)>,
    size: usize,
    cfg: Config,
    cfg_tracker: Tracker<Config>,
    loads: Arc<ThreadLoadPool>,
}

impl BatchMessageBuffer {
    fn new(
        global_cfg_track: &Arc<VersionTrack<Config>>,
        loads: Arc<ThreadLoadPool>,
    ) -> BatchMessageBuffer {
        let cfg_tracker = Arc::clone(global_cfg_track).tracker("raft-client-buffer".into());
        let cfg = global_cfg_track.value().clone();
        BatchMessageBuffer {
            batch: Vec::with_capacity(cfg.raft_msg_max_batch_size),
            overflowing: None,
            size: 0,
            cfg,
            cfg_tracker,
            loads,
        }
    }

    #[inline]
    fn message_size(msg: &RaftMessage) -> usize {
        let mut msg_size = msg.start_key.len()
            + msg.end_key.len()
            + msg.get_message().context.len()
            + msg.extra_ctx.len()
            // index: 3, term: 2, data tag and size: 3, entry tag and size: 3
            + 11 * msg.get_message().get_entries().len();
        for entry in msg.get_message().get_entries() {
            msg_size += entry.data.len();
        }
        msg_size
    }

    #[cfg(test)]
    #[allow(unused)]
    fn clear(&mut self) {
        self.batch.clear();
        self.size = 0;
        self.overflowing = None;
    }
}

impl Buffer for BatchMessageBuffer {
    type OutputMessage = BatchRaftMessage;

    #[inline]
    fn full(&self) -> bool {
        self.overflowing.is_some()
    }

    #[inline]
    fn push(&mut self, msg_with_time: (RaftMessage, Instant)) {
        let msg_size = Self::message_size(&msg_with_time.0);
        // try refresh config before check
        if let Some(new_cfg) = self.cfg_tracker.any_new() {
            self.cfg = new_cfg.clone();
        }
        // To avoid building too large batch, we limit each batch's size. Since
        // `msg_size` is estimated, `GRPC_SEND_MSG_BUF` is reserved for errors.
        if self.size > 0
            && (self.size + msg_size + self.cfg.raft_client_grpc_send_msg_buffer
                >= self.cfg.max_grpc_send_msg_len as usize
                || self.batch.len() >= self.cfg.raft_msg_max_batch_size)
        {
            self.overflowing = Some(msg_with_time);
            return;
        }
        self.size += msg_size;
        self.batch.push(msg_with_time);
    }

    #[inline]
    fn empty(&self) -> bool {
        self.batch.is_empty()
    }

    #[inline]
    fn flush(&mut self, sender: &mut ClientCStreamSender<BatchRaftMessage>) -> grpcio::Result<()> {
        let mut batch_msgs = BatchRaftMessage::default();
        self.batch.drain(..).for_each(|(msg, time)| {
            RAFT_MESSAGE_DURATION
                .send_wait
                .observe(time.saturating_elapsed().as_secs_f64());
            batch_msgs.msgs.push(msg);
        });
        let now = SystemTime::now()
            .duration_since(UNIX_EPOCH)
            .unwrap()
            .as_nanos() as u64;
        batch_msgs.last_observed_time = now;

        let res = Pin::new(sender).start_send((
            batch_msgs,
            WriteFlags::default().buffer_hint(self.overflowing.is_some()),
        ));

        self.size = 0;
        if let Some(more) = self.overflowing.take() {
            self.push(more);
        }
        res
    }

    #[inline]
    fn wait_hint(&mut self) -> Option<Duration> {
        let wait_dur = self.cfg.heavy_load_wait_duration();
        if !wait_dur.is_zero() {
            if self.loads.current_thread_in_heavy_load() {
                Some(wait_dur)
            } else {
                None
            }
        } else {
            None
        }
    }
}

/// A buffer for non-batch RaftMessage.
struct MessageBuffer {
    batch: VecDeque<RaftMessage>,
}

impl MessageBuffer {
    fn new() -> MessageBuffer {
        MessageBuffer {
            batch: VecDeque::with_capacity(2),
        }
    }
}

impl Buffer for MessageBuffer {
    type OutputMessage = RaftMessage;

    #[inline]
    fn full(&self) -> bool {
        self.batch.len() >= 2
    }

    #[inline]
    fn push(&mut self, msg_with_time: (RaftMessage, Instant)) {
        self.batch.push_back(msg_with_time.0);
    }

    #[inline]
    fn empty(&self) -> bool {
        self.batch.is_empty()
    }

    #[inline]
    fn flush(&mut self, sender: &mut ClientCStreamSender<RaftMessage>) -> grpcio::Result<()> {
        if let Some(msg) = self.batch.pop_front() {
            Pin::new(sender).start_send((
                msg,
                WriteFlags::default().buffer_hint(!self.batch.is_empty()),
            ))
        } else {
            Ok(())
        }
    }
}

fn report_unreachable<R>(router: &R, msg: &RaftMessage)
where
    R: RaftStoreRouter,
{
    let to_peer = msg.get_to_peer();
    router.report_unreachable(msg.region_id, to_peer.store_id);
}

fn grpc_error_is_unimplemented(e: &grpcio::Error) -> bool {
    if let grpcio::Error::RpcFailure(ref status) = e {
        status.code() == RpcStatusCode::UNIMPLEMENTED
    } else {
        false
    }
}

/// Struct tracks the lifetime of a `raft` or `batch_raft` RPC.
struct AsyncRaftSender<M, B> {
    sender: ClientCStreamSender<M>,
    queue: Arc<Queue>,
    buffer: B,
    addr: String,
    flush_timeout: Option<Delay>,
}

impl<M, B> AsyncRaftSender<M, B>
where
    B: Buffer<OutputMessage = M>,
{
    fn fill_msg(&mut self, ctx: &Context<'_>) {
        while !self.buffer.full() {
            let msg_with_time = match self.queue.pop(ctx) {
                Some(msg_with_time) => msg_with_time,
                None => return,
            };
            self.buffer.push(msg_with_time);
        }
    }
}

impl<M, B> Future for AsyncRaftSender<M, B>
where
    B: Buffer<OutputMessage = M> + Unpin,
{
    type Output = grpcio::Result<()>;

    fn poll(mut self: Pin<&mut Self>, ctx: &mut Context<'_>) -> Poll<grpcio::Result<()>> {
        let s = &mut *self;
        loop {
            s.fill_msg(ctx);
            if !s.buffer.empty() {
                // Then it's the first time visit this block since last flush.
                if s.flush_timeout.is_none() {
                    ready!(Pin::new(&mut s.sender).poll_ready(ctx))?;
                }
                // Only set up a timer if buffer is not full.
                if !s.buffer.full() {
                    if s.flush_timeout.is_none() {
                        // Only set up a timer if necessary.
                        if let Some(wait_time) = s.buffer.wait_hint() {
                            s.flush_timeout = Some(Delay::new(wait_time));
                        }
                    }

                    // It will be woken up again when the timer fires or new messages are enqueued.
                    if let Some(timeout) = &mut s.flush_timeout {
                        if Pin::new(timeout).poll(ctx).is_pending() {
                            return Poll::Pending;
                        } else {
                            RAFT_MESSAGE_FLUSH_COUNTER.delay.inc_by(1);
                        }
                    } else {
                        RAFT_MESSAGE_FLUSH_COUNTER.eof.inc_by(1);
                    }
                } else if s.flush_timeout.is_some() {
                    RAFT_MESSAGE_FLUSH_COUNTER.full_after_delay.inc_by(1);
                } else {
                    RAFT_MESSAGE_FLUSH_COUNTER.full.inc_by(1);
                }

                // So either enough messages are batched up or don't need to wait or wait
                // timeouts.
                s.flush_timeout.take();
                ready!(Poll::Ready(s.buffer.flush(&mut s.sender)))?;
                continue;
            }

            if let Poll::Ready(Err(e)) = Pin::new(&mut s.sender).poll_flush(ctx) {
                return Poll::Ready(Err(e));
            }
            return Poll::Pending;
        }
    }
}

#[derive(PartialEq)]
enum RaftCallRes {
    // the call is not supported, probably due to visiting to older version TiKV
    Fallback,
    // the connection is aborted or closed
    Disconnected,
}

struct RaftCall<M, B> {
    sender: AsyncRaftSender<M, B>,
    receiver: ClientCStreamReceiver<Done>,
    lifetime: Option<oneshot::Sender<RaftCallRes>>,
    store_id: u64,
}

impl<M, B> RaftCall<M, B>
where
    B: Buffer<OutputMessage = M> + Unpin,
{
    async fn poll(&mut self) {
        let res = futures::join!(&mut self.sender, &mut self.receiver);
        if let (Ok(()), Ok(Done { .. })) = res {
            info!("connection close"; "store_id" => self.store_id, "addr" => %self.sender.addr);
            if let Some(tx) = self.lifetime.take() {
                let _ = tx.send(RaftCallRes::Disconnected);
            }
            return;
        }

        let (sink_err, recv_err) = (res.0.err(), res.1.err());
        error!("connection aborted"; "store_id" => self.store_id, "sink_error" => ?sink_err, "receiver_err" => ?recv_err, "addr" => %self.sender.addr);
        if let Some(tx) = self.lifetime.take() {
            let should_fallback = [sink_err, recv_err]
                .iter()
                .any(|e| e.as_ref().is_some_and(grpc_error_is_unimplemented));

            let res = if should_fallback {
                // Asks backend to fallback.
                RaftCallRes::Fallback
            } else {
                RaftCallRes::Disconnected
            };
            let _ = tx.send(res);
        }
    }
}

#[derive(Clone)]
pub struct ConnectionBuilder<S, R> {
    env: Arc<Environment>,
    cfg: Arc<VersionTrack<Config>>,
    security_mgr: Arc<SecurityManager>,
    resolver: S,
    router: R,
    loads: Arc<ThreadLoadPool>,
}

impl<S, R> ConnectionBuilder<S, R> {
    pub fn new(
        env: Arc<Environment>,
        cfg: Arc<VersionTrack<Config>>,
        security_mgr: Arc<SecurityManager>,
        resolver: S,
        router: R,
        loads: Arc<ThreadLoadPool>,
    ) -> ConnectionBuilder<S, R> {
        ConnectionBuilder {
            env,
            cfg,
            security_mgr,
            resolver,
            router,
            loads,
        }
    }
}

/// StreamBackEnd watches lifetime of a connection and handles reconnecting,
/// spawn new RPC.
struct StreamBackEnd<S, R> {
    store_id: u64,
    queue: Arc<Queue>,
    builder: ConnectionBuilder<S, R>,
}

impl<S, R> StreamBackEnd<S, R>
where
    S: StoreAddrResolver,
    R: RaftStoreRouter + Unpin + 'static,
{
    fn resolve(&self) -> impl Future<Output = server::Result<String>> {
        let (tx, rx) = oneshot::channel();
        let store_id = self.store_id;
        let res = self.builder.resolver.resolve(
            store_id,
            #[allow(unused_mut)]
            Box::new(move |mut addr| {
                {
                    // Wrapping the fail point in a closure, so we can modify
                    // local variables without return.
                    let mut transport_on_resolve_fp = || {
                        fail_point!(_ON_RESOLVE_FP, |sid| if let Some(sid) = sid {
                            let sid: u64 = sid.parse().unwrap();
                            if sid == store_id {
                                addr = Err(box_err!("injected failure"));
                                // Sleep some time to avoid race between enqueuing message and
                                // resolving address.
                                std::thread::sleep(std::time::Duration::from_millis(10));
                            }
                        })
                    };
                    transport_on_resolve_fp();
                }
                let _ = tx.send(addr);
            }),
        );
        async move {
            res?;
            match rx.await {
                Ok(a) => a,
                Err(_) => Err(server::Error::Other(
                    "failed to receive resolve result".into(),
                )),
            }
        }
    }

    fn clear_pending_message(&self, reason: &str) {
        let len = self.queue.len();
        for _ in 0..len {
            let msg_with_time = self.queue.try_pop().unwrap();
            report_unreachable(&self.builder.router, &msg_with_time.0)
        }
        REPORT_FAILURE_MSG_COUNTER
            .with_label_values(&[reason, &self.store_id.to_string()])
            .inc_by(len as u64);
    }

    fn connect(&self, addr: &str) -> Channel {
        info!("server: new connection with tikv endpoint"; "addr" => addr, "store_id" => self.store_id);

        let cfg = self.builder.cfg.value();
        let cb = ChannelBuilder::new(self.builder.env.clone())
            .stream_initial_window_size(cfg.grpc_stream_initial_window_size.0 as i32)
            .keepalive_time(cfg.grpc_keepalive_time.0)
            .keepalive_timeout(cfg.grpc_keepalive_timeout.0)
            .default_compression_algorithm(cfg.grpc_compression_algorithm(true))
            .default_gzip_compression_level(cfg.grpc_gzip_compression_level)
            .default_grpc_min_message_size_to_compress(cfg.grpc_min_message_size_to_compress)
            .max_reconnect_backoff(cfg.raft_client_max_backoff.0)
            .initial_reconnect_backoff(cfg.raft_client_initial_reconnect_backoff.0)
            // hack: so it's different args, grpc will always create a new connection.
            .raw_cfg_int(
                CString::new("random id").unwrap(),
                CONN_ID.fetch_add(1, Ordering::SeqCst),
            );
        self.builder.security_mgr.connect(cb, addr)
    }

    fn batch_call(&self, client: &TikvClient, addr: String) -> oneshot::Receiver<RaftCallRes> {
        let (batch_sink, batch_stream) = client.batch_raft().unwrap();
        let (tx, rx) = oneshot::channel();
        let mut call = RaftCall {
            sender: AsyncRaftSender {
                sender: batch_sink,
                queue: self.queue.clone(),
                buffer: BatchMessageBuffer::new(&self.builder.cfg, self.builder.loads.clone()),
                addr,
                flush_timeout: None,
            },
            receiver: batch_stream,
            lifetime: Some(tx),
            store_id: self.store_id,
        };
        // TODO: verify it will be notified if client is dropped while env still alive.
        client.spawn(async move {
            call.poll().await;
        });
        rx
    }

    fn call(&self, client: &TikvClient, addr: String) -> oneshot::Receiver<RaftCallRes> {
        let (sink, stream) = client.raft().unwrap();
        let (tx, rx) = oneshot::channel();
        let mut call = RaftCall {
            sender: AsyncRaftSender {
                sender: sink,
                queue: self.queue.clone(),
                buffer: MessageBuffer::new(),
                addr,
                flush_timeout: None,
            },
            receiver: stream,
            lifetime: Some(tx),
            store_id: self.store_id,
        };
        client.spawn(async move {
            call.poll().await;
        });
        rx
    }
}

async fn maybe_backoff(backoff: Duration, last_wake_time: &mut Option<Instant>) {
    let now = Instant::now();
    if let Some(last) = *last_wake_time {
        if last + backoff < now {
            // We have spent long enough time in last retry, no need to backoff again.
            *last_wake_time = Some(now);
            return;
        }
    } else {
        *last_wake_time = Some(now);
        return;
    }

    if let Err(e) = GLOBAL_TIMER_HANDLE.delay(now + backoff).compat().await {
        error_unknown!(?e; "failed to backoff");
    }
    *last_wake_time = Some(Instant::now());
}

/// A future that drives the life cycle of a connection.
///
/// The general progress of connection is:
///
/// 1. resolve address
/// 2. connect
/// 3. make batch call
/// 4. fallback to legacy API if incompatible
///
/// Every failure during the process should trigger retry automatically.
async fn start<S, R>(
    back_end: StreamBackEnd<S, R>,
    conn_id: usize,
    pool: Arc<Mutex<ConnectionPool>>,
) where
    S: StoreAddrResolver + Send,
    R: RaftStoreRouter + Unpin + Send + 'static,
{
    let mut last_wake_time = None;
    let mut first_time = true;
    let backoff_duration = back_end.builder.cfg.value().raft_client_max_backoff.0;
    let mut addr_channel = None;
    loop {
        maybe_backoff(backoff_duration, &mut last_wake_time).await;
        let f = back_end.resolve();
        let addr = match f.await {
            Ok(addr) => {
                RESOLVE_STORE_COUNTER.with_label_values(&["success"]).inc();
                info!("resolve store address ok"; "store_id" => back_end.store_id, "addr" => %addr);
                addr
            }
            Err(e) => {
                RESOLVE_STORE_COUNTER.with_label_values(&["failed"]).inc();
                back_end.clear_pending_message("resolve");
                error_unknown!(?e; "resolve store address failed"; "store_id" => back_end.store_id,);
                // TOMBSTONE
                if format!("{}", e).contains("has been removed") {
                    let mut pool = pool.lock().unwrap();
                    if let Some(s) = pool.connections.remove(&(back_end.store_id, conn_id)) {
                        s.set_conn_state(ConnState::Disconnected);
                    }
                    pool.tombstone_stores.insert(back_end.store_id);
                    return;
                }
                continue;
            }
        };

        // reuse channel if the address is the same.
        if addr_channel
            .as_ref()
            .is_none_or(|(_, prev_addr)| prev_addr != &addr)
        {
            addr_channel = Some((back_end.connect(&addr), addr.clone()));
        }
        let channel = addr_channel.as_ref().unwrap().0.clone();

        debug!("connecting to store"; "store_id" => back_end.store_id, "addr" => %addr);
        if !channel.wait_for_connected(backoff_duration).await {
            error!("wait connect timeout"; "store_id" => back_end.store_id, "addr" => addr);

            // Clears pending messages to avoid consuming high memory when one node is
            // shutdown.
            back_end.clear_pending_message("unreachable");

            // broadcast is time consuming operation which would blocks raftstore, so report
            // unreachable only once until being connected again.
            if first_time {
                first_time = false;
                back_end
                    .builder
                    .router
                    .broadcast_unreachable(back_end.store_id);
            }
            continue;
        }
        debug!("connection established"; "store_id" => back_end.store_id, "addr" => %addr);

        let client = TikvClient::new(channel);
        let f = back_end.batch_call(&client, addr.clone());
        let mut res = f.await; // block here until the stream call is closed or aborted.
        if res == Ok(RaftCallRes::Fallback) {
            // If the call is setup successfully, it will never finish. Returning
            // `UnImplemented` means the batch_call is not supported, we are probably
            // connect to an old version of TiKV. So we need to fallback to use
            // legacy API.
            let f = back_end.call(&client, addr.clone());
            res = f.await;
        }
        match res {
            Ok(RaftCallRes::Fallback) => {
                error!("connection fail"; "store_id" => back_end.store_id, "addr" => addr, "err" => "require fallback even with legacy API");
            }
            // Err(_) should be tx is dropped
            Ok(RaftCallRes::Disconnected) | Err(_) => {
                error!("connection abort"; "store_id" => back_end.store_id, "addr" => addr);
                REPORT_FAILURE_MSG_COUNTER
                    .with_label_values(&["unreachable", &back_end.store_id.to_string()])
                    .inc_by(1);
                back_end
                    .builder
                    .router
                    .broadcast_unreachable(back_end.store_id);
                addr_channel = None;
                first_time = false;
            }
        }
    }
}

/// A global connection pool.
///
/// All valid connections should be stored as a record. Once it's removed
/// from the struct, all cache clone should also remove it at some time.
#[derive(Default)]
struct ConnectionPool {
    connections: HashMap<(u64, usize), Arc<Queue>>,
    tombstone_stores: HashSet<u64>,
    store_allowlist: Vec<u64>,
}

impl ConnectionPool {
    fn set_store_allowlist(&mut self, stores: Vec<u64>) {
        self.store_allowlist = stores;
        for (&(store_id, _), q) in self.connections.iter() {
            let mut state = ConnState::Established;
            if self.need_pause(store_id) {
                state = ConnState::Paused;
            }
            q.set_conn_state(state);
        }
    }

    fn need_pause(&self, store_id: u64) -> bool {
        !self.store_allowlist.is_empty() && !self.store_allowlist.contains(&store_id)
    }
}

/// Queue in cache.
struct CachedQueue {
    queue: Arc<Queue>,
    /// If a msg is enqueued, but the queue has not been notified for polling,
    /// it will be marked to true. And all dirty queues are expected to be
    /// notified during flushing.
    dirty: bool,
    /// Mark if the connection is full.
    full: bool,
}

/// A raft client that can manages connections correctly.
///
/// A correct usage of raft client is:
///
/// ```text
/// for m in msgs {
///     if !raft_client.send(m) {
///         // handle error.
///     }
/// }
/// raft_client.flush();
/// ```
pub struct RaftClient<S, R> {
    pool: Arc<Mutex<ConnectionPool>>,
    cache: LruCache<(u64, usize), CachedQueue>,
    need_flush: Vec<(u64, usize)>,
    full_stores: Vec<(u64, usize)>,
    future_pool: Arc<ThreadPool<TaskCell>>,
    builder: ConnectionBuilder<S, R>,
    last_hash: (u64, u64),
}

impl<S, R> RaftClient<S, R>
where
    S: StoreAddrResolver + Send + 'static,
    R: RaftStoreRouter + Unpin + Send + 'static,
{
    pub fn new(builder: ConnectionBuilder<S, R>) -> RaftClient<S, R> {
        let future_pool = Arc::new(
            yatp::Builder::new(thd_name!("raft-stream"))
                .max_thread_count(1)
                .build_future_pool(),
        );
        RaftClient {
            pool: Arc::default(),
            cache: LruCache::with_capacity_and_sample(0, 7),
            need_flush: vec![],
            full_stores: vec![],
            future_pool,
            builder,
            last_hash: (0, 0),
        }
    }

    /// Loads connection from pool.
    ///
    /// Creates it if it doesn't exist. `false` is returned if such connection
    /// can't be established.
    fn load_stream(&mut self, store_id: u64, conn_id: usize) -> bool {
        let (s, pool_len) = {
            let mut pool = self.pool.lock().unwrap();
            if pool.tombstone_stores.contains(&store_id) {
                let pool_len = pool.connections.len();
                drop(pool);
                self.cache.resize(pool_len);
                return false;
            }
            let need_pause = pool.need_pause(store_id);
            let conn = pool
                .connections
                .entry((store_id, conn_id))
                .or_insert_with(|| {
                    let queue = Arc::new(Queue::with_capacity(
                        self.builder.cfg.value().raft_client_queue_size,
                    ));
                    if need_pause {
                        queue.set_conn_state(ConnState::Paused);
                    }
                    let back_end = StreamBackEnd {
                        store_id,
                        queue: queue.clone(),
                        builder: self.builder.clone(),
                    };
                    self.future_pool
                        .spawn(start(back_end, conn_id, self.pool.clone()));
                    queue
                })
                .clone();
            (conn, pool.connections.len())
        };
        self.cache.resize(pool_len);
        self.cache.insert(
            (store_id, conn_id),
            CachedQueue {
                queue: s,
                dirty: false,
                full: false,
            },
        );
        true
    }

    /// Sends a message.
    ///
    /// If the message fails to be sent, false is returned. Returning true means
    /// the message is enqueued to buffer. Caller is expected to call
    /// `flush` to ensure all buffered messages are sent out.
    pub fn send(&mut self, msg: RaftMessage) -> result::Result<(), DiscardReason> {
        let wait_send_start = Instant::now();
        let store_id = msg.get_to_peer().store_id;
        let grpc_raft_conn_num = self.builder.cfg.value().grpc_raft_conn_num as u64;
        let conn_id = if grpc_raft_conn_num == 1 {
            0
        } else {
            if self.last_hash.0 == 0 || msg.region_id != self.last_hash.0 {
                self.last_hash = (
                    msg.region_id,
                    seahash::hash(&msg.region_id.to_ne_bytes()) % grpc_raft_conn_num,
                );
            };
            self.last_hash.1 as usize
        };

        #[allow(unused_mut)]
        let mut transport_on_send_store_fp = || {
            fail_point!(
                "transport_on_send_snapshot",
                msg.get_message().get_msg_type() == raft::eraftpb::MessageType::MsgSnapshot,
                |sid| if let Some(sid) = sid {
                    let sid: u64 = sid.parse().unwrap();
                    if sid == store_id {
                        // Forbid building new connections.
                        fail::cfg(_ON_RESOLVE_FP, &format!("1*return({})", sid)).unwrap();
                        self.cache.remove(&(store_id, conn_id));
                        self.pool
                            .lock()
                            .unwrap()
                            .connections
                            .remove(&(store_id, conn_id));
                    }
                }
            )
        };
        transport_on_send_store_fp();
        loop {
            if let Some(s) = self.cache.get_mut(&(store_id, conn_id)) {
                match s.queue.push((msg, wait_send_start)) {
                    Ok(_) => {
                        if !s.dirty {
                            s.dirty = true;
                            self.need_flush.push((store_id, conn_id));
                        }
                        return Ok(());
                    }
                    Err(DiscardReason::Full) => {
                        s.queue.notify();
                        s.dirty = false;
                        if !s.full {
                            s.full = true;
                            self.full_stores.push((store_id, conn_id));
                        }
                        return Err(DiscardReason::Full);
                    }
                    Err(DiscardReason::Disconnected) => break,
                    Err(DiscardReason::Paused) => return Err(DiscardReason::Paused),
                    Err(DiscardReason::Filtered) => return Err(DiscardReason::Filtered),
                }
            }
            if !self.load_stream(store_id, conn_id) {
                return Err(DiscardReason::Disconnected);
            }
        }
        self.cache.remove(&(store_id, conn_id));
        Err(DiscardReason::Disconnected)
    }

    pub fn need_flush(&self) -> bool {
        !self.need_flush.is_empty() || !self.full_stores.is_empty()
    }

    fn flush_full_metrics(&mut self) {
        if self.full_stores.is_empty() {
            return;
        }

        for id in &self.full_stores {
            if let Some(s) = self.cache.get_mut(id) {
                s.full = false;
            }
            REPORT_FAILURE_MSG_COUNTER
                .with_label_values(&["full", &id.0.to_string()])
                .inc();
        }
        self.full_stores.clear();
        if self.full_stores.capacity() > 2048 {
            self.full_stores.shrink_to(512);
        }
    }

    /// Flushes all buffered messages.
    pub fn flush(&mut self) {
        self.flush_full_metrics();
        if self.need_flush.is_empty() {
            return;
        }
        let mut counter = 0;
        for id in &self.need_flush {
            if let Some(s) = self.cache.get_mut(id) {
                if s.dirty {
                    s.dirty = false;
                    counter += 1;
                    s.queue.notify();
                }
                continue;
            }
            let l = self.pool.lock().unwrap();
            if let Some(q) = l.connections.get(id) {
                counter += 1;
                q.notify();
            }
        }
        self.need_flush.clear();
        if self.need_flush.capacity() > 2048 {
            self.need_flush.shrink_to(512);
        }
        RAFT_MESSAGE_FLUSH_COUNTER.wake.inc_by(counter);
    }

    pub fn set_store_allowlist(&mut self, stores: Vec<u64>) {
        let mut p = self.pool.lock().unwrap();
        p.set_store_allowlist(stores);
    }
}

impl<S, R> Clone for RaftClient<S, R>
where
    S: Clone,
    R: Clone,
{
    fn clone(&self) -> Self {
        RaftClient {
            pool: self.pool.clone(),
            cache: LruCache::with_capacity_and_sample(0, 7),
            need_flush: vec![],
            full_stores: vec![],
            future_pool: self.future_pool.clone(),
            builder: self.builder.clone(),
            last_hash: (0, 0),
        }
    }
}