proxy.rs

use std::time::{Duration, Instant};

use anyhow::{anyhow, bail};
use bytes::BytesMut;
use futures::StreamExt;
use log::{debug, error};
use tokio::sync::mpsc;
use tokio::sync::mpsc::{Receiver, Sender};

use redis_codec_core::resp_decoder::ResFramedData;
use redis_command::CommandFlags;
use redis_command_gen::CmdType;
use redis_proxy_common::command::utils::{has_flag, has_key};
use redis_proxy_common::ReqPkt;

use crate::config::WriteMode;
use crate::connection::Connection;
use crate::double_writer::DoubleWriter;
use crate::filter_trait::Filter;
use crate::handler::get_handler;
use crate::prometheus::{METRICS, RESP_FAILED, SEND_ASYNC};
use crate::session::{Session, SessionFlags};

const TASK_BUFFER_SIZE: usize = 16;

pub struct RedisProxy<P> {
    pub filter: P,
    pub upstream_addr: String,
    pub server_name: String,
    pub keepalive_interval: u64,
    pub double_writer: DoubleWriter,
}

impl<P> RedisProxy<P> {
    fn double_write_mode(&self, session: &Session, header_frame: &&ReqPkt) -> WriteMode {
        // double write is disabled
        if self.double_writer.write_mode == WriteMode::DISABLED {
            return WriteMode::DISABLED;
        }
        //current session is in transaction, should not double write
        if session.contains_client_flags(SessionFlags::InTrasactions) {
            return WriteMode::DISABLED;
        }
        // ping command usually used to check the connection's status, 
        // so it should double write anyway
        // for example, jedis use ping command to "test on borrowed" connection 
        // this will keep both upstream and double write conn available
        if header_frame.cmd_type == CmdType::PING
            || header_frame.cmd_type == CmdType::HELLO  // to use resp3 or resp2, will call hello command
            || header_frame.cmd_type == CmdType::CLIENT_SETINFO //set client info
            || header_frame.cmd_type == CmdType::CLIENT_SETNAME { //set client name
            return self.double_writer.write_mode;
        }
        // not a write command, should not double write
        if !has_flag(&header_frame.cmd_type, CommandFlags::Write) {
            return WriteMode::DISABLED;
        }
        // write command and has no key in command
        // only FLUSHALL、FLUSHDB、FUNCTION DELETE、FUNCTION FLUSH、FUNCTION LOAD、FUNCTION RESTORE、SWAPDB
        // will double write anyway
        if !has_key(&header_frame.cmd_type) {
            return self.double_writer.write_mode;
        }
        // double write if the key is in the double write key list
        if self.double_writer.should_double_write(header_frame) {
            return self.double_writer.write_mode;
        }
        WriteMode::DISABLED
    }
}

