ecksport_net/

worker.rs

//! Concurrent connection worker that returns an event queue of new channel
//! handles and notifications, operating in the background.
// TODO make this also send and respond to ping messages

use std::collections::HashMap;
use std::sync::Arc;

use ecksport_core::frame::MsgFlags;
use ecksport_core::peer::PeerData;
use futures::{future, pin_mut};
use tokio::sync::{mpsc, oneshot};
use tracing::*;

use ecksport_core::topic;
use ecksport_core::traits::{AsyncRecvFrame, AsyncSendFrame};

use crate::channel::InbMsg;
use crate::channel_state::Creator;
use crate::event::{InbEvent, OpenChanCmd, PushFlags, WorkerCommand};
use crate::shared_state::{ChanSharedState, ConnSharedState};
use crate::{channel, connection::Connection, errors::Error};

/// Event about the connection that whatever is holding the worker handle would
/// be interested in.
pub enum WorkerEvent {
    /// New channel that's been initiated by the remote.
    NewChan(channel::ChannelHandle),

    /// Oneshot notification sent by the remote.
    Notification(topic::Topic, Vec<u8>),
}

/// Handle to a connection worker that accepts new channels as channel handles
/// and relays messages.
pub struct ConnectionHandle {
    shared: Arc<ConnSharedState>,
    event_rx: mpsc::Receiver<Result<WorkerEvent, Error>>,
    cmd_tx: mpsc::Sender<WorkerCommand>,
}

impl ConnectionHandle {
    /// Returns the protocol that was negotiated in the client.
    pub fn protocol(&self) -> topic::Topic {
        self.shared.protocol()
    }

    /// Returns the peer data structure.
    pub fn peer(&self) -> &PeerData {
        self.shared.peer_data()
    }

    /// The party that initiated the underlying connection.
    pub fn initiator(&self) -> Creator {
        self.shared.initiator()
    }

    async fn submit_command(&self, cmd: WorkerCommand) -> Result<(), Error> {
        if self.cmd_tx.send(cmd).await.is_err() {
            return Err(Error::ConnWorkerExit);
        }

        Ok(())
    }

    /// Opens a channel by queueing a command and waiting for it to be executed
    /// by the worker task.
    pub async fn open_channel(
        &self,
        topic: topic::Topic,
        init_msg: Vec<u8>,
        flags: MsgFlags,
    ) -> Result<channel::ChannelHandle, Error> {
        // Create and send the open channel command.
        let (chh_tx, chh_rx) = oneshot::channel();
        let cmd = OpenChanCmd {
            topic,
            init_msg,
            flags,
            chh_tx,
        };

        self.submit_command(WorkerCommand::OpenChannel(cmd)).await?;

        // Now wait for a response.
        match chh_rx.await {
            Ok(chh) => Ok(chh),
            // TODO more expressive handling
            Err(_) => Err(Error::ConnWorkerExit),
        }
    }

    /// Queues a notification to be sent.
    pub async fn send_notification(&self, topic: topic::Topic, msg: Vec<u8>) -> Result<(), Error> {
        self.submit_command(WorkerCommand::SendNotification(topic, msg))
            .await?;
        Ok(())
    }

    /// Returns if there are any events in the queue.  This might be an error
    /// message.
    pub fn has_event(&mut self) -> bool {
        !self.event_rx.is_empty()
    }

    /// Waits for an event, asynchronously.
    pub async fn wait_event(&mut self) -> Result<Option<WorkerEvent>, Error> {
        self.event_rx.recv().await.transpose()
    }

    /// Waits for an event, blockingly.
    pub fn wait_event_blocking(&mut self) -> Result<Option<WorkerEvent>, Error> {
        self.event_rx.blocking_recv().transpose()
    }
}

impl Drop for ConnectionHandle {
    fn drop(&mut self) {
        self.shared.set_dropped();
    }
}

/// Encapsulates worker IO so that we can select across all the different inputs
/// easily.
struct WorkerIo<'c, T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin + 'static> {
    /// Low-level connection type that manages the low-level protocol state machine.
    conn: &'c mut Connection<T>,

