[broadcast] Add a result type to the broadcast (#804)

BrendanChou · web-flow · commit 7ca3160c9565 · 2025-04-23T08:29:28.000-07:00
diff --git a/broadcast/src/buffered/engine.rs b/broadcast/src/buffered/engine.rs
@@ -61,7 +61,7 @@ pub struct Engine<
     // Messaging
     ////////////////////////////////////////
     /// The mailbox for receiving messages.
-    mailbox_receiver: mpsc::Receiver<Message<D, M>>,
+    mailbox_receiver: mpsc::Receiver<Message<P, D, M>>,
 
     /// Pending requests from the application.
     waiters: HashMap<D, Vec<oneshot::Sender<M>>>,
@@ -103,9 +103,9 @@ impl<
 {
     /// Creates a new engine with the given context and configuration.
     /// Returns the engine and a mailbox for sending messages to the engine.
-    pub fn new(context: E, cfg: Config<Cfg, P>) -> (Self, Mailbox<D, M>) {
+    pub fn new(context: E, cfg: Config<Cfg, P>) -> (Self, Mailbox<P, D, M>) {
         let (mailbox_sender, mailbox_receiver) = mpsc::channel(cfg.mailbox_size);
-        let mailbox = Mailbox::<D, M>::new(mailbox_sender);
+        let mailbox = Mailbox::<P, D, M>::new(mailbox_sender);
         let metrics = metrics::Metrics::init(context.clone());
 
         let result = Self {
@@ -154,9 +154,9 @@ impl<
                         break;
                     };
                     match msg {
-                        Message::Broadcast{ message } => {
+                        Message::Broadcast{ message, responder } => {
                             trace!("mailbox: broadcast");
-                            self.handle_broadcast(&mut net_sender, message).await;
+                            self.handle_broadcast(&mut net_sender, message, responder).await;
                         }
                         Message::Get{ digest, responder } => {
                             trace!("mailbox: get");
@@ -199,16 +199,26 @@ impl<
     ////////////////////////////////////////
 
     /// Handles a `broadcast` request from the application.
-    async fn handle_broadcast(&mut self, net_sender: &mut NetS, msg: M) {
+    async fn handle_broadcast(
+        &mut self,
+        net_sender: &mut NetS,
+        msg: M,
+        responder: oneshot::Sender<Vec<P>>,
+    ) {
         // Store the message, continue even if it was already stored
         let _ = self.insert_message(self.public_key.clone(), msg.clone());
 
         // Broadcast the message to the network
         let recipients = Recipients::All;
         let msg = Bytes::from(msg.encode());
-        if let Err(err) = net_sender.send(recipients, msg, self.priority).await {
-            warn!(?err, "failed to send message");
-        }
+        let sent_to = net_sender
+            .send(recipients, msg, self.priority)
+            .await
+            .unwrap_or_else(|err| {
+                error!(?err, "failed to send message");
+                vec![]
+            });
+        let _ = responder.send(sent_to);
     }
 
     /// Handles a `get` request from the application.
diff --git a/broadcast/src/buffered/ingress.rs b/broadcast/src/buffered/ingress.rs
@@ -1,15 +1,21 @@
 use crate::Broadcaster;
 use commonware_codec::{Codec, Config};
 use commonware_cryptography::{Digest, Digestible};
+use commonware_utils::Array;
 use futures::{
     channel::{mpsc, oneshot},
     SinkExt,
 };
 
 /// Message types that can be sent to the `Mailbox`
-pub enum Message<D, M> {
+pub enum Message<P: Array, D: Digest, M: Digestible<D>> {
     /// Broadcast a [`Message`](crate::Broadcaster::Message) to the network.
-    Broadcast { message: M },
+    ///
+    /// The responder will be sent a list of peers that received the message.
+    Broadcast {
+        message: M,
+        responder: oneshot::Sender<Vec<P>>,
+    },
 
     /// Get a message by digest.
     ///
@@ -23,17 +29,17 @@ pub enum Message<D, M> {
 
 /// Ingress mailbox for [`Engine`](super::Engine).
 #[derive(Clone)]
-pub struct Mailbox<D: Digest, M: Digestible<D>> {
-    sender: mpsc::Sender<Message<D, M>>,
+pub struct Mailbox<P: Array, D: Digest, M: Digestible<D>> {
+    sender: mpsc::Sender<Message<P, D, M>>,
 }
 
-impl<D: Digest, M: Digestible<D>> Mailbox<D, M> {
-    pub(super) fn new(sender: mpsc::Sender<Message<D, M>>) -> Self {
+impl<P: Array, D: Digest, M: Digestible<D>> Mailbox<P, D, M> {
+    pub(super) fn new(sender: mpsc::Sender<Message<P, D, M>>) -> Self {
         Self { sender }
     }
 }
 
-impl<D: Digest, M: Digestible<D>> Mailbox<D, M> {
+impl<P: Array, D: Digest, M: Digestible<D>> Mailbox<P, D, M> {
     /// Get a message by digest.
     pub async fn get(&mut self, digest: D) -> oneshot::Receiver<M> {
         let (sender, receiver) = oneshot::channel();
@@ -48,13 +54,21 @@ impl<D: Digest, M: Digestible<D>> Mailbox<D, M> {
     }
 }
 
-impl<Cfg: Config, D: Digest, M: Codec<Cfg> + Digestible<D>> Broadcaster<Cfg> for Mailbox<D, M> {
+impl<Cfg: Config, P: Array, D: Digest, M: Codec<Cfg> + Digestible<D>> Broadcaster<Cfg>
+    for Mailbox<P, D, M>
+{
     type Message = M;
+    type Response = Vec<P>;
 
-    async fn broadcast(&mut self, message: Self::Message) {
+    async fn broadcast(&mut self, message: Self::Message) -> oneshot::Receiver<Vec<P>> {
+        let (sender, receiver) = oneshot::channel();
         self.sender
-            .send(Message::Broadcast { message })
+            .send(Message::Broadcast {
+                message,
+                responder: sender,
+            })
             .await
             .expect("mailbox closed");
+        receiver
     }
 }
diff --git a/broadcast/src/buffered/mod.rs b/broadcast/src/buffered/mod.rs
@@ -108,8 +108,9 @@ mod tests {
     fn spawn_peer_engines(
         context: deterministic::Context,
         registrations: &mut Registrations,
-    ) -> BTreeMap<PublicKey, Mailbox<Sha256Digest, TestMessage>> {
-        let mut mailboxes = BTreeMap::<PublicKey, Mailbox<Sha256Digest, TestMessage>>::new();
+    ) -> BTreeMap<PublicKey, Mailbox<PublicKey, Sha256Digest, TestMessage>> {
+        let mut mailboxes =
+            BTreeMap::<PublicKey, Mailbox<PublicKey, Sha256Digest, TestMessage>>::new();
         while let Some((peer, network)) = registrations.pop_first() {
             let context = context.with_label(&peer.to_string());
             let config = Config {
@@ -141,7 +142,7 @@ mod tests {
             // Send a single broadcast message from the first peer
             let message = TestMessage::new(b"hello world test message");
             let mut first_mailbox = mailboxes.get(peers.first().unwrap()).unwrap().clone();
-            first_mailbox.broadcast(message.clone()).await;
+            let result = first_mailbox.broadcast(message.clone()).await;
 
             // Allow time for propagation
             context.sleep(Duration::from_secs(1)).await;
@@ -154,6 +155,7 @@ mod tests {
                 let received_message = receiver.await.ok();
                 assert_eq!(received_message.unwrap(), message);
             }
+            assert_eq!(result.await.unwrap().len(), peers.len() - 1);
         });
     }
 
@@ -177,7 +179,8 @@ mod tests {
             let receiver_before = mailbox_a.get(digest_m1).await;
 
             // Broadcast the message
-            mailbox_a.broadcast(m1.clone()).await;
+            let result = mailbox_a.broadcast(m1.clone()).await;
+            assert_eq!(result.await.unwrap().len(), peers.len() - 1);
 
             // Wait for the pre-broadcast retrieval to complete
             let msg_before = receiver_before
@@ -219,7 +222,7 @@ mod tests {
             let digest = message.digest();
             for i in 0..100 {
                 // Broadcast the message
-                first_mailbox.broadcast(message.clone()).await;
+                let result = first_mailbox.broadcast(message.clone()).await;
                 context.sleep(NETWORK_SPEED_WITH_BUFFER).await;
 
                 // Check if all peers received the message
@@ -233,6 +236,7 @@ mod tests {
                     };
                     all_received &= has;
                 }
+                assert_eq!(result.await.unwrap().len(), peers.len() - 1);
 
                 // If all received, we're done
                 if all_received {
@@ -255,7 +259,8 @@ mod tests {
             // Broadcast a message
             let message = TestMessage::new(b"cached message");
             let mut first_mailbox = mailboxes.get(peers.first().unwrap()).unwrap().clone();
-            first_mailbox.broadcast(message.clone()).await;
+            let result = first_mailbox.broadcast(message.clone()).await;
+            assert_eq!(result.await.unwrap().len(), peers.len() - 1);
 
             // Wait for propagation
             context.sleep(NETWORK_SPEED_WITH_BUFFER).await;
@@ -294,7 +299,8 @@ mod tests {
             drop(dummy2);
 
             // Broadcast the message
-            mailbox1.broadcast(message.clone()).await;
+            let result = mailbox1.broadcast(message.clone()).await;
+            assert_eq!(result.await.unwrap().len(), peers.len() - 1);
 
             // Wait for propagation
             context.sleep(NETWORK_SPEED_WITH_BUFFER).await;
@@ -320,7 +326,8 @@ mod tests {
                 messages.push(TestMessage::new(format!("message {}", i).as_bytes()));
             }
             for message in messages.iter() {
-                mailbox.broadcast(message.clone()).await;
+                let result = mailbox.broadcast(message.clone()).await;
+                assert_eq!(result.await.unwrap().len(), peers.len() - 1);
             }
 
             // Wait for propagation
@@ -359,11 +366,13 @@ mod tests {
             // Create and broadcast message M1 from A
             let m1 = TestMessage::new(b"message M1");
             let digest_m1 = m1.digest();
-            mailbox_a.broadcast(m1.clone()).await;
+            let result = mailbox_a.broadcast(m1.clone()).await;
+            assert_eq!(result.await.unwrap().len(), peers.len() - 1);
             context.sleep(NETWORK_SPEED_WITH_BUFFER).await;
 
             // Broadcast M1 from C
-            mailbox_c.broadcast(m1.clone()).await;
+            let result = mailbox_c.broadcast(m1.clone()).await;
+            assert_eq!(result.await.unwrap().len(), peers.len() - 1);
             context.sleep(NETWORK_SPEED_WITH_BUFFER).await;
 
             // M1 is now in A's and C's deques in B's engine
@@ -374,7 +383,8 @@ mod tests {
                 new_messages_a.push(TestMessage::new(format!("A{}", i).as_bytes()));
             }
             for msg in &new_messages_a {
-                mailbox_a.broadcast(msg.clone()).await;
+                let result = mailbox_a.broadcast(msg.clone()).await;
+                assert_eq!(result.await.unwrap().len(), peers.len() - 1);
             }
             context.sleep(NETWORK_SPEED_WITH_BUFFER).await;
 
@@ -389,7 +399,8 @@ mod tests {
                 new_messages_c.push(TestMessage::new(format!("C{}", i).as_bytes()));
             }
             for msg in &new_messages_c {
-                mailbox_c.broadcast(msg.clone()).await;
+                let result = mailbox_c.broadcast(msg.clone()).await;
+                assert_eq!(result.await.unwrap().len(), peers.len() - 1);
             }
             context.sleep(NETWORK_SPEED_WITH_BUFFER).await;
 
diff --git a/broadcast/src/lib.rs b/broadcast/src/lib.rs
@@ -1,6 +1,7 @@
 //! Disseminate data over a wide-area network.
 
 use commonware_codec::{Codec, Config};
+use futures::channel::oneshot;
 use std::future::Future;
 
 pub mod buffered;
@@ -14,6 +15,14 @@ pub trait Broadcaster<Cfg: Config>: Clone + Send + 'static {
     /// - deserialized upon reception
     type Message: Codec<Cfg> + Clone + Send + 'static;
 
+    /// Response is the type of data that is returned once the message is broadcasted.
+    ///
+    /// It may also indicate the success or failure of the broadcast attempt.
+    type Response: Clone + Send + 'static;
+
     /// Attempt to broadcast a message to the network.
-    fn broadcast(&mut self, message: Self::Message) -> impl Future<Output = ()> + Send;
+    fn broadcast(
+        &mut self,
+        message: Self::Message,
+    ) -> impl Future<Output = oneshot::Receiver<Self::Response>> + Send;
 }
