use futures-channel for mailbox

Author: lherman-cs

Makefile

@@ -1,37 +1,38 @@
 TARGET_DIR = target/profiling
 BINARY = $(TARGET_DIR)/pulsebeam
 
-# Rust and perf settings
+# Rust settings
 CARGO = cargo
 RUSTFLAGS = -C force-frame-pointers=yes
-PERF = perf
-PERF_RECORD_FLAGS = -F 4999 --call-graph dwarf
 
-# Default: Build and capture perf data
+# Default: Build and capture flamegraph
 .PHONY: all
-all: build profile
+all: build flamegraph
 
 # Build with profiling profile
 .PHONY: build
 build:
 	$(CARGO) build --profile profiling 
 
-# Capture perf data
-.PHONY: profile
-profile: $(BINARY)
-	sudo $(PERF) record $(PERF_RECORD_FLAGS) -- $(BINARY)
+# Capture flamegraph
+.PHONY: flamegraph
+flamegraph: build
+	taskset -c 2-5 cargo flamegraph --profile profiling -p pulsebeam --bin pulsebeam
 
 brew-deps:
 	brew install git-cliff axodotdev/tap/cargo-dist
 
 cargo-deps:
 	cargo install cargo-smart-release --features allow-emoji
+	cargo install flamegraph
 
 release:
 	cargo smart-release --execute
 
-# Clean build artifacts and perf data
+# Clean build artifacts and flamegraph data
 .PHONY: clean
 clean:
 	$(CARGO) clean
 	rm -f perf.data
+	rm -f flamegraph.svg
+

pulsebeam-runtime/Cargo.toml

@@ -19,10 +19,10 @@ tokio = { workspace = true }
 tracing = { workspace = true }
 rand = { workspace = true }
 bytes = { workspace = true }
+futures.workspace = true
 tokio-rustls = "0.26.2"
 rustls = "0.23.31"
 turmoil = "0.6.6"
-futures.workspace = true
 quinn-udp = "0.5.14"
 socket2 = "0.6.0"
 

pulsebeam-runtime/src/actor.rs

