connection.rs

//! All Sink transforms use SinkConnection for their outgoing connections.

use crate::codec::{CodecBuilder, CodecReadError, CodecWriteError};
use crate::frame::Frame;
use crate::message::{Message, MessageId, Messages};
use crate::tcp;
use crate::tls::{TlsConnector, ToHostname};
use futures::{SinkExt, StreamExt};
use std::io::ErrorKind;
use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::Duration;
use tokio::io::{AsyncRead, AsyncWrite, split};
use tokio::net::ToSocketAddrs;
use tokio::sync::mpsc::error::TryRecvError;
use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender};
use tokio::sync::{Notify, mpsc};
use tokio_util::codec::{FramedRead, FramedWrite};
use tracing::Instrument;
use tracing::error;

pub struct SinkConnection {
    in_rx: mpsc::Receiver<Vec<Message>>,
    out_tx: mpsc::UnboundedSender<Vec<Message>>,
    connection_closed_rx: mpsc::Receiver<ConnectionError>,
    error: Option<ConnectionError>,
    dummy_response_inserter: DummyResponseInserter,
}

impl SinkConnection {
    pub async fn new<A: ToSocketAddrs + ToHostname + std::fmt::Debug, C: CodecBuilder + 'static>(
        host: A,
        codec_builder: C,
        tls: &Option<TlsConnector>,
        connect_timeout: Duration,
        force_run_chain: Arc<Notify>,
        read_timeout: Option<Duration>,
    ) -> anyhow::Result<Self> {
        let destination = tokio::net::lookup_host(&host).await?.next().unwrap();
        let (in_tx, in_rx) = mpsc::channel::<Messages>(10_000);
        let (out_tx, out_rx) = mpsc::unbounded_channel::<Messages>();
        let (connection_closed_tx, connection_closed_rx) = mpsc::channel(1);

        if let Some(tls) = tls.as_ref() {
            let tls_stream = tls.connect(connect_timeout, host).await?;
            let (rx, tx) = split(tls_stream);
            spawn_read_write_tasks(
                codec_builder,
                rx,
                tx,
                in_tx,
                out_rx,
                out_tx.clone(),
                force_run_chain,
                connection_closed_tx,
                read_timeout,
            );
        } else {
            let tcp_stream = tcp::tcp_stream(connect_timeout, destination).await?;
            let (rx, tx) = tcp_stream.into_split();
            spawn_read_write_tasks(
                codec_builder,
                rx,
                tx,
                in_tx,
                out_rx,
                out_tx.clone(),
                force_run_chain,
                connection_closed_tx,
                read_timeout,
            );
        }

        let dummy_response_inserter = DummyResponseInserter::new();

        Ok(SinkConnection {
            in_rx,
            out_tx,
            connection_closed_rx,
            error: None,
            dummy_response_inserter,
        })
    }

    /// Number of requests waiting on a response.
    /// The count includes requests that will have a dummy response generated by shotover.
    pub fn pending_requests_count(&self) -> usize {
        self.dummy_response_inserter.pending_requests_count()
    }

    /// This method must only be called when the read or write tasks have closed their `in_` or `out_` channel.
    /// In this case it is gauranteed that the `connection_closed_` channel will
    /// have an error sent to it before the closing of `in_` or `out_`.
    fn get_error_for_close(&mut self) -> ConnectionError {
        self.error = Some(self.connection_closed_rx.try_recv().unwrap());
        self.error.clone().unwrap()
    }

    /// This method can be called at any time.
    /// If the connection has hit an error that error will be returned.
    pub fn get_error(&mut self) -> Option<ConnectionError> {
        if self.error.is_none() {
            self.error = self.connection_closed_rx.try_recv().ok();
        }
        self.error.clone()
    }

    /// Send messages.
    /// If there is a problem with the connection an error is returned.
    pub fn send(&mut self, mut messages: Vec<Message>) -> Result<(), ConnectionError> {
        self.dummy_response_inserter.process_requests(&mut messages);

        if let Some(error) = &self.error {
            Err(error.clone())
        } else {
            self.out_tx
                .send(messages)
                .map_err(|_| self.get_error_for_close())
        }
    }

    /// Receives messages, if there are no messages available it awaits until there are messages.
    /// If there is a problem with the connection an error is returned.
    pub async fn recv(&mut self) -> Result<Vec<Message>, ConnectionError> {
        let mut result = vec![];
        self.recv_into(&mut result).await?;
        Ok(result)
    }

    /// Receives messages, if there are no messages available it awaits until there are messages.
    /// If there is a problem with the connection an error is returned.
    pub async fn recv_into(&mut self, responses: &mut Vec<Message>) -> Result<(), ConnectionError> {
        let initial_count = responses.len();
        // exhaust channel and ensure dummy messages are inserted by running try_recv_into
        self.try_recv_into(responses)?;

        // If we didnt receive any messages from the channel or the dummy inserter then await more messages
        if initial_count == responses.len() {
            match self.in_rx.recv().await {
                Some(mut new) => {
                    self.dummy_response_inserter.process_responses(&mut new, 0);

                    // If there is no allocation behind the vec then just directly use the new vec.
                    if responses.capacity() == 0 {
                        *responses = new;
                    } else {
                        responses.extend(new)
                    }
                }
                None => {
                    return Err(self.get_error_for_close());
                }
            }
        }

        Ok(())
    }

    /// Attempts to receive messages, if there are no messages available it immediately returns an empty vec.
    /// If there is a problem with the connection an error is returned.
    pub fn try_recv_into(&mut self, responses: &mut Vec<Message>) -> Result<(), ConnectionError> {
        if let Some(error) = &self.error {
            Err(error.clone())
        } else {
            let initial_count = responses.len();

            // Even if >1 message batches are in the channel force_run_chain will only notify us once.
            // So we need to ensure we fully exhaust the channel to prevent messages being stuck when there is a backlog.
            loop {
                match self.in_rx.try_recv() {
                    Ok(new) => {
                        // If there is no allocation behind the vec then just directly use the new vec.
                        if responses.capacity() == 0 {
                            *responses = new;
                        } else {
                            responses.extend(new)
                        }
                    }
                    Err(TryRecvError::Disconnected) => {
                        self.dummy_response_inserter
                            .process_responses(responses, initial_count);
                        return self.handle_disconnect_on_recv(responses, initial_count);
                    }
                    Err(TryRecvError::Empty) => {
                        // We need to ensure this runs even when we receive nothing from the channel.
                        // this will ensure that the dummy message at index 0 is inserted if it exists
                        self.dummy_response_inserter
                            .process_responses(responses, initial_count);
                        return Ok(());
                    }
                }
            }
        }
    }

    fn handle_disconnect_on_recv(
        &mut self,
        responses: &[Message],
        initial_count: usize,
    ) -> Result<(), ConnectionError> {
        // call this first to ensure the next send call will have an error
        let err = self.get_error_for_close();

        if responses.len() == initial_count {
            // We failed to get any messages return the error.
            Err(err)
        } else {
            // We got at least some messages, so consider this a success and let the error be returned when the user next calls send()
            Ok(())
        }
    }
}