impl<P> RedisProxy<P> where P: Filter + Send + Sync {
    pub async fn handle_new_request(&self, mut session: Session) -> Option<Session> {
        debug!("handle new request");

        let req_pkt = session.read_req_pkt().await?.ok()?;
        session.init_from_req(&req_pkt);

        let response_sent =
            self.filter.on_request(&mut session, &req_pkt).await.ok()?;
        if response_sent {
            return Some(session);
        }

        match self.handle_request_inner(&mut session, &req_pkt).await {
            Ok(_) => {
                self.filter.on_request_done(&mut session, req_pkt.cmd_type, None).await;
                return Some(session);
            }
            Err(err) => {
                debug!("handle request error: {:?}", err);
                session.upstream_conn = None;
                session.dw_conn = None;
                self.filter.on_request_done(&mut session, req_pkt.cmd_type, Some(&err)).await;
            }
        }
        None
    }
    pub async fn handle_request_inner(&self, session: &mut Session, req_pkt: &ReqPkt) -> anyhow::Result<()> {
        let start = Instant::now();
        let write_mode = self.double_write_mode(session, &req_pkt);
        match write_mode {
            WriteMode::SYNC => {
                match (&session.upstream_conn.is_none(), &session.dw_conn.is_none()) {
                    (false, false) => {
                        session.upstream_conn_elapsed = Duration::ZERO;
                    }
                    (true, false) => {
                        bail!("should not happen!!! upstream connection is none, but dw_conn is not none")
                    }
                    (true, true) => {
                        let (r1, r2) = tokio::join!(
                            Connection::connect(
                                &self.server_name,
                                &self.upstream_addr,
                                self.keepalive_interval,
                                false
                            ),
                            Connection::connect(
                                &self.server_name,
                                &self.double_writer.double_write_addr,
                                self.double_writer.keepalive_interval,
                                true
                            ),
                        );
                        session.upstream_conn_elapsed = start.elapsed();
                        match (r1, r2) {
                            (Ok(c), Ok(d)) => {
                                session.upstream_conn = Some(c);
                                session.dw_conn = Some(d);
                            }
                            _ => bail!("get connection from pool error")
                        };
                    }
                    (false, true) => {
                        let conn =
                            Connection::connect(
                                &self.server_name,
                                &self.double_writer.double_write_addr,
                                self.double_writer.keepalive_interval,
                                true
                            ).await.map_err(|e| anyhow!("get connection from dw pool error: {:?}", e))?;
                        session.upstream_conn_elapsed = start.elapsed();
                        session.dw_conn = Some(conn);
                    }
                }

                let mut conn = session.upstream_conn.take().unwrap();
                let mut dw_conn = session.dw_conn.take().unwrap();
                self.handle_double_write(session, &req_pkt, &mut conn, &mut dw_conn).await?;
                session.upstream_conn = Some(conn);
                session.dw_conn = Some(dw_conn);
                Ok(())
            }
            WriteMode::ASYNC | WriteMode::DISABLED => {
                if session.upstream_conn.is_none() {
                    let conn =
                        Connection::connect(
                            &self.server_name,
                            &self.upstream_addr,
                            self.keepalive_interval,
                            false
                        ).await.map_err(|e| anyhow!("get connection from upstream pool error: {:?}", e))?;
                    session.upstream_conn_elapsed = start.elapsed();
                    session.upstream_conn = Some(conn);
                } else {
                    session.upstream_conn_elapsed = Duration::ZERO;
                }

                let mut conn = session.upstream_conn.take().unwrap();

                if req_pkt.is_subscribe() {
                    //only support "not double write request" in subscribe mode
                    //such as psubscribe、ssubscribe、subscribe
                    //see https://redis.io/docs/reference/protocol-spec/
                    self.handle_subscribe_request(session, &req_pkt, &mut conn).await?;
                } else {
                    self.handle_normal_request(session, &req_pkt, &mut conn).await?;
                }
                session.upstream_conn = Some(conn);
                if write_mode == WriteMode::ASYNC {
                    if let Err(err) = self.double_writer.async_write(req_pkt) {
                        error!("async send retry error: {:?}, will lost it forever", err);
                        METRICS.request_status
                            .with_label_values(&[
                                &self.server_name,
                                req_pkt.cmd_type.as_ref(),
                                SEND_ASYNC,
                                RESP_FAILED
                            ]).inc();
                    }
                }
                Ok(())
            }
        }
    }

    async fn handle_subscribe_request(&self, session: &mut Session, req_pkt: &ReqPkt,
                                      conn: &mut Connection) -> anyhow::Result<()> {
        let (tx_upstream, rx_upstream) = mpsc::channel::<(Vec<ResFramedData>, usize)>(TASK_BUFFER_SIZE);
        let (tx_downstream, rx_downstream) = mpsc::channel::<ReqPkt>(TASK_BUFFER_SIZE);

        session.upstream_start = Instant::now();

        //send first command
        conn.send_bytes_vectored(req_pkt).await?;

        // handle subscribe until one side is done
        let join_result = tokio::try_join!(
            self.proxy_handle_subscribe_downstream(session, tx_downstream, rx_upstream),
            self.proxy_handle_subscribe_upstream(conn, tx_upstream, rx_downstream)
        );

        // in subscribe req, upstream_elapsed is the whole subscribe time
        session.upstream_elapsed = session.upstream_start.elapsed();
        if let Err(err) = join_result {
            session.res_is_ok = false;
            return Err(err);
        }
        session.res_is_ok = true;
        Ok(())
    }

