commonware_p2p/utils/
mux.rs

1//! This utility wraps a [Sender] and [Receiver], providing lightweight sub-channels keyed by
2//! [Channel].
3//!
4//! Usage:
5//! - Call [Muxer::new] to obtain a ([Muxer], [MuxHandle]) pair.
6//! - Call [Muxer::start] or run [Muxer::run] in a background task to demux incoming messages into
7//!   per-subchannel queues.
8//! - Call [MuxHandle::register] to obtain a ([SubSender], [SubReceiver]) pair for that subchannel,
9//!   even if the muxer is already running.
10
11use crate::{Channel, CheckedSender, LimitedSender, Message, Receiver, Recipients, Sender};
12use commonware_actor::{Feedback, Unreliable};
13use commonware_codec::{Encode, Error as CodecError, ReadExt, varint::UInt};
14use commonware_macros::select_loop;
15use commonware_runtime::{ContextCell, Handle, IoBuf, IoBufs, Spawner, spawn_cell};
16use commonware_utils::channel::{
17    fallible::FallibleExt,
18    mpsc::{self, error::TrySendError},
19    oneshot,
20};
21use std::{collections::HashMap, fmt::Debug, time::SystemTime};
22use thiserror::Error;
23use tracing::debug;
24
25/// Errors that can occur when interacting with a [SubReceiver] or [MuxHandle].
26#[derive(Error, Debug)]
27pub enum Error {
28    #[error("subchannel already registered: {0}")]
29    AlreadyRegistered(Channel),
30    #[error("muxer is closed")]
31    Closed,
32    #[error("recv failed")]
33    RecvFailed,
34}
35
36/// Parse a muxed message into its subchannel and payload.
37pub fn parse(mut buf: IoBuf) -> Result<(Channel, IoBuf), CodecError> {
38    let subchannel: Channel = UInt::read(&mut buf)?.into();
39    Ok((subchannel, buf))
40}
41
42/// Control messages for the [Muxer].
43enum Control<R: Receiver> {
44    Register {
45        subchannel: Channel,
46        sender: oneshot::Sender<mpsc::Receiver<Message<R::PublicKey>>>,
47    },
48    Deregister {
49        subchannel: Channel,
50    },
51}
52
53/// Thread-safe routing table mapping each [Channel] to the [mpsc::Sender] for [`Message<P>`].
54type Routes<P> = HashMap<Channel, mpsc::Sender<Message<P>>>;
55
56/// A backup channel response, with a [SubSender] to respond, the [Channel] that wasn't registered,
57/// and the [Message] received.
58type BackupResponse<P> = (Channel, Message<P>);
59
60/// A multiplexer of p2p channels into subchannels.
61pub struct Muxer<E: Spawner, S: Sender, R: Receiver> {
62    context: ContextCell<E>,
63    sender: S,
64    receiver: R,
65    mailbox_size: usize,
66    control_rx: mpsc::UnboundedReceiver<Control<R>>,
67    routes: Routes<R::PublicKey>,
68    backup: Option<mpsc::Sender<BackupResponse<R::PublicKey>>>,
69}
70
71impl<E: Spawner, S: Sender, R: Receiver> Muxer<E, S, R> {
72    /// Create a multiplexed wrapper around a [Sender] and [Receiver] pair, and return a ([Muxer],
73    /// [MuxHandle]) pair that can be used to register routes dynamically.
74    pub fn new(context: E, sender: S, receiver: R, mailbox_size: usize) -> (Self, MuxHandle<S, R>) {
75        Self::builder(context, sender, receiver, mailbox_size).build()
76    }
77
78    /// Creates a [MuxerBuilder] that can be used to configure and build a [Muxer].
79    pub fn builder(
80        context: E,
81        sender: S,
82        receiver: R,
83        mailbox_size: usize,
84    ) -> MuxerBuilder<E, S, R> {
85        let (control_tx, control_rx) = mpsc::unbounded_channel();
86        let mux = Self {
87            context: ContextCell::new(context),
88            sender,
89            receiver,
90            mailbox_size,
91            control_rx,
92            routes: HashMap::new(),
93            backup: None,
94        };
95
96        let mux_handle = MuxHandle {
97            sender: mux.sender.clone(),
98            control_tx,
99        };
100
101        MuxerBuilder { mux, mux_handle }
102    }
103
104    /// Start the demuxer using the given spawner.
105    pub fn start(mut self) -> Handle<Result<(), R::Error>> {
106        spawn_cell!(self.context, self.run())
107    }
108
109    /// Drive demultiplexing of messages into per-subchannel receivers.
110    ///
111    /// Callers should run this in a background task for as long as the underlying `Receiver` is
112    /// expected to receive traffic.
113    pub async fn run(mut self) -> Result<(), R::Error> {
114        select_loop! {
115            self.context,
116            on_stopped => {
117                debug!("context shutdown, stopping muxer");
118            },
119            // Prefer control messages because network messages will
120            // already block when full (providing backpressure).
121            Some(control) = self.control_rx.recv() else {
122                // If the control channel is closed, we can shut down since there must
123                // be no more registrations, and all receivers must have been dropped.
124                return Ok(());
125            } => match control {
126                Control::Register { subchannel, sender } => {
127                    // If the subchannel is already registered, drop the sender.
128                    if self.routes.contains_key(&subchannel) {
129                        continue;
130                    }
131
132                    // Otherwise, create a new subchannel and send the receiver to the caller.
133                    let (tx, rx) = mpsc::channel(self.mailbox_size);
134                    self.routes.insert(subchannel, tx);
135                    let _ = sender.send(rx);
136                }
137                Control::Deregister { subchannel } => {
138                    // Remove the route.
139                    self.routes.remove(&subchannel);
140                }
141            },
142            // Process network messages.
143            message = self.receiver.recv() => {
144                // Decode the message.
145                let (pk, bytes) = message?;
146                let (subchannel, bytes) = match parse(bytes) {
147                    Ok(parsed) => parsed,
148                    Err(_) => {
149                        debug!(?pk, "invalid message: missing subchannel");
150                        continue;
151                    }
152                };
153
154                // Get the route for the subchannel.
155                let Some(sender) = self.routes.get_mut(&subchannel) else {
156                    // Attempt to use the backup channel if available.
157                    if let Some(backup) = &mut self.backup
158                        && let Err(e) = backup.try_send((subchannel, (pk, bytes)))
159                    {
160                        debug!(?subchannel, ?e, "failed to send message to backup channel");
161                    }
162
163                    // Drops the message if the subchannel is not found or the backup
164                    // channel was not used.
165                    continue;
166                };
167
168                // Send the message to the subchannel using non-blocking try_send
169                // to avoid head-of-line blocking when one subchannel is slow.
170                if let Err(e) = sender.try_send((pk, bytes)) {
171                    // Check if the channel is disconnected (receiver dropped)
172                    if matches!(e, TrySendError::Closed(_)) {
173                        // Remove the route for the subchannel.
174                        self.routes.remove(&subchannel);
175                        debug!(?subchannel, "subchannel receiver dropped, removing route");
176                    } else {
177                        // Channel is full, drop the message
178                        debug!(?subchannel, "subchannel full, dropping message");
179                    }
180                }
181            },
182        }
183
184        Ok(())
185    }
186}
187
188/// A clonable handle that allows registering routes at any time, even after the [Muxer] is running.
189#[derive(Clone)]
190pub struct MuxHandle<S: Sender, R: Receiver> {
191    sender: S,
192    control_tx: mpsc::UnboundedSender<Control<R>>,
193}
194
195impl<S: Sender, R: Receiver> MuxHandle<S, R> {
196    /// Open a `subchannel`. Returns a ([SubSender], [SubReceiver]) pair that can be used to send
197    /// and receive messages for that subchannel.
198    ///
199    /// Panics if the subchannel is already registered at any point.
200    pub async fn register(
201        &mut self,
202        subchannel: Channel,
203    ) -> Result<(SubSender<S>, SubReceiver<R>), Error> {
204        let (tx, rx) = oneshot::channel();
205        self.control_tx
206            .send(Control::Register {
207                subchannel,
208                sender: tx,
209            })
210            .map_err(|_| Error::Closed)?;
211        let receiver = rx.await.map_err(|_| Error::AlreadyRegistered(subchannel))?;
212
213        Ok((
214            SubSender {
215                subchannel,
216                inner: GlobalSender::new(self.sender.clone()),
217            },
218            SubReceiver {
219                receiver,
220                control_tx: Some(self.control_tx.clone()),
221                subchannel,
222            },
223        ))
224    }
225}
226
227/// Sender that routes messages to the `subchannel`.
228#[derive(Clone, Debug)]
229pub struct SubSender<S: Sender> {
230    inner: GlobalSender<S>,
231    subchannel: Channel,
232}
233
234impl<S: Sender> LimitedSender for SubSender<S> {
235    type PublicKey = S::PublicKey;
236    type Checked<'a> = CheckedGlobalSender<'a, S>;
237
238    fn check(
239        &mut self,
240        recipients: Recipients<Self::PublicKey>,
241    ) -> Result<Self::Checked<'_>, SystemTime> {
242        self.inner
243            .check(recipients)
244            .map(|checked| checked.with_subchannel(self.subchannel))
245    }
246}
247
248/// Receiver that yields messages for a specific subchannel.
249pub struct SubReceiver<R: Receiver> {
250    receiver: mpsc::Receiver<Message<R::PublicKey>>,
251    control_tx: Option<mpsc::UnboundedSender<Control<R>>>,
252    subchannel: Channel,
253}
254
255impl<R: Receiver> Receiver for SubReceiver<R> {
256    type Error = Error;
257    type PublicKey = R::PublicKey;
258
259    async fn recv(&mut self) -> Result<Message<Self::PublicKey>, Self::Error> {
260        self.receiver.recv().await.ok_or(Error::RecvFailed)
261    }
262}
263
264impl<R: Receiver> Debug for SubReceiver<R> {
265    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
266        write!(f, "SubReceiver({})", self.subchannel)
267    }
268}
269
270impl<R: Receiver> Drop for SubReceiver<R> {
271    fn drop(&mut self) {
272        // Take the control channel to avoid cloning.
273        let control_tx = self
274            .control_tx
275            .take()
276            .expect("SubReceiver::drop called twice");
277
278        // Deregister the subchannel immediately.
279        control_tx.send_lossy(Control::Deregister {
280            subchannel: self.subchannel,
281        });
282    }
283}
284
285/// Sender that can send messages over any sub [Channel].
286#[derive(Clone, Debug)]
287pub struct GlobalSender<S: Sender> {
288    inner: S,
289}
290
291impl<S: Sender> GlobalSender<S> {
292    /// Create a new [GlobalSender] wrapping the given [Sender].
293    pub const fn new(inner: S) -> Self {
294        Self { inner }
295    }
296
297    /// Send a message over the given `subchannel`.
298    pub fn send(
299        &mut self,
300        subchannel: Channel,
301        recipients: Recipients<S::PublicKey>,
302        payload: impl Into<IoBufs> + Send,
303        priority: bool,
304    ) -> Unreliable<Feedback> {
305        self.check(recipients).map_or_else(
306            |_| Unreliable::Rejected,
307            |checked| checked.with_subchannel(subchannel).send(payload, priority),
308        )
309    }
310}
311
312impl<S: Sender> LimitedSender for GlobalSender<S> {
313    type PublicKey = S::PublicKey;
314    type Checked<'a> = CheckedGlobalSender<'a, S>;
315
316    fn check(
317        &mut self,
318        recipients: Recipients<Self::PublicKey>,
319    ) -> Result<Self::Checked<'_>, SystemTime> {
320        self.inner
321            .check(recipients)
322            .map(|checked| CheckedGlobalSender {
323                subchannel: None,
324                inner: checked,
325            })
326    }
327}
328
329/// A checked sender for a [GlobalSender].
330pub struct CheckedGlobalSender<'a, S: Sender> {
331    subchannel: Option<Channel>,
332    inner: S::Checked<'a>,
333}
334
335impl<'a, S: Sender> CheckedGlobalSender<'a, S> {
336    /// Set the subchannel for this sender.
337    pub const fn with_subchannel(mut self, subchannel: Channel) -> Self {
338        self.subchannel = Some(subchannel);
339        self
340    }
341}
342
343impl<'a, S: Sender> CheckedSender for CheckedGlobalSender<'a, S> {
344    type PublicKey = S::PublicKey;
345
346    fn recipients(&self) -> Vec<Self::PublicKey> {
347        self.inner.recipients()
348    }
349
350    fn send(self, message: impl Into<IoBufs> + Send, priority: bool) -> Unreliable<Feedback> {
351        let subchannel = UInt(self.subchannel.expect("subchannel not set"));
352        let mut message = message.into();
353        message.prepend(subchannel.encode().into());
354        self.inner.send(message, priority)
355    }
356}
357
358/// A generic builder interface.
359pub trait Builder {
360    /// The output type produced by the builder.
361    type Output;
362
363    /// Builds the output type, consuming `self`.
364    fn build(self) -> Self::Output;
365}
366
367/// A builder that constructs a [Muxer].
368pub struct MuxerBuilder<E: Spawner, S: Sender, R: Receiver> {
369    mux: Muxer<E, S, R>,
370    mux_handle: MuxHandle<S, R>,
371}
372
373impl<E: Spawner, S: Sender, R: Receiver> Builder for MuxerBuilder<E, S, R> {
374    type Output = (Muxer<E, S, R>, MuxHandle<S, R>);
375
376    fn build(self) -> Self::Output {
377        (self.mux, self.mux_handle)
378    }
379}
380
381impl<E: Spawner, S: Sender, R: Receiver> MuxerBuilder<E, S, R> {
382    /// Registers a backup channel with the muxer.
383    pub fn with_backup(mut self) -> MuxerBuilderWithBackup<E, S, R> {
384        let (tx, rx) = mpsc::channel(self.mux.mailbox_size);
385        self.mux.backup = Some(tx);
386
387        MuxerBuilderWithBackup {
388            mux: self.mux,
389            mux_handle: self.mux_handle,
390            backup_rx: rx,
391        }
392    }
393
394    /// Registers a global sender with the muxer.
395    pub fn with_global_sender(self) -> MuxerBuilderWithGlobalSender<E, S, R> {
396        let global_sender = GlobalSender::new(self.mux.sender.clone());
397
398        MuxerBuilderWithGlobalSender {
399            mux: self.mux,
400            mux_handle: self.mux_handle,
401            global_sender,
402        }
403    }
404}
405
406/// A builder that constructs a [Muxer] with a backup channel.
407pub struct MuxerBuilderWithBackup<E: Spawner, S: Sender, R: Receiver> {
408    mux: Muxer<E, S, R>,
409    mux_handle: MuxHandle<S, R>,
410    backup_rx: mpsc::Receiver<BackupResponse<R::PublicKey>>,
411}
412
413impl<E: Spawner, S: Sender, R: Receiver> MuxerBuilderWithBackup<E, S, R> {
414    /// Registers a global sender with the muxer.
415    pub fn with_global_sender(self) -> MuxerBuilderAllOpts<E, S, R> {
416        let global_sender = GlobalSender::new(self.mux.sender.clone());
417
418        MuxerBuilderAllOpts {
419            mux: self.mux,
420            mux_handle: self.mux_handle,
421            backup_rx: self.backup_rx,
422            global_sender,
423        }
424    }
425}
426
427impl<E: Spawner, S: Sender, R: Receiver> Builder for MuxerBuilderWithBackup<E, S, R> {
428    type Output = (
429        Muxer<E, S, R>,
430        MuxHandle<S, R>,
431        mpsc::Receiver<BackupResponse<R::PublicKey>>,
432    );
433
434    fn build(self) -> Self::Output {
435        (self.mux, self.mux_handle, self.backup_rx)
436    }
437}
438
439/// A builder that constructs a [Muxer] with a [GlobalSender].
440pub struct MuxerBuilderWithGlobalSender<E: Spawner, S: Sender, R: Receiver> {
441    mux: Muxer<E, S, R>,
442    mux_handle: MuxHandle<S, R>,
443    global_sender: GlobalSender<S>,
444}
445
446impl<E: Spawner, S: Sender, R: Receiver> MuxerBuilderWithGlobalSender<E, S, R> {
447    /// Registers a backup channel with the muxer.
448    pub fn with_backup(mut self) -> MuxerBuilderAllOpts<E, S, R> {
449        let (tx, rx) = mpsc::channel(self.mux.mailbox_size);
450        self.mux.backup = Some(tx);
451
452        MuxerBuilderAllOpts {
453            mux: self.mux,
454            mux_handle: self.mux_handle,
455            backup_rx: rx,
456            global_sender: self.global_sender,
457        }
458    }
459}
460
461impl<E: Spawner, S: Sender, R: Receiver> Builder for MuxerBuilderWithGlobalSender<E, S, R> {
462    type Output = (Muxer<E, S, R>, MuxHandle<S, R>, GlobalSender<S>);
463
464    fn build(self) -> Self::Output {
465        (self.mux, self.mux_handle, self.global_sender)
466    }
467}
468
469/// A builder that constructs a [Muxer] with a [GlobalSender] and backup channel.
470pub struct MuxerBuilderAllOpts<E: Spawner, S: Sender, R: Receiver> {
471    mux: Muxer<E, S, R>,
472    mux_handle: MuxHandle<S, R>,
473    backup_rx: mpsc::Receiver<BackupResponse<R::PublicKey>>,
474    global_sender: GlobalSender<S>,
475}
476
477impl<E: Spawner, S: Sender, R: Receiver> Builder for MuxerBuilderAllOpts<E, S, R> {
478    type Output = (
479        Muxer<E, S, R>,
480        MuxHandle<S, R>,
481        mpsc::Receiver<BackupResponse<R::PublicKey>>,
482        GlobalSender<S>,
483    );
484
485    fn build(self) -> Self::Output {
486        (
487            self.mux,
488            self.mux_handle,
489            self.backup_rx,
490            self.global_sender,
491        )
492    }
493}
494
495#[cfg(test)]
496mod tests {
497    use super::*;
498    use crate::{
499        Manager as _, Provider as _, Recipients,
500        simulated::{self, Link, Network, Oracle},
501    };
502    use commonware_cryptography::{
503        Signer,
504        ed25519::{PrivateKey, PublicKey},
505    };
506    use commonware_macros::{select, test_traced};
507    use commonware_runtime::{IoBuf, Quota, Runner, Supervisor as _, deterministic};
508    use commonware_utils::{NZUsize, ordered::Set, probability};
509    use std::{
510        num::NonZeroU32,
511        time::{Duration, SystemTime},
512    };
513
514    const LINK: Link = Link {
515        latency: Duration::from_millis(0),
516        jitter: Duration::from_millis(0),
517        success_rate: probability!(1.0),
518    };
519    const CAPACITY: usize = 5usize;
520
521    /// Default rate limit set high enough to not interfere with normal operation
522    const TEST_QUOTA: Quota = Quota::per_second(NonZeroU32::MAX);
523
524    /// Start the network and return the oracle.
525    fn start_network(context: deterministic::Context) -> Oracle<PublicKey, deterministic::Context> {
526        let (network, oracle) = Network::new(
527            context.child("network"),
528            simulated::Config {
529                max_size: 1024 * 1024,
530                max_peers_per_set: NZUsize!(2),
531                disconnect_on_block: true,
532                tracked_peer_sets: NZUsize!(1),
533            },
534        );
535        network.start();
536        oracle
537    }
538
539    /// Create a public key from a seed.
540    fn pk(seed: u64) -> PublicKey {
541        PrivateKey::from_seed(seed).public_key()
542    }
543
544    /// Link two peers bidirectionally.
545    async fn link_bidirectional(
546        oracle: &mut Oracle<PublicKey, deterministic::Context>,
547        a: PublicKey,
548        b: PublicKey,
549    ) {
550        let mut manager = oracle.manager();
551        let peers = manager.peer_set(0).await.unwrap_or_default();
552        manager.track(
553            0,
554            Set::from_iter_dedup(peers.primary.iter().cloned().chain([a.clone(), b.clone()])),
555        );
556        oracle.add_link(a.clone(), b.clone(), LINK).await.unwrap();
557        oracle.add_link(b, a, LINK).await.unwrap();
558    }
559
560    /// Create a peer and register it with the oracle.
561    async fn create_peer(
562        context: &deterministic::Context,
563        oracle: &mut Oracle<PublicKey, deterministic::Context>,
564        seed: u64,
565    ) -> (
566        PublicKey,
567        MuxHandle<
568            impl Sender<PublicKey = PublicKey> + use<>,
569            impl Receiver<PublicKey = PublicKey> + use<>,
570        >,
571    ) {
572        let pubkey = pk(seed);
573        let (sender, receiver) = oracle
574            .control(pubkey.clone())
575            .register(0, TEST_QUOTA)
576            .await
577            .unwrap();
578        let (mux, handle) = Muxer::new(context.child("mux"), sender, receiver, CAPACITY);
579        mux.start();
580        (pubkey, handle)
581    }
582
583    /// Create a peer and register it with the oracle.
584    async fn create_peer_with_backup_and_global_sender(
585        context: &deterministic::Context,
586        oracle: &mut Oracle<PublicKey, deterministic::Context>,
587        seed: u64,
588    ) -> (
589        PublicKey,
590        MuxHandle<
591            impl Sender<PublicKey = PublicKey> + use<>,
592            impl Receiver<PublicKey = PublicKey> + use<>,
593        >,
594        mpsc::Receiver<BackupResponse<PublicKey>>,
595        GlobalSender<simulated::Sender<PublicKey, deterministic::Context>>,
596    ) {
597        let pubkey = pk(seed);
598        let (sender, receiver) = oracle
599            .control(pubkey.clone())
600            .register(0, TEST_QUOTA)
601            .await
602            .unwrap();
603        let (mux, handle, backup, global_sender) =
604            Muxer::builder(context.child("mux"), sender, receiver, CAPACITY)
605                .with_backup()
606                .with_global_sender()
607                .build();
608        mux.start();
609        (pubkey, handle, backup, global_sender)
610    }
611
612    /// Send a burst of messages to a list of senders.
613    fn send_burst<S: Sender>(txs: &mut [SubSender<S>], count: usize) {
614        for i in 0..count {
615            let payload = IoBuf::from(vec![i as u8]);
616            for tx in txs.iter_mut() {
617                tx.send(Recipients::All, payload.clone(), false);
618            }
619        }
620    }
621
622    /// Wait for `n` messages to be received on the receiver.
623    async fn expect_n_messages(
624        rx: &mut SubReceiver<impl Receiver<PublicKey = PublicKey>>,
625        n: usize,
626    ) {
627        let mut count = 0;
628        loop {
629            select! {
630                res = rx.recv() => {
631                    res.expect("should have received message");
632                    count += 1;
633                },
634            }
635
636            if count >= n {
637                break;
638            }
639        }
640        assert_eq!(n, count);
641    }
642
643    /// Wait for `n` messages to be received on the receiver + backup receiver.
644    async fn expect_n_messages_with_backup(
645        rx: &mut SubReceiver<impl Receiver<PublicKey = PublicKey>>,
646        backup_rx: &mut mpsc::Receiver<BackupResponse<PublicKey>>,
647        n: usize,
648        n_backup: usize,
649    ) {
650        let mut count_std = 0;
651        let mut count_backup = 0;
652        loop {
653            select! {
654                res = rx.recv() => {
655                    res.expect("should have received message");
656                    count_std += 1;
657                },
658                res = backup_rx.recv() => {
659                    res.expect("should have received message");
660                    count_backup += 1;
661                },
662            }
663
664            if count_std >= n && count_backup >= n_backup {
665                break;
666            }
667        }
668        assert_eq!(n, count_std);
669        assert_eq!(n_backup, count_backup);
670    }
671
672    #[derive(Clone)]
673    struct RateLimitedSender;
674
675    struct UnusedCheckedSender;
676
677    impl CheckedSender for UnusedCheckedSender {
678        type PublicKey = PublicKey;
679
680        fn recipients(&self) -> Vec<Self::PublicKey> {
681            unreachable!("rate-limited sender should not produce a checked sender");
682        }
683
684        fn send(self, _: impl Into<IoBufs> + Send, _: bool) -> Unreliable<Feedback> {
685            unreachable!("rate-limited sender should not send");
686        }
687    }
688
689    impl LimitedSender for RateLimitedSender {
690        type PublicKey = PublicKey;
691        type Checked<'a> = UnusedCheckedSender;
692
693        fn check(
694            &mut self,
695            _: Recipients<Self::PublicKey>,
696        ) -> Result<Self::Checked<'_>, SystemTime> {
697            Err(SystemTime::UNIX_EPOCH)
698        }
699    }
700
701    #[test]
702    fn test_global_sender_rate_limited_send_rejected() {
703        let mut sender = GlobalSender::new(RateLimitedSender);
704        let feedback = sender.send(0, Recipients::One(pk(0)), b"rate-limited", false);
705        assert_eq!(feedback, Unreliable::Rejected);
706        assert!(!feedback.accepted());
707    }
708
709    #[test]
710    fn test_basic_routing() {
711        // Can register a subchannel and send messages to it.
712        let executor = deterministic::Runner::default();
713        executor.start(|context| async move {
714            let mut oracle = start_network(context.child("network"));
715
716            let (pk1, mut handle1) = create_peer(&context, &mut oracle, 0).await;
717            let (pk2, mut handle2) = create_peer(&context, &mut oracle, 1).await;
718            link_bidirectional(&mut oracle, pk1.clone(), pk2.clone()).await;
719
720            let (_, mut sub_rx1) = handle1.register(7).await.unwrap();
721            let (mut sub_tx2, _) = handle2.register(7).await.unwrap();
722
723            // Send and receive
724            let payload = IoBuf::from(b"hello");
725            let _ = sub_tx2.send(Recipients::One(pk1.clone()), payload.clone(), false);
726            let (from, bytes) = sub_rx1.recv().await.unwrap();
727            assert_eq!(from, pk2);
728            assert_eq!(bytes, payload);
729        });
730    }
731
732    #[test]
733    fn test_multiple_routes() {
734        // Can register multiple subchannels and send messages to each.
735        let executor = deterministic::Runner::default();
736        executor.start(|context| async move {
737            let mut oracle = start_network(context.child("network"));
738
739            let (pk1, mut handle1) = create_peer(&context, &mut oracle, 0).await;
740            let (pk2, mut handle2) = create_peer(&context, &mut oracle, 1).await;
741            link_bidirectional(&mut oracle, pk1.clone(), pk2.clone()).await;
742
743            let (_, mut rx_a) = handle1.register(10).await.unwrap();
744            let (_, mut rx_b) = handle1.register(20).await.unwrap();
745
746            let (mut tx2_a, _) = handle2.register(10).await.unwrap();
747            let (mut tx2_b, _) = handle2.register(20).await.unwrap();
748
749            let payload_a = IoBuf::from(b"A");
750            let payload_b = IoBuf::from(b"B");
751            let _ = tx2_a.send(Recipients::One(pk1.clone()), payload_a.clone(), false);
752            let _ = tx2_b.send(Recipients::One(pk1.clone()), payload_b.clone(), false);
753
754            let (from_a, bytes_a) = rx_a.recv().await.unwrap();
755            assert_eq!(from_a, pk2);
756            assert_eq!(bytes_a, payload_a);
757
758            let (from_b, bytes_b) = rx_b.recv().await.unwrap();
759            assert_eq!(from_b, pk2);
760            assert_eq!(bytes_b, payload_b);
761        });
762    }
763
764    #[test_traced]
765    fn test_mailbox_capacity_drops_when_full() {
766        // Messages are dropped (not blocked) when a subchannel buffer is full.
767        // This prevents head-of-line blocking where one slow subchannel blocks all others.
768        let executor = deterministic::Runner::default();
769        executor.start(|context| async move {
770            let mut oracle = start_network(context.child("network"));
771
772            let (pk1, mut handle1) = create_peer(&context, &mut oracle, 0).await;
773            let (pk2, mut handle2) = create_peer(&context, &mut oracle, 1).await;
774            link_bidirectional(&mut oracle, pk1.clone(), pk2.clone()).await;
775
776            // Register the subchannels.
777            let (tx1, _) = handle1.register(99).await.unwrap();
778            let (tx2, _) = handle1.register(100).await.unwrap();
779            let (_, mut rx1) = handle2.register(99).await.unwrap();
780            let (_, mut rx2) = handle2.register(100).await.unwrap();
781
782            // Send 10 messages to each subchannel from pk1 to pk2.
783            // With buffer size of CAPACITY=5, messages beyond that are dropped.
784            send_burst(&mut [tx1, tx2], CAPACITY * 2);
785
786            // Each subchannel should receive up to CAPACITY messages (the rest are dropped).
787            expect_n_messages(&mut rx1, CAPACITY).await;
788            expect_n_messages(&mut rx2, CAPACITY).await;
789        });
790    }
791
792    #[test]
793    fn test_drop_subchannel_receiver_deregisters_route() {
794        // Dropping a subchannel receiver deregisters the route, and subsequent
795        // messages to that subchannel are dropped.
796        let executor = deterministic::Runner::default();
797        executor.start(|context| async move {
798            let mut oracle = start_network(context.child("network"));
799
800            let (pk1, mut handle1) = create_peer(&context, &mut oracle, 0).await;
801            let (pk2, mut handle2) = create_peer(&context, &mut oracle, 1).await;
802            link_bidirectional(&mut oracle, pk1.clone(), pk2.clone()).await;
803
804            // Register the subchannels.
805            let (tx1, _) = handle1.register(99).await.unwrap();
806            let (tx2, _) = handle1.register(100).await.unwrap();
807            let (_, rx1) = handle2.register(99).await.unwrap();
808            let (_, mut rx2) = handle2.register(100).await.unwrap();
809
810            // Drop rx1 before any messages are sent - its route is now deregistered.
811            drop(rx1);
812
813            // Send messages to both subchannels. Messages to subchannel 99 will be dropped
814            // since its receiver was dropped.
815            send_burst(&mut [tx1, tx2], CAPACITY);
816
817            // rx2 should receive all CAPACITY messages sent to subchannel 100.
818            expect_n_messages(&mut rx2, CAPACITY).await;
819        });
820    }
821
822    #[test]
823    fn test_drop_messages_for_unregistered_subchannel() {
824        // Messages are dropped if the subchannel they are for is not registered.
825        // The unregistered subchannel does not affect the registered one.
826        let executor = deterministic::Runner::default();
827        executor.start(|context| async move {
828            let mut oracle = start_network(context.child("network"));
829
830            let (pk1, mut handle1) = create_peer(&context, &mut oracle, 0).await;
831            let (pk2, mut handle2) = create_peer(&context, &mut oracle, 1).await;
832            link_bidirectional(&mut oracle, pk1.clone(), pk2.clone()).await;
833
834            // Register the subchannels.
835            let (tx1, _) = handle1.register(1).await.unwrap();
836            let (tx2, _) = handle1.register(2).await.unwrap();
837            // Do not register the first subchannel on the second peer.
838            let (_, mut rx2) = handle2.register(2).await.unwrap();
839
840            // Send CAPACITY messages to each subchannel.
841            // Messages to subchannel 1 are dropped (unregistered).
842            // Messages to subchannel 2 fill the buffer.
843            send_burst(&mut [tx1, tx2], CAPACITY);
844
845            // Receive messages from subchannel 2.
846            expect_n_messages(&mut rx2, CAPACITY).await;
847        });
848    }
849
850    #[test]
851    fn test_backup_for_unregistered_subchannel() {
852        // Messages are forwarded to the backup channel if the subchannel they are for
853        // is not registered.
854        let executor = deterministic::Runner::default();
855        executor.start(|context| async move {
856            let mut oracle = start_network(context.child("network"));
857
858            let (pk1, mut handle1) = create_peer(&context, &mut oracle, 0).await;
859            let (pk2, mut handle2, mut backup2, _) =
860                create_peer_with_backup_and_global_sender(&context, &mut oracle, 1).await;
861            link_bidirectional(&mut oracle, pk1.clone(), pk2.clone()).await;
862
863            // Register the subchannels.
864            let (tx1, _) = handle1.register(1).await.unwrap();
865            let (tx2, _) = handle1.register(2).await.unwrap();
866            // Do not register the first subchannel on the second peer.
867            let (_, mut rx2) = handle2.register(2).await.unwrap();
868
869            // Send CAPACITY messages to each subchannel.
870            // Subchannel 1 messages go to backup, subchannel 2 messages go to rx2.
871            send_burst(&mut [tx1, tx2], CAPACITY);
872
873            // Both channels should receive CAPACITY messages each.
874            expect_n_messages_with_backup(&mut rx2, &mut backup2, CAPACITY, CAPACITY).await;
875        });
876    }
877
878    #[test]
879    fn test_backup_for_unregistered_subchannel_response() {
880        // Messages are forwarded to the backup channel if the subchannel they are for
881        // is not registered.
882        let executor = deterministic::Runner::default();
883        executor.start(|context| async move {
884            let mut oracle = start_network(context.child("network"));
885
886            let (pk1, mut handle1) = create_peer(&context, &mut oracle, 0).await;
887            let (pk2, _handle2, mut backup2, mut global_sender2) =
888                create_peer_with_backup_and_global_sender(&context, &mut oracle, 1).await;
889            link_bidirectional(&mut oracle, pk1.clone(), pk2.clone()).await;
890
891            // Register the subchannels.
892            let (mut tx1, mut rx1) = handle1.register(1).await.unwrap();
893            // Do not register any subchannels on the second peer.
894
895            // Send a message from pk1 to pk2.
896            tx1.send(Recipients::One(pk2.clone()), b"REQUEST", false);
897
898            // Get the message from pk2's backup channel and respond.
899            let (subchannel, (from, _)) = backup2.recv().await.unwrap();
900            assert_eq!(subchannel, 1);
901            assert_eq!(from, pk1);
902            let checked = global_sender2.check(Recipients::One(pk1.clone())).unwrap();
903            assert_eq!(checked.recipients(), vec![pk1]);
904            checked.with_subchannel(subchannel).send(b"TEST", true);
905
906            // Receive the response with pk1's receiver.
907            let (from, bytes) = rx1.recv().await.unwrap();
908            assert_eq!(from, pk2);
909            assert_eq!(bytes, b"TEST");
910        });
911    }
912
913    #[test]
914    fn test_message_dropped_for_closed_subchannel() {
915        // Messages are dropped if the subchannel they are for is registered, but has been closed.
916        //
917        // NOTE: This case should be exceedingly rare in practice due to `SubReceiver` deregistering
918        // the subchannel on drop, but is included for completeness.
919        let executor = deterministic::Runner::default();
920        executor.start(|context| async move {
921            let mut oracle = start_network(context.child("network"));
922
923            let (pk1, mut handle1) = create_peer(&context, &mut oracle, 0).await;
924            let (pk2, mut handle2) = create_peer(&context, &mut oracle, 1).await;
925            link_bidirectional(&mut oracle, pk1.clone(), pk2.clone()).await;
926
927            // Register the subchannels.
928            let (mut tx1, _) = handle1.register(1).await.unwrap();
929            let (mut tx2, _) = handle1.register(2).await.unwrap();
930            let (_, rx1) = handle2.register(1).await.unwrap();
931            let (_, mut rx2) = handle2.register(2).await.unwrap();
932
933            // Explicitly close the first subchannel.
934            drop(rx1);
935
936            // Send to the closed subchannel, then verify it does not block
937            // messages to another subchannel.
938            tx1.send(Recipients::One(pk2.clone()), b"closed", false);
939            tx2.send(Recipients::One(pk2.clone()), b"open", false);
940
941            // Subchannel 2 should still receive messages.
942            expect_n_messages(&mut rx2, 1).await;
943        });
944    }
945
946    #[test]
947    fn test_dropped_backup_channel_doesnt_block() {
948        // Dropping the backup receiver doesn't block message processing.
949        // Messages to unregistered subchannels are simply dropped.
950        let executor = deterministic::Runner::default();
951        executor.start(|context| async move {
952            let mut oracle = start_network(context.child("network"));
953
954            let (pk1, mut handle1) = create_peer(&context, &mut oracle, 0).await;
955            let (pk2, mut handle2, backup2, _) =
956                create_peer_with_backup_and_global_sender(&context, &mut oracle, 1).await;
957            link_bidirectional(&mut oracle, pk1.clone(), pk2.clone()).await;
958
959            // Explicitly drop the backup receiver.
960            drop(backup2);
961
962            // Register the subchannels.
963            let (tx1, _) = handle1.register(1).await.unwrap();
964            let (tx2, _) = handle1.register(2).await.unwrap();
965            // Do not register the first subchannel on the second peer.
966            let (_, mut rx2) = handle2.register(2).await.unwrap();
967
968            // Send CAPACITY messages to each subchannel.
969            // Subchannel 1 messages are dropped (backup is closed).
970            // Subchannel 2 messages go to rx2.
971            send_burst(&mut [tx1, tx2], CAPACITY);
972
973            // rx2 should receive all CAPACITY messages.
974            expect_n_messages(&mut rx2, CAPACITY).await;
975        });
976    }
977
978    #[test]
979    fn test_duplicate_registration() {
980        // Returns an error if the subchannel is already registered.
981        let executor = deterministic::Runner::default();
982        executor.start(|context| async move {
983            let mut oracle = start_network(context.child("network"));
984
985            let (_pk1, mut handle1) = create_peer(&context, &mut oracle, 0).await;
986
987            // Register the subchannel.
988            let (_, _rx) = handle1.register(7).await.unwrap();
989
990            // Registering again should return an error.
991            assert!(matches!(
992                handle1.register(7).await,
993                Err(Error::AlreadyRegistered(_))
994            ));
995        });
996    }
997
998    #[test]
999    fn test_register_after_deregister() {
1000        // Can register a channel after it has been deregistered.
1001        let executor = deterministic::Runner::default();
1002        executor.start(|context| async move {
1003            let mut oracle = start_network(context.child("network"));
1004
1005            let (_, mut handle) = create_peer(&context, &mut oracle, 0).await;
1006            let (_, rx) = handle.register(7).await.unwrap();
1007            drop(rx);
1008
1009            // Registering again should not return an error.
1010            handle.register(7).await.unwrap();
1011        });
1012    }
1013}
commonware_p2p/utils/mux.rs

commonware_p2p/utils/
mux.rs
