[collector] Minimal Panic Fix (#1560)

patrick-ogrady · web-flow · commit c654d6da73d6 · 2025-09-08T10:45:11.000-07:00
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -48,6 +48,7 @@ commonware-runtime = { version = "0.0.60", path = "runtime" }
 commonware-storage = { version = "0.0.60", path = "storage" }
 commonware-stream = { version = "0.0.60", path = "stream" }
 commonware-utils = { version = "0.0.60", path = "utils", default-features = false }
+anyhow = { version = "1.0.99", default-features = false }
 thiserror = { version = "2.0.12", default-features = false }
 bytes = { version = "1.7.1", default-features = false }
 sha2 = { version = "0.10.8", default-features = false }
diff --git a/collector/Cargo.toml b/collector/Cargo.toml
@@ -25,6 +25,7 @@ tracing = { workspace = true }
 governor = { workspace = true }
 prometheus-client = { workspace = true }
 thiserror = { workspace = true }
+anyhow = { workspace = true }
 
 [dev-dependencies]
 tracing-subscriber = { workspace = true }
diff --git a/collector/src/lib.rs b/collector/src/lib.rs
@@ -15,9 +15,19 @@ use commonware_cryptography::{Committable, Digestible, PublicKey};
 use commonware_p2p::Recipients;
 use futures::channel::oneshot;
 use std::future::Future;
+use thiserror::Error;
 
 pub mod p2p;
 
+/// Errors that can occur when interacting with a [Originator].
+#[derive(Error, Debug)]
+pub enum Error {
+    #[error("send failed: {0}")]
+    SendFailed(anyhow::Error),
+    #[error("canceled")]
+    Canceled,
+}
+
 /// An [Originator] sends requests out to a set of [Handler]s and collects replies.
 pub trait Originator: Clone + Send + 'static {
     /// The [PublicKey] of a recipient.
@@ -32,7 +42,7 @@ pub trait Originator: Clone + Send + 'static {
         &mut self,
         recipients: Recipients<Self::PublicKey>,
         request: Self::Request,
-    ) -> impl Future<Output = Vec<Self::PublicKey>> + Send;
+    ) -> impl Future<Output = Result<Vec<Self::PublicKey>, Error>> + Send;
 
     /// Cancel a request by `commitment`, ignoring any future responses.
     ///
diff --git a/collector/src/p2p/engine.rs b/collector/src/p2p/engine.rs
@@ -2,7 +2,10 @@ use super::{
     ingress::{Mailbox, Message},
     Config,
 };
-use crate::p2p::{Handler, Monitor};
+use crate::{
+    p2p::{Handler, Monitor},
+    Error,
+};
 use commonware_codec::Codec;
 use commonware_cryptography::{Committable, Digestible, PublicKey};
 use commonware_macros::select;
@@ -143,13 +146,12 @@ where
                                     self.priority_request
                                 ).await {
                                     Ok(recipients) => {
-                                        for peer in &recipients {
-                                            entry.0.insert(peer.clone());
-                                        }
-                                        let _ = responder.send(recipients);
+                                        entry.0.extend(recipients.iter().cloned());
+                                        let _ = responder.send(Ok(recipients));
                                     }
                                     Err(err) => {
                                         error!(?err, ?commitment, "failed to send message");
+                                        let _ = responder.send(Err(Error::SendFailed(err.into())));
                                     }
                                 }
                             },
diff --git a/collector/src/p2p/ingress.rs b/collector/src/p2p/ingress.rs
@@ -1,4 +1,4 @@
-use crate::Originator;
+use crate::{Error, Originator};
 use commonware_codec::Codec;
 use commonware_cryptography::{Committable, Digestible, PublicKey};
 use commonware_p2p::Recipients;
