init mpsc

Author: lherman-cs

pulsebeam-runtime/src/sync/mod.rs

@@ -1,3 +1,4 @@
+pub mod mpsc;
 pub mod spmc;
 
 #[cfg(not(feature = "loom"))]

pulsebeam-runtime/src/sync/mpsc.rs

@@ -0,0 +1,293 @@
+use event_listener::{Event, EventListener};
+use futures_lite::Stream;
+use parking_lot::Mutex;
+use std::pin::Pin;
+use std::sync::Arc;
+use std::sync::atomic::{AtomicU64, Ordering};
+use std::task::{Context, Poll, ready};
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum RecvError {
+    Lagged(u64),
+    Closed,
+}
+
+#[derive(Debug, Clone, Copy, PartialEq, Eq)]
+pub enum StreamRecvError {
+    Lagged(u64),
+}
+
+#[derive(Debug)]
+struct Slot<T> {
+    seq: u64,
+    val: Option<T>,
+}
+
+#[derive(Debug)]
+struct Ring<T> {
+    slots: Vec<Mutex<Slot<T>>>, // per-slot mutex
+    mask: usize,
+    head: AtomicU64, // next free sequence number
+    event: Event,
+    closed: AtomicU64,
+}
+
+impl<T> Ring<T> {
+    fn new(mut capacity: usize) -> Arc<Self> {
+        if capacity == 0 {
+            capacity = 1;
+        } else if !capacity.is_power_of_two() {
+            capacity = capacity.next_power_of_two();
+        }
+
+        let mut slots = Vec::with_capacity(capacity);
+        for _ in 0..capacity {
+            slots.push(Mutex::new(Slot { seq: 0, val: None }));
+        }
+
+        Arc::new(Self {
+            slots,
+            mask: capacity - 1,
+            head: AtomicU64::new(0),
+            event: Event::new(),
+            closed: AtomicU64::new(0),
+        })
+    }
+}
+
+#[derive(Debug)]
+pub enum TrySendError<T> {
+    Closed(T),
+}
+
+#[derive(Debug, Clone)]
+pub struct Sender<T> {
+    ring: Arc<Ring<T>>,
+}
+
+impl<T> Sender<T> {
+    pub fn try_send(&self, val: T) -> Result<(), TrySendError<T>> {
+        if self.ring.closed.load(Ordering::Relaxed) == 1 {
+            return Err(TrySendError::Closed(val));
+        }
+
+        // Atomically claim slot
+        let seq = self.ring.head.fetch_add(1, Ordering::AcqRel);
+        let idx = (seq as usize) & self.ring.mask;
+
+        let mut slot = self.ring.slots[idx].lock();
+        slot.val = Some(val);
+        slot.seq = seq;
+
+        // there's only 1 consumer
+        self.ring.event.notify(1);
+        Ok(())
+    }
+}
+
+impl<T> Drop for Sender<T> {
+    fn drop(&mut self) {
+        // TODO: there's no notification for receiver when no senders left.
+        // self.ring.closed.store(1, Ordering::Release);
+        // self.ring.event.notify(usize::MAX);
+    }
+}
+
+#[derive(Debug)]
+pub struct Receiver<T> {
+    ring: Arc<Ring<T>>,
+    next_seq: u64,
+    local_head: u64,
+    listener: Option<EventListener>,
+}
+
+impl<T> Drop for Receiver<T> {
+    fn drop(&mut self) {
+        self.ring.closed.store(1, Ordering::Release);
+    }
+}
+
+impl<T: Clone> Receiver<T> {
+    pub async fn recv(&mut self) -> Result<T, RecvError> {
+        std::future::poll_fn(|cx| self.poll_recv(cx)).await
+    }
+
+    pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Result<T, RecvError>> {
+        loop {
+            let coop = std::task::ready!(tokio::task::coop::poll_proceed(cx));
+
+            // Refresh local snapshot of head
+            if self.next_seq == self.local_head {
+                self.local_head = self.ring.head.load(Ordering::Acquire);
+            }
+
+            let capacity = (self.ring.mask as u64) + 1;
+            let earliest = self.local_head.saturating_sub(capacity);
+
+            // Closed + nothing left
+            if self.ring.closed.load(Ordering::Acquire) == 1 && self.next_seq >= self.local_head {
+                return Poll::Ready(Err(RecvError::Closed));
+            }
+
+            // No new items
+            if self.next_seq >= self.local_head {
+                match &mut self.listener {
+                    Some(l) => {
+                        if Pin::new(l).poll(cx).is_pending() {
+                            return Poll::Pending;
+                        }
+                        self.listener = None;
+                        continue;
+                    }
+                    None => {
+                        self.listener = Some(self.ring.event.listen());
+                        continue;
+                    }
+                }
+            }
+
+            let idx = (self.next_seq as usize) & self.ring.mask;
+            let slot = self.ring.slots[idx].lock();
+            let slot_seq = slot.seq;
+
+            if slot_seq < earliest {
+                self.next_seq = self.local_head;
+                return Poll::Ready(Err(RecvError::Lagged(self.local_head)));
+            }
+
+            if slot_seq != self.next_seq {
+                self.next_seq = self.local_head;
+                return Poll::Ready(Err(RecvError::Lagged(self.local_head)));
+            }
+
+            if let Some(ref v) = slot.val {
+                coop.made_progress();
+                let out = v.clone();
+                self.next_seq += 1;
+                return Poll::Ready(Ok(out));
+            }
+
+            self.next_seq = self.local_head;
+            return Poll::Ready(Err(RecvError::Lagged(self.local_head)));
+        }
+    }
+}
+
+impl<T: Clone> Stream for Receiver<T> {
+    type Item = Result<T, StreamRecvError>;
+    fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        let this = self.get_mut();
+        let res = match ready!(this.poll_recv(cx)) {
+            Ok(item) => Some(Ok(item)),
+            Err(RecvError::Lagged(n)) => Some(Err(StreamRecvError::Lagged(n))),
+            Err(RecvError::Closed) => None,
+        };
+        Poll::Ready(res)
+    }
+}
+
+impl<T: Clone> Clone for Receiver<T> {
+    fn clone(&self) -> Self {
+        Self {
+            ring: self.ring.clone(),
+            next_seq: self.next_seq,
+            local_head: self.local_head,
+            listener: None,
+        }
+    }
+}
+
+pub fn channel<T: Send + Sync + Clone + 'static>(capacity: usize) -> (Sender<T>, Receiver<T>) {
+    let ring = Ring::new(capacity);
+    (
+        Sender { ring: ring.clone() },
+        Receiver {
+            ring,
+            next_seq: 0,
+            local_head: 0,
+            listener: None,
+        },
+    )
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio::time::Duration;
+
+    #[tokio::test]
+    async fn basic_send_recv() {
+        let (tx, mut rx) = channel::<u64>(8);
+
+        tx.try_send(42).unwrap();
+        assert_eq!(rx.recv().await, Ok(42));
+
+        tx.try_send(123).unwrap();
+        assert_eq!(rx.recv().await, Ok(123));
+    }
+
+    #[tokio::test]
+    async fn lagging_consumer_jumps_to_head() {
+        let (tx, mut rx) = channel::<u64>(4);
+
+        for i in 0..6 {
+            tx.try_send(i).unwrap();
+        }
+
+        match rx.recv().await {
+            Err(RecvError::Lagged(seq)) => assert_eq!(seq, 6),
+            _ => panic!("Expected lag error"),
+        }
+
+        tx.try_send(6).unwrap();
+        assert_eq!(rx.recv().await, Ok(6));
+    }
+
+    #[tokio::test]
+    async fn close_signal_drains_then_stops() {
+        let (tx, mut rx) = channel::<u64>(4);
+
+        tx.try_send(1).unwrap();
+        tx.try_send(2).unwrap();
+
+        drop(tx);
+
+        assert_eq!(rx.recv().await, Ok(1));
+        assert_eq!(rx.recv().await, Ok(2));
+        assert_eq!(rx.recv().await, Err(RecvError::Closed));
+        assert_eq!(rx.recv().await, Err(RecvError::Closed));
+    }
+
+    #[tokio::test]
+    async fn async_waker_notification() {
+        let (tx, mut rx) = channel::<u64>(4);
+
+        let handle = tokio::spawn(async move { rx.recv().await });
+
+        tokio::time::sleep(Duration::from_millis(10)).await;
+        tx.try_send(99).unwrap();
+
+        let result = handle.await.unwrap();
+        assert_eq!(result, Ok(99));
+    }
+
+    #[tokio::test]
+    async fn buffer_wrap_around_behavior() {
+        let (tx, mut rx) = channel::<u64>(4);
+
+        for i in 0..4 {
+            tx.try_send(i).unwrap();
+        }
+
+        assert_eq!(rx.recv().await, Ok(0));
+        assert_eq!(rx.recv().await, Ok(1));
+
+        tx.try_send(4).unwrap();
+        tx.try_send(5).unwrap();
+
+        assert_eq!(rx.recv().await, Ok(2));
+        assert_eq!(rx.recv().await, Ok(3));
+        assert_eq!(rx.recv().await, Ok(4));
+        assert_eq!(rx.recv().await, Ok(5));
+    }
+}