    /// Queue that commands are sent to the worker from consumer code holding
    /// channels and whatnot.
    cmd_rx: mpsc::Receiver<WorkerCommand>,
}

impl<'c, T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin + 'static> WorkerIo<'c, T> {
    pub fn new(conn: &'c mut Connection<T>, cmd_rx: mpsc::Receiver<WorkerCommand>) -> Self {
        Self { conn, cmd_rx }
    }

    /// Selects across the IO inputs and returns a signal to act on.
    async fn wait_for_signal(&mut self) -> Result<Signal, Error> {
        let ev_fut = self.conn.next_event();
        let cmd_fut = self.cmd_rx.recv();
        pin_mut!(ev_fut);
        pin_mut!(cmd_fut);

        match future::select(ev_fut, cmd_fut).await {
            future::Either::Left((ev, _)) => match ev? {
                Some(ev) => Ok(Signal::ConnEvent(ev)),
                None => Ok(Signal::RemoteClosed),
            },

            future::Either::Right((cmd, _)) => match cmd {
                Some(cmd) => Ok(Signal::Command(cmd)),
                // All cleaned up, I guess?
                None => Ok(Signal::Shutdown),
            },
        }
    }
}

/// Signal from IO handles to be processed by the connection worker.
enum Signal {
    ConnEvent(InbEvent),
    Command(WorkerCommand),
    RemoteClosed,
    Shutdown,
}

/// Wraps the persistent worker bookkeeping state we need to share in different
/// places.
struct WorkerState {
    /// Protocol the worker is operating on.
    protocol: topic::Topic,

    /// State shared between the worker and channels, used to track when
    /// channels are ungracefully dropped so we can still clean them up.
    shared_state: Arc<ConnSharedState>,

    /// Inbound queues we maintain for each channel.  When the remote closes
    /// their side of the channel, we remove these queues to signal to the
    /// consumer that the channel has been closed.
    chan_inb_tbl: HashMap<u32, mpsc::Sender<Result<channel::InbMsg, Error>>>,

    /// Event queue sent to whatever is monitoring the worker.
    event_tx: mpsc::Sender<Result<WorkerEvent, Error>>,

    /// Queue cloned when constructing new channels so they can send commands
    /// back to the worker.
    cmd_tx: mpsc::Sender<WorkerCommand>,
}

impl WorkerState {
    fn new(
        protocol: topic::Topic,
        shared_state: Arc<ConnSharedState>,
        event_tx: mpsc::Sender<Result<WorkerEvent, Error>>,
        cmd_tx: mpsc::Sender<WorkerCommand>,
    ) -> Self {
        Self {
            protocol,
            shared_state,
            chan_inb_tbl: HashMap::new(),
            event_tx,
            cmd_tx,
        }
    }

    /// Relays an event to the worker handle.
    async fn relay_event(&self, ev: WorkerEvent) -> Result<(), Error> {
        if self.event_tx.send(Ok(ev)).await.is_err() {
            // TODO should we do something to handle this more?
            return Err(Error::ConnRecvDropped);
        }
        Ok(())
    }

    /// Relays an error message to the receiver.
    async fn relay_err(&self, e: Error) -> Result<(), Error> {
        if self.event_tx.send(Err(e)).await.is_err() {
            return Err(Error::ConnRecvDropped);
        }
        Ok(())
    }

    /// Closes the inbound queue for a channel, in response to the remote cleaning it up.
    fn close_chan_inb(&mut self, id: u32) {
        assert!(self.chan_inb_tbl.remove(&id).is_some());
    }

    /// After we've seen that a channel has been closed at the connection level,
    /// we should clean it up here.
    async fn cleanup_chan(&mut self, id: u32) {
        if self.chan_inb_tbl.contains_key(&id) {
            warn!(%id, "cleaning up channel that we still have inbound queue open");
            self.chan_inb_tbl.remove(&id);
        }

        let mut states = self.shared_state.chan_shared.write().await;
        assert!(states.remove(&id).is_some());
    }

