[p2p/utils] (Optionally) Wrap Channel with Codec (#830)

Author: patrick-ogrady

consensus/src/ordered_broadcast/engine.rs

@@ -13,13 +13,15 @@ use super::{
     AckManager, Config, TipManager,
 };
 use crate::{Automaton, Monitor, Relay, Reporter, Supervisor, ThresholdSupervisor};
-use commonware_codec::{DecodeExt, Encode};
 use commonware_cryptography::{
     bls12381::primitives::{group, poly},
     Digest, Scheme,
 };
 use commonware_macros::select;
-use commonware_p2p::{Receiver, Recipients, Sender};
+use commonware_p2p::{
+    utils::codec::{wrap, WrappedSender},
+    Receiver, Recipients, Sender,
+};
 use commonware_runtime::{
     telemetry::metrics::{
         histogram,
@@ -259,8 +261,8 @@ impl<
 
     /// Inner run loop called by `start`.
     async fn run(mut self, chunk_network: (NetS, NetR), ack_network: (NetS, NetR)) {
-        let (mut node_sender, mut node_receiver) = chunk_network;
-        let (mut ack_sender, mut ack_receiver) = ack_network;
+        let (mut node_sender, mut node_receiver) = wrap((), chunk_network.0, chunk_network.1);
+        let (mut ack_sender, mut ack_receiver) = wrap((), ack_network.0, ack_network.1);
         let mut shutdown = self.context.stopped();
 
         // Tracks if there is an outstanding proposal request to the automaton.
@@ -361,7 +363,7 @@ impl<
                         }
                     };
                     let mut guard = self.metrics.nodes.guard(Status::Invalid);
-                    let node = match Node::decode(msg) {
+                    let node = match msg {
                         Ok(node) => node,
                         Err(err) => {
                             warn!(?err, ?sender, "node decode failed");
@@ -405,7 +407,7 @@ impl<
                         }
                     };
                     let mut guard = self.metrics.acks.guard(Status::Invalid);
-                    let ack = match Ack::decode(msg) {
+                    let ack = match msg {
                         Ok(ack) => ack,
                         Err(err) => {
                             warn!(?err, ?sender, "ack decode failed");
@@ -469,7 +471,7 @@ impl<
         &mut self,
         context: &Context<C::PublicKey>,
         payload: &D,
-        ack_sender: &mut NetS,
+        ack_sender: &mut WrappedSender<NetS, (), Ack<C::PublicKey, D>>,
     ) -> Result<(), Error> {
         // Get the tip
         let Some(tip) = self.tip_manager.get(&context.sequencer) else {
@@ -513,19 +515,15 @@ impl<
             recipients
         };
 
+        // Handle the ack internally
+        self.handle_ack(&ack).await?;
+
         // Send the ack to the network
         ack_sender
-            .send(
-                Recipients::Some(recipients),
-                ack.encode().into(),
-                self.priority_acks,
-            )
+            .send(Recipients::Some(recipients), ack, self.priority_acks)
             .await
             .map_err(|_| Error::UnableToSendMessage)?;
 
-        // Handle the ack internally
-        self.handle_ack(&ack).await?;
-
         // Emit the activity
         self.reporter
             .report(Activity::Proposal(Proposal::new(
@@ -669,7 +667,7 @@ impl<
         &mut self,
         context: Context<C::PublicKey>,
         payload: D,
-        node_sender: &mut NetS,
+        node_sender: &mut WrappedSender<NetS, (), Node<C, D>>,
     ) -> Result<(), Error> {
         let mut guard = self.metrics.propose.guard(Status::Dropped);
         let me = self.crypto.public_key();
@@ -718,7 +716,7 @@ impl<
         self.propose_timer = Some(self.metrics.e2e_duration.timer());
 
         // Broadcast to network
-        if let Err(err) = self.broadcast(&node, node_sender, self.epoch).await {
+        if let Err(err) = self.broadcast(node, node_sender, self.epoch).await {
             guard.set(Status::Failure);
             return Err(err);
         };
@@ -734,7 +732,10 @@ impl<
     /// - this instance is the sequencer for the current epoch.
     /// - this instance has a chunk to rebroadcast.
     /// - this instance has not yet collected the threshold signature for the chunk.
-    async fn rebroadcast(&mut self, node_sender: &mut NetS) -> Result<(), Error> {
+    async fn rebroadcast(
+        &mut self,
+        node_sender: &mut WrappedSender<NetS, (), Node<C, D>>,
+    ) -> Result<(), Error> {
         let mut guard = self.metrics.rebroadcast.guard(Status::Dropped);
 
         // Unset the rebroadcast deadline
@@ -762,16 +763,16 @@ impl<
 
         // Broadcast the message, which resets the rebroadcast deadline
         guard.set(Status::Failure);
-        self.broadcast(&tip, node_sender, self.epoch).await?;
+        self.broadcast(tip, node_sender, self.epoch).await?;
         guard.set(Status::Success);
         Ok(())
     }
 
     /// Send a  `Node` message to all validators in the given epoch.
     async fn broadcast(
         &mut self,
-        node: &Node<C, D>,
-        node_sender: &mut NetS,
+        node: Node<C, D>,
+        node_sender: &mut WrappedSender<NetS, (), Node<C, D>>,
         epoch: Epoch,
     ) -> Result<(), Error> {
         // Get the validators for the epoch
@@ -786,7 +787,7 @@ impl<
         node_sender
             .send(
                 Recipients::Some(validators.clone()),
-                node.encode().into(),
+                node,
                 self.priority_proposals,
             )
             .await

consensus/src/simplex/actors/resolver/actor.rs

@@ -9,10 +9,15 @@ use crate::{
     },
     Supervisor,
 };
-use commonware_codec::{Decode, Encode};
 use commonware_cryptography::{Digest, Scheme};
 use commonware_macros::select;
-use commonware_p2p::{utils::requester, Receiver, Recipients, Sender};
+use commonware_p2p::{
+    utils::{
+        codec::{wrap, WrappedSender},
+        requester,
+    },
+    Receiver, Recipients, Sender,
+};
 use commonware_runtime::{Clock, Handle, Metrics, Spawner};
 use futures::{channel::mpsc, future::Either, StreamExt};
 use governor::clock::Clock as GClock;
@@ -191,7 +196,11 @@ impl<
     }
 
     /// Concurrent indicates whether we should send a new request (only if we see a request for the first time)
-    async fn send(&mut self, shuffle: bool, sender: &mut impl Sender<PublicKey = C::PublicKey>) {
+    async fn send<Sr: Sender<PublicKey = C::PublicKey>>(
+        &mut self,
+        shuffle: bool,
+        sender: &mut WrappedSender<Sr, (usize, usize), Backfiller<C::Signature, D>>,
+    ) {
         // Clear retry
         self.retry = None;
 
@@ -255,9 +264,7 @@ impl<
 
                 // Create new message
                 msg.id = request;
-                let encoded = Backfiller::Request::<C::Signature, D>(msg.clone())
-                    .encode()
-                    .into();
+                let encoded = Backfiller::Request(msg.clone());
 
                 // Try to send
                 if sender
@@ -298,9 +305,16 @@ impl<
     async fn run(
         mut self,
         mut voter: voter::Mailbox<C::Signature, D>,
-        mut sender: impl Sender<PublicKey = C::PublicKey>,
-        mut receiver: impl Receiver<PublicKey = C::PublicKey>,
+        sender: impl Sender<PublicKey = C::PublicKey>,
+        receiver: impl Receiver<PublicKey = C::PublicKey>,
     ) {
+        // Wrap channel
+        let (mut sender, mut receiver) = wrap(
+            (self.max_fetch_count, self.max_participants),
+            sender,
+            receiver,
+        );
+
         // Wait for an event
         let mut current_view = 0;
         let mut finalized_view = 0;
@@ -423,14 +437,19 @@ impl<
                     }
                 },
                 network = receiver.recv() => {
-                    let (s, msg) = network.unwrap();
-                    let msg = match Backfiller::decode_cfg(msg, &(self.max_fetch_count, self.max_participants)) {
+                    // Break if there is an internal error
+                    let Ok((s, msg)) = network else {
+                        break;
+                    };
+
+                    // Skip if there is a decoding error
+                    let msg = match msg {
                         Ok(msg) => msg,
                         Err(err) => {
                             warn!(?err, sender = ?s, "failed to decode message");
                             self.requester.block(s);
                             continue;
-                        },
+                        }
                     };
                     match msg {
                         Backfiller::Request(request) => {
@@ -465,13 +484,11 @@ impl<
 
                             // Send response
                             debug!(sender = ?s, ?notarizations, ?missing_notarizations, ?nullifications, ?missing_nullifications, "sending response");
-                            let msg = Backfiller::Response::<C::Signature, D>(Response::new(
+                            let msg = Backfiller::Response(Response::new(
                                 request.id,
                                 notarizations_found,
                                 nullifications_found,
-                            ))
-                            .encode()
-                            .into();
+                            ));
                             sender
                                 .send(Recipients::One(s), msg, false)
                                 .await

consensus/src/simplex/actors/voter/actor.rs

@@ -11,10 +11,12 @@ use crate::{
     },
     Automaton, Relay, Reporter, Supervisor, LATENCY,
 };
-use commonware_codec::{Decode, Encode};
 use commonware_cryptography::{Digest, Scheme};
 use commonware_macros::select;
-use commonware_p2p::{Receiver, Recipients, Sender};
+use commonware_p2p::{
+    utils::codec::{wrap, WrappedSender},
+    Receiver, Recipients, Sender,
+};
 use commonware_runtime::{Clock, Handle, Metrics, Spawner, Storage};
 use commonware_storage::journal::variable::{Config as JConfig, Journal};
 use commonware_utils::quorum;
@@ -598,7 +600,10 @@ impl<
         null_retry
     }
 
-    async fn timeout(&mut self, sender: &mut impl Sender) {
+    async fn timeout<Sr: Sender>(
+        &mut self,
+        sender: &mut WrappedSender<Sr, usize, Voter<C::Signature, D>>,
+    ) {
         // Set timeout fired
         let round = self.views.get_mut(&self.view).unwrap();
         let mut retry = false;
@@ -616,16 +621,14 @@ impl<
         let past_view = self.view - 1;
         if retry && past_view > 0 {
             if let Some(notarization) = self.construct_notarization(past_view, true) {
-                let msg = Voter::Notarization(notarization).encode().into();
+                let msg = Voter::Notarization(notarization);
                 sender.send(Recipients::All, msg, true).await.unwrap();
                 self.broadcast_messages
                     .get_or_create(&metrics::NOTARIZATION)
                     .inc();
                 debug!(view = past_view, "rebroadcast entry notarization");
             } else if let Some(nullification) = self.construct_nullification(past_view, true) {
-                let msg = Voter::Nullification::<C::Signature, D>(nullification)
-                    .encode()
-                    .into();
+                let msg = Voter::Nullification(nullification);
                 sender.send(Recipients::All, msg, true).await.unwrap();
                 self.broadcast_messages
                     .get_or_create(&metrics::NULLIFICATION)
@@ -663,7 +666,7 @@ impl<
             .expect("unable to sync journal");
 
         // Broadcast nullify
-        let msg = Voter::Nullify::<C::Signature, D>(nullify).encode().into();
+        let msg = Voter::Nullify(nullify);
         sender.send(Recipients::All, msg, true).await.unwrap();
         self.broadcast_messages
             .get_or_create(&metrics::NULLIFY)
@@ -1444,10 +1447,10 @@ impl<
         Some(Finalization::new(proposal, signatures))
     }
 
-    async fn notify(
+    async fn notify<Sr: Sender>(
         &mut self,
         backfiller: &mut resolver::Mailbox<C::Signature, D>,
-        sender: &mut impl Sender,
+        sender: &mut WrappedSender<Sr, usize, Voter<C::Signature, D>>,
         view: u64,
     ) {
         // Attempt to notarize
@@ -1464,7 +1467,7 @@ impl<
                 .expect("unable to sync journal");
 
             // Broadcast the notarize
-            let msg = Voter::Notarize(notarize).encode().into();
+            let msg = Voter::Notarize(notarize);
             sender.send(Recipients::All, msg, true).await.unwrap();
             self.broadcast_messages
                 .get_or_create(&metrics::NOTARIZE)
@@ -1500,7 +1503,7 @@ impl<
                 .await;
 
             // Broadcast the notarization
-            let msg = Voter::Notarization(notarization).encode().into();
+            let msg = Voter::Notarization(notarization);
             sender.send(Recipients::All, msg, true).await.unwrap();
             self.broadcast_messages
                 .get_or_create(&metrics::NOTARIZATION)
@@ -1531,9 +1534,7 @@ impl<
                 .await;
 
             // Broadcast the nullification
-            let msg = Voter::Nullification::<C::Signature, D>(nullification)
-                .encode()
-                .into();
+            let msg = Voter::Nullification(nullification);
             sender.send(Recipients::All, msg, true).await.unwrap();
             self.broadcast_messages
                 .get_or_create(&metrics::NULLIFICATION)
@@ -1573,7 +1574,7 @@ impl<
                     );
                     let finalization = self.construct_finalization(self.last_finalized, true);
                     if let Some(finalization) = finalization {
-                        let msg = Voter::Finalization(finalization).encode().into();
+                        let msg = Voter::Finalization(finalization);
                         sender
                             .send(Recipients::All, msg, true)
                             .await
@@ -1605,7 +1606,7 @@ impl<
                 .expect("unable to sync journal");
 
             // Broadcast the finalize
-            let msg = Voter::Finalize(finalize).encode().into();
+            let msg = Voter::Finalize(finalize);
             sender.send(Recipients::All, msg, true).await.unwrap();
             self.broadcast_messages
                 .get_or_create(&metrics::FINALIZE)
@@ -1641,7 +1642,7 @@ impl<
                 .await;
 
             // Broadcast the finalization
-            let msg = Voter::Finalization(finalization).encode().into();
+            let msg = Voter::Finalization(finalization);
             sender.send(Recipients::All, msg, true).await.unwrap();
             self.broadcast_messages
                 .get_or_create(&metrics::FINALIZATION)
@@ -1661,9 +1662,12 @@ impl<
     async fn run(
         mut self,
         mut backfiller: resolver::Mailbox<C::Signature, D>,
-        mut sender: impl Sender,
-        mut receiver: impl Receiver,
+        sender: impl Sender,
+        receiver: impl Receiver,
     ) {
+        // Wrap channel
+        let (mut sender, mut receiver) = wrap(self.max_participants, sender, receiver);
+
         // Compute genesis
         let genesis = self.automaton.genesis().await;
         self.genesis = Some(genesis);
@@ -1907,11 +1911,13 @@ impl<
                     }
                 },
                 msg = receiver.recv() => {
-                    // Parse message
+                    // Break if there is an internal error
                     let Ok((s, msg)) = msg else {
                         break;
                     };
-                    let Ok(msg) = Voter::decode_cfg(msg, &self.max_participants) else {
+
+                    // Skip if there is a decoding error
+                    let Ok(msg) = msg else {
                         continue;
                     };