/// This represents an unrecoverable error to the connection.
/// The connection is no longer usable after this error is received.
#[derive(thiserror::Error, Debug, Clone)]
pub enum ConnectionError {
    #[error("The other side of this connection closed the connection")]
    OtherSideClosed,
    #[error("Shotover closed the connection due to protocol requirements")]
    ShotoverClosed,
    #[error("Message decode error {0}")]
    MessageDecode(Arc<anyhow::Error>),
    #[error("Message encode error {0}")]
    MessageEncode(Arc<anyhow::Error>),
    #[error("IO error {0}")]
    Io(Arc<std::io::Error>),
    #[error("The other side of this connection did not respond within {0:?}")]
    ReadTimeout(Duration),
}

struct RequestPending {
    pub notify: Notify,
    count: AtomicU64,
}

impl RequestPending {
    fn add(&self, value: u64) {
        self.count.fetch_add(value, Ordering::SeqCst);
        self.notify.notify_one();
    }

    fn sub(&self, value: u64) {
        self.count.fetch_sub(value, Ordering::SeqCst);
    }

    fn get(&self) -> u64 {
        self.count.load(Ordering::SeqCst)
    }
}

#[expect(clippy::too_many_arguments)]
fn spawn_read_write_tasks<
    C: CodecBuilder + 'static,
    R: AsyncRead + Unpin + Send + 'static,
    W: AsyncWrite + Unpin + Send + 'static,
