async_nng/

context.rs

//! Types and traits used for performing asynchronous socket operations that utilize contexts in
//! conjunction with 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 [`Context`](nng::Context) to enable async-await send and receive
/// operations.
///
/// This type allows for asynchronous context-based send and receive operations. One trouble with
/// `AsyncSocket` and `Socket`s in general is that managing concurrent sends and receives can be
/// troublesome depending on which scalability protocol your socket implements. For example,
/// consider the following program:
///
/// ```ignore
/// use async_nng::AsyncSocket;
/// # use std::{io::Write, time::Duration};
///
/// struct MyMessageHandler {
///     // Assume a Rep0 socket
///     socket: AsyncSocket,
/// }
///
/// impl MyMessageHandler {
///     async fn some_api_call(&self, msg: nng::Message) -> Result<(), nng::Error> {
///         self.socket.send(msg, None).await;
///
///         // a really long operation
///         smol::Timer::after(Duration::from_millis(3000)).await;
///
///         let reply = self.socket.receive(None).await;
///         Ok(())
///     }
/// }
/// ```
///
/// We have some message handling type, which may or may not break at one of the await points
/// above. If it breaks ono the middle await point (our really long operation), then we might have
/// an issue depending on the protocol being used. Above, we're using req-rep, which has the
/// restriction that generally (unless you're using sockets in raw mode) that you cannot send two
/// different requests before the reply is received.
///
/// Therefore, if one tries to run multiple message handlers or API calls from different tasks or
/// threads concurrently they are at risk to run into an issue where the req-rep cycle invariants
/// are violated because of long-awaiting tasks in between the send and receive.
///
/// This is exclusive because `AsyncSocket::send` and `AsyncSocket::receive` take `&mut self`
/// instead of `&self`. The underlying `nng` crate does not actually require this, but is done to
/// prevent multiple accesses of the same underlying AIO object in Rust. This still runs into an
/// issue. If users still want to concurrently construct API calls such as the above, they will
/// then choose to wrap the async socket inside of something like an `Arc<Mutex<_>>`. If that mutex
/// is asynchronous and releases across await points, then you still have the issue that concurrent
/// calls to these APIs can fail because e.g. two sends were called before a receive was.
///
/// ## Contexts
///
/// The way around this is to use contexts. Each context can be constructed locally on the stack,
/// and then indepedently manage sending and receiving in a safe way, without requiring a lock on
/// the socket as a whole:
///
/// ```no_run
/// use async_nng::AsyncContext;
/// use nng::Socket;
/// # use std::{io::Write, time::Duration};
///
/// struct MyMessageHandler {
///     // Assume a Rep0 socket
///     socket: Socket,
/// }
///
/// impl MyMessageHandler {
///     async fn some_api_call(&self, msg: nng::Message) -> Result<(), nng::Error> {
///         // Notice that we use a regular borrow on `&self.socket`, not `&mut self.socket`.
///         //
///         // Contexts defined locally will not conflict with one another, which makes writing
///         // asynchronous, concurrent programs easier.
///         let context = AsyncContext::try_from(&self.socket)?;
///
///         context.send(msg, None).await;
///
///         // a really long operation
///         smol::Timer::after(Duration::from_millis(3000)).await;
///
///         let reply = context.receive(None).await;
///         Ok(())
///     }
/// }
/// ```
///
/// Contexts cannot be used with raw-mode sockets; however, by borrowing the underlying socket they
/// are able to create an object that can concurrently and independently operate without the
/// aforementioned bugs that can arise at runtime due to trying to use e.g. a req-rep socket
/// concurrently.
///
/// In almost all cases, one should prefer to use the `AsyncContext` type when possible. It
/// provides better guarantees about concurrent access, which for asynchronous code is necessarily
/// a concern.
#[derive(Debug)]
pub struct AsyncContext<'a> {
    /// The underlying context to run async operations on.
    context: nng::Context,

    /// 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>,

    /// The socket that the context is bound to.
    ///
    /// This is pretty much entirely unused but is here to tie the lifetime of the context to the
    /// socket. This is particularly useful because `nng::Context` doesn't forward this lifetime,
    /// and as a result you can end up getting `nng::Error::Closed` on operations. By tying this
    /// lifetime here artificially (because the context is closed if the socket is closed, but not
    /// vice-versa), we can guarantee that contexts do not outlive their sockets using the borrow
    /// checker instead of fishing for `nng::Error::Closed`.
    ///
    /// See <https://doc.rust-lang.org/nomicon/phantom-data.html#table-of-phantomdata-patterns> for
    /// more info on variance of the phantom. The reference to the socket was intentionally not
    /// chosen to be held here because contexts do not generally appreciate if one touches the
    /// socket (e.g. call `nng::Socket::close`) while borrowed, and this is an easy way to enforce
    /// the lifetime bound artificially without providing a footgun that is a reference to the
    /// actual socket.
    _phantom: std::marker::PhantomData<&'a nng::Socket>,
}

impl<'a> TryFrom<&'a nng::Socket> for AsyncContext<'a> {
    type Error = nng::Error;

    #[instrument]
    fn try_from(socket: &'a nng::Socket) -> Result<Self, Self::Error> {
        let context = nng::Context::new(socket)?;
        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 {
            context,
            aio,
            receiver,
            _phantom: std::marker::PhantomData,
        })
    }
}

impl<'a> Cancel for AsyncContext<'a> {
    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<'a> AsyncContext<'a> {
    /// 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.
    /// - `Closed`: The context or socket has been closed and future operations will not work.
    #[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.context.send(&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 AsyncContext::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 receive 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.
    /// - `Closed`: The context or socket has been closed and future operations will not work.
    #[instrument]
    pub async fn receive(&mut self, timeout: Option<Duration>) -> Result<nng::Message, nng::Error> {
        self.aio.set_timeout(timeout)?;
        self.context.recv(&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 AsyncContext::receive.");
                unreachable!();
            }
        };
        std::mem::forget(guard);
        res
    }
}

#[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() -> (nng::Socket, nng::Socket) {
        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.")
        };

        (req, rep)
    }

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

        let mut req_ctx = AsyncContext::try_from(&req).expect("Make async context for req.");
        let mut rep_ctx = AsyncContext::try_from(&rep).expect("Make async context for rep.");

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

        let request = rep_ctx
            .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_ctx
            .send(pong, None)
            .await
            .expect("Reply should send successfully.");

        let response = req_ctx
            .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 (req, rep) = make_req_rep();

        let mut req_ctx = AsyncContext::try_from(&req).expect("Make async context for req.");
        let rep_ctx = AsyncContext::try_from(&rep).expect("Make async context for rep.");

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

        std::mem::drop(rep_ctx);

        let err = req_ctx
            .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 context_send_is_cancel_safe() {
        let (req, rep) = make_req_rep();

        let mut req_ctx = AsyncContext::try_from(&req).expect("Make async context for req.");
        let mut rep_ctx = AsyncContext::try_from(&rep).expect("Make async context for rep.");

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

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

        let err = rep_ctx
            .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_ctx
            .send(ping, None)
            .await
            .expect("Request should send successfully");

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