pulsebeam-runtime/src/sync/spmc.rs

@@ -168,30 +168,26 @@ impl<T: Clone> Receiver<T> {
 
             // Slot overwritten before we got here
             if slot_seq < earliest {
-                drop(slot);
                 self.next_seq = self.local_head;
                 return Poll::Ready(Err(RecvError::Lagged(self.local_head)));
             }
 
             // Seq mismatch — producer overwrote after head snapshot
             if slot_seq != self.next_seq {
-                drop(slot);
                 self.next_seq = self.local_head;
                 return Poll::Ready(Err(RecvError::Lagged(self.local_head)));
             }
 
             // Valid message
             if let Some(v) = &slot.val {
                 let out = v.clone();
-                drop(slot);
                 coop.made_progress();
                 self.next_seq += 1;
                 return Poll::Ready(Ok(out));
             }
 
             // This shouldn't never happen, but just in case..
             // Seq was correct but value missing — treat as lag
-            drop(slot);
             self.next_seq = self.local_head;
             return Poll::Ready(Err(RecvError::Lagged(self.local_head)));
         }

pulsebeam/src/gateway/actor.rs

@@ -1,13 +1,16 @@
 use crate::gateway::demux::Demuxer;
 use pulsebeam_runtime::actor::{ActorKind, ActorStatus, RunnerConfig};
 use pulsebeam_runtime::prelude::*;