>(
    codec: C,
    rx: R,
    tx: W,
    in_tx: mpsc::Sender<Messages>,
    mut out_rx: UnboundedReceiver<Messages>,
    out_tx: UnboundedSender<Messages>,
    force_run_chain: Arc<Notify>,
    connection_closed_tx: mpsc::Sender<ConnectionError>,
    read_timeout: Option<Duration>,
) {
    let (decoder, encoder) = codec.build();
    let reader = FramedRead::new(rx, decoder);
    let writer = FramedWrite::new(tx, encoder);

    let request_pending = Arc::new(RequestPending {
        notify: Notify::new(),
        count: 0.into(),
    });

    // Shutdown flows
    //
    // The Connection is dropped:
    // 1. The Connection is dropped, dropping in_rx and the first out_tx
    // 2. The reader task detects that in_rx has dropped and terminates, the last out_tx instance is dropped
    // 3. The writer task detects that the last out_tx is dropped by out_rx returning None and terminates
    //
    // Destination closes connection (proactive cleanup via force_run_chain):
    // 1.   The reader task detects that the destination has closed the connection via reader returning None
    // 1.1. connection_closed_tx is sent `ConnectionError::OtherSideClosed`
    // 1.2. in_tx is dropped, closing the channel
    // 1.3. force_run_chain.notify_one() is called to wake up the transform chain
    // 2.   The transform chain wakes up and runs with empty requests
    // 2.1. The transform is responsible for detecting the closed connection and cleaning it up
    // 3.   When the Connection is dropped, out_tx is dropped
    // 4.   The writer task detects that out_tx is dropped by out_rx returning None and terminates
    // 4.1. The writer task could also terminate early by detecting a BrokenPipe/ConnectionReset error when writing
    //
    // Destination closes connection and then shotover tries to send:
    // if a send or recv has not been attempted yet the send will appear to have succeeded.
    // if a recv was already attempted, then the logic is the same as the above example.
    // if a send was already attempted, then the following logic occurs:
    // 1.  Connection::send sends a message to the writer task via out_tx.
    // 2.  The writer task attempts to send the message to the writer but it returns a BrokenPipe or ConnectionReset error.
    // 3.1 The writer task sends an OtherSideClosed error to the Connection.
    // 3.2 The writer task terminates.
    // 4.  Connection::send sends a message to the writer task via out_tx but detects the writer task terminated due to out_tx returning None.
    // 4.1 Connection::send checks connection_closed_rx for the error, stores it and returns it to the caller.

    let connection_closed_tx2 = connection_closed_tx.clone();
    let request_pending2 = request_pending.clone();
    let force_run_chain2 = force_run_chain.clone();
    tokio::spawn(
        async move {
            match reader_task::<C, _>(
                reader,
                &in_tx,
                out_tx,
                force_run_chain,
                request_pending2,
                read_timeout,
            )
            .await
            {
                Ok(()) => {}
                Err(err) => {
                    connection_closed_tx2.try_send(err).ok();
                    // Drop in_tx only after sending the error message.
                    // This ensures the handle side logic will always have an
                    // error available to consult as to why the `in_` channel was closed.
                    std::mem::drop(in_tx);
                    // Notify the transform chain about the error so it can run recv_responses()
                    // and detect the closed connection. This must happen after dropping in_tx
                    // so that when the transform wakes up, the channel is already closed.
                    force_run_chain2.notify_one();
                }
            }
        }
        .in_current_span(),
    );

    tokio::spawn(
        async move {
            match writer_task::<C, _>(writer, &mut out_rx, request_pending).await {
                Ok(()) => {}
                Err(err) => {
                    connection_closed_tx.try_send(err).ok();
                    // Drop out_rx only after sending the error message.
                    // This ensures the handle side logic will always have an
                    // error available to consult as to why the `out_` channel was closed.
                    std::mem::drop(out_rx);
                }
            }
        }
        .in_current_span(),
    );
}