    /// Updates the internal bookkeeping and creates the appropriate channels
    /// to construct a new channel handle.
    async fn create_chan(
        &mut self,
        new_id: u32,
        topic: topic::Topic,
    ) -> Result<channel::ChannelHandle, Error> {
        // TODO make configurable
        let (inb_tx, inb_rx) = mpsc::channel(64);
        assert!(self.chan_inb_tbl.insert(new_id, inb_tx).is_none());

        let css = Arc::new(ChanSharedState::new(self.protocol, topic));
        {
            let mut states = self.shared_state.chan_shared.write().await;
            assert!(!states.contains_key(&new_id));
            states.insert(new_id, css.clone());
        }

        let cmd_tx = self.cmd_tx.clone();
        let pd = self.shared_state.peer_data().clone();
        let handle = channel::ChannelHandle::new(new_id, pd, css, inb_rx, cmd_tx);
        Ok(handle)
    }

    /// Relays an inbound message by sending it on the channel to the handle.
    async fn relay_inb_msg(
        &mut self,
        id: u32,
        flags: PushFlags,
        payload: Vec<u8>,
    ) -> Result<(), Error> {
        let ch_inb_tx = self.chan_inb_tbl.get(&id).ok_or(Error::RecvOnUnkChan(id))?;
        if ch_inb_tx
            .send(Ok(InbMsg::new(flags, payload)))
            .await
            .is_err()
        {
            // This means the channel handle was dropped.  We'll get around to
            // checking for this on the next loop around.
            warn!(%id, "channel dropped without being explicitly closed");
        }

        Ok(())
    }

    async fn relay_inb_err(&mut self, id: u32, err: Error) -> Result<(), Error> {
        let ch_inb_tx = self.chan_inb_tbl.get(&id).ok_or(Error::RecvOnUnkChan(id))?;
        if ch_inb_tx.send(Err(err)).await.is_err() {
            // This means the channel handle was dropped.  We'll get around to
            // checking for this on the next loop around.
            warn!(%id, "channel dropped without being explicitly closed");
        }

        Ok(())
    }

    /// Relays a shutdown error message to all channels that can receive it, and drops our
    /// end of the channel.
    async fn shutdown_channels(&mut self) -> Result<(), Error> {
        // Calling .drain() here will drop the channels.
        for (id, ch) in self.chan_inb_tbl.drain() {
            if ch.send(Err(Error::ConnWorkerExit)).await.is_err() {
                warn!(%id, "channel dropped without being explicitly closed");
            }
        }

        Ok(())
    }
}

/// Spawns a connection worker and returns a [``WorkerHandle``] to it.
pub async fn spawn_connection_worker<
    T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin + 'static,
>(
    conn: Connection<T>,
) -> ConnectionHandle {
    // Create the channels.
    // TODO make these configurable
    let (event_tx, event_rx) = mpsc::channel(64);
    let (cmd_tx, cmd_rx) = mpsc::channel(256);

    // Copy some stuff into a new shared state instance.
    let proto = conn.protocol();
    let pd = conn.peer_data().clone();
    let peer = pd.location().clone();
    let initer = conn.initiator();
    let shared = Arc::new(ConnSharedState::new(proto, pd, initer));

    let worker_span = debug_span!("conn", %peer);
    debug!(parent: &worker_span, "spawning worker task");

    // Actually spawn the task!  This took fucking long enough to get to!
    tokio::spawn(
        conn_worker_task(conn, shared.clone(), event_tx, cmd_rx, cmd_tx.clone())
            .instrument(worker_span),
    );

    ConnectionHandle {
        shared,
        cmd_tx,
        event_rx,
    }
}

pub async fn conn_worker_task<
    T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin + 'static,
>(
    mut conn: Connection<T>,
    shared: Arc<ConnSharedState>,
    event_tx: mpsc::Sender<Result<WorkerEvent, Error>>,
    cmd_rx: mpsc::Receiver<WorkerCommand>,
    cmd_tx: mpsc::Sender<WorkerCommand>,
) {
    let proto = conn.protocol();
    let wio = WorkerIo::new(&mut conn, cmd_rx);
    let mut wstate = WorkerState::new(proto, shared, event_tx, cmd_tx);
    if let Err(e) = do_worker(wio, &mut wstate).await {
        warn!(err = %e, "connection worker task exited");
    }
}

