feat: volatile emits

crates/engineioxide/Cargo.toml

@@ -28,7 +28,7 @@ http-body.workspace = true
 serde.workspace = true
 serde_json.workspace = true
 thiserror.workspace = true
-tokio = { workspace = true, features = ["rt", "time"] }
+tokio = { workspace = true, features = ["rt", "time", "macros"] }
 tokio-util.workspace = true
 tower-service.workspace = true
 tower-layer.workspace = true

crates/engineioxide/src/engine.rs

@@ -99,7 +99,7 @@ impl<H: EngineIoHandler> EngineIo<H> {
             // channel so that `Socket::closed` resolves and no more packets can be buffered.
             // When no transport holds the lock (e.g. polling between requests) this resolves
             // immediately.
-            socket.internal_rx.lock().await.rx.close();
+            socket.internal_rx.lock().await.buffered_rx.close();
 
             #[cfg(feature = "tracing")]
             tracing::debug!(sid = %socket.id, "session closed, notifying handler");

crates/engineioxide/src/socket.rs

@@ -57,7 +57,6 @@
 use std::{
     collections::VecDeque,
     fmt::Debug,
-    ops::{Deref, DerefMut},
     sync::{
         Arc,
         atomic::{AtomicBool, AtomicU8, Ordering},
@@ -74,6 +73,7 @@ use smallvec::{SmallVec, smallvec};
 use tokio::sync::{
     Mutex,
     mpsc::{self, Receiver, error::TrySendError},
+    watch,
 };
 
 pub use engineioxide_core::Sid;
@@ -156,24 +156,21 @@ impl Permit<'_> {
 
 #[derive(Debug)]
 pub(crate) struct InternalRx {
-    pub rx: Receiver<PacketBuf>,
-    pub peeked: Option<PacketBuf>,
+    pub buffered_rx: Receiver<PacketBuf>,
+    pub peeked_packet: Option<PacketBuf>,
+    pub volatile_rx: watch::Receiver<Option<PacketBuf>>,
 }
-impl InternalRx {
-    fn new(rx: Receiver<PacketBuf>) -> Self {
-        Self { rx, peeked: None }
-    }
-}
-impl Deref for InternalRx {
-    type Target = Receiver<PacketBuf>;
 
-    fn deref(&self) -> &Self::Target {
-        &self.rx
-    }
-}
-impl DerefMut for InternalRx {
-    fn deref_mut(&mut self) -> &mut Self::Target {
-        &mut self.rx
+impl InternalRx {
+    fn new(
+        buffered_rx: Receiver<PacketBuf>,
+        volatile_rx: watch::Receiver<Option<PacketBuf>>,
+    ) -> Self {
+        Self {
+            buffered_rx,
+            peeked_packet: None,
+            volatile_rx,
+        }
     }
 }
 
@@ -219,6 +216,11 @@ where
     /// Channel to send [PacketBuf] to the internal connection
     pub(crate) internal_tx: mpsc::Sender<PacketBuf>,
 
+    /// Channel to send volatile [PacketBuf]s that bypass the internal buffer.
+    /// Uses a [`watch`](tokio::sync::watch) channel so only the latest volatile
+    /// message is retained; subsequent volatile sends overwrite the previous one.
+    volatile_tx: watch::Sender<Option<PacketBuf>>,
+
     /// Internal channel to receive Pong [`Packets`](Packet) (v4 protocol) or Ping (v3 protocol) in the heartbeat job
     /// which is running in a separate task
     heartbeat_rx: Mutex<Receiver<()>>,
@@ -255,16 +257,19 @@ where
     ) -> Self {
         let (internal_tx, internal_rx) = mpsc::channel(config.max_buffer_size);
         let (heartbeat_tx, heartbeat_rx) = mpsc::channel(1);
+        let (volatile_tx, volatile_rx) = watch::channel(None);
 
         Self {
             id: Sid::new(),
             protocol,
             transport: AtomicU8::new(transport as u8),
             upgrading: AtomicBool::new(false),
 
-            internal_rx: Mutex::new(InternalRx::new(internal_rx)),
+            internal_rx: Mutex::new(InternalRx::new(internal_rx, volatile_rx)),
             internal_tx,
 
+            volatile_tx,
+
             heartbeat_rx: Mutex::new(heartbeat_rx),
             heartbeat_tx,
             cancellation_token: CancellationToken::new(),
@@ -457,6 +462,7 @@ where
 
     /// Immediately closes the socket and the underlying connection.
     /// The socket will be removed from the `Engine` and the [`Handler`](crate::handler::EngineIoHandler) will be notified.
+    #[cfg_attr(feature = "tracing", tracing::instrument(skip(self)))]
     pub fn close(&self, reason: DisconnectReason) {
         // Try to send a close packet is the connection is still operational.
         self.send(Packet::Close).ok();
@@ -493,6 +499,55 @@ where
             TrySendError::Closed(p) => TrySendError::Closed(p.into_binary()),
         })
     }
+
+    /// Try to send a volatile message bypassing the internal buffer channel.
+    /// Volatile messages may be dropped if the transport is not ready to
+    /// receive them.
+    ///
+    /// Because volatile messages bypass the main mpsc buffer queue, they may
+    /// arrive out of order relative to regular messages.
+    ///
+    /// Returns `true` if the message was queued for sending, `false` if it
+    /// was dropped (channel full or transport shutting down).
+    #[inline]
+    pub fn emit_volatile(&self, msg: impl Into<Str>) -> bool {
+        self.send_volatile(smallvec![Packet::Message(msg.into())])
+    }
+
+    /// Try to send a volatile binary message bypassing the internal buffer channel.
+    /// Volatile messages may be dropped if the transport is not ready to
+    /// receive them.
+    ///
+    /// Returns `true` if the message was queued for sending, `false` if it
+    /// was dropped.
+    #[inline]
+    pub fn emit_binary_volatile<B: Into<Bytes>>(&self, data: B) -> bool {
+        if self.protocol == ProtocolVersion::V3 {
+            self.send_volatile(smallvec![Packet::BinaryV3(data.into())])
+        } else {
+            self.send_volatile(smallvec![Packet::Binary(data.into())])
+        }
+    }
+
+    /// Try to send a volatile message with multiple adjacent binary payloads.
+    /// The message and all binary payloads are sent atomically as a single
+    /// volatile write.
+    ///
+    /// Returns `true` if the message was queued for sending, `false` if it
+    /// was dropped.
+    #[inline]
+    pub fn emit_many_volatile(&self, msg: Str, data: VecDeque<Bytes>) -> bool {
+        let mut packets = SmallVec::with_capacity(1 + data.len());
+        packets.push(Packet::Message(msg));
+        for bin in data {
+            packets.push(Packet::Binary(bin));
+        }
+        self.send_volatile(packets)
+    }
+
+    pub(crate) fn send_volatile(&self, packets: PacketBuf) -> bool {
+        self.volatile_tx.send(Some(packets)).is_ok()
+    }
 }
 
 impl<D: Default + Send + Sync + 'static> std::fmt::Debug for Socket<D> {
@@ -553,16 +608,19 @@ where
     ) -> (Arc<Socket<D>>, tokio::sync::mpsc::Receiver<Packet>) {
         let (internal_tx, internal_rx) = mpsc::channel(buffer_size);
         let (heartbeat_tx, heartbeat_rx) = mpsc::channel(1);
+        let (volatile_tx, volatile_rx) = watch::channel(None);
 
         let sock = Self {
             id: sid,
             protocol: ProtocolVersion::V4,
             transport: AtomicU8::new(TransportType::Websocket as u8),
             upgrading: AtomicBool::new(false),
 
-            internal_rx: Mutex::new(InternalRx::new(internal_rx)),
+            internal_rx: Mutex::new(InternalRx::new(internal_rx, volatile_rx)),
             internal_tx,
 
+            volatile_tx,
+
             heartbeat_rx: Mutex::new(heartbeat_rx),
             heartbeat_tx,
             cancellation_token: CancellationToken::new(),
@@ -579,7 +637,7 @@ where
         let sock_clone = sock.clone();
         tokio::spawn(async move {
             let mut internal_rx = sock_clone.internal_rx.try_lock().unwrap();
-            while let Some(packets) = internal_rx.recv().await {
+            while let Some(packets) = internal_rx.buffered_rx.recv().await {
                 for packet in packets {
                     tx.send(packet).await.unwrap();
                 }

crates/engineioxide/src/transport/polling.rs

@@ -13,7 +13,7 @@ use engineioxide_core::{Packet, ProtocolVersion, Sid, TransportType};
 
 use crate::{
     DisconnectReason, body::ResponseBody, engine::EngineIo, errors::Error,
-    handler::EngineIoHandler, transport::make_open_packet,
+    handler::EngineIoHandler, socket::InternalRx, transport::make_open_packet,
 };
 
 /// Create a response for http request
@@ -116,16 +116,24 @@ where
     tracing::debug!(%sid, %protocol, supports_binary = socket.supports_binary, "polling request");
 
     let max_payload = engine.config.max_payload;
+    let payload = {
+        let InternalRx {
+            buffered_rx,
+            volatile_rx,
+            peeked_packet,
+        } = &mut *rx;
 
-    // Prepend the packet peeked during the previous
-    // polling request so it is encoded first, ahead of the newly received packets.
-    let rx_stream = stream::iter(rx.peeked.take()).chain(rx_stream::ReceiverStream::new(&mut rx));
+        // Prepend the packet peeked during the previous
+        // polling request so it is encoded first, ahead of the newly received packets.
+        let rx_stream = stream::iter(peeked_packet.take())
+            .chain(rx_stream::EncoderStream::new(buffered_rx, volatile_rx));
 
-    // Stop waiting for the next packet as soon as the session is closed. The combinator is biased
-    // towards the encoder completion, so any buffered close packet is still flushed to the client.
-    let payload = payload::encoder(rx_stream, protocol, socket.supports_binary, max_payload)
-        .with_cancellation_token(&socket.cancellation_token)
-        .await;
+        // Stop waiting for the next packet as soon as the session is closed. The combinator is biased
+        // towards the encoder completion, so any buffered close packet is still flushed to the client.
+        payload::encoder(rx_stream, protocol, socket.supports_binary, max_payload)
+            .with_cancellation_token(&socket.cancellation_token)
+            .await
+    };
 
     let Some(Payload {
         data,
@@ -140,7 +148,7 @@ where
 
     // set back the peeked packet so it can be read again
     // on the next polling request
-    rx.peeked = peeked;
+    rx.peeked_packet = peeked;
 
     #[cfg(feature = "tracing")]
     tracing::trace!(%sid, %protocol, supports_binary = socket.supports_binary, "sending data: {:?}", data);
@@ -214,20 +222,22 @@ where
 mod rx_stream {
     use std::{
         pin::Pin,
-        task::{Context, Poll},
+        task::{Context, Poll, ready},
     };
 
     use futures_core::Stream;
-    use tokio::sync::mpsc::Receiver;
+    use pin_project_lite::pin_project;
+    use tokio::sync::{mpsc, watch};
+    use tokio_util::sync::ReusableBoxFuture;
 
-    /// [`ReceiverStream`] is a stream that wraps a tokio::sync::mpsc::Receiver by reference.
+    /// [`ReceiverStream`] is a stream that wraps a [`mpsc::Receiver`] by reference.
     /// Allowing to use it as a stream even if it is behind a mutex.
     pub struct ReceiverStream<'a, T> {
-        inner: &'a mut Receiver<T>,
+        inner: &'a mut mpsc::Receiver<T>,
     }
 
     impl<'a, T> ReceiverStream<'a, T> {
-        pub fn new(inner: &'a mut Receiver<T>) -> Self {
+        pub fn new(inner: &'a mut mpsc::Receiver<T>) -> Self {
             Self { inner }
         }
     }
@@ -239,4 +249,91 @@ mod rx_stream {
             self.inner.poll_recv(cx)
         }
     }
+
+    type WatchFutOutput<'a, T> = (
+        Result<(), watch::error::RecvError>,
+        &'a mut watch::Receiver<T>,
+    );
+
+    /// Wraps a [`watch::Receiver`] into a [`Stream`] by reference.
+    /// Allowing to use it as a stream even if it is behind a mutex.
+    ///
+    /// Inspired by <https://docs.rs/tokio-stream/latest/tokio_stream/wrappers/struct.WatchStream.html>
+    pub struct WatchStream<'a, T> {
+        inner: ReusableBoxFuture<'a, WatchFutOutput<'a, T>>,
+    }
+    impl<'a, T: Send + Sync + 'static> WatchStream<'a, T> {
+        pub fn new(rx: &'a mut watch::Receiver<T>) -> Self {
+            Self {
+                inner: ReusableBoxFuture::new(make_future(rx)),
+            }
+        }
+    }
+
+    async fn make_future<T>(
+        rx: &mut watch::Receiver<T>,
+    ) -> (Result<(), watch::error::RecvError>, &mut watch::Receiver<T>) {
+        let result = rx.changed().await;
+        (result, rx)
+    }
+
+    impl<'a, T: Clone + Send + Sync + 'static> Stream for WatchStream<'a, T> {
+        type Item = T;
+
+        fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+            let (result, rx) = ready!(self.inner.poll(cx));
+            match result {
+                Ok(_) => {
+                    let received = (*rx.borrow_and_update()).clone();
+                    self.inner.set(make_future(rx));
+                    Poll::Ready(Some(received))
+                }
+                Err(_) => {
+                    self.inner.set(make_future(rx));
+                    Poll::Ready(None)
+                }
+            }
+        }
+    }
+
+    impl<T> Unpin for WatchStream<'_, T> {}
+
+    pin_project! {
+        /// Combines a [`WatchStream`] and a [`ReceiverStream`] into a single [`Stream`].
+        pub struct EncoderStream<'a, T> {
+            #[pin]
+            watch: WatchStream<'a, Option<T>>,
+            #[pin]
+            rx: ReceiverStream<'a, T>,
+        }
+    }
+
+    impl<'a, T: Clone + Send + Sync + 'static> EncoderStream<'a, T> {
+        pub fn new(
+            rx: &'a mut mpsc::Receiver<T>,
+            watch: &'a mut watch::Receiver<Option<T>>,
+        ) -> Self {
+            Self {
+                rx: ReceiverStream::new(rx),
+                watch: WatchStream::new(watch),
+            }
+        }
+    }
+
+    impl<'a, T: Clone + Send + Sync + 'static> Stream for EncoderStream<'a, T> {
+        type Item = T;
+
+        fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+            let this = self.project();
+            // low priority: poll the watch stream first, then the rx stream.
+            // same than with ws transport
+            match this.rx.poll_next(cx) {
+                Poll::Ready(v) => Poll::Ready(v),
+                Poll::Pending => match this.watch.poll_next(cx) {
+                    Poll::Ready(Some(v)) => Poll::Ready(v),
+                    Poll::Ready(None) | Poll::Pending => Poll::Pending,
+                },
+            }
+        }
+    }
 }