async fn reader_task<C: CodecBuilder + 'static, R: AsyncRead + Unpin + Send + 'static>(
    mut reader: FramedRead<R, <C as CodecBuilder>::Decoder>,
    in_tx: &mpsc::Sender<Messages>,
    out_tx: UnboundedSender<Messages>,
    force_run_chain: Arc<Notify>,
    request_pending: Arc<RequestPending>,
    read_timeout: Option<Duration>,
) -> Result<(), ConnectionError> {
    loop {
        let read_timeout = if request_pending.get() == 0 {
            // There are no requests pending so we should not trigger a timeout.
            // To achieve this, sleep forever.
            None
        } else {
            // There are requests pending so we need to timeout after the configure timeout elapses.
            read_timeout
        };
        tokio::select! {
            biased;
            _ = in_tx.closed() => {
                // shotover is no longer listening for responses, this task is no longer needed
                return Ok(());
            }

            // reader.next is supposedly cancel safe: https://github.com/tokio-rs/tokio/discussions/4416#discussioncomment-2023884
            result = reader.next() => {
                if let Some(messages) = result {
                    match messages {
                        Ok(messages) => {
                            let count = messages.iter().filter(|x| x.request_id.is_some()).count();
                            request_pending.sub(count as u64);

                            if in_tx.send(messages).await.is_err() {
                                // main task has shutdown, this task is no longer needed
                                return Ok(());
                            }

                            force_run_chain.notify_one();
                        }
                        Err(CodecReadError::RespondAndThenCloseConnection(messages)) => {
                            if out_tx.send(messages).is_err() {
                                error!("Failed to send RespondAndThenCloseConnection message");
                            }
                            return Err(ConnectionError::ShotoverClosed);
                        }
                        Err(CodecReadError::Parser(err)) => {
                            return Err(ConnectionError::MessageDecode(Arc::new(err)));
                        }
                        Err(CodecReadError::Io(err)) => {
                            return Err(ConnectionError::Io(Arc::new(err)));
                        }
                    }
                } else {
                    return Err(ConnectionError::OtherSideClosed);
                }
            }
            // The timeout logic can occur after read.next(), if we succesfully read we dont need the timeout at all.
            _ = request_pending.notify.notified() => {
                continue;
            }
            _ = sleep_for_duration_or_forever(read_timeout) => {
                return Err(ConnectionError::ReadTimeout(read_timeout.unwrap()));
            }
        }
    }
}

async fn sleep_for_duration_or_forever(duration: Option<Duration>) {
    if let Some(duration) = duration {
        tokio::time::sleep(duration).await
    } else {
        std::future::pending().await
    }
}

async fn writer_task<C: CodecBuilder + 'static, W: AsyncWrite + Unpin + Send + 'static>(
    mut writer: FramedWrite<W, <C as CodecBuilder>::Encoder>,
    out_rx: &mut UnboundedReceiver<Messages>,
    request_pending: Arc<RequestPending>,
) -> Result<(), ConnectionError> {
    loop {
        if let Some(messages) = out_rx.recv().await {
            request_pending.add(messages.len() as u64);
            match writer.send(messages).await {
                Err(CodecWriteError::Encoder(err)) => {
                    return Err(ConnectionError::MessageEncode(Arc::new(err)));
                }
                Err(CodecWriteError::Io(err)) => {
                    if matches!(
                        err.kind(),
                        ErrorKind::BrokenPipe | ErrorKind::ConnectionReset
                    ) {
                        return Err(ConnectionError::OtherSideClosed);
                    } else {
                        return Err(ConnectionError::Io(Arc::new(err)));
                    }
                }
                Ok(()) => {}
            }
        } else {
            // shotover is no longer sending responses, this task is no longer needed
            return Ok(());
        }
    }
}

