async_nng/

socket.rs

//! Types and traits used for performing asynchronous socket operations based on NNG's AIO.

use super::cancel_guard::{Cancel, CancelGuard};
use async_channel::{bounded, Receiver};
use std::time::Duration;
use tracing::instrument;

/// A wrapper type around [`Socket`](nng::Socket) to enable async-await send and receive
/// operations.
///
/// This type allows for getting a single future (send or receive) from a socket at a time. This is
/// useful for raw sockets and scalability protocols which cannot leverage
/// [`Context`](nng::Context), and thus need to perform send and receive operations on an owned
/// socket directly.
///
/// If you have set up your socket such that it can be used with contexts, it is instead
/// recommended to use [`AsyncContext`](super::context::AsyncContext), which only needs to borrow
/// the underlying socket. Contexts are more useful for concurrent operations where
/// of futures for each independent operation on a socket.
#[derive(Debug)]
pub struct AsyncSocket {
    /// The underlying socket
    socket: nng::Socket,

    /// AIO object for managing async operations on the socket
    aio: nng::Aio,

    /// Receiving end of a bounded channel for getting results back from the AIO object.
    receiver: Receiver<nng::AioResult>,
}

impl TryFrom<nng::Socket> for AsyncSocket {
    type Error = nng::Error;

    #[instrument]
    fn try_from(socket: nng::Socket) -> Result<Self, Self::Error> {
        let (sender, receiver) = bounded(1);
        let aio = nng::Aio::new(move |aio, result| {
            if let Err(err) = sender.try_send(result) {
                tracing::warn!(?err, "Failed to forward IO event from nng::Aio.");
                aio.cancel();
            }
        })?;
        Ok(Self {
            socket,
            aio,
            receiver,
        })
    }
}

impl Cancel for AsyncSocket {
    fn cancel(&mut self) {
        self.aio.cancel();
        // Blocking call which isn't great but because we've cancelled above this will either:
        //
        // 1. Return NNG_ECANCELED very quickly
        // 2. Return the actual IO event
        //
        // Either way, the result isn't useful to us because it has been canceled. One trade-off of
        // this approach is that we have to accept that we are "blocking" but realistically we
        // aren't expecting to be blocking long, because the call to cancel the AIO object should
        // not take excessively long aside from potential context switching from the OS.
        let _ = self.receiver.recv_blocking();
    }
}

impl AsyncSocket {
    /// Sends a [`Message`](nng::Message) to the socket asynchronously.
    ///
    /// # Errors
    ///
    /// - `IncorrectState` if the internal `Aio` is already running an operation, or the socket
    /// cannot send messages in its current state.
    /// - `MessageTooLarge`: The message is too large.
    /// - `NotSupported`: The protocol does not support sending messages.
    /// - `OutOfMemory`: Insufficient memory available.
    /// - `TimedOut`: The operation timed out.
    #[instrument(skip(msg))]
    pub async fn send<M>(
        &mut self,
        msg: M,
        timeout: Option<Duration>,
    ) -> Result<(), (nng::Message, nng::Error)>
    where
        M: Into<nng::Message>,
    {
        if let Err(err) = self.aio.set_timeout(timeout) {
            return Err((msg.into(), err));
        }

        self.socket.send_async(&self.aio, msg.into())?;
        let receiver = self.receiver.clone();

        let guard = CancelGuard::new(self);
        // Unwrap is safe here because we know that the sender cannot be dropped unless `self.aio`
        // is dropped.
        let res = match receiver.recv().await.unwrap() {
            nng::AioResult::Send(res) => res,
            _ => {
                tracing::warn!("Reached invalid state in AsyncSocket::send.");
                unreachable!();
            }
        };
        std::mem::forget(guard);
        res
    }

    /// Receives a [`Message`](nng::Message) from the socket asynchronously.
    ///
    /// # Errors
    ///
    /// - `IncorrectState` if the internal `Aio` is already running an operation, or the socket
    /// cannot send messages in its current state.
    /// - `MessageTooLarge`: The message is too large.
    /// - `NotSupported`: The protocol does not support sending messages.
    /// - `OutOfMemory`: Insufficient memory available.
    /// - `TimedOut`: The operation timed out.
    #[instrument]
    pub async fn receive(&mut self, timeout: Option<Duration>) -> Result<nng::Message, nng::Error> {
        self.aio.set_timeout(timeout)?;
        self.socket.recv_async(&self.aio)?;
        let receiver = self.receiver.clone();

        let guard = CancelGuard::new(self);
        // Unwrap is safe here because we know that the sender cannot be dropped unless `self.aio`
        // is dropped.
        let res = match receiver.recv().await.unwrap() {
            nng::AioResult::Recv(res) => res,
            _ => {
                tracing::warn!("Reached invalid state in AsyncSocket::receive.");
                unreachable!();
            }
        };
        std::mem::forget(guard);
        res
    }