    async fn handle_normal_request(&self, session: &mut Session, req_pkt: &ReqPkt,
                                   conn: &mut Connection) -> anyhow::Result<()> {
        let (tx_upstream, rx_upstream) = mpsc::channel::<ResFramedData>(TASK_BUFFER_SIZE);

        session.upstream_start = Instant::now();

        conn.send_bytes_vectored(req_pkt).await?;

        // bi-directional proxy
        tokio::try_join!(
            self.proxy_handle_downstream(session, req_pkt.cmd_type, rx_upstream),
            self.proxy_handle_upstream(conn, tx_upstream)
        )?;
        Ok(())
    }

    async fn handle_double_write(&self, session: &mut Session, req_pkt: &ReqPkt,
                                 conn: &mut Connection, dw_conn: &mut Connection) -> anyhow::Result<()> {
        session.upstream_start = Instant::now();
        let join_result = tokio::join!(
            conn.send_bytes_vectored_and_wait_resp(&req_pkt),
            dw_conn.send_bytes_vectored_and_wait_resp(&req_pkt)
        );

        debug!("double write result: {:?}", join_result);

        session.upstream_elapsed = session.upstream_start.elapsed();

        return match join_result {
            (Ok((upstream_resp_ok, upstream_pkts, upstream_resp_size)),
                Ok((dw_resp_ok, dw_pkts, _))) => {
                get_handler(req_pkt.cmd_type).map(|h|
                    h.handle_session_on_resp(session, upstream_resp_ok));
                session.res_size += upstream_resp_size;
                session.res_is_ok = upstream_resp_ok && dw_resp_ok;

                if dw_resp_ok {
                    session.write_downstream_batch(upstream_pkts).await?;
                } else {
                    // if double write is not ok, send upstream response to downstream
                    session.write_downstream_batch(dw_pkts).await?;
                }
                Ok(())
            }
            (r1, r2) => {
                bail!("double write error: r1: {:?}, r2:{:?}", r1, r2)
            }
        };
    }

    async fn proxy_handle_downstream(&self, session: &mut Session, cmd_type: CmdType, mut rx_upstream: Receiver<ResFramedData>) -> anyhow::Result<()> {
        while let Some(res_framed_data) = rx_upstream.recv().await {
            if res_framed_data.is_done {
                session.upstream_elapsed = session.upstream_start.elapsed();
                session.res_is_ok = res_framed_data.res_is_ok;
            }

            session.res_size += res_framed_data.data.len();

            if res_framed_data.is_done {
                get_handler(cmd_type).map(|h|
                    h.handle_session_on_resp(session, res_framed_data.res_is_ok));
            }

            session.write_downstream(&res_framed_data.data).await?;

            if res_framed_data.is_done {
                debug!("upstream result: {:?}", res_framed_data);
                return Ok(());
            }
        }
        bail!("rx_upstream channel closed")
    }

    async fn proxy_handle_upstream(&self, conn: &mut Connection,
                                   tx_upstream: Sender<ResFramedData>) -> anyhow::Result<()> {
        while let Some(Ok(data)) = conn.r.next().await {
            let is_done = data.is_done;
            tx_upstream.send(data).await?;
            if is_done {
                return Ok(());
            }
        }
        bail!("upstream read eof or error")
    }

    async fn proxy_handle_subscribe_downstream(&self, session: &mut Session,
                                               tx_downstream: Sender<ReqPkt>,
                                               mut rx_upstream: Receiver<(Vec<ResFramedData>, usize)>) -> anyhow::Result<()> {
        loop {
            tokio::select! {
                res_framed_data = rx_upstream.recv() => {
                    match res_framed_data {
                        None => {
                            break; //upstream closed
                        }
                        Some((res_framed_data, size)) => {
                            //handle session here, so there is no need to call handler's handle_session_on_resp method at all.
                            handle_session_on_pubsub_resp(session, &res_framed_data, size);

                            session.res_size += size;
                            session.write_downstream_batch(res_framed_data).await?;
                            if !session.in_subscribe() { //if not in subscribe mode after this resp, break
                                debug!("subscribe end");
                                break;
                            }
                        }
                    }
                }
                downstream_data = session.read_req_pkt() => {
                    match downstream_data {
                        None => {
                            drop(tx_downstream);
                            bail!("subscribe downstream read eof");
                        }
                        Some(Ok(req_pkt)) => {
                            session.req_size += req_pkt.bytes_total;
                            tx_downstream.send(req_pkt).await?
                        }
                        Some(Err(e)) => {
                            drop(tx_downstream);
                            bail!("subscribe downstream read error: {:?}", e);
                        }
                    }
                }
            }
        }
        Ok(())
    }