/// Keeps track of all dummy requests that pass through this connection and inserts a dummy response at the same index as the request.
struct DummyResponseInserter {
    dummy_requests: Vec<DummyRequest>,
    pending_requests_count: usize,
}

#[derive(Debug)]
struct DummyRequest {
    request_id: MessageId,
    request_index: usize,
}

impl DummyResponseInserter {
    fn new() -> Self {
        DummyResponseInserter {
            dummy_requests: vec![],
            pending_requests_count: 0,
        }
    }

    /// All requests must be passed through this method so that DummyResponseInserter can record all requests that need a dummy response generated.
    /// The requests will not be modified.
    fn process_requests(&mut self, requests: &mut [Message]) {
        for (i, request) in requests.iter_mut().enumerate() {
            if request.response_is_dummy() {
                self.dummy_requests.push(DummyRequest {
                    request_id: request.id(),
                    request_index: self.pending_requests_count + i,
                });
            }
        }
        self.pending_requests_count += requests.len();
    }

    /// Insert dummy responses into the list of responses.
    /// All elements before the element at index `start_at` is ignored,
    /// those elements should have been already processed by a previous call to process_responses.
    fn process_responses(&mut self, responses: &mut Vec<Message>, start_at: usize) {
        let mut len = responses.len() - start_at;
        // responses with no request will invalidate our indexes, so we need to fix them up here.
        for (response_i, response) in responses[start_at..].iter().enumerate() {
            if response.request_id().is_none() {
                for (dummy_request_i, dummy_request) in
                    &mut self.dummy_requests.iter_mut().enumerate()
                {
                    // Either of `<` or `<=` could work here.
                    // If its `<` then the dummy response comes before the unrequested response
                    // If its `<=` then the dummy response comes after the unrequested response
                    if response_i <= dummy_request.request_index + dummy_request_i {
                        dummy_request.request_index += 1;
                    }
                }
                self.pending_requests_count += 1;
            }
        }

        // It is important that retain_mut iterates in the order of the vec.
        // This is because it is only once the previous insert_index is inserted that the following insert_index becomes valid again.
        self.dummy_requests.retain_mut(|dummy_request| {
            if dummy_request.request_index <= len {
                let mut dummy = Message::from_frame(Frame::Dummy);
                dummy.set_request_id(dummy_request.request_id);
                responses.insert(dummy_request.request_index + start_at, dummy);
                len += 1;
                false
            } else {
                true
            }
        });

        // Decrement indexes so that they will be offset from 0 the next time responses come in.
        for dummy_request in &mut self.dummy_requests {
            dummy_request.request_index -= len;
        }
        self.pending_requests_count -= len;
    }

    fn pending_requests_count(&self) -> usize {
        self.pending_requests_count
    }
}

#[cfg(all(test, feature = "valkey"))]
mod tests {
    use super::DummyResponseInserter;
    use crate::frame::{Frame, ValkeyFrame};
    use crate::message::Message;
    use pretty_assertions::assert_eq;

    fn dummy() -> Message {
        Message::from_frame(Frame::Dummy)
    }

    fn valkey_request() -> Message {
        Message::from_frame(Frame::Valkey(ValkeyFrame::Null))
    }

    fn valkey_response(request: &Message) -> Message {
        let mut message = Message::from_frame(Frame::Valkey(ValkeyFrame::Null));
        message.set_request_id(request.id());
        message
    }

    fn valkey_response_unsolicited() -> Message {
        Message::from_frame(Frame::Valkey(ValkeyFrame::Null))
    }