// TODO add some provision so that we can clean up dropped channel handles
// without having to receive a message on them first, maybe using an atomicbool
async fn do_worker<'c, T: AsyncRecvFrame + AsyncSendFrame + Sync + Send + Unpin + 'static>(
    mut wio: WorkerIo<'c, T>,
    wstate: &mut WorkerState,
) -> Result<(), Error> {
    loop {
        let signal = match wio.wait_for_signal().await {
            Ok(s) => s,
            Err(e) => {
                // Shut down all the channels with an error, just so the
                // receivers can know what happened.
                // TODO maybe make this more explicit about being an error
                // condition so the channels can take appropriate action
                if let Err(e) = wstate.shutdown_channels().await {
                    warn!(err = %e, "encountered error while shutting down channels");
                }

                return Err(e.into());
            }
        };

        match signal {
            Signal::ConnEvent(ev) => handle_conn_event(ev, wstate).await?,

            Signal::Command(cmd) => match cmd {
                WorkerCommand::OpenChannel(occ) => {
                    let topic = occ.topic;
                    let flags = occ.flags;

                    // Send the initial opening so we can get the chan ID.
                    let id = wio.conn.open_channel(topic, occ.init_msg, flags).await?;

                    let ch_handle = wstate.create_chan(id, topic).await?;

                    if occ.chh_tx.send(ch_handle).is_err() {
                        warn!(%topic, %id, "channel sendback closed before open completed");
                        // TODO should we drop the rest of this?
                    }
                }

                WorkerCommand::SendMsg(msg) => {
                    let id = msg.id();
                    let flags = *msg.flags();
                    let still_open = wio.conn.send_message(id, msg.into_payload(), flags).await?;
                    if !still_open {
                        wstate.cleanup_chan(id).await;
                    }
                }

                WorkerCommand::CloseChannel(id) => {
                    let still_open = wio.conn.close_channel(id).await?;
                    if !still_open {
                        wstate.cleanup_chan(id).await;
                    }
                }

                WorkerCommand::SendNotification(topic, notif) => {
                    wio.conn.send_notification(topic, notif).await?;
                }
            },

            Signal::RemoteClosed => {
                if !wstate.chan_inb_tbl.is_empty() {
                    // Is this even a warning?
                    let channels = wstate.chan_inb_tbl.len();
                    warn!(%channels, "remote side closed with channels still open");
                }

                return Ok(());
            }

            Signal::Shutdown => {
                // TODO Is there any other cleanup we should do?
                wstate.shutdown_channels().await?;
                return Ok(());
            }
        }
    }
}

async fn handle_conn_event(conn_ev: InbEvent, wstate: &mut WorkerState) -> Result<(), Error> {
    match conn_ev {
        InbEvent::NewChannel(id, topic, flags, payload) => {
            // Create the channel and write the payload to it.
            let ch_handle = wstate.create_chan(id, topic).await?;
            wstate.relay_inb_msg(id, flags, payload).await?;

            // Send off the new channel event.
            wstate.relay_event(WorkerEvent::NewChan(ch_handle)).await?;
            Ok(())
        }

        InbEvent::PushChannel(id, flags, payload) => {
            wstate.relay_inb_msg(id, flags, payload).await?;
            Ok(())
        }

        InbEvent::CloseChannel(id, still_alive) => {
            // Close the queue.
            wstate.close_chan_inb(id);

            // In this case we just remove the channel channel.
            if !still_alive {
                wstate.cleanup_chan(id).await;
            }

            Ok(())
        }

        InbEvent::Notification(topic, payload) => {
            // Send the notification, exit if we have to.
            wstate
                .relay_event(WorkerEvent::Notification(topic, payload))
                .await?;
            Ok(())
        }
    }
}