-use pulsebeam_runtime::{actor, mailbox, net};
+use pulsebeam_runtime::{actor, net};
 use std::{io, sync::Arc};
 use tokio::task::JoinSet;
 
 #[derive(Clone)]
 pub enum GatewayControlMessage {
-    AddParticipant(String, mailbox::Sender<net::RecvPacketBatch>),
+    AddParticipant(
+        String,
+        pulsebeam_runtime::sync::mpsc::Sender<net::RecvPacketBatch>,
+    ),
     RemoveParticipant(String),
 }
 
@@ -171,7 +174,7 @@ impl GatewayWorkerActor {
                         };
 
                         let src = batch.src;
-                        if !self.demuxer.demux(&mut self.socket, batch).await {
+                        if !self.demuxer.demux(&mut self.socket, batch) {
                             // In case there's a malicious actor, close immediately as there's no
                             // associated participant.
                             self.socket.close_peer(&src);

pulsebeam/src/gateway/demux.rs

@@ -1,11 +1,11 @@
+use pulsebeam_runtime::net;
 use pulsebeam_runtime::net::UnifiedSocketReader;
-use pulsebeam_runtime::{mailbox, net};
 
 use crate::gateway::ice;
 use std::collections::HashMap;
 use std::net::SocketAddr;
 
-pub type ParticipantHandle = mailbox::Sender<net::RecvPacketBatch>;
+pub type ParticipantHandle = pulsebeam_runtime::sync::mpsc::Sender<net::RecvPacketBatch>;
 
 /// A UDP demuxer that maps packets to participants based on source address and STUN ufrag.
 ///
@@ -57,11 +57,7 @@ impl Demuxer {
 
     /// Routes a packet to the correct participant.
     /// Returns `true` if sent, `false` if dropped
-    pub async fn demux(
-        &mut self,
-        socket: &mut UnifiedSocketReader,
-        batch: net::RecvPacketBatch,
-    ) -> bool {
+    pub fn demux(&mut self, socket: &mut UnifiedSocketReader, batch: net::RecvPacketBatch) -> bool {
         let src = batch.src;
 
         let handle = if let Some(h) = self.addr_map.get_mut(&src) {
@@ -83,7 +79,7 @@ impl Demuxer {
             return false;
         };
 
-        if let Err(_) = handle.send(batch).await {
+        if let Err(_) = handle.try_send(batch) {
             // Handle is closed! Clean up everything related to this participant.
             if let Some(ufrag) = self.addr_to_ufrag.get(&src).cloned() {
                 tracing::info!("Participant handle closed, cleaning up ufrag: {:?}", ufrag);