    #[test]
    fn dummy_response_inserter() {
        let mut inserter = DummyResponseInserter::new();

        // send an empty list of requests
        {
            let mut requests = vec![];
            inserter.process_requests(&mut requests);
            let mut responses = vec![];
            inserter.process_responses(&mut responses, 0);
            assert_eq!(responses, []);
        }

        // send one dummy request
        {
            let mut requests = vec![dummy()];
            inserter.process_requests(&mut requests);
            let mut responses = vec![];
            inserter.process_responses(&mut responses, 0);
            assert_eq!(responses, vec![dummy()]);
            inserter.process_responses(&mut responses, 1);
            assert_eq!(responses, vec![dummy()]);
            assert_eq!(inserter.pending_requests_count(), 0);
        }

        // send one valkey request
        {
            let mut requests = vec![valkey_request()];
            inserter.process_requests(&mut requests);
            let mut responses = vec![];
            inserter.process_responses(&mut responses, 0);
            assert_eq!(responses, []);

            // received valkey response
            assert_eq!(inserter.pending_requests_count(), 1);
            responses.insert(0, valkey_response(&requests[0]));
            inserter.process_responses(&mut responses, 0);
            assert_eq!(responses, vec![valkey_response(&requests[0])]);
            assert_eq!(inserter.pending_requests_count(), 0);
        }

        // send one dummy request and then one valkey request
        {
            let mut requests = vec![dummy(), valkey_request()];
            inserter.process_requests(&mut requests);
            let mut responses = vec![];
            inserter.process_responses(&mut responses, 0);
            assert_eq!(responses, vec![dummy()]);

            // received valkey response
            assert_eq!(inserter.pending_requests_count(), 1);
            responses.insert(1, valkey_response(&requests[1]));
            inserter.process_responses(&mut responses, 1);
            assert_eq!(responses, vec![dummy(), valkey_response(&requests[1])]);
            assert_eq!(inserter.pending_requests_count(), 0);

            // try_recv with no responses
            responses.clear();
            inserter.process_responses(&mut responses, 0);
            assert_eq!(inserter.pending_requests_count(), 0);
        }

        // send one valkey request and then one dummy request
        {
            let mut requests = vec![valkey_request(), dummy()];
            inserter.process_requests(&mut requests);
            let mut responses = vec![];
            inserter.process_responses(&mut responses, 0);
            assert_eq!(responses, vec![]);

            // received valkey response
            assert_eq!(inserter.pending_requests_count(), 2);
            responses.insert(0, valkey_response(&requests[0]));
            inserter.process_responses(&mut responses, 0);
            assert_eq!(responses, vec![valkey_response(&requests[0]), dummy()]);
            assert_eq!(inserter.pending_requests_count(), 0);
            inserter.process_responses(&mut responses, 2);
            assert_eq!(responses, vec![valkey_response(&requests[0]), dummy()]);
            assert_eq!(inserter.pending_requests_count(), 0);

            // try_recv with no responses
            responses.clear();
            inserter.process_responses(&mut responses, 0);
            assert_eq!(inserter.pending_requests_count(), 0);
        }

        // send one valkey request and then one dummy request
        // an unsolicited response is received before and after the valkey response
        {
            let mut requests = vec![valkey_request(), dummy()];
            inserter.process_requests(&mut requests);
            let mut responses = vec![];
            inserter.process_responses(&mut responses, 0);
            assert_eq!(responses, vec![]);

            // received valkey response surrounded by unsolicted responses
            assert_eq!(inserter.pending_requests_count(), 2);
            responses.insert(0, valkey_response_unsolicited());
            responses.insert(0, valkey_response(&requests[1]));
            responses.insert(0, valkey_response_unsolicited());
            inserter.process_responses(&mut responses, 0);
            assert_eq!(
                responses,
                vec![
                    valkey_response_unsolicited(),
                    valkey_response(&requests[0]),
                    valkey_response_unsolicited(),
                    dummy()
                ]
            );
            assert_eq!(inserter.pending_requests_count(), 0);
            inserter.process_responses(&mut responses, 4);
            assert_eq!(
                responses,
                vec![
                    valkey_response_unsolicited(),
                    valkey_response(&requests[0]),
                    valkey_response_unsolicited(),
                    dummy()
                ]
            );
            assert_eq!(inserter.pending_requests_count(), 0);

            // try_recv with no responses
            responses.clear();
            inserter.process_responses(&mut responses, 0);
            assert_eq!(inserter.pending_requests_count(), 0);
        }
    }
}