    /// Grabs the inner [`Socket`](nng::Socket).
    #[instrument]
    pub fn into_inner(self) -> nng::Socket {
        self.socket
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use macro_rules_attribute::apply;
    use nng::options::Options;
    use std::{io::Write, time::Duration};
    use tracing_test::traced_test;

    /// Bytes to send on the request
    const PING_BYTES: &[u8] = b"ping";

    /// Bytes to send on the reply
    const PONG_BYTES: &[u8] = b"pong";

    /// Helper function for getting a configured pair of Req/Rep socket pairs.
    #[instrument]
    fn make_req_rep() -> (AsyncSocket, AsyncSocket) {
        let addr = nng::SocketAddr::InProc(random_string::generate(128, "abcdef1234567890"));
        tracing::info!(addr = addr.to_string(), "Local address");

        let rep = nng::Socket::new(nng::Protocol::Rep0).expect("Construct Rep0");
        let _listener = {
            let builder = nng::ListenerBuilder::new(&rep, addr.to_string().as_str())
                .expect("Listener create succeeds.");
            // 1MiB RECVMAXSZ
            builder
                .set_opt::<nng::options::RecvMaxSize>(1024 * 1024)
                .expect("Set opt on listener builder.");
            builder.start().expect("Start listener.")
        };

        let req = nng::Socket::new(nng::Protocol::Req0).expect("Construct Req0");
        let _dialer = {
            let builder = nng::DialerBuilder::new(&req, addr.to_string().as_str())
                .expect("Dial create succeeds.");
            // 1MiB RECVMAXSZ
            builder
                .set_opt::<nng::options::RecvMaxSize>(1024 * 1024)
                .expect("Set opt on dialer builder.");
            builder.start(false).expect("Start dialer.")
        };

        let req_async = AsyncSocket::try_from(req).expect("Async req.");
        let rep_async = AsyncSocket::try_from(rep).expect("Async rep.");

        (req_async, rep_async)
    }

    #[traced_test]
    #[apply(smol_macros::test)]
    async fn ping_pong_with_req_rep() {
        let (mut req, mut rep) = make_req_rep();

        let mut ping = nng::Message::new();
        ping.write_all(PING_BYTES).expect("All bytes written.");
        req.send(ping, None)
            .await
            .expect("Request should send successfully");

        let request = rep
            .receive(None)
            .await
            .expect("Request should be received successfully.");
        assert_eq!(PING_BYTES, request.as_slice());

        let mut pong = nng::Message::new();
        pong.write_all(PONG_BYTES).expect("All bytes written.");

        rep.send(pong, None)
            .await
            .expect("Reply should send successfully.");

        let response = req
            .receive(None)
            .await
            .expect("Response should be received successfully.");
        assert_eq!(PONG_BYTES, response.as_slice());
    }

    #[traced_test]
    #[apply(smol_macros::test)]
    async fn drop_rep_after_request_sent_fails() {
        let (mut req, rep) = make_req_rep();

        let mut ping = nng::Message::new();
        ping.write_all(PING_BYTES).expect("All bytes written.");
        req.send(ping, None)
            .await
            .expect("Request should send successfully");

        std::mem::drop(rep);

        let err = req
            .receive(Some(Duration::from_millis(20)))
            .await
            .expect_err("Request should timeout because the response will never come.");

        assert_eq!(err, nng::Error::TimedOut);
    }

    #[traced_test]
    #[apply(smol_macros::test)]
    async fn send_is_cancel_safe() {
        let (mut req, mut rep) = make_req_rep();

        let mut ping = nng::Message::new();
        ping.write_all(PING_BYTES).expect("All bytes written.");

        smol::future::or(smol::future::ready(Ok(())), req.send(ping, None))
            .await
            .expect("Immediate future should return Ok.");

        let err = rep
            .receive(Some(Duration::from_millis(20)))
            .await
            .expect_err("Request should timeout because nothing was sent.");

        assert_eq!(err, nng::Error::TimedOut);

        // Now send it for real. The state of the req socket should be good to send again.
        let mut ping = nng::Message::new();
        ping.write_all(PING_BYTES).expect("All bytes written.");

        req.send(ping, None)
            .await
            .expect("Request should send successfully");

        let request = rep
            .receive(None)
            .await
            .expect("Request should be received successfully.");
        assert_eq!(PING_BYTES, request.as_slice());
    }
}