refactor: [#1389] change channel in UDP core from mpsc to broadcast

Author: josecelano

packages/udp-tracker-core/benches/helpers/utils.rs

@@ -3,7 +3,7 @@ use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use bittorrent_udp_tracker_core::statistics;
 use futures::future::BoxFuture;
 use mockall::mock;
-use tokio::sync::mpsc::error::SendError;
+use tokio::sync::broadcast::error::SendError;
 
 pub(crate) fn sample_ipv4_remote_addr() -> SocketAddr {
     sample_ipv4_socket_address()
@@ -20,6 +20,6 @@ pub(crate) fn sample_issue_time() -> f64 {
 mock! {
     pub(crate) UdpCoreStatsEventSender {}
     impl statistics::event::sender::Sender for UdpCoreStatsEventSender {
-         fn send_event(&self, event: statistics::event::Event) -> BoxFuture<'static,Option<Result<(),SendError<statistics::event::Event> > > > ;
+         fn send_event(&self, event: statistics::event::Event) -> BoxFuture<'static,Option<Result<usize,SendError<statistics::event::Event> > > > ;
     }
 }

packages/udp-tracker-core/src/services/connect.rs

@@ -142,7 +142,7 @@ mod tests {
                     context: ConnectionContext::new(client_socket_addr, server_socket_addr),
                 }))
                 .times(1)
-                .returning(|_| Box::pin(future::ready(Some(Ok(())))));
+                .returning(|_| Box::pin(future::ready(Some(Ok(1)))));
             let opt_udp_stats_event_sender: Arc<Option<Box<dyn statistics::event::sender::Sender>>> =
                 Arc::new(Some(Box::new(udp_stats_event_sender_mock)));
 
@@ -165,7 +165,7 @@ mod tests {
                     context: ConnectionContext::new(client_socket_addr, server_socket_addr),
                 }))
                 .times(1)
-                .returning(|_| Box::pin(future::ready(Some(Ok(())))));
+                .returning(|_| Box::pin(future::ready(Some(Ok(1)))));
             let opt_udp_stats_event_sender: Arc<Option<Box<dyn statistics::event::sender::Sender>>> =
                 Arc::new(Some(Box::new(udp_stats_event_sender_mock)));
 

packages/udp-tracker-core/src/services/mod.rs

