1use crossbeam_channel::{bounded, Receiver, Sender, TrySendError};
2use std::sync::atomic::{AtomicU64, Ordering};
3use std::sync::Arc;
4
5#[derive(Debug)]
7pub struct Bus<M> {
8 sender: Sender<M>,
9 receiver: Receiver<M>,
10 capacity: usize,
11 dropped: Arc<AtomicU64>,
12}
13
14impl<M: Send + 'static> Bus<M> {
15 pub fn new(capacity: usize) -> Self {
16 let (sender, receiver) = bounded(capacity.max(1));
17 Self {
18 sender,
19 receiver,
20 capacity: capacity.max(1),
21 dropped: Arc::new(AtomicU64::new(0)),
22 }
23 }
24
25 pub fn sender(&self) -> BusSender<M> {
26 BusSender {
27 inner: self.sender.clone(),
28 dropped: Arc::clone(&self.dropped),
29 }
30 }
31
32 pub fn receiver(&self) -> &Receiver<M> {
33 &self.receiver
34 }
35
36 pub fn capacity(&self) -> usize {
37 self.capacity
38 }
39
40 pub fn dropped_count(&self) -> u64 {
41 self.dropped.load(Ordering::Relaxed)
42 }
43}
44
45#[derive(Debug)]
46pub struct BusSender<M> {
47 inner: Sender<M>,
48 dropped: Arc<AtomicU64>,
49}
50
51impl<M> Clone for BusSender<M> {
52 fn clone(&self) -> Self {
53 Self {
54 inner: self.inner.clone(),
55 dropped: Arc::clone(&self.dropped),
56 }
57 }
58}
59
60impl<M> BusSender<M> {
61 pub fn send(&self, msg: M) -> Result<(), M> {
62 match self.inner.try_send(msg) {
63 Ok(()) => Ok(()),
64 Err(TrySendError::Full(msg)) => {
65 self.dropped.fetch_add(1, Ordering::Relaxed);
66 Err(msg)
67 }
68 Err(TrySendError::Disconnected(msg)) => Err(msg),
69 }
70 }
71
72 pub fn send_blocking(&self, msg: M) -> Result<(), M> {
73 self.inner.send(msg).map_err(|e| e.into_inner())
74 }
75}
76
77#[cfg(test)]
78mod tests {
79 use super::*;
80
81 #[test]
82 fn backpressure_counts_drops() {
83 let bus = Bus::<u32>::new(1);
84 let tx = bus.sender();
85 assert!(tx.send(1).is_ok());
86 let dropped = tx.send(2).unwrap_err();
87 assert_eq!(dropped, 2);
88 assert_eq!(bus.dropped_count(), 1);
89 }
90}