@@ -12,7 +12,7 @@ pub enum Message<P: PublicKey, R: Committable + Digestible + Codec> {
     Send {
         request: R,
         recipients: Recipients<P>,
-        responder: oneshot::Sender<Vec<P>>,
+        responder: oneshot::Sender<Result<Vec<P>, Error>>,
     },
     Cancel {
         commitment: R::Commitment,
@@ -36,7 +36,7 @@ impl<P: PublicKey, R: Committable + Digestible + Codec> Originator for Mailbox<P
     type Request = R;
     type PublicKey = P;
 
-    async fn send(&mut self, recipients: Recipients<P>, request: R) -> Vec<P> {
+    async fn send(&mut self, recipients: Recipients<P>, request: R) -> Result<Vec<P>, Error> {
         let (tx, rx) = oneshot::channel();
         let _ = self
             .sender
@@ -46,7 +46,7 @@ impl<P: PublicKey, R: Committable + Digestible + Codec> Originator for Mailbox<P
                 responder: tx,
             })
             .await;
-        rx.await.unwrap()
+        rx.await.map_err(|_| Error::Canceled)?
     }
 
     async fn cancel(&mut self, commitment: R::Commitment) {
diff --git a/collector/src/p2p/mocks/mod.rs b/collector/src/p2p/mocks/mod.rs
@@ -2,4 +2,5 @@
 
 pub mod handler;
 pub mod monitor;
+pub mod sender;
 pub mod types;
diff --git a/collector/src/p2p/mocks/sender.rs b/collector/src/p2p/mocks/sender.rs
@@ -0,0 +1,42 @@
+//! Mock sender implementations for testing.
+
+use bytes::Bytes;
+use commonware_cryptography::PublicKey;
+use commonware_p2p::{Recipients, Sender};
+use thiserror::Error;
+
+/// Errors that can be returned by [Failing].
+#[derive(Debug, Error)]
+pub enum Error {
+    #[error("send failed")]
+    Failed,
+}
+
+/// A sender that always fails with [Error::Canceled].
+#[derive(Clone, Debug)]
+pub struct Failing<P: PublicKey> {
+    _phantom: std::marker::PhantomData<P>,
+}
+
+impl<P: PublicKey> Failing<P> {
+    /// Creates a new failing sender.
+    pub fn new() -> Self {
+        Self {
+            _phantom: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<P: PublicKey> Sender for Failing<P> {
+    type PublicKey = P;
+    type Error = Error;
+
+    async fn send(
+        &mut self,
+        _recipients: Recipients<P>,
+        _message: Bytes,
+        _priority: bool,
+    ) -> Result<Vec<P>, Self::Error> {
+        Err(Error::Failed)
+    }
+}
diff --git a/collector/src/p2p/mod.rs b/collector/src/p2p/mod.rs
@@ -49,7 +49,7 @@ mod tests {
         },
         Config, Engine, Mailbox,
     };
-    use crate::{Handler, Monitor, Originator};
+    use crate::{Error, Handler, Monitor, Originator};
     use commonware_codec::Encode;
     use commonware_cryptography::{
         ed25519::{PrivateKey, PublicKey},
@@ -208,7 +208,8 @@ mod tests {
             let request = Request { id: 1, data: 1 };
             let recipients = mailbox1
                 .send(Recipients::One(peers[1].clone()), request.clone())
-                .await;
+                .await
+                .expect("send failed");
             assert_eq!(recipients, vec![peers[1].clone()]);
 
             // Verify peer 2 received the request
@@ -275,7 +276,8 @@ mod tests {
             let commitment = request.commitment();
             let recipients = mailbox
                 .send(Recipients::One(peers[1].clone()), request.clone())
-                .await;
+                .await
+                .expect("send failed");
             assert_eq!(recipients, vec![peers[1].clone()]);
 
             // Cancel immediately
@@ -356,7 +358,10 @@ mod tests {
 
             // Broadcast request
             let request = Request { id: 3, data: 3 };
-            let recipients = mailbox1.send(Recipients::All, request.clone()).await;
+            let recipients = mailbox1
+                .send(Recipients::All, request.clone())
+                .await
+                .expect("send failed");
             assert_eq!(recipients.len(), 2);
             assert!(recipients.contains(&peers[1]));
             assert!(recipients.contains(&peers[2]));
@@ -434,7 +439,8 @@ mod tests {
             for _ in 0..3 {
                 let recipients = mailbox1
                     .send(Recipients::One(peers[1].clone()), request.clone())
-                    .await;
+                    .await
+                    .expect("send failed");
                 assert_eq!(recipients, vec![peers[1].clone()]);
             }
 
@@ -507,10 +513,12 @@ mod tests {
             let request2 = Request { id: 20, data: 20 };
             mailbox1
                 .send(Recipients::One(peers[1].clone()), request1)
-                .await;
+                .await
+                .expect("send failed");
             mailbox1
                 .send(Recipients::One(peers[1].clone()), request2)
-                .await;
+                .await
+                .expect("send failed");
 
             // Collect both responses
             let mut response10_received = false;
@@ -585,7 +593,8 @@ mod tests {
             let request = Request { id: 100, data: 100 };
             let recipients = mailbox1
                 .send(Recipients::One(peers[1].clone()), request.clone())
-                .await;
+                .await
+                .expect("send failed");
             assert_eq!(recipients, vec![peers[1].clone()]);
 
             // Verify handler received request but didn't respond
@@ -632,7 +641,10 @@ mod tests {
 
             // Send request with empty recipients list
             let request = Request { id: 1, data: 1 };
-            let recipients = mailbox.send(Recipients::All, request.clone()).await;
+            let recipients = mailbox
+                .send(Recipients::All, request.clone())
+                .await
+                .expect("send failed");
             assert_eq!(recipients, Vec::<PublicKey>::new());
 
             // Verify no responses collected
@@ -647,6 +659,92 @@ mod tests {
         });
     }
 
+    #[test_traced]
+    fn test_send_fails_with_network_error() {
+        let executor = deterministic::Runner::timed(Duration::from_secs(10));
+        executor.start(|context| async move {
+            let (oracle, schemes, peers, connections) =
+                setup_network_and_peers(&context, &[0, 1]).await;
+            let mut schemes = schemes.into_iter();
+            let mut connections = connections.into_iter();
+
+            // Setup peer 1 with a failing sender
+            let scheme = schemes.next().unwrap();
+            let conn = connections.next().unwrap();
+            let (_, receiver1) = conn.0; // Request channel
+            let sender1 = super::mocks::sender::Failing::<PublicKey>::new();
+            let (sender2, receiver2) = conn.1; // Response channel
+            let (engine, mut mailbox) = Engine::new(
+                context.with_label(&format!("engine_{}", scheme.public_key())),
+                Config {
+                    blocker: oracle.control(scheme.public_key()),
+                    monitor: MockMonitor::dummy(),
+                    handler: MockHandler::dummy(),
+                    mailbox_size: MAILBOX_SIZE,
+                    priority_request: false,
+                    request_codec: (),
+                    priority_response: false,
+                    response_codec: (),
+                },
+            );
+
+            // Start engine
+            engine.start((sender1, receiver1), (sender2, receiver2));
+
+            // Send request
+            let request = Request { id: 1, data: 1 };
+            let err = mailbox
+                .send(Recipients::One(peers[1].clone()), request.clone())
+                .await
+                .unwrap_err();
+            assert!(matches!(err, Error::SendFailed(_)));
+        });
+    }
+
+    #[test_traced]
+    fn test_send_fails_with_canceled() {
+        let executor = deterministic::Runner::timed(Duration::from_secs(10));
+        executor.start(|context| async move {
+            let (oracle, schemes, peers, connections) =
+                setup_network_and_peers(&context, &[0, 1]).await;
+            let mut schemes = schemes.into_iter();
+            let mut connections = connections.into_iter();
+
+            // Setup peer 1 with a failing sender
+            let scheme = schemes.next().unwrap();
+            let conn = connections.next().unwrap();
+            let (sender1, receiver1) = conn.0; // Request channel
+            let (sender2, receiver2) = conn.1; // Response channel
+            let (engine, mut mailbox) = Engine::new(
+                context.with_label(&format!("engine_{}", scheme.public_key())),
+                Config {
+                    blocker: oracle.control(scheme.public_key()),
+                    monitor: MockMonitor::dummy(),
+                    handler: MockHandler::dummy(),
+                    mailbox_size: MAILBOX_SIZE,
+                    priority_request: false,
+                    request_codec: (),
+                    priority_response: false,
+                    response_codec: (),
+                },
+            );
+
+            // Start engine
+            let handle = engine.start((sender1, receiver1), (sender2, receiver2));
+
+            // Stop the engine (which will result in all further requests being canceled)
+            handle.abort();
+
+            // Send request (will return Error::Canceled instead of Error::SendFailed)
+            let request = Request { id: 1, data: 1 };
+            let err = mailbox
+                .send(Recipients::One(peers[1].clone()), request.clone())
+                .await
+                .unwrap_err();
+            assert!(matches!(err, Error::Canceled));
+        });
+    }
+
     #[test_traced]
     fn test_response_from_unknown_peer() {
         let executor = deterministic::Runner::timed(Duration::from_secs(10));
@@ -701,7 +799,8 @@ mod tests {
             let request_to_peer2 = Request { id: 42, data: 42 };
             let recipients = mailbox1
                 .send(Recipients::One(peers[1].clone()), request_to_peer2.clone())
-                .await;
+                .await
+                .expect("send failed");
             assert_eq!(recipients, vec![peers[1].clone()]);
 
             // Send a response from peer 3 to peer 1
diff --git a/p2p/src/lib.rs b/p2p/src/lib.rs
@@ -38,7 +38,7 @@ pub enum Recipients<P: PublicKey> {
 /// Interface for sending messages to a set of recipients.
 pub trait Sender: Clone + Debug + Send + 'static {
     /// Error that can occur when sending a message.
-    type Error: Debug + StdError + Send;
+    type Error: Debug + StdError + Send + Sync;
 
     /// Public key type used to identify recipients.
     type PublicKey: PublicKey;
@@ -55,7 +55,7 @@ pub trait Sender: Clone + Debug + Send + 'static {
 /// Interface for receiving messages from arbitrary recipients.
 pub trait Receiver: Debug + Send + 'static {
     /// Error that can occur when receiving a message.
-    type Error: Debug + StdError + Send;
+    type Error: Debug + StdError + Send + Sync;
 
     /// Public key type used to identify recipients.
     type PublicKey: PublicKey;
diff --git a/utils/src/lib.rs b/utils/src/lib.rs
@@ -13,7 +13,10 @@ extern crate alloc;
 use alloc::{string::String, vec::Vec};
 use bytes::{BufMut, BytesMut};
 use commonware_codec::{EncodeSize, Write};
-use core::{fmt::Write as FmtWrite, time::Duration};
+use core::{
+    fmt::{Debug, Write as FmtWrite},
+    time::Duration,
+};
 
 pub mod sequence;
 pub use sequence::{Array, Span};