@@ -10,7 +10,7 @@ pub(crate) mod tests {
 
     use futures::future::BoxFuture;
     use mockall::mock;
-    use tokio::sync::mpsc::error::SendError;
+    use tokio::sync::broadcast::error::SendError;
 
     use crate::connection_cookie::gen_remote_fingerprint;
     use crate::statistics;
@@ -46,7 +46,7 @@ pub(crate) mod tests {
     mock! {
         pub(crate) UdpCoreStatsEventSender {}
         impl statistics::event::sender::Sender for UdpCoreStatsEventSender {
-             fn send_event(&self, event: statistics::event::Event) -> BoxFuture<'static,Option<Result<(),SendError<statistics::event::Event> > > > ;
+             fn send_event(&self, event: statistics::event::Event) -> BoxFuture<'static,Option<Result<usize,SendError<statistics::event::Event> > > > ;
         }
     }
 }

packages/udp-tracker-core/src/statistics/event/listener.rs

@@ -1,11 +1,17 @@
-use tokio::sync::mpsc;
+use tokio::sync::broadcast;
 
 use super::handler::handle_event;
 use super::Event;
 use crate::statistics::repository::Repository;
 
-pub async fn dispatch_events(mut receiver: mpsc::Receiver<Event>, stats_repository: Repository) {
-    while let Some(event) = receiver.recv().await {
-        handle_event(event, &stats_repository).await;
+pub async fn dispatch_events(mut receiver: broadcast::Receiver<Event>, stats_repository: Repository) {
+    loop {
+        match receiver.recv().await {
+            Ok(event) => handle_event(event, &stats_repository).await,
+            Err(e) => {
+                tracing::error!("Error receiving udp tracker core event: {:?}", e);
+                break;
+            }
+        }
     }
 }

packages/udp-tracker-core/src/statistics/event/mod.rs

@@ -8,14 +8,14 @@ pub mod sender;
 ///
 /// - `Udp` prefix means the event was triggered by the UDP tracker.
 /// - The event suffix is the type of request: `announce`, `scrape` or `connection`.
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub enum Event {
     UdpConnect { context: ConnectionContext },
     UdpAnnounce { context: ConnectionContext },
     UdpScrape { context: ConnectionContext },
 }
 
-#[derive(Debug, PartialEq, Eq)]
+#[derive(Debug, PartialEq, Eq, Clone)]
 pub struct ConnectionContext {
     client_socket_addr: SocketAddr,
     server_socket_addr: SocketAddr,

packages/udp-tracker-core/src/statistics/event/sender.rs

@@ -2,15 +2,15 @@ use futures::future::BoxFuture;
 use futures::FutureExt;
 #[cfg(test)]
 use mockall::{automock, predicate::str};
-use tokio::sync::mpsc;
-use tokio::sync::mpsc::error::SendError;
+use tokio::sync::broadcast;
+use tokio::sync::broadcast::error::SendError;
 
 use super::Event;
 
 /// A trait to allow sending statistics events
 #[cfg_attr(test, automock)]
 pub trait Sender: Sync + Send {
-    fn send_event(&self, event: Event) -> BoxFuture<'_, Option<Result<(), SendError<Event>>>>;
+    fn send_event(&self, event: Event) -> BoxFuture<'_, Option<Result<usize, SendError<Event>>>>;
 }
 
 /// An [`statistics::EventSender`](crate::statistics::event::sender::Sender) implementation.
@@ -19,11 +19,11 @@ pub trait Sender: Sync + Send {
 /// [`statistics::Keeper`](crate::statistics::keeper::Keeper)
 #[allow(clippy::module_name_repetitions)]
 pub struct ChannelSender {
-    pub(crate) sender: mpsc::Sender<Event>,
+    pub(crate) sender: broadcast::Sender<Event>,
 }
 
 impl Sender for ChannelSender {
-    fn send_event(&self, event: Event) -> BoxFuture<'_, Option<Result<(), SendError<Event>>>> {
-        async move { Some(self.sender.send(event).await) }.boxed()
+    fn send_event(&self, event: Event) -> BoxFuture<'_, Option<Result<usize, SendError<Event>>>> {
+        async move { Some(self.sender.send(event)) }.boxed()
     }
 }

packages/udp-tracker-core/src/statistics/keeper.rs

@@ -1,12 +1,9 @@
-use tokio::sync::mpsc;
+use tokio::sync::broadcast::Receiver;
 
 use super::event::listener::dispatch_events;
-use super::event::sender::{ChannelSender, Sender};
 use super::event::Event;
 use super::repository::Repository;
 
-const CHANNEL_BUFFER_SIZE: usize = 65_535;
-
 /// The service responsible for keeping tracker metrics (listening to statistics events and handle them).
 ///
 /// It actively listen to new statistics events. When it receives a new event
@@ -29,31 +26,15 @@ impl Keeper {
         }
     }
 
-    #[must_use]
-    pub fn new_active_instance() -> (Box<dyn Sender>, Repository) {
-        let mut stats_tracker = Self::new();
-
-        let stats_event_sender = stats_tracker.run_event_listener();
-
-        (stats_event_sender, stats_tracker.repository)
-    }
-
-    pub fn run_event_listener(&mut self) -> Box<dyn Sender> {
-        let (sender, receiver) = mpsc::channel::<Event>(CHANNEL_BUFFER_SIZE);
-
+    pub fn run_event_listener(&mut self, receiver: Receiver<Event>) {
         let stats_repository = self.repository.clone();
 
         tokio::spawn(async move { dispatch_events(receiver, stats_repository).await });
-
-        Box::new(ChannelSender { sender })
     }
 }
 
 #[cfg(test)]
 mod tests {
-    use std::net::{IpAddr, Ipv4Addr, SocketAddr};
-
-    use crate::statistics::event::{ConnectionContext, Event};
     use crate::statistics::keeper::Keeper;
     use crate::statistics::metrics::Metrics;
 
@@ -65,22 +46,4 @@ mod tests {
 
         assert_eq!(stats.udp4_announces_handled, Metrics::default().udp4_announces_handled);
     }
-
-    #[tokio::test]
-    async fn should_create_an_event_sender_to_send_statistical_events() {
-        let mut stats_tracker = Keeper::new();
-
-        let event_sender = stats_tracker.run_event_listener();
-
-        let result = event_sender
-            .send_event(Event::UdpConnect {
-                context: ConnectionContext::new(
-                    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 195)), 8080),
-                    SocketAddr::new(IpAddr::V4(Ipv4Addr::new(203, 0, 113, 196)), 6969),
-                ),
-            })
-            .await;
-
-        assert!(result.is_some());
-    }
 }