    async fn proxy_handle_subscribe_upstream(&self, conn: &mut Connection,
                                             tx_upstream: Sender<(Vec<ResFramedData>, usize)>,
                                             mut rx_downstream: Receiver<ReqPkt>) -> anyhow::Result<()> {
        let mut result: Vec<ResFramedData> = Vec::with_capacity(1);
        let mut size: usize = 0;
        loop {
            tokio::select! {
                req_pkt = rx_downstream.recv() => {
                    match req_pkt {
                        None => {
                            break; //downstream closed
                        }
                        Some(req_pkt) => {
                            conn.send_bytes_vectored(&req_pkt).await?;
                        }
                    };
                }
                upstream_data = conn.r.next() => {
                    match upstream_data {
                        None => {
                            drop(tx_upstream);
                            bail!("subscribe upstream read eof");
                        }
                        Some(Ok(data)) => {
                            let is_done = data.is_done;
                            size += data.data.len();
                            result.push(data);
                            if is_done {
                                //send to downstream
                                tx_upstream.send((result, size)).await?;
                                // reset result and it's size
                                result = Vec::with_capacity(1);
                                size = 0;
                            }
                        }
                        Some(Err(e)) => {
                            drop(tx_upstream);
                            bail!("subscribe upstream read error: {:?}", e);
                        }
                    }
                }
            }
        }
        Ok(())
    }
}

fn handle_session_on_pubsub_resp(session: &mut Session, result: &Vec<ResFramedData>, size: usize) {
    if result.is_empty() {
        return;
    }
    let first_resp = result.first().unwrap();
    match first_resp.data[0] {
        b'*' | b'>' => { //in resp3 is >, in resp2 is *
            let resp_bytes = gather_resp_bytes(&result, size);
            if is_pub_sub_resp(&resp_bytes) {
                session.attrs.sub = get_remaining_channels(resp_bytes);
            } else if is_spub_ssub_resp(&resp_bytes) {
                session.attrs.ssub = get_remaining_channels(resp_bytes);
            }
        }
        b'+' => { // on reset command
            let resp_bytes = gather_resp_bytes(&result, size);
            if &resp_bytes[..] == b"+RESET\r\n" {
                session.attrs.reset()
            }
        }
        _ => {
            // do nothing
        }
    }
}

fn is_pub_sub_resp(resp_bytes: &BytesMut) -> bool {
    let resp_len = resp_bytes.len();
    return (resp_len > 19 && &resp_bytes[1..19] == b"3\r\n$9\r\nsubscribe\r\n")
        || (resp_len > 21 && &resp_bytes[1..21] == b"3\r\n$10\r\npsubscribe\r\n")
        || (resp_len > 22 && &resp_bytes[1..22] == b"3\r\n$11\r\nunsubscribe\r\n")
        || (resp_len > 23 && &resp_bytes[1..23] == b"3\r\n$12\r\npunsubscribe\r\n");
}

fn is_spub_ssub_resp(resp_bytes: &BytesMut) -> bool {
    let resp_len = resp_bytes.len();
    return (resp_len > 21 && &resp_bytes[1..21] == b"3\r\n$10\r\nssubscribe\r\n")
        || (resp_len > 23 && &resp_bytes[1..23] == b"3\r\n$12\r\nsunsubscribe\r\n");
}

