1use crate::{Feedback, Unreliable};
44use commonware_runtime::{
45 Metrics,
46 telemetry::metrics::{Counter, MetricsExt as _},
47};
48use std::{
49 collections::VecDeque,
50 fmt,
51 future::poll_fn,
52 marker::PhantomData,
53 num::NonZeroUsize,
54 sync::mpsc::TryRecvError,
55 task::{Context, Poll},
56};
57
58pub trait Overflow<T>: Default {
60 fn is_empty(&self) -> bool;
62
63 fn drain<F>(&mut self, push: F)
69 where
70 F: FnMut(T) -> Option<T>;
71}
72
73impl<T> Overflow<T> for VecDeque<T> {
74 fn is_empty(&self) -> bool {
75 self.is_empty()
76 }
77
78 fn drain<F>(&mut self, mut push: F)
79 where
80 F: FnMut(T) -> Option<T>,
81 {
82 while let Some(message) = self.pop_front() {
83 if let Some(message) = push(message) {
84 self.push_front(message);
85 break;
86 }
87 }
88 }
89}
90
91pub trait Policy: Sized {
93 type Overflow: Overflow<Self>;
95
96 fn handle(overflow: &mut Self::Overflow, message: Self);
112}
113
114pub trait UnreliablePolicy: Sized {
116 type Overflow: Overflow<Self>;
118
119 fn handle(overflow: &mut Self::Overflow, message: Self) -> bool;
140}
141
142mod mode {
144 pub(super) struct Reliable;
146
147 pub(super) struct Unreliable;
149}
150
151trait Mode<T>: Sized {
152 type Overflow: Overflow<T>;
154 type Feedback;
156
157 fn handle(overflow: &mut Self::Overflow, message: T) -> bool;
159 fn ready_feedback(feedback: Feedback) -> Self::Feedback;
161 fn overflow_feedback(handled: bool) -> Self::Feedback;
163 fn is_backoff(feedback: &Self::Feedback) -> bool;
165 fn is_closed(feedback: &Self::Feedback) -> bool;
167}
168
169impl<T: Policy> Mode<T> for mode::Reliable {
170 type Overflow = T::Overflow;
171 type Feedback = Feedback;
172
173 fn handle(overflow: &mut Self::Overflow, message: T) -> bool {
174 T::handle(overflow, message);
175 true
176 }
177
178 fn ready_feedback(feedback: Feedback) -> Self::Feedback {
179 feedback
180 }
181
182 fn overflow_feedback(_handled: bool) -> Self::Feedback {
183 Feedback::Backoff
184 }
185
186 fn is_backoff(feedback: &Self::Feedback) -> bool {
187 *feedback == Feedback::Backoff
188 }
189
190 fn is_closed(feedback: &Self::Feedback) -> bool {
191 *feedback == Feedback::Closed
192 }
193}
194
195impl<T: UnreliablePolicy> Mode<T> for mode::Unreliable {
196 type Overflow = T::Overflow;
197 type Feedback = Unreliable<Feedback>;
198
199 fn handle(overflow: &mut Self::Overflow, message: T) -> bool {
200 T::handle(overflow, message)
201 }
202
203 fn ready_feedback(feedback: Feedback) -> Self::Feedback {
204 Unreliable::new(feedback)
205 }
206
207 fn overflow_feedback(handled: bool) -> Self::Feedback {
208 if handled {
209 Unreliable::new(Feedback::Backoff)
210 } else {
211 Unreliable::Rejected
212 }
213 }
214
215 fn is_backoff(feedback: &Self::Feedback) -> bool {
216 *feedback == Unreliable::new(Feedback::Backoff)
217 }
218
219 fn is_closed(feedback: &Self::Feedback) -> bool {
220 *feedback == Unreliable::new(Feedback::Closed)
221 }
222}
223
224pub struct Sender<T: Policy> {
226 state: Arc<State<T, mode::Reliable>>,
227}
228
229pub struct UnreliableSender<T: UnreliablePolicy> {
231 state: Arc<State<T, mode::Unreliable>>,
232}
233
234impl<T: Policy> Clone for Sender<T> {
235 fn clone(&self) -> Self {
236 Self {
237 state: clone_sender_state(&self.state),
238 }
239 }
240}
241
242impl<T: UnreliablePolicy> Clone for UnreliableSender<T> {
243 fn clone(&self) -> Self {
244 Self {
245 state: clone_sender_state(&self.state),
246 }
247 }
248}
249
250impl<T: Policy> Drop for Sender<T> {
251 fn drop(&mut self) {
252 drop_sender_state(&self.state);
253 }
254}
255
256impl<T: UnreliablePolicy> Drop for UnreliableSender<T> {
257 fn drop(&mut self) {
258 drop_sender_state(&self.state);
259 }
260}
261
262impl<T: Policy> fmt::Debug for Sender<T> {
263 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
264 fmt_sender_state("Sender", &self.state, f)
265 }
266}
267
268impl<T: UnreliablePolicy> fmt::Debug for UnreliableSender<T> {
269 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
270 fmt_sender_state("UnreliableSender", &self.state, f)
271 }
272}
273
274impl<T: Policy> Sender<T> {
275 #[must_use = "caller must handle enqueue feedback"]
277 pub fn enqueue(&self, message: T) -> Feedback {
278 self.state.enqueue(message)
279 }
280}
281
282impl<T: UnreliablePolicy> UnreliableSender<T> {
283 #[must_use = "caller must handle enqueue feedback"]
285 pub fn enqueue(&self, message: T) -> Unreliable<Feedback> {
286 self.state.enqueue(message)
287 }
288}
289
290pub struct Receiver<T: Policy> {
297 state: Arc<State<T, mode::Reliable>>,
298}
299
300pub struct UnreliableReceiver<T: UnreliablePolicy> {
307 state: Arc<State<T, mode::Unreliable>>,
308}
309
310impl<T: Policy> Receiver<T> {
311 pub async fn recv(&mut self) -> Option<T> {
316 recv_from(&self.state).await
317 }
318
319 pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
324 try_recv_from(&self.state)
325 }
326}
327
328impl<T: UnreliablePolicy> UnreliableReceiver<T> {
329 pub async fn recv(&mut self) -> Option<T> {
334 recv_from(&self.state).await
335 }
336
337 pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
342 try_recv_from(&self.state)
343 }
344}
345
346impl<T: Policy> Drop for Receiver<T> {
347 fn drop(&mut self) {
348 self.state.close();
349 }
350}
351
352impl<T: UnreliablePolicy> Drop for UnreliableReceiver<T> {
353 fn drop(&mut self) {
354 self.state.close();
355 }
356}
357
358pub fn new<T: Policy>(metrics: impl Metrics, capacity: NonZeroUsize) -> (Sender<T>, Receiver<T>) {
360 let state = new_state(metrics, capacity);
361 (
362 Sender {
363 state: state.clone(),
364 },
365 Receiver { state },
366 )
367}
368
369pub fn new_unreliable<T: UnreliablePolicy>(
371 metrics: impl Metrics,
372 capacity: NonZeroUsize,
373) -> (UnreliableSender<T>, UnreliableReceiver<T>) {
374 let state = new_state(metrics, capacity);
375 (
376 UnreliableSender {
377 state: state.clone(),
378 },
379 UnreliableReceiver { state },
380 )
381}
382
383const OVERFLOW_HAS_MESSAGES: usize = 1;
409const OVERFLOW_MUTATION: usize = 2;
410
411cfg_if::cfg_if! {
412 if #[cfg(feature = "loom")] {
413 use loom::{
414 future::AtomicWaker,
415 sync::{
416 Arc, Mutex, MutexGuard,
417 atomic::{AtomicBool, AtomicUsize, Ordering},
418 },
419 };
420
421 fn register_waker(waker: &AtomicWaker, task: &std::task::Waker) {
422 waker.register_by_ref(task);
423 }
424
425 fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
426 mutex.lock().unwrap()
427 }
428
429 struct ReadyState<T> {
430 published: VecDeque<T>,
431 reserved: usize,
432 }
433
434 struct Ready<T> {
435 state: Mutex<ReadyState<T>>,
436 capacity: usize,
437 }
438
439 impl<T> Ready<T> {
440 fn new(capacity: usize) -> Self {
441 Self {
442 state: Mutex::new(ReadyState {
443 published: VecDeque::new(),
444 reserved: 0,
445 }),
446 capacity,
447 }
448 }
449
450 const fn capacity(&self) -> usize {
451 self.capacity
452 }
453
454 fn push(&self, message: T) -> Result<(), T> {
455 {
456 let mut state = lock(&self.state);
457 if state.published.len() + state.reserved >= self.capacity {
458 return Err(message);
459 }
460 state.reserved += 1;
461 }
462
463 loom::thread::yield_now();
464
465 let mut state = lock(&self.state);
466 state.reserved -= 1;
467 state.published.push_back(message);
468 Ok(())
469 }
470
471 fn pop(&self) -> Option<T> {
472 loop {
473 let mut state = lock(&self.state);
474 if let Some(message) = state.published.pop_front() {
475 return Some(message);
476 }
477 if state.reserved == 0 {
478 return None;
479 }
480 drop(state);
481 loom::thread::yield_now();
482 }
483 }
484 }
485 } else {
486 use crossbeam_queue::ArrayQueue;
487 use futures_util::task::AtomicWaker;
488 use parking_lot::{Mutex, MutexGuard};
489 use std::sync::{
490 Arc,
491 atomic::{AtomicBool, AtomicUsize, Ordering},
492 };
493
494 fn register_waker(waker: &AtomicWaker, task: &std::task::Waker) {
495 waker.register(task);
496 }
497
498 fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
499 mutex.lock()
500 }
501
502 struct Ready<T> {
503 queue: ArrayQueue<T>,
504 }
505
506 impl<T> Ready<T> {
507 fn new(capacity: usize) -> Self {
508 Self {
509 queue: ArrayQueue::new(capacity),
510 }
511 }
512
513 fn capacity(&self) -> usize {
514 self.queue.capacity()
515 }
516
517 fn push(&self, message: T) -> Result<(), T> {
518 self.queue.push(message)
519 }
520
521 fn pop(&self) -> Option<T> {
522 self.queue.pop()
523 }
524 }
525 }
526}
527
528struct OverflowState<T, M: Mode<T>> {
529 queue: Mutex<M::Overflow>,
530 activity: AtomicUsize,
531 _phantom: PhantomData<fn() -> T>,
532}
533
534impl<T, M: Mode<T>> OverflowState<T, M> {
535 #[allow(clippy::missing_const_for_fn)]
536 fn new() -> Self {
537 Self {
538 queue: Mutex::new(M::Overflow::default()),
539 activity: AtomicUsize::new(0),
540 _phantom: PhantomData,
541 }
542 }
543
544 fn try_ready(&self, ready: &Ready<T>, message: T) -> Result<(), T> {
545 if self.activity.load(Ordering::Relaxed) != 0 {
547 return Err(message);
548 }
549 ready.push(message)
550 }
551
552 fn enqueue_overflow(
553 &self,
554 ready: &Ready<T>,
555 message: T,
556 is_closed: impl Fn() -> bool,
557 ) -> M::Feedback {
558 let mutation = Mutation::begin(&self.activity);
560 let mut queue = lock(&self.queue);
561 if is_closed() {
562 mutation.publish(queue.is_empty());
563 return M::ready_feedback(Feedback::Closed);
564 }
565
566 let message = if queue.is_empty() {
570 match ready.push(message) {
571 Ok(()) => {
572 mutation.publish(queue.is_empty());
573 return M::ready_feedback(Feedback::Ok);
574 }
575 Err(message) => message,
576 }
577 } else {
578 message
579 };
580
581 let handled = M::handle(&mut queue, message);
583 mutation.publish(queue.is_empty());
584 M::overflow_feedback(handled)
585 }
586
587 fn refill(&self, ready: &Ready<T>) {
588 if self.activity.load(Ordering::Relaxed) & OVERFLOW_HAS_MESSAGES == 0 {
590 return;
591 }
592
593 let mutation = Mutation::begin(&self.activity);
594 let mut queue = lock(&self.queue);
595 queue.drain(|message| ready.push(message).err());
596 mutation.publish(queue.is_empty());
597 }
598
599 fn drain(&self, ready: &Ready<T>) {
600 let mutation = Mutation::begin(&self.activity);
602 while ready.pop().is_some() {}
603
604 let mut drained = Vec::new();
607 let mut queue = lock(&self.queue);
608 queue.drain(|message| {
609 drained.push(message);
610 None
611 });
612 mutation.publish(queue.is_empty());
613 drop(queue);
614 drop(drained);
615
616 while ready.pop().is_some() {}
619 }
620}
621
622struct Mutation<'a> {
623 activity: &'a AtomicUsize,
624}
625
626impl<'a> Mutation<'a> {
627 fn begin(activity: &'a AtomicUsize) -> Self {
628 activity.fetch_add(OVERFLOW_MUTATION, Ordering::Relaxed);
629 Self { activity }
630 }
631
632 fn publish(&self, is_empty: bool) {
633 if is_empty {
634 self.activity
635 .fetch_and(!OVERFLOW_HAS_MESSAGES, Ordering::Relaxed);
636 } else {
637 self.activity
638 .fetch_or(OVERFLOW_HAS_MESSAGES, Ordering::Relaxed);
639 }
640 }
641}
642
643impl Drop for Mutation<'_> {
644 fn drop(&mut self) {
645 let previous = self
646 .activity
647 .fetch_sub(OVERFLOW_MUTATION, Ordering::Relaxed);
648 assert!(previous >= OVERFLOW_MUTATION);
649 }
650}
651
652struct State<T, M: Mode<T>> {
653 ready: Ready<T>,
654 overflow: OverflowState<T, M>,
655 backoff: Counter,
656 closed: AtomicBool,
657 senders: AtomicUsize,
658 waker: AtomicWaker,
659}
660
661impl<T, M: Mode<T>> State<T, M> {
662 fn enqueue(&self, message: T) -> M::Feedback {
663 if self.closed.load(Ordering::Acquire) {
665 return M::ready_feedback(Feedback::Closed);
666 }
667
668 let message = match self.overflow.try_ready(&self.ready, message) {
670 Ok(()) => {
671 if self.closed.load(Ordering::Acquire) {
672 self.overflow.drain(&self.ready);
673 return M::ready_feedback(Feedback::Closed);
674 }
675 self.waker.wake();
676 return M::ready_feedback(Feedback::Ok);
677 }
678 Err(message) => message,
679 };
680
681 let feedback = self
683 .overflow
684 .enqueue_overflow(&self.ready, message, || self.closed.load(Ordering::Acquire));
685
686 if M::is_backoff(&feedback) {
688 self.backoff.inc();
689 }
690
691 if !M::is_closed(&feedback) {
696 self.waker.wake();
697 }
698 feedback
699 }
700
701 fn poll_recv(&self, cx: &mut Context<'_>) -> Poll<Option<T>> {
702 if let Some(message) = self.pop() {
704 return Poll::Ready(Some(message));
705 }
706
707 if self.is_disconnected() {
708 return Poll::Ready(self.pop());
709 }
710
711 register_waker(&self.waker, cx.waker());
712
713 if let Some(message) = self.pop() {
716 return Poll::Ready(Some(message));
717 }
718
719 if self.is_disconnected() {
720 Poll::Ready(self.pop())
721 } else {
722 Poll::Pending
723 }
724 }
725
726 fn pop(&self) -> Option<T> {
727 if let Some(message) = self.ready.pop() {
728 self.overflow.refill(&self.ready);
730 return Some(message);
731 }
732
733 self.overflow.refill(&self.ready);
736 self.ready.pop()
737 }
738
739 fn is_disconnected(&self) -> bool {
740 self.closed.load(Ordering::Acquire) || self.senders.load(Ordering::Acquire) == 0
741 }
742
743 fn close(&self) {
744 self.closed.store(true, Ordering::Release);
745 self.overflow.drain(&self.ready);
746 }
747}
748
749fn new_state<T, M: Mode<T>>(metrics: impl Metrics, capacity: NonZeroUsize) -> Arc<State<T, M>> {
750 Arc::new(State {
751 ready: Ready::new(capacity.get()),
752 overflow: OverflowState::new(),
753 backoff: metrics.counter("backoff", "number of enqueue calls that requested backoff"),
754 closed: AtomicBool::new(false),
755 senders: AtomicUsize::new(1),
756 waker: AtomicWaker::new(),
757 })
758}
759
760fn clone_sender_state<T, M: Mode<T>>(state: &Arc<State<T, M>>) -> Arc<State<T, M>> {
761 state.senders.fetch_add(1, Ordering::Relaxed);
763 state.clone()
764}
765
766fn drop_sender_state<T, M: Mode<T>>(state: &State<T, M>) {
767 let previous = state.senders.fetch_sub(1, Ordering::AcqRel);
768 assert!(previous > 0);
769 if previous == 1 {
771 state.waker.wake();
772 }
773}
774
775fn fmt_sender_state<T, M: Mode<T>>(
776 name: &str,
777 state: &State<T, M>,
778 f: &mut fmt::Formatter<'_>,
779) -> fmt::Result {
780 f.debug_struct(name)
781 .field("capacity", &state.ready.capacity())
782 .field("closed", &state.closed.load(Ordering::Acquire))
783 .finish()
784}
785
786async fn recv_from<T, M: Mode<T>>(state: &State<T, M>) -> Option<T> {
787 poll_fn(|cx| state.poll_recv(cx)).await
788}
789
790fn try_recv_from<T, M: Mode<T>>(state: &State<T, M>) -> Result<T, TryRecvError> {
791 if let Some(message) = state.pop() {
792 return Ok(message);
793 }
794 if state.is_disconnected() {
795 return state.pop().ok_or(TryRecvError::Disconnected);
796 }
797 Err(TryRecvError::Empty)
798}
799
800#[cfg(test)]
801mod mocks {
802 use commonware_runtime::{
803 Metrics as RuntimeMetrics, Name, Supervisor,
804 telemetry::metrics::{Metric, Registered, Registration},
805 };
806 use std::fmt;
807
808 #[derive(Clone, Copy, Debug, Default)]
809 pub(super) struct Metrics;
810
811 impl Supervisor for Metrics {
812 fn name(&self) -> Name {
813 Name::default()
814 }
815
816 fn child(&self, _label: &'static str) -> Self {
817 Self
818 }
819
820 fn with_attribute(self, _key: &'static str, _value: impl fmt::Display) -> Self {
821 self
822 }
823 }
824
825 impl RuntimeMetrics for Metrics {
826 fn register<N: Into<String>, H: Into<String>, M: Metric>(
827 &self,
828 _name: N,
829 _help: H,
830 metric: M,
831 ) -> Registered<M> {
832 Registered::with_registration(metric, Registration::from(()))
833 }
834
835 fn encode(&self) -> String {
836 String::new()
837 }
838 }
839}
840
841#[cfg(all(test, not(feature = "loom")))]
842mod tests {
843 use super::{mocks, *};
844 use commonware_macros::test_async;
845 use commonware_runtime::{Runner as _, Supervisor, deterministic};
846 use commonware_utils::{NZUsize, channel::oneshot};
847 use futures::{
848 FutureExt, pin_mut,
849 task::{ArcWake, waker_ref},
850 };
851 use std::sync::{
852 Arc,
853 atomic::{AtomicUsize, Ordering},
854 mpsc::TryRecvError,
855 };
856
857 fn new<T: Policy>(capacity: NonZeroUsize) -> (Sender<T>, Receiver<T>) {
858 super::new(mocks::Metrics, capacity)
859 }
860
861 fn new_unreliable<T: UnreliablePolicy>(
862 capacity: NonZeroUsize,
863 ) -> (UnreliableSender<T>, UnreliableReceiver<T>) {
864 super::new_unreliable(mocks::Metrics, capacity)
865 }
866
867 #[derive(Debug, PartialEq, Eq)]
868 enum Message {
869 Update(u64),
870 Vote(u64),
871 Required(u64),
872 Buffered(u64),
873 Hint(u64),
874 }
875
876 impl UnreliablePolicy for Message {
877 type Overflow = VecDeque<Self>;
878
879 fn handle(overflow: &mut VecDeque<Self>, message: Self) -> bool {
880 match message {
881 Self::Update(value) => {
882 if let Some(index) = overflow
883 .iter()
884 .rposition(|pending| matches!(pending, Self::Update(_)))
885 {
886 overflow.remove(index);
887 }
888 overflow.push_back(Self::Update(value));
889 true
890 }
891 Self::Required(_) | Self::Buffered(_) => {
892 overflow.push_back(message);
893 true
894 }
895 Self::Hint(value) => {
896 let Some(index) = overflow
897 .iter()
898 .rposition(|pending| matches!(pending, Self::Update(_)))
899 else {
900 return true;
901 };
902 overflow.remove(index);
903 overflow.push_back(Self::Hint(value));
904 true
905 }
906 Self::Vote(_) => false,
907 }
908 }
909 }
910
911 struct Ack {
912 _sender: oneshot::Sender<()>,
913 }
914
915 impl Policy for Ack {
916 type Overflow = VecDeque<Self>;
917
918 fn handle(overflow: &mut VecDeque<Self>, message: Self) {
919 overflow.push_back(message);
920 }
921 }
922
923 #[derive(Default)]
924 struct WakeCounter {
925 wakes: AtomicUsize,
926 }
927
928 impl WakeCounter {
929 fn count(&self) -> usize {
930 self.wakes.load(Ordering::Acquire)
931 }
932 }
933
934 impl ArcWake for WakeCounter {
935 fn wake_by_ref(arc_self: &Arc<Self>) {
936 arc_self.wakes.fetch_add(1, Ordering::AcqRel);
937 }
938 }
939
940 #[test]
941 fn vecdeque_overflow_drain_stops_after_rejected_message() {
942 let mut overflow = VecDeque::from([Message::Vote(1), Message::Vote(2), Message::Vote(3)]);
943 let mut drained = VecDeque::new();
944
945 Overflow::drain(&mut overflow, |message| {
946 drained.push_back(message);
947 if drained.len() == 2 {
948 drained.pop_back()
949 } else {
950 None
951 }
952 });
953
954 assert_eq!(drained, VecDeque::from([Message::Vote(1)]));
955 assert_eq!(
956 overflow,
957 VecDeque::from([Message::Vote(2), Message::Vote(3)])
958 );
959 }
960
961 #[test_async]
962 async fn full_inbox_replaces_stale_overflow_message() {
963 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
964 assert_eq!(
965 sender.enqueue(Message::Update(1)),
966 Unreliable::new(Feedback::Ok)
967 );
968 assert_eq!(
969 sender.enqueue(Message::Update(2)),
970 Unreliable::new(Feedback::Backoff)
971 );
972 assert_eq!(
973 sender.enqueue(Message::Update(3)),
974 Unreliable::new(Feedback::Backoff)
975 );
976
977 assert_eq!(receiver.recv().await, Some(Message::Update(1)));
978 assert_eq!(receiver.recv().await, Some(Message::Update(3)));
979 }
980
981 #[test_async]
982 async fn policy_can_replace_stale_overflow_at_back() {
983 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
984 assert_eq!(
985 sender.enqueue(Message::Vote(1)),
986 Unreliable::new(Feedback::Ok)
987 );
988 assert_eq!(
989 sender.enqueue(Message::Update(2)),
990 Unreliable::new(Feedback::Backoff)
991 );
992 assert_eq!(
993 sender.enqueue(Message::Required(3)),
994 Unreliable::new(Feedback::Backoff)
995 );
996 assert_eq!(
997 sender.enqueue(Message::Update(4)),
998 Unreliable::new(Feedback::Backoff)
999 );
1000
1001 assert_eq!(receiver.recv().await, Some(Message::Vote(1)));
1002 assert_eq!(receiver.recv().await, Some(Message::Required(3)));
1003 assert_eq!(receiver.recv().await, Some(Message::Update(4)));
1004 }
1005
1006 #[test_async]
1007 async fn full_inbox_rejects_non_replaceable_message() {
1008 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
1009 assert_eq!(
1010 sender.enqueue(Message::Vote(1)),
1011 Unreliable::new(Feedback::Ok)
1012 );
1013 assert_eq!(sender.enqueue(Message::Vote(2)), Unreliable::Rejected);
1014
1015 assert_eq!(receiver.recv().await, Some(Message::Vote(1)));
1016 }
1017
1018 #[test_async]
1019 async fn full_inbox_retains_required_message() {
1020 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
1021 assert_eq!(
1022 sender.enqueue(Message::Vote(1)),
1023 Unreliable::new(Feedback::Ok)
1024 );
1025 assert_eq!(
1026 sender.enqueue(Message::Buffered(2)),
1027 Unreliable::new(Feedback::Backoff)
1028 );
1029
1030 assert_eq!(receiver.recv().await, Some(Message::Vote(1)));
1031 assert_eq!(receiver.recv().await, Some(Message::Buffered(2)));
1032 }
1033
1034 #[test]
1035 fn try_recv_refills_from_overflow() {
1036 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
1037 assert_eq!(
1038 sender.enqueue(Message::Vote(1)),
1039 Unreliable::new(Feedback::Ok)
1040 );
1041 assert_eq!(
1042 sender.enqueue(Message::Buffered(2)),
1043 Unreliable::new(Feedback::Backoff)
1044 );
1045
1046 assert_eq!(receiver.try_recv(), Ok(Message::Vote(1)));
1047 assert_eq!(receiver.try_recv(), Ok(Message::Buffered(2)));
1048 }
1049
1050 #[test]
1051 fn backoff_metric_counts_backoff_feedback() {
1052 let executor = deterministic::Runner::default();
1053 executor.start(|context| async move {
1054 let (sender, _receiver) = super::new_unreliable(context.child("mailbox"), NZUsize!(1));
1055 assert_eq!(
1056 sender.enqueue(Message::Vote(1)),
1057 Unreliable::new(Feedback::Ok)
1058 );
1059 assert_eq!(
1060 sender.enqueue(Message::Buffered(2)),
1061 Unreliable::new(Feedback::Backoff)
1062 );
1063 assert_eq!(
1064 sender.enqueue(Message::Buffered(3)),
1065 Unreliable::new(Feedback::Backoff)
1066 );
1067
1068 let buffer = context.encode();
1069 assert!(
1070 buffer.contains("mailbox_backoff_total 2"),
1071 "missing backoff count in metrics: {buffer}"
1072 );
1073 });
1074 }
1075
1076 #[test]
1077 fn unreliable_rejected_feedback_is_not_accepted_or_counted_as_backoff() {
1078 let executor = deterministic::Runner::default();
1079 executor.start(|context| async move {
1080 let (sender, _receiver) = super::new_unreliable(context.child("mailbox"), NZUsize!(1));
1081 assert_eq!(
1082 sender.enqueue(Message::Vote(1)),
1083 Unreliable::new(Feedback::Ok)
1084 );
1085 let feedback = sender.enqueue(Message::Vote(2));
1086
1087 assert_eq!(feedback, Unreliable::Rejected);
1088 assert!(!feedback.accepted());
1089
1090 let buffer = context.encode();
1091 assert!(
1092 buffer.contains("mailbox_backoff_total 0"),
1093 "unexpected backoff count in metrics: {buffer}"
1094 );
1095 });
1096 }
1097
1098 #[test]
1099 fn try_recv_drains_buffered_messages_after_senders_drop() {
1100 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
1101 assert_eq!(
1102 sender.enqueue(Message::Vote(1)),
1103 Unreliable::new(Feedback::Ok)
1104 );
1105 assert_eq!(
1106 sender.enqueue(Message::Buffered(2)),
1107 Unreliable::new(Feedback::Backoff)
1108 );
1109 drop(sender);
1110
1111 assert_eq!(receiver.try_recv(), Ok(Message::Vote(1)));
1112 assert_eq!(receiver.try_recv(), Ok(Message::Buffered(2)));
1113 assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected));
1114 }
1115
1116 #[test]
1117 fn poll_recv_drains_buffered_messages_after_senders_drop() {
1118 let (sender, receiver) = new_unreliable(NZUsize!(1));
1119 let wakes = Arc::new(WakeCounter::default());
1120 let waker = waker_ref(&wakes);
1121 let mut cx = Context::from_waker(&waker);
1122
1123 assert_eq!(
1124 sender.enqueue(Message::Vote(1)),
1125 Unreliable::new(Feedback::Ok)
1126 );
1127 assert_eq!(
1128 sender.enqueue(Message::Buffered(2)),
1129 Unreliable::new(Feedback::Backoff)
1130 );
1131 drop(sender);
1132
1133 assert_eq!(
1134 receiver.state.poll_recv(&mut cx),
1135 Poll::Ready(Some(Message::Vote(1)))
1136 );
1137 assert_eq!(
1138 receiver.state.poll_recv(&mut cx),
1139 Poll::Ready(Some(Message::Buffered(2)))
1140 );
1141 assert_eq!(receiver.state.poll_recv(&mut cx), Poll::Ready(None));
1142 }
1143
1144 #[test]
1145 fn enqueue_uses_ready_capacity_after_partial_drain() {
1146 let (sender, mut receiver) = new_unreliable(NZUsize!(2));
1147 assert_eq!(
1148 sender.enqueue(Message::Vote(1)),
1149 Unreliable::new(Feedback::Ok)
1150 );
1151 assert_eq!(
1152 sender.enqueue(Message::Vote(2)),
1153 Unreliable::new(Feedback::Ok)
1154 );
1155 assert_eq!(
1156 sender.enqueue(Message::Required(3)),
1157 Unreliable::new(Feedback::Backoff)
1158 );
1159
1160 assert_eq!(receiver.try_recv(), Ok(Message::Vote(1)));
1161 assert_eq!(receiver.try_recv(), Ok(Message::Vote(2)));
1162
1163 assert_eq!(
1164 sender.enqueue(Message::Vote(4)),
1165 Unreliable::new(Feedback::Ok)
1166 );
1167 assert_eq!(receiver.try_recv(), Ok(Message::Required(3)));
1168 assert_eq!(receiver.try_recv(), Ok(Message::Vote(4)));
1169 }
1170
1171 #[test]
1172 fn receiver_refills_overflow_after_partial_drain() {
1173 let (sender, mut receiver) = new_unreliable(NZUsize!(3));
1174 assert_eq!(
1175 sender.enqueue(Message::Vote(1)),
1176 Unreliable::new(Feedback::Ok)
1177 );
1178 assert_eq!(
1179 sender.enqueue(Message::Vote(2)),
1180 Unreliable::new(Feedback::Ok)
1181 );
1182 assert_eq!(
1183 sender.enqueue(Message::Vote(3)),
1184 Unreliable::new(Feedback::Ok)
1185 );
1186 assert_eq!(
1187 sender.enqueue(Message::Required(4)),
1188 Unreliable::new(Feedback::Backoff)
1189 );
1190
1191 assert_eq!(receiver.try_recv(), Ok(Message::Vote(1)));
1192 assert_eq!(receiver.try_recv(), Ok(Message::Vote(2)));
1193
1194 assert_eq!(
1195 sender.enqueue(Message::Vote(5)),
1196 Unreliable::new(Feedback::Ok)
1197 );
1198 assert_eq!(receiver.try_recv(), Ok(Message::Vote(3)));
1199 assert_eq!(receiver.try_recv(), Ok(Message::Required(4)));
1200 assert_eq!(receiver.try_recv(), Ok(Message::Vote(5)));
1201 }
1202
1203 #[test_async]
1204 async fn full_inbox_retains_unmatched_replaceable_message() {
1205 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
1206 assert_eq!(
1207 sender.enqueue(Message::Vote(1)),
1208 Unreliable::new(Feedback::Ok)
1209 );
1210 assert_eq!(
1211 sender.enqueue(Message::Required(2)),
1212 Unreliable::new(Feedback::Backoff)
1213 );
1214
1215 assert_eq!(receiver.recv().await, Some(Message::Vote(1)));
1216 assert_eq!(receiver.recv().await, Some(Message::Required(2)));
1217 }
1218
1219 #[test_async]
1220 async fn full_inbox_replaces_stale_overflow_after_ready_fills() {
1221 let (sender, mut receiver) = new_unreliable(NZUsize!(2));
1222 assert_eq!(
1223 sender.enqueue(Message::Vote(1)),
1224 Unreliable::new(Feedback::Ok)
1225 );
1226 assert_eq!(
1227 sender.enqueue(Message::Update(2)),
1228 Unreliable::new(Feedback::Ok)
1229 );
1230 assert_eq!(
1231 sender.enqueue(Message::Update(3)),
1232 Unreliable::new(Feedback::Backoff)
1233 );
1234 assert_eq!(
1235 sender.enqueue(Message::Update(4)),
1236 Unreliable::new(Feedback::Backoff)
1237 );
1238
1239 assert_eq!(receiver.recv().await, Some(Message::Vote(1)));
1240 assert_eq!(receiver.recv().await, Some(Message::Update(2)));
1241 assert_eq!(receiver.recv().await, Some(Message::Update(4)));
1242 }
1243
1244 #[test_async]
1245 async fn mailbox_capacity_is_soft_limit_for_required_messages() {
1246 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
1247 assert_eq!(
1248 sender.enqueue(Message::Vote(1)),
1249 Unreliable::new(Feedback::Ok)
1250 );
1251 assert_eq!(
1252 sender.enqueue(Message::Required(2)),
1253 Unreliable::new(Feedback::Backoff)
1254 );
1255 assert_eq!(
1256 sender.enqueue(Message::Required(3)),
1257 Unreliable::new(Feedback::Backoff)
1258 );
1259
1260 assert_eq!(receiver.recv().await, Some(Message::Vote(1)));
1261 assert_eq!(receiver.recv().await, Some(Message::Required(2)));
1262 assert_eq!(receiver.recv().await, Some(Message::Required(3)));
1263 }
1264
1265 #[test_async]
1266 async fn full_inbox_rejects_hint() {
1267 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
1268 assert_eq!(
1269 sender.enqueue(Message::Vote(1)),
1270 Unreliable::new(Feedback::Ok)
1271 );
1272 assert_eq!(
1273 sender.enqueue(Message::Hint(2)),
1274 Unreliable::new(Feedback::Backoff)
1275 );
1276
1277 assert_eq!(receiver.recv().await, Some(Message::Vote(1)));
1278 }
1279
1280 #[test_async]
1281 async fn full_inbox_can_replace_or_drop_by_message() {
1282 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
1283 assert_eq!(
1284 sender.enqueue(Message::Vote(1)),
1285 Unreliable::new(Feedback::Ok)
1286 );
1287 assert_eq!(
1288 sender.enqueue(Message::Update(2)),
1289 Unreliable::new(Feedback::Backoff)
1290 );
1291 assert_eq!(
1292 sender.enqueue(Message::Hint(3)),
1293 Unreliable::new(Feedback::Backoff)
1294 );
1295
1296 assert_eq!(receiver.recv().await, Some(Message::Vote(1)));
1297 assert_eq!(receiver.recv().await, Some(Message::Hint(3)));
1298 }
1299
1300 #[test_async]
1301 async fn empty_inbox_wakes_on_enqueue() {
1302 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
1303
1304 let next = receiver.recv();
1305 pin_mut!(next);
1306 assert!(next.as_mut().now_or_never().is_none());
1307
1308 assert_eq!(
1309 sender.enqueue(Message::Vote(1)),
1310 Unreliable::new(Feedback::Ok)
1311 );
1312 assert_eq!(next.await, Some(Message::Vote(1)));
1313 }
1314
1315 #[test]
1316 fn pending_recv_wakes_when_senders_drop() {
1317 let (sender, receiver) = new_unreliable::<Message>(NZUsize!(1));
1318 let wakes = Arc::new(WakeCounter::default());
1319 let waker = waker_ref(&wakes);
1320 let mut cx = Context::from_waker(&waker);
1321
1322 assert_eq!(receiver.state.poll_recv(&mut cx), Poll::Pending);
1323 assert_eq!(wakes.count(), 0);
1324
1325 drop(sender);
1326
1327 assert_eq!(wakes.count(), 1);
1328 assert_eq!(receiver.state.poll_recv(&mut cx), Poll::Ready(None));
1329 }
1330
1331 #[test]
1332 fn pending_recv_wakes_on_handled_overflow_enqueue() {
1333 let (sender, mut receiver) = new_unreliable(NZUsize!(1));
1334 let wakes = Arc::new(WakeCounter::default());
1335 let waker = waker_ref(&wakes);
1336 let mut cx = Context::from_waker(&waker);
1337
1338 assert_eq!(receiver.state.poll_recv(&mut cx), Poll::Pending);
1339 assert_eq!(wakes.count(), 0);
1340
1341 assert_eq!(sender.state.ready.push(Message::Vote(1)), Ok(()));
1343 assert_eq!(
1344 sender.enqueue(Message::Buffered(2)),
1345 Unreliable::new(Feedback::Backoff)
1346 );
1347
1348 assert_eq!(wakes.count(), 1);
1349 assert_eq!(receiver.try_recv(), Ok(Message::Vote(1)));
1350 assert_eq!(receiver.try_recv(), Ok(Message::Buffered(2)));
1351 }
1352
1353 #[test]
1354 fn receiver_drop_blocks_ready_fast_path_feedback() {
1355 let (sender, receiver) = new_unreliable(NZUsize!(1));
1356 let wakes = Arc::new(WakeCounter::default());
1357 let waker = waker_ref(&wakes);
1358 let mut cx = Context::from_waker(&waker);
1359
1360 assert_eq!(receiver.state.poll_recv(&mut cx), Poll::Pending);
1361 drop(receiver);
1362
1363 assert_eq!(
1364 sender.enqueue(Message::Vote(1)),
1365 Unreliable::new(Feedback::Closed)
1366 );
1367 assert_eq!(wakes.count(), 0);
1368 }
1369
1370 #[test_async]
1371 async fn empty_inbox_closes_when_senders_drop() {
1372 let (sender, mut receiver) = new_unreliable::<Message>(NZUsize!(1));
1373 drop(sender);
1374
1375 assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected));
1376 assert_eq!(receiver.recv().await, None);
1377 }
1378
1379 #[test]
1380 fn enqueue_after_receiver_drop_returns_closed() {
1381 let (sender, receiver) = new_unreliable(NZUsize!(1));
1382 drop(receiver);
1383
1384 assert_eq!(
1385 sender.enqueue(Message::Vote(1)),
1386 Unreliable::new(Feedback::Closed)
1387 );
1388 }
1389
1390 #[test_async]
1391 async fn receiver_drop_cancels_buffered_responders() {
1392 let (sender, receiver) = new(NZUsize!(1));
1393 let (ready_tx, ready_rx) = oneshot::channel();
1394 let (overflow_tx, overflow_rx) = oneshot::channel();
1395
1396 assert_eq!(sender.enqueue(Ack { _sender: ready_tx }), Feedback::Ok);
1397 assert_eq!(
1398 sender.enqueue(Ack {
1399 _sender: overflow_tx
1400 }),
1401 Feedback::Backoff
1402 );
1403 drop(receiver);
1404
1405 assert!(ready_rx.await.is_err());
1406 assert!(overflow_rx.await.is_err());
1407 }
1408
1409 #[derive(Debug, PartialEq, Eq)]
1410 enum ClearingMessage {
1411 FillReady,
1412 ClearOverflow,
1413 }
1414
1415 impl Policy for ClearingMessage {
1416 type Overflow = VecDeque<Self>;
1417
1418 fn handle(overflow: &mut VecDeque<Self>, message: Self) {
1419 overflow.push_back(message);
1420 overflow.clear();
1421 }
1422 }
1423
1424 #[test]
1425 fn policy_can_clear_overflow_and_request_backoff() {
1426 let (sender, mut receiver) = new(NZUsize!(1));
1427 assert_eq!(sender.enqueue(ClearingMessage::FillReady), Feedback::Ok);
1428 assert_eq!(
1429 sender.enqueue(ClearingMessage::ClearOverflow),
1430 Feedback::Backoff
1431 );
1432
1433 assert!(matches!(
1434 receiver.try_recv(),
1435 Ok(ClearingMessage::FillReady)
1436 ));
1437 assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty));
1438 }
1439
1440 #[derive(Debug, PartialEq, Eq)]
1441 enum SpillMessage {
1442 FillReady,
1443 Spill,
1444 }
1445
1446 impl Policy for SpillMessage {
1447 type Overflow = VecDeque<Self>;
1448
1449 fn handle(overflow: &mut VecDeque<Self>, message: Self) {
1450 overflow.push_back(message);
1451 }
1452 }
1453
1454 #[test]
1455 fn pending_recv_wakes_when_policy_spills() {
1456 let (sender, mut receiver) = new(NZUsize!(1));
1457 let wakes = Arc::new(WakeCounter::default());
1458 let waker = waker_ref(&wakes);
1459 let mut cx = Context::from_waker(&waker);
1460
1461 assert_eq!(receiver.state.poll_recv(&mut cx), Poll::Pending);
1462 assert_eq!(wakes.count(), 0);
1463
1464 assert_eq!(sender.state.ready.push(SpillMessage::FillReady), Ok(()));
1465 assert_eq!(sender.enqueue(SpillMessage::Spill), Feedback::Backoff);
1466
1467 assert_eq!(wakes.count(), 1);
1468 assert_eq!(receiver.try_recv(), Ok(SpillMessage::FillReady));
1469 assert_eq!(receiver.try_recv(), Ok(SpillMessage::Spill));
1470 }
1471}
1472
1473#[cfg(all(test, feature = "loom"))]
1474mod loom_tests {
1475 use super::{mocks, *};
1476 use commonware_utils::NZUsize;
1477 use futures::pin_mut;
1478 use loom::{
1479 sync::{
1480 Arc,
1481 atomic::{AtomicUsize, Ordering},
1482 },
1483 thread,
1484 };
1485 use std::{
1486 future::Future,
1487 task::{RawWaker, RawWakerVTable, Waker},
1488 };
1489
1490 fn new<T: Policy>(capacity: NonZeroUsize) -> (Sender<T>, Receiver<T>) {
1491 super::new(mocks::Metrics, capacity)
1492 }
1493
1494 fn new_unreliable<T: UnreliablePolicy>(
1495 capacity: NonZeroUsize,
1496 ) -> (UnreliableSender<T>, UnreliableReceiver<T>) {
1497 super::new_unreliable(mocks::Metrics, capacity)
1498 }
1499
1500 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
1501 enum Message {
1502 Drop(u8),
1503 Spill(u8),
1504 }
1505
1506 #[derive(Clone, Debug)]
1507 enum OrderedMessage {
1508 Item(u8),
1509 Coordinated(u8, Arc<AtomicUsize>),
1510 }
1511
1512 #[derive(Clone, Copy, Debug, PartialEq, Eq)]
1513 enum ReplacingMessage {
1514 FillReady,
1515 Replace(u8),
1516 }
1517
1518 struct TrackedMessage {
1519 drops: Arc<AtomicUsize>,
1520 }
1521
1522 struct CyclicMessage {
1523 _sender: Sender<Self>,
1524 drops: Arc<AtomicUsize>,
1525 }
1526
1527 impl TrackedMessage {
1528 const fn new(drops: Arc<AtomicUsize>) -> Self {
1529 Self { drops }
1530 }
1531 }
1532
1533 impl Drop for TrackedMessage {
1534 fn drop(&mut self) {
1535 self.drops.fetch_add(1, Ordering::AcqRel);
1536 }
1537 }
1538
1539 impl Drop for CyclicMessage {
1540 fn drop(&mut self) {
1541 self.drops.fetch_add(1, Ordering::AcqRel);
1542 }
1543 }
1544
1545 impl UnreliablePolicy for Message {
1546 type Overflow = VecDeque<Self>;
1547
1548 fn handle(overflow: &mut VecDeque<Self>, message: Self) -> bool {
1549 match message {
1550 Self::Drop(_) => false,
1551 Self::Spill(_) => {
1552 overflow.push_back(message);
1553 true
1554 }
1555 }
1556 }
1557 }
1558
1559 impl Policy for OrderedMessage {
1560 type Overflow = VecDeque<Self>;
1561
1562 fn handle(overflow: &mut VecDeque<Self>, message: Self) {
1563 let gate = match &message {
1564 Self::Item(_) => None,
1565 Self::Coordinated(_, gate) => Some(gate.clone()),
1566 };
1567 overflow.push_back(message);
1568 if let Some(gate) = gate {
1569 gate.store(1, Ordering::Release);
1570 while gate.load(Ordering::Acquire) == 1 {
1571 thread::yield_now();
1572 }
1573 }
1574 }
1575 }
1576
1577 impl UnreliablePolicy for ReplacingMessage {
1578 type Overflow = VecDeque<Self>;
1579
1580 fn handle(overflow: &mut VecDeque<Self>, message: Self) -> bool {
1581 match message {
1582 Self::FillReady => false,
1583 Self::Replace(_) => {
1584 if let Some(pending) = overflow
1585 .iter_mut()
1586 .rev()
1587 .find(|pending| matches!(pending, Self::Replace(_)))
1588 {
1589 *pending = message;
1590 } else {
1591 overflow.push_back(message);
1592 }
1593 true
1594 }
1595 }
1596 }
1597 }
1598
1599 impl Policy for TrackedMessage {
1600 type Overflow = VecDeque<Self>;
1601
1602 fn handle(overflow: &mut VecDeque<Self>, message: Self) {
1603 overflow.push_back(message);
1604 }
1605 }
1606
1607 impl Policy for CyclicMessage {
1608 type Overflow = VecDeque<Self>;
1609
1610 fn handle(overflow: &mut VecDeque<Self>, message: Self) {
1611 overflow.push_back(message);
1612 }
1613 }
1614
1615 fn record(seen: &AtomicUsize, message: Message) {
1616 let value = match message {
1617 Message::Drop(value) | Message::Spill(value) => value,
1618 };
1619 seen.fetch_or(1usize << usize::from(value), Ordering::AcqRel);
1620 }
1621
1622 fn value(message: OrderedMessage) -> u8 {
1623 match message {
1624 OrderedMessage::Item(value) | OrderedMessage::Coordinated(value, _) => value,
1625 }
1626 }
1627
1628 const fn replacement_value(message: ReplacingMessage) -> Option<u8> {
1629 match message {
1630 ReplacingMessage::FillReady => None,
1631 ReplacingMessage::Replace(value) => Some(value),
1632 }
1633 }
1634
1635 unsafe fn clone_counter(data: *const ()) -> RawWaker {
1636 let wakes = unsafe { Arc::<AtomicUsize>::from_raw(data.cast()) };
1639 let cloned = wakes.clone();
1640 let _ = Arc::into_raw(wakes);
1641 RawWaker::new(Arc::into_raw(cloned).cast(), &COUNTER_WAKER_VTABLE)
1642 }
1643
1644 unsafe fn wake_counter(data: *const ()) {
1645 let wakes = unsafe { Arc::<AtomicUsize>::from_raw(data.cast()) };
1648 wakes.fetch_add(1, Ordering::AcqRel);
1649 }
1650
1651 unsafe fn wake_counter_by_ref(data: *const ()) {
1652 let wakes = unsafe { Arc::<AtomicUsize>::from_raw(data.cast()) };
1655 wakes.fetch_add(1, Ordering::AcqRel);
1656 let _ = Arc::into_raw(wakes);
1657 }
1658
1659 unsafe fn drop_counter(data: *const ()) {
1660 unsafe {
1663 drop(Arc::<AtomicUsize>::from_raw(data.cast()));
1664 }
1665 }
1666
1667 static COUNTER_WAKER_VTABLE: RawWakerVTable = RawWakerVTable::new(
1668 clone_counter,
1669 wake_counter,
1670 wake_counter_by_ref,
1671 drop_counter,
1672 );
1673
1674 fn counting_waker(wakes: Arc<AtomicUsize>) -> Waker {
1675 let raw = RawWaker::new(Arc::into_raw(wakes).cast(), &COUNTER_WAKER_VTABLE);
1676 unsafe { Waker::from_raw(raw) }
1679 }
1680
1681 #[test]
1682 fn sender_drop_racing_waker_registration_wakes_or_disconnects() {
1683 loom::model(|| {
1684 let (sender, receiver) = new_unreliable::<Message>(NZUsize!(1));
1685 let wakes = Arc::new(AtomicUsize::new(0));
1686 let waker = counting_waker(wakes.clone());
1687 let mut cx = Context::from_waker(&waker);
1688
1689 let close = thread::spawn(move || {
1690 drop(sender);
1691 });
1692
1693 let poll = receiver.state.poll_recv(&mut cx);
1694 close.join().unwrap();
1695
1696 match poll {
1697 Poll::Ready(None) => {}
1698 Poll::Pending => {
1699 assert!(wakes.load(Ordering::Acquire) > 0);
1700 assert_eq!(receiver.state.poll_recv(&mut cx), Poll::Ready(None));
1701 }
1702 Poll::Ready(Some(_)) => panic!("unexpected message"),
1703 }
1704 });
1705 }
1706
1707 #[test]
1708 fn sender_enqueue_then_drop_racing_poll_recv_drains_message() {
1709 loom::model(|| {
1710 let (sender, receiver) = new_unreliable::<Message>(NZUsize!(1));
1711 let wakes = Arc::new(AtomicUsize::new(0));
1712 let waker = counting_waker(wakes.clone());
1713 let mut cx = Context::from_waker(&waker);
1714
1715 let enqueue = thread::spawn(move || {
1716 assert_eq!(
1717 sender.enqueue(Message::Spill(0)),
1718 Unreliable::new(Feedback::Ok)
1719 );
1720 });
1721
1722 let poll = receiver.state.poll_recv(&mut cx);
1723 enqueue.join().unwrap();
1724
1725 match poll {
1726 Poll::Ready(Some(Message::Spill(0))) => {}
1727 Poll::Pending => {
1728 assert!(wakes.load(Ordering::Acquire) > 0);
1729 assert_eq!(
1730 receiver.state.poll_recv(&mut cx),
1731 Poll::Ready(Some(Message::Spill(0)))
1732 );
1733 }
1734 Poll::Ready(None) => panic!("disconnected before draining message"),
1735 Poll::Ready(Some(message)) => panic!("unexpected message: {message:?}"),
1736 }
1737
1738 assert_eq!(receiver.state.poll_recv(&mut cx), Poll::Ready(None));
1739 });
1740 }
1741
1742 #[test]
1743 fn sender_enqueue_then_drop_racing_try_recv_drains_message() {
1744 loom::model(|| {
1745 let (sender, mut receiver) = new_unreliable::<Message>(NZUsize!(1));
1746
1747 let enqueue = thread::spawn(move || {
1748 assert_eq!(
1749 sender.enqueue(Message::Spill(0)),
1750 Unreliable::new(Feedback::Ok)
1751 );
1752 });
1753
1754 let result = receiver.try_recv();
1755 enqueue.join().unwrap();
1756
1757 match result {
1758 Ok(Message::Spill(0)) => {}
1759 Err(TryRecvError::Empty) => {
1760 assert_eq!(receiver.try_recv(), Ok(Message::Spill(0)));
1761 }
1762 Err(TryRecvError::Disconnected) => {
1763 panic!("disconnected before draining message");
1764 }
1765 Ok(message) => panic!("unexpected message: {message:?}"),
1766 }
1767
1768 assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected));
1769 });
1770 }
1771
1772 #[test]
1773 fn handled_enqueue_wakes_registered_receiver() {
1774 loom::model(|| {
1775 let (sender, mut receiver) = new_unreliable::<Message>(NZUsize!(1));
1776 let wakes = Arc::new(AtomicUsize::new(0));
1777 let waker = counting_waker(wakes.clone());
1778 let mut cx = Context::from_waker(&waker);
1779
1780 let next = receiver.recv();
1781 pin_mut!(next);
1782 assert!(matches!(next.as_mut().poll(&mut cx), Poll::Pending));
1783 assert_eq!(
1784 sender.enqueue(Message::Spill(0)),
1785 Unreliable::new(Feedback::Ok)
1786 );
1787
1788 assert_eq!(wakes.load(Ordering::Acquire), 1);
1789 assert_eq!(
1790 next.as_mut().poll(&mut cx),
1791 Poll::Ready(Some(Message::Spill(0)))
1792 );
1793 });
1794 }
1795
1796 #[test]
1797 fn receiver_drop_racing_ready_fast_path_feedback_wakes_if_ready() {
1798 loom::model(|| {
1799 let (sender, receiver) = new_unreliable::<Message>(NZUsize!(1));
1800 let wakes = Arc::new(AtomicUsize::new(0));
1801 let waker = counting_waker(wakes.clone());
1802 let mut cx = Context::from_waker(&waker);
1803
1804 assert_eq!(receiver.state.poll_recv(&mut cx), Poll::Pending);
1805
1806 let close = thread::spawn(move || {
1807 drop(receiver);
1808 });
1809 let feedback = sender.enqueue(Message::Spill(0));
1810 close.join().unwrap();
1811
1812 if feedback.accepted() {
1813 assert!(wakes.load(Ordering::Acquire) > 0);
1814 } else {
1815 assert_eq!(feedback, Unreliable::new(Feedback::Closed));
1816 }
1817 assert_eq!(
1818 sender.enqueue(Message::Spill(1)),
1819 Unreliable::new(Feedback::Closed)
1820 );
1821 });
1822 }
1823
1824 #[test]
1825 fn receiver_drop_racing_ready_enqueue_drops_message() {
1826 loom::model(|| {
1827 let (sender, receiver) = new::<TrackedMessage>(NZUsize!(1));
1828 let drops = Arc::new(AtomicUsize::new(0));
1829
1830 let close = thread::spawn(move || {
1831 drop(receiver);
1832 });
1833 let _ = sender.enqueue(TrackedMessage::new(drops.clone()));
1834 close.join().unwrap();
1835
1836 assert_eq!(drops.load(Ordering::Acquire), 1);
1837 });
1838 }
1839
1840 #[test]
1841 fn receiver_drop_racing_overflow_enqueue_drops_messages() {
1842 loom::model(|| {
1843 let (sender, receiver) = new::<TrackedMessage>(NZUsize!(1));
1844 let ready_drops = Arc::new(AtomicUsize::new(0));
1845 let overflow_drops = Arc::new(AtomicUsize::new(0));
1846
1847 assert_eq!(
1848 sender.enqueue(TrackedMessage::new(ready_drops.clone())),
1849 Feedback::Ok
1850 );
1851 let close = thread::spawn(move || {
1852 drop(receiver);
1853 });
1854 let _ = sender.enqueue(TrackedMessage::new(overflow_drops.clone()));
1855 close.join().unwrap();
1856
1857 assert_eq!(ready_drops.load(Ordering::Acquire), 1);
1858 assert_eq!(overflow_drops.load(Ordering::Acquire), 1);
1859 });
1860 }
1861
1862 #[test]
1863 fn receiver_drop_drains_ready_message_published_under_overflow_lock() {
1864 loom::model(|| {
1865 let (sender, receiver) = new::<TrackedMessage>(NZUsize!(1));
1866 let drops = Arc::new(AtomicUsize::new(0));
1867 let mutation = Mutation::begin(&sender.state.overflow.activity);
1868 let queue = lock(&sender.state.overflow.queue);
1869
1870 let close = thread::spawn(move || {
1871 drop(receiver);
1872 });
1873
1874 assert!(
1875 sender
1876 .state
1877 .ready
1878 .push(TrackedMessage::new(drops.clone()))
1879 .is_ok()
1880 );
1881 mutation.publish(queue.is_empty());
1882 drop(queue);
1883 drop(mutation);
1884 close.join().unwrap();
1885
1886 assert_eq!(drops.load(Ordering::Acquire), 1);
1887 });
1888 }
1889
1890 #[test]
1891 fn receiver_drop_drains_overflow_message_published_under_overflow_lock() {
1892 loom::model(|| {
1893 let (sender, receiver) = new::<TrackedMessage>(NZUsize!(1));
1894 let ready_drops = Arc::new(AtomicUsize::new(0));
1895 let overflow_drops = Arc::new(AtomicUsize::new(0));
1896
1897 assert_eq!(
1898 sender.enqueue(TrackedMessage::new(ready_drops.clone())),
1899 Feedback::Ok
1900 );
1901
1902 let mutation = Mutation::begin(&sender.state.overflow.activity);
1903 let mut queue = lock(&sender.state.overflow.queue);
1904 let close = thread::spawn(move || {
1905 drop(receiver);
1906 });
1907
1908 queue.push_back(TrackedMessage::new(overflow_drops.clone()));
1909 mutation.publish(queue.is_empty());
1910 drop(queue);
1911 drop(mutation);
1912 close.join().unwrap();
1913
1914 assert_eq!(ready_drops.load(Ordering::Acquire), 1);
1915 assert_eq!(overflow_drops.load(Ordering::Acquire), 1);
1916 });
1917 }
1918
1919 #[test]
1920 fn receiver_drop_breaks_message_sender_cycle() {
1921 loom::model(|| {
1922 let (sender, receiver) = new::<CyclicMessage>(NZUsize!(1));
1923 let drops = Arc::new(AtomicUsize::new(0));
1924
1925 assert_eq!(
1926 sender.enqueue(CyclicMessage {
1927 _sender: sender.clone(),
1928 drops: drops.clone(),
1929 }),
1930 Feedback::Ok
1931 );
1932 assert_eq!(
1933 sender.enqueue(CyclicMessage {
1934 _sender: sender.clone(),
1935 drops: drops.clone(),
1936 }),
1937 Feedback::Backoff
1938 );
1939
1940 drop(receiver);
1941
1942 assert_eq!(drops.load(Ordering::Acquire), 2);
1943 assert_eq!(
1944 sender.enqueue(CyclicMessage {
1945 _sender: sender.clone(),
1946 drops,
1947 }),
1948 Feedback::Closed
1949 );
1950 });
1951 }
1952
1953 #[test]
1954 fn concurrent_close_and_ready_enqueue_remains_closed() {
1955 loom::model(|| {
1956 let (sender, receiver) = new_unreliable::<Message>(NZUsize!(1));
1957
1958 let enqueue_sender = sender.clone();
1959 let enqueue = thread::spawn(move || {
1960 let _ = enqueue_sender.enqueue(Message::Spill(1));
1961 });
1962
1963 let close = thread::spawn(move || {
1964 drop(receiver);
1965 });
1966
1967 enqueue.join().unwrap();
1968 close.join().unwrap();
1969 assert_eq!(
1970 sender.enqueue(Message::Spill(2)),
1971 Unreliable::new(Feedback::Closed)
1972 );
1973 });
1974 }
1975
1976 #[test]
1977 fn concurrent_close_and_overflow_enqueue_remains_closed() {
1978 loom::model(|| {
1979 let (sender, receiver) = new_unreliable::<Message>(NZUsize!(1));
1980 assert_eq!(
1981 sender.enqueue(Message::Drop(0)),
1982 Unreliable::new(Feedback::Ok)
1983 );
1984
1985 let enqueue_sender = sender.clone();
1986 let enqueue = thread::spawn(move || {
1987 let _ = enqueue_sender.enqueue(Message::Spill(1));
1988 });
1989
1990 let close = thread::spawn(move || {
1991 drop(receiver);
1992 });
1993
1994 enqueue.join().unwrap();
1995 close.join().unwrap();
1996 assert_eq!(
1997 sender.enqueue(Message::Spill(2)),
1998 Unreliable::new(Feedback::Closed)
1999 );
2000 });
2001 }
2002
2003 #[test]
2004 fn concurrent_spill_and_refill_preserves_messages() {
2005 loom::model(|| {
2006 let (sender, mut receiver) = new_unreliable::<Message>(NZUsize!(1));
2007 let idle_sender = sender.clone();
2008 assert_eq!(
2009 sender.enqueue(Message::Spill(0)),
2010 Unreliable::new(Feedback::Ok)
2011 );
2012
2013 let seen = Arc::new(AtomicUsize::new(0));
2014 let enqueue = thread::spawn(move || {
2015 let feedback = sender.enqueue(Message::Spill(1));
2016 assert!(feedback.accepted());
2017 });
2018
2019 let seen_by_receiver = seen.clone();
2020 let recv = thread::spawn(move || {
2021 if let Ok(message) = receiver.try_recv() {
2022 record(&seen_by_receiver, message);
2023 }
2024 receiver
2025 });
2026
2027 enqueue.join().unwrap();
2028 let mut receiver = recv.join().unwrap();
2029
2030 while let Ok(message) = receiver.try_recv() {
2031 record(&seen, message);
2032 }
2033 assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty));
2034 drop(idle_sender);
2035 assert_eq!(seen.load(Ordering::Acquire), 0b11);
2036 });
2037 }
2038
2039 #[test]
2040 fn concurrent_spill_senders_preserve_messages() {
2041 loom::model(|| {
2042 let (sender, mut receiver) = new_unreliable::<Message>(NZUsize!(1));
2043 let idle_sender = sender.clone();
2044 assert_eq!(
2045 sender.enqueue(Message::Spill(0)),
2046 Unreliable::new(Feedback::Ok)
2047 );
2048
2049 let sender_1 = sender.clone();
2050 let enqueue_1 = thread::spawn(move || sender_1.enqueue(Message::Spill(1)));
2051 let enqueue_2 = thread::spawn(move || sender.enqueue(Message::Spill(2)));
2052
2053 let seen = Arc::new(AtomicUsize::new(0));
2054
2055 assert!(enqueue_1.join().unwrap().accepted());
2056 assert!(enqueue_2.join().unwrap().accepted());
2057
2058 while let Ok(message) = receiver.try_recv() {
2059 record(&seen, message);
2060 }
2061 assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty));
2062 drop(idle_sender);
2063 assert_eq!(seen.load(Ordering::Acquire), 0b111);
2064 });
2065 }
2066
2067 #[test]
2068 fn concurrent_replace_keeps_one_overflow_message() {
2069 loom::model(|| {
2070 let (sender, mut receiver) = new_unreliable::<ReplacingMessage>(NZUsize!(1));
2071 let idle_sender = sender.clone();
2072 assert_eq!(
2073 sender.enqueue(ReplacingMessage::FillReady),
2074 Unreliable::new(Feedback::Ok)
2075 );
2076 assert_eq!(
2077 sender.enqueue(ReplacingMessage::Replace(1)),
2078 Unreliable::new(Feedback::Backoff)
2079 );
2080
2081 let sender_1 = sender.clone();
2082 let replace_1 = thread::spawn(move || sender_1.enqueue(ReplacingMessage::Replace(2)));
2083 let replace_2 = thread::spawn(move || sender.enqueue(ReplacingMessage::Replace(3)));
2084
2085 assert_eq!(
2086 replace_1.join().unwrap(),
2087 Unreliable::new(Feedback::Backoff)
2088 );
2089 assert_eq!(
2090 replace_2.join().unwrap(),
2091 Unreliable::new(Feedback::Backoff)
2092 );
2093 assert_eq!(receiver.try_recv(), Ok(ReplacingMessage::FillReady));
2094
2095 let retained = replacement_value(receiver.try_recv().unwrap()).unwrap();
2096 assert!(retained == 2 || retained == 3);
2097 assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty));
2098 drop(idle_sender);
2099 });
2100 }
2101
2102 #[test]
2103 fn stale_overflow_hint_retries_ready_before_policy() {
2104 loom::model(|| {
2105 let (sender, mut receiver) = new_unreliable::<Message>(NZUsize!(2));
2106 assert_eq!(
2107 sender.enqueue(Message::Drop(0)),
2108 Unreliable::new(Feedback::Ok)
2109 );
2110 assert_eq!(
2111 sender.enqueue(Message::Drop(1)),
2112 Unreliable::new(Feedback::Ok)
2113 );
2114 assert_eq!(
2115 sender.enqueue(Message::Spill(2)),
2116 Unreliable::new(Feedback::Backoff)
2117 );
2118
2119 assert_eq!(receiver.try_recv(), Ok(Message::Drop(0)));
2120 assert_eq!(receiver.try_recv(), Ok(Message::Drop(1)));
2121
2122 assert_eq!(
2123 sender.enqueue(Message::Drop(3)),
2124 Unreliable::new(Feedback::Ok)
2125 );
2126 assert_eq!(receiver.try_recv(), Ok(Message::Spill(2)));
2127 assert_eq!(receiver.try_recv(), Ok(Message::Drop(3)));
2128 });
2129 }
2130
2131 #[test]
2132 fn concurrent_overflow_cannot_be_bypassed_by_ready_fast_path() {
2133 loom::model(|| {
2134 let (sender, mut receiver) = new::<OrderedMessage>(NZUsize!(2));
2135 assert_eq!(sender.enqueue(OrderedMessage::Item(0)), Feedback::Ok);
2136 assert_eq!(sender.enqueue(OrderedMessage::Item(1)), Feedback::Ok);
2137
2138 let gate = Arc::new(AtomicUsize::new(0));
2139 let overflow_sender = sender.clone();
2140 let overflow_gate = gate.clone();
2141 let overflow = thread::spawn(move || {
2142 assert_eq!(
2143 overflow_sender.enqueue(OrderedMessage::Coordinated(2, overflow_gate)),
2144 Feedback::Backoff
2145 );
2146 });
2147
2148 while gate.load(Ordering::Acquire) == 0 {
2149 thread::yield_now();
2150 }
2151
2152 let mut observed = vec![value(receiver.try_recv().unwrap())];
2155 gate.store(2, Ordering::Release);
2156 let feedback = sender.enqueue(OrderedMessage::Item(3));
2157 assert!(feedback.accepted());
2158
2159 overflow.join().unwrap();
2160 while let Ok(message) = receiver.try_recv() {
2161 observed.push(value(message));
2162 }
2163
2164 assert_eq!(observed, vec![0, 1, 2, 3]);
2165 });
2166 }
2167
2168 #[test]
2169 fn concurrent_overflow_mutation_does_not_hide_published_overflow() {
2170 loom::model(|| {
2171 let (sender, mut receiver) = new::<OrderedMessage>(NZUsize!(1));
2172 assert_eq!(sender.enqueue(OrderedMessage::Item(0)), Feedback::Ok);
2173 assert_eq!(sender.enqueue(OrderedMessage::Item(1)), Feedback::Backoff);
2174
2175 let gate = Arc::new(AtomicUsize::new(0));
2176 let overflow_gate = gate.clone();
2177 let overflow = thread::spawn(move || {
2178 sender.enqueue(OrderedMessage::Coordinated(2, overflow_gate))
2179 });
2180
2181 while gate.load(Ordering::Acquire) == 0 {
2182 thread::yield_now();
2183 }
2184
2185 let release_gate = gate;
2186 let release = thread::spawn(move || {
2187 release_gate.store(2, Ordering::Release);
2188 });
2189
2190 let receive = thread::spawn(move || {
2191 assert_eq!(receiver.try_recv().map(value), Ok(0));
2192 assert_eq!(receiver.try_recv().map(value), Ok(1));
2193 receiver
2194 });
2195
2196 release.join().unwrap();
2197 let mut receiver = receive.join().unwrap();
2198 assert_eq!(overflow.join().unwrap(), Feedback::Backoff);
2199 assert_eq!(receiver.try_recv().map(value), Ok(2));
2200 });
2201 }
2202
2203 #[test]
2204 fn published_overflow_wakes_pending_receiver() {
2205 loom::model(|| {
2206 let (sender, mut receiver) = new::<OrderedMessage>(NZUsize!(1));
2207 let wakes = Arc::new(AtomicUsize::new(0));
2208 let waker = counting_waker(wakes.clone());
2209 let mut cx = Context::from_waker(&waker);
2210
2211 let gate = Arc::new(AtomicUsize::new(0));
2212 let overflow = {
2213 let next = receiver.recv();
2214 pin_mut!(next);
2215 assert!(matches!(next.as_mut().poll(&mut cx), Poll::Pending));
2216
2217 assert_eq!(sender.enqueue(OrderedMessage::Item(0)), Feedback::Ok);
2218 while wakes.load(Ordering::Acquire) == 0 {
2219 thread::yield_now();
2220 }
2221
2222 let overflow_gate = gate.clone();
2223 let overflow = thread::spawn(move || {
2224 sender.enqueue(OrderedMessage::Coordinated(1, overflow_gate))
2225 });
2226
2227 while gate.load(Ordering::Acquire) == 0 {
2228 thread::yield_now();
2229 }
2230
2231 assert_eq!(
2232 next.as_mut()
2233 .poll(&mut cx)
2234 .map(|message| message.map(value)),
2235 Poll::Ready(Some(0))
2236 );
2237 overflow
2238 };
2239
2240 {
2241 let next = receiver.recv();
2242 pin_mut!(next);
2243 assert!(matches!(next.as_mut().poll(&mut cx), Poll::Pending));
2244 assert_eq!(wakes.load(Ordering::Acquire), 1);
2245
2246 gate.store(2, Ordering::Release);
2247 while wakes.load(Ordering::Acquire) < 2 {
2248 thread::yield_now();
2249 }
2250
2251 assert_eq!(
2252 next.as_mut()
2253 .poll(&mut cx)
2254 .map(|message| message.map(value)),
2255 Poll::Ready(Some(1))
2256 );
2257 }
2258 assert_eq!(overflow.join().unwrap(), Feedback::Backoff);
2259 });
2260 }
2261
2262 #[test]
2263 fn concurrent_refill_and_enqueue_preserves_overflow_order() {
2264 loom::model(|| {
2265 let (sender, mut receiver) = new::<OrderedMessage>(NZUsize!(1));
2266 assert_eq!(sender.enqueue(OrderedMessage::Item(0)), Feedback::Ok);
2267 assert_eq!(sender.enqueue(OrderedMessage::Item(1)), Feedback::Backoff);
2268
2269 let enqueue = thread::spawn(move || sender.enqueue(OrderedMessage::Item(2)));
2270 let receive = thread::spawn(move || {
2271 assert_eq!(receiver.try_recv().map(value), Ok(0));
2272 receiver
2273 });
2274
2275 let mut receiver = receive.join().unwrap();
2276 assert_eq!(enqueue.join().unwrap(), Feedback::Backoff);
2277 assert_eq!(receiver.try_recv().map(value), Ok(1));
2278 assert_eq!(receiver.try_recv().map(value), Ok(2));
2279 });
2280 }
2281}