feat: use broadcast channel for event listeners (#8193)

Co-authored-by: Emilia Hane <elsaemiliaevahane@gmail.com>

crates/tokio-util/Cargo.toml

@@ -12,7 +12,11 @@ description = "Additional utilities for working with Tokio in reth."
 workspace = true
 
 [dependencies]
+tracing.workspace = true
 
 # async
 tokio = { workspace = true, features = ["sync"] }
 tokio-stream = { workspace = true, features = ["sync"] }
+
+[dev-dependencies]
+tokio = { workspace = true, features = ["full", "macros"] }

crates/tokio-util/src/event_listeners.rs

@@ -1,46 +0,0 @@
-use tokio::sync::mpsc;
-use tokio_stream::wrappers::UnboundedReceiverStream;
-
-/// A collection of event listeners for a task.
-#[derive(Clone, Debug)]
-pub struct EventListeners<T> {
-    /// All listeners for events
-    listeners: Vec<mpsc::UnboundedSender<T>>,
-}
-
-impl<T> Default for EventListeners<T> {
-    fn default() -> Self {
-        Self { listeners: Vec::new() }
-    }
-}
-
-impl<T: Clone> EventListeners<T> {
-    /// Send an event to all listeners.
-    ///
-    /// Channels that were closed are removed.
-    pub fn notify(&mut self, event: T) {
-        self.listeners.retain(|listener| listener.send(event.clone()).is_ok())
-    }
-
-    /// Add a new event listener.
-    pub fn new_listener(&mut self) -> UnboundedReceiverStream<T> {
-        let (sender, receiver) = mpsc::unbounded_channel();
-        self.listeners.push(sender);
-        UnboundedReceiverStream::new(receiver)
-    }
-
-    /// Push new event listener.
-    pub fn push_listener(&mut self, listener: mpsc::UnboundedSender<T>) {
-        self.listeners.push(listener);
-    }
-
-    /// Returns the number of registered listeners.
-    pub fn len(&self) -> usize {
-        self.listeners.len()
-    }
-
-    /// Returns true if there are no registered listeners.
-    pub fn is_empty(&self) -> bool {
-        self.listeners.is_empty()
-    }
-}

crates/tokio-util/src/event_sender.rs

@@ -0,0 +1,42 @@
+use crate::EventStream;
+use tokio::sync::broadcast::{self, Sender};
+use tracing::error;
+
+const DEFAULT_SIZE_BROADCAST_CHANNEL: usize = 2000;
+
+/// A bounded broadcast channel for a task.
+#[derive(Debug, Clone)]
+pub struct EventSender<T> {
+    /// The sender part of the broadcast channel
+    sender: Sender<T>,
+}
+
+impl<T> Default for EventSender<T>
+where
+    T: Clone + Send + Sync + 'static,
+{
+    fn default() -> Self {
+        Self::new(DEFAULT_SIZE_BROADCAST_CHANNEL)
+    }
+}
+
+impl<T: Clone + Send + Sync + 'static> EventSender<T> {
+    /// Creates a new `EventSender`.
+    pub fn new(events_channel_size: usize) -> Self {
+        let (sender, _) = broadcast::channel(events_channel_size);
+        Self { sender }
+    }
+
+    /// Broadcasts an event to all listeners.
+    pub fn notify(&self, event: T) {
+        if self.sender.send(event).is_err() {
+            error!("channel closed");
+        }
+    }
+
+    /// Creates a new event stream with a subscriber to the sender as the
+    /// receiver.
+    pub fn new_listener(&self) -> EventStream<T> {
+        EventStream::new(self.sender.subscribe())
+    }
+}

crates/tokio-util/src/event_stream.rs

@@ -0,0 +1,92 @@
+//! Event streams related functionality.
+
+use std::{
+    pin::Pin,
+    task::{Context, Poll},
+};
+use tokio_stream::Stream;
+use tracing::warn;
+
+/// Thin wrapper around tokio's BroadcastStream to allow skipping broadcast errors.
+#[derive(Debug)]
+pub struct EventStream<T> {
+    inner: tokio_stream::wrappers::BroadcastStream<T>,
+}
+
+impl<T> EventStream<T>
+where
+    T: Clone + Send + 'static,
+{
+    /// Creates a new `EventStream`.
+    pub fn new(receiver: tokio::sync::broadcast::Receiver<T>) -> Self {
+        let inner = tokio_stream::wrappers::BroadcastStream::new(receiver);
+        EventStream { inner }
+    }
+}
+
+impl<T> Stream for EventStream<T>
+where
+    T: Clone + Send + 'static,
+{
+    type Item = T;
+
+    fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
+        loop {
+            match Pin::new(&mut self.inner).poll_next(cx) {
+                Poll::Ready(Some(Ok(item))) => return Poll::Ready(Some(item)),
+                Poll::Ready(Some(Err(e))) => {
+                    warn!("BroadcastStream lagged: {e:?}");
+                    continue;
+                }
+                Poll::Ready(None) => return Poll::Ready(None),
+                Poll::Pending => return Poll::Pending,
+            }
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use tokio::sync::broadcast;
+    use tokio_stream::StreamExt;
+
+    #[tokio::test]
+    async fn test_event_stream_yields_items() {
+        let (tx, _) = broadcast::channel(16);
+        let my_stream = EventStream::new(tx.subscribe());
+
+        tx.send(1).unwrap();
+        tx.send(2).unwrap();
+        tx.send(3).unwrap();
+
+        // drop the sender to terminate the stream and allow collect to work.
+        drop(tx);
+
+        let items: Vec<i32> = my_stream.collect().await;
+
+        assert_eq!(items, vec![1, 2, 3]);
+    }
+
+    #[tokio::test]
+    async fn test_event_stream_skips_lag_errors() {
+        let (tx, _) = broadcast::channel(2);
+        let my_stream = EventStream::new(tx.subscribe());
+
+        let mut _rx2 = tx.subscribe();
+        let mut _rx3 = tx.subscribe();
+
+        tx.send(1).unwrap();
+        tx.send(2).unwrap();
+        tx.send(3).unwrap();
+        tx.send(4).unwrap(); // This will cause lag for the first subscriber
+
+        // drop the sender to terminate the stream and allow collect to work.
+        drop(tx);
+
+        // Ensure lag errors are skipped and only valid items are collected
+        let items: Vec<i32> = my_stream.collect().await;
+
+        assert_eq!(items, vec![3, 4]);
+    }
+}

crates/tokio-util/src/lib.rs

@@ -8,5 +8,7 @@
 #![cfg_attr(not(test), warn(unused_crate_dependencies))]
 #![cfg_attr(docsrs, feature(doc_cfg, doc_auto_cfg))]
 
-mod event_listeners;
-pub use event_listeners::EventListeners;
+mod event_sender;
+mod event_stream;
+pub use event_sender::EventSender;
+pub use event_stream::EventStream;