fn get_remaining_channels(resp_bytes: BytesMut) -> usize {
    // example:
    // ...
    // *3\r\n
    // $11\r\n
    // unsubscribe\r\n
    // $4\r\n
    // chan\r\n
    // :0\r\n <------------------- this is the remaining number of channels in this session
    let mut remaining = 0_usize;
    let mut i = resp_bytes.len() - 3;
    let mut base = 1_usize;
    while resp_bytes[i] != b':' {
        remaining += (resp_bytes[i] - b'0') as usize * base;
        base *= 10;
        i -= 1;
    }
    remaining
}

fn gather_resp_bytes(resp: &Vec<ResFramedData>, size: usize) -> BytesMut {
    let mut bytes = BytesMut::with_capacity(size);
    for res in resp {
        bytes.extend_from_slice(&res.data);
    }
    bytes
}

#[cfg(test)]
mod tests {
    use bytes::Bytes;

    use super::*;

    #[test]
    fn test_gather_resp_bytes() {
        let mut resp = Vec::new();
        resp.push(ResFramedData {
            data: Bytes::from("hello".as_bytes()),
            is_done: false,
            res_is_ok: true,
        });
        resp.push(ResFramedData {
            data: Bytes::from("world".as_bytes()),
            is_done: true,
            res_is_ok: true,
        });
        let size = 10;
        let bytes = gather_resp_bytes(&resp, size);
        assert_eq!(bytes, "helloworld".as_bytes());
    }

    #[test]
    fn test_get_remaining_channels() {
        let resp = BytesMut::from("*3\r\n$11\r\nunsubscribe\r\n$4\r\nchan\r\n:0\r\n".as_bytes());
        assert_eq!(get_remaining_channels(resp), 0);

        let resp = BytesMut::from("*3\r\n$11\r\nunsubscribe\r\n$4\r\nchan\r\n:1\r\n".as_bytes());
        assert_eq!(get_remaining_channels(resp), 1);

        let resp = BytesMut::from("*3\r\n$11\r\nunsubscribe\r\n$4\r\nchan\r\n:111\r\n".as_bytes());
        assert_eq!(get_remaining_channels(resp), 111);
    }

    #[test]
    fn test_is_pub_sub_resp() {
        let resp = BytesMut::from("*3\r\n$9\r\nsubscribe\r\n$4\r\nchan\r\n:0\r\n".as_bytes());
        assert_eq!(is_pub_sub_resp(&resp), true);

        let resp = BytesMut::from("*3\r\n$10\r\npsubscribe\r\n$4\r\nchan\r\n:0\r\n".as_bytes());
        assert_eq!(is_pub_sub_resp(&resp), true);

        let resp = BytesMut::from("*3\r\n$11\r\nunsubscribe\r\n$4\r\nchan\r\n:0\r\n".as_bytes());
        assert_eq!(is_pub_sub_resp(&resp), true);

        let resp = BytesMut::from("*3\r\n$12\r\npunsubscribe\r\n$4\r\nchan\r\n:0\r\n".as_bytes());
        assert_eq!(is_pub_sub_resp(&resp), true);

        let resp = BytesMut::from("*3\r\n$7\r\nmessage\r\n$4\r\nchan\r\n$4\r\ntest\r\n".as_bytes());
        assert_eq!(is_pub_sub_resp(&resp), false);

        let resp = BytesMut::from("*4\r\n$8\r\npmessage\r\n$4\r\nchan\r\n$4\r\nchan\r\n$4\r\ntest\r\n".as_bytes());
        assert_eq!(is_pub_sub_resp(&resp), false);
    }

    #[test]
    fn test_is_spub_ssub_resp() {
        let resp = BytesMut::from("*3\r\n$10\r\nssubscribe\r\n$4\r\nchan\r\n:0\r\n".as_bytes());
        assert_eq!(is_spub_ssub_resp(&resp), true);

        let resp = BytesMut::from("*3\r\n$12\r\nsunsubscribe\r\n$4\r\nchan\r\n:0\r\n".as_bytes());
        assert_eq!(is_spub_ssub_resp(&resp), true);

        let resp = BytesMut::from("*3\r\n$8\r\nsmessage\r\n$4\r\nchan\r\n$4\r\ntest\r\n".as_bytes());
        assert_eq!(is_pub_sub_resp(&resp), false);
    }
}