packages/udp-tracker-core/src/statistics/setup.rs

@@ -1,8 +1,13 @@
 //! Setup for the tracker statistics.
 //!
 //! The [`factory`] function builds the structs needed for handling the tracker metrics.
+use tokio::sync::broadcast;
+
+use super::event::sender::ChannelSender;
 use crate::statistics;
 
+const CHANNEL_CAPACITY: usize = 1024;
+
 /// It builds the structs needed for handling the tracker metrics.
 ///
 /// It returns:
@@ -19,15 +24,21 @@ pub fn factory(
     Option<Box<dyn statistics::event::sender::Sender>>,
     statistics::repository::Repository,
 ) {
-    let mut stats_event_sender = None;
+    let mut stats_event_sender: Option<Box<dyn statistics::event::sender::Sender>> = None;
 
-    let mut stats_tracker = statistics::keeper::Keeper::new();
+    let mut keeper = statistics::keeper::Keeper::new();
 
     if tracker_usage_statistics {
-        stats_event_sender = Some(stats_tracker.run_event_listener());
+        let (sender, _) = broadcast::channel(CHANNEL_CAPACITY);
+
+        let receiver = sender.subscribe();
+
+        stats_event_sender = Some(Box::new(ChannelSender { sender }));
+
+        keeper.run_event_listener(receiver);
     }
 
-    (stats_event_sender, stats_tracker.repository)
+    (stats_event_sender, keeper.repository)
 }
 
 #[cfg(test)]

packages/udp-tracker-server/src/handlers/announce.rs

@@ -854,7 +854,7 @@ mod tests {
                             context: core_statistics::event::ConnectionContext::new(client_socket_addr, server_socket_addr),
                         }))
                         .times(1)
-                        .returning(|_| Box::pin(future::ready(Some(Ok(())))));
+                        .returning(|_| Box::pin(future::ready(Some(Ok(1)))));
                     let udp_core_stats_event_sender: Arc<Option<Box<dyn core_statistics::event::sender::Sender>>> =
                         Arc::new(Some(Box::new(udp_core_stats_event_sender_mock)));
 

packages/udp-tracker-server/src/handlers/connect.rs

@@ -196,7 +196,7 @@ mod tests {
                     context: core_statistics::event::ConnectionContext::new(client_socket_addr, server_socket_addr),
                 }))
                 .times(1)
-                .returning(|_| Box::pin(future::ready(Some(Ok(())))));
+                .returning(|_| Box::pin(future::ready(Some(Ok(1)))));
             let udp_core_stats_event_sender: Arc<Option<Box<dyn core_statistics::event::sender::Sender>>> =
                 Arc::new(Some(Box::new(udp_core_stats_event_sender_mock)));
 
@@ -237,7 +237,7 @@ mod tests {
                     context: core_statistics::event::ConnectionContext::new(client_socket_addr, server_socket_addr),
                 }))
                 .times(1)
-                .returning(|_| Box::pin(future::ready(Some(Ok(())))));
+                .returning(|_| Box::pin(future::ready(Some(Ok(1)))));
             let udp_core_stats_event_sender: Arc<Option<Box<dyn core_statistics::event::sender::Sender>>> =
                 Arc::new(Some(Box::new(udp_core_stats_event_sender_mock)));