@@ -214,9 +214,9 @@ impl<A: Actor> ActorHandle<A> {
     /// Returns an error if the actor's mailbox is closed.
     #[inline]
     pub async fn send_high(
-        &self,
+        &mut self,
         message: A::HighPriorityMsg,
-    ) -> Result<(), mailbox::SendError<A::HighPriorityMsg>> {
+    ) -> Result<(), mailbox::SendError> {
         self.hi_tx.send(message).await
     }
 
@@ -225,7 +225,7 @@ impl<A: Actor> ActorHandle<A> {
     /// Returns an error if the mailbox is full or closed.
     #[inline]
     pub fn try_send_high(
-        &self,
+        &mut self,
         message: A::HighPriorityMsg,
     ) -> Result<(), mailbox::TrySendError<A::HighPriorityMsg>> {
         self.hi_tx.try_send(message)
@@ -235,10 +235,7 @@ impl<A: Actor> ActorHandle<A> {
     ///
     /// Returns an error if the actor's mailbox is closed.
     #[inline]
-    pub async fn send_low(
-        &self,
-        message: A::LowPriorityMsg,
-    ) -> Result<(), mailbox::SendError<A::LowPriorityMsg>> {
+    pub async fn send_low(&mut self, message: A::LowPriorityMsg) -> Result<(), mailbox::SendError> {
         self.lo_tx.send(message).await
     }
 
@@ -247,13 +244,13 @@ impl<A: Actor> ActorHandle<A> {
     /// Returns an error if the mailbox is full or closed.
     #[inline]
     pub fn try_send_low(
-        &self,
+        &mut self,
         message: A::LowPriorityMsg,
     ) -> Result<(), mailbox::TrySendError<A::LowPriorityMsg>> {
         self.lo_tx.try_send(message)
     }
 
-    pub async fn get_state(&self) -> Result<A::ObservableState, ActorError> {
+    pub async fn get_state(&mut self) -> Result<A::ObservableState, ActorError> {
         let (tx, rx) = tokio::sync::oneshot::channel();
         let msg = SystemMsg::GetState(tx);
 
@@ -267,7 +264,7 @@ impl<A: Actor> ActorHandle<A> {
         rx.await.map_err(|_| ActorError::SystemError)
     }
 
-    pub async fn terminate(&self) -> Result<(), ActorError> {
+    pub async fn terminate(&mut self) -> Result<(), ActorError> {
         self.sys_tx
             .send(SystemMsg::Terminate)
             .await

pulsebeam-runtime/src/mailbox.rs

@@ -1,22 +1,23 @@
+use futures::channel::mpsc;
+use futures::{SinkExt, StreamExt};
 use std::error::Error;
 use std::fmt;
-use tokio::sync::mpsc;
 
 /// An error returned when sending on a closed mailbox.
 ///
 /// This error is returned by the asynchronous `send` method. It contains
 /// the message that could not be sent.
 #[derive(Debug, PartialEq, Eq)]
-pub struct SendError<T>(pub T);
+pub enum SendError {
+    Closed,
+}
 
-impl<T> fmt::Display for SendError<T> {
+impl fmt::Display for SendError {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         write!(f, "sending on a closed mailbox")
     }
 }
 
-impl<T: fmt::Debug> Error for SendError<T> {}
-
 /// An error returned from the `try_send` method on a `Sender`.
 #[derive(Debug, PartialEq, Eq)]
 pub enum TrySendError<T> {
@@ -37,14 +38,6 @@ impl<T> fmt::Display for TrySendError<T> {
 
 impl<T: fmt::Debug> Error for TrySendError<T> {}
 
-// For convenience, we can convert a SendError into a TrySendError,
-// as a 'Closed' state is a subset of 'TrySend' failures.
-impl<T> From<SendError<T>> for TrySendError<T> {
-    fn from(err: SendError<T>) -> Self {
-        TrySendError::Closed(err.0)
-    }
-}
-
 /// A handle to send messages to an actor's mailbox.
 pub struct Sender<T> {
     sender: mpsc::Sender<T>,
@@ -62,22 +55,25 @@ impl<T> Sender<T> {
     /// Sends a message asynchronously, waiting if the mailbox is full.
     ///
     /// Returns a `SendError` only if the receiving actor has terminated.
-    pub async fn send(&self, message: T) -> Result<(), SendError<T>> {
-        self.sender.send(message).await.map_err(|e| SendError(e.0))
+    pub async fn send(&mut self, message: T) -> Result<(), SendError> {
+        self.sender
+            .send(message)
+            .await
+            .map_err(|_e| SendError::Closed)
     }
 
     /// Attempts to immediately send a message.
     ///
     /// Returns a `TrySendError` if the mailbox is full or if the
     /// receiving actor has terminated.
-    pub fn try_send(&self, message: T) -> Result<(), TrySendError<T>> {
-        use mpsc::error::TrySendError as TokioTrySendError;
-        self.sender.try_send(message).map_err(|e| match e {
-            TokioTrySendError::Full(m) => {
+    pub fn try_send(&mut self, message: T) -> Result<(), TrySendError<T>> {
+        self.sender.try_send(message).map_err(|e| {
+            if e.is_disconnected() {
+                TrySendError::Closed(e.into_inner())
+            } else {
                 tracing::info!("try_send dropped a packet due to full queue");
-                TrySendError::Full(m)
+                TrySendError::Full(e.into_inner())
             }
-            TokioTrySendError::Closed(m) => TrySendError::Closed(m),
         })
     }
 }
@@ -90,7 +86,7 @@ pub struct Receiver<T> {
 impl<T> Receiver<T> {
     /// Receives the next message from the mailbox.
     pub async fn recv(&mut self) -> Option<T> {
-        self.receiver.recv().await
+        self.receiver.next().await
     }
 }
 
@@ -107,7 +103,7 @@ mod tests {
 
     #[tokio::test]
     async fn send_and_recv_single_message() {
-        let (sender, mut mailbox) = mailbox::new(10);
+        let (mut sender, mut mailbox) = mailbox::new(10);
         let message = "hello".to_string();
 
         sender.send(message.clone()).await.unwrap();
@@ -118,7 +114,7 @@ mod tests {
 
     #[tokio::test]
     async fn recv_returns_none_when_all_senders_are_dropped() {
-        let (sender, mut mailbox) = mailbox::new::<i32>(10);
+        let (mut sender, mut mailbox) = mailbox::new::<i32>(10);
         sender.send(1).await.unwrap();
         sender.send(2).await.unwrap();
 
@@ -135,7 +131,7 @@ mod tests {
 
     #[tokio::test]
     async fn async_send_fails_when_receiver_is_dropped() {
-        let (sender, mailbox) = mailbox::new::<String>(10);
+        let (mut sender, mailbox) = mailbox::new::<String>(10);
 
         // Drop the mailbox immediately
         drop(mailbox);
@@ -144,27 +140,27 @@ mod tests {
         assert!(result.is_err());
 
         // Check that the error is our custom SendError and we can get the message back
-        if let Err(mailbox::SendError(msg)) = result {
-            assert_eq!(msg, "should fail");
+        if let Err(mailbox::SendError::Closed) = result {
         } else {
             panic!("Expected a SendError");
         }
     }
 
     #[test]
     fn try_send_success_on_capacity() {
-        let (sender, _mailbox) = mailbox::new::<i32>(1);
+        let (mut sender, _mailbox) = mailbox::new::<i32>(1);
         let result = sender.try_send(123);
         assert!(result.is_ok());
     }
 
     #[test]
     fn try_send_fails_when_full() {
         // Create a mailbox with a buffer of 1
-        let (sender, _mailbox) = mailbox::new::<i32>(1);
+        let (mut sender, _mailbox) = mailbox::new::<i32>(1);
 
         // The first send succeeds and fills the buffer
         sender.try_send(1).unwrap();
+        sender.try_send(1).unwrap();
 
         // The second send should fail because the buffer is full
         let result = sender.try_send(2);
@@ -179,7 +175,7 @@ mod tests {
 
     #[test]
     fn try_send_fails_when_receiver_is_dropped() {
-        let (sender, mailbox) = mailbox::new::<i32>(1);
+        let (mut sender, mailbox) = mailbox::new::<i32>(1);
 
         // Drop the receiver
         drop(mailbox);
@@ -196,14 +192,14 @@ mod tests {
 
     #[tokio::test]
     async fn async_send_waits_when_full() {
-        let (sender, mut mailbox) = mailbox::new::<i32>(1);
+        let (mut sender, mut mailbox) = mailbox::new::<i32>(1);
 
         // Fill the buffer
         sender.send(1).await.unwrap();
 
         // This send should wait. We spawn it in a separate task.
         let send_task = tokio::spawn({
-            let sender = sender.clone();
+            let mut sender = sender.clone();
             async move {
                 sender.send(2).await.unwrap();
             }
@@ -227,8 +223,8 @@ mod tests {
 
     #[tokio::test]
     async fn cloned_sender_works_and_channel_stays_open() {
-        let (sender1, mut mailbox) = mailbox::new::<i32>(10);
-        let sender2 = sender1.clone();
+        let (mut sender1, mut mailbox) = mailbox::new::<i32>(10);
+        let mut sender2 = sender1.clone();
 
         sender1.send(1).await.unwrap();
         sender2.send(2).await.unwrap();

pulsebeam/src/controller.rs

@@ -83,7 +83,7 @@ impl actor::Actor for ControllerActor {
             }
 
             ControllerMessage::RemoveParticipant(room_id, participant_id) => {
-                if let Some(room_handle) = self.rooms.get(&room_id) {
+                if let Some(room_handle) = self.rooms.get_mut(&room_id) {
                     // if the room has exited, the participants have already cleaned up too.
                     let _ = room_handle
                         .send_high(room::RoomMessage::RemoveParticipant(participant_id))
@@ -141,7 +141,7 @@ impl ControllerActor {
             .accept_offer(offer)
             .map_err(ControllerError::OfferRejected)?;
 
-        let room_handle = self.get_or_create_room(room_id);
+        let mut room_handle = self.get_or_create_room(room_id);
 
         // TODO: probably retry? Or, let the client to retry instead?
         // Each room will always have a graceful timeout before closing.

pulsebeam/src/gateway.rs

@@ -127,7 +127,8 @@ impl GatewayActor {
 
         tracing::trace!("received {count} packets from socket");
         for packet in self.recv_batch.drain(..) {
-            let participant_handle = if let Some(participant_handle) = self.mapping.get(&packet.src)
+            let mut participant_handle = if let Some(participant_handle) =
+                self.mapping.get(&packet.src)
             {
                 tracing::trace!(
                     "found connection from mapping: {} -> {:?}",

pulsebeam/src/participant.rs

@@ -487,12 +487,12 @@ impl ParticipantActor {
                 }
             }
             Event::MediaData(e) => {
-                if let Some(track) = self.published_video_tracks.get(&e.mid) {
+                if let Some(track) = self.published_video_tracks.get_mut(&e.mid) {
                     tracing::debug!("forwarded media: participant -> track, {track:?}");
                     let _ = track
                         .send_low(track::TrackDataMessage::ForwardMedia(Arc::new(e)))
                         .await;
-                } else if let Some(track) = self.published_audio_tracks.get(&e.mid) {
+                } else if let Some(track) = self.published_audio_tracks.get_mut(&e.mid) {
                     let _ = track
                         .send_low(track::TrackDataMessage::ForwardMedia(Arc::new(e)))
                         .await;
@@ -515,7 +515,7 @@ impl ParticipantActor {
             return;
         };
 
-        let Some(track) = self.available_video_tracks.get(&track_id.internal) else {
+        let Some(track) = self.available_video_tracks.get_mut(&track_id.internal) else {
             return;
         };