use async_broadcast::{InactiveReceiver, Sender as InnerSender};
pub mod error {
#[derive(Debug, thiserror::Error)]
pub enum RecvError {
#[error("broadcast channel closed")]
Closed,
#[error("broadcast receiver lagged behind by {0} messages")]
Lagged(u64),
}
}
#[derive(Debug, thiserror::Error)]
#[error("broadcast channel has no receivers")]
pub struct SendError<T>(pub T);
#[derive(Debug)]
pub struct Sender<T: Clone> {
inner: InnerSender<T>,
_liveness: InactiveReceiver<T>,
}
impl<T: Clone> Clone for Sender<T> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
_liveness: self._liveness.clone(),
}
}
}
impl<T: Clone> Sender<T> {
pub fn send(&self, value: T) -> Result<usize, SendError<T>> {
if self.inner.receiver_count() == 0 {
return Err(SendError(value));
}
match self.inner.try_broadcast(value) {
Ok(_displaced) => Ok(self.inner.receiver_count()),
Err(async_broadcast::TrySendError::Closed(value)) => Err(SendError(value)),
Err(async_broadcast::TrySendError::Inactive(value)) => Err(SendError(value)),
Err(async_broadcast::TrySendError::Full(value)) => {
Err(SendError(value))
}
}
}
pub fn subscribe(&self) -> Receiver<T> {
Receiver {
inner: self.inner.new_receiver(),
}
}
pub fn receiver_count(&self) -> usize {
self.inner.receiver_count()
}
}
#[derive(Debug)]
pub struct Receiver<T: Clone> {
inner: async_broadcast::Receiver<T>,
}
impl<T: Clone> Receiver<T> {
pub async fn recv(&mut self) -> Result<T, error::RecvError> {
match self.inner.recv().await {
Ok(value) => Ok(value),
Err(async_broadcast::RecvError::Closed) => Err(error::RecvError::Closed),
Err(async_broadcast::RecvError::Overflowed(skipped)) => {
Err(error::RecvError::Lagged(skipped))
}
}
}
}
impl<T: Clone> Clone for Receiver<T> {
fn clone(&self) -> Self {
Self {
inner: self.inner.clone(),
}
}
}
pub fn channel<T: Clone>(capacity: usize) -> (Sender<T>, Receiver<T>) {
let (mut sender, receiver) = async_broadcast::broadcast(capacity.max(1));
sender.set_overflow(true);
sender.set_await_active(false);
let liveness = receiver.clone().deactivate();
(
Sender {
inner: sender,
_liveness: liveness,
},
Receiver { inner: receiver },
)
}
#[cfg(all(test, not(loom)))]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use futures::executor::block_on;
#[test]
fn send_to_active_receiver_roundtrips() {
let (sender, mut receiver) = channel::<u32>(4);
block_on(async move {
sender.send(7).expect("send");
assert_eq!(receiver.recv().await.expect("recv"), 7);
});
}
#[test]
fn second_subscriber_only_sees_messages_sent_after_subscribe() {
let (sender, mut alpha) = channel::<u32>(4);
sender.send(1).expect("send 1");
let mut beta = sender.subscribe();
sender.send(2).expect("send 2");
block_on(async move {
assert_eq!(alpha.recv().await.expect("alpha first"), 1);
assert_eq!(alpha.recv().await.expect("alpha second"), 2);
assert_eq!(beta.recv().await.expect("beta"), 2);
});
}
#[test]
fn send_succeeds_when_no_receivers_but_inactive_holds_open() {
let (sender, receiver) = channel::<u32>(4);
drop(receiver);
match sender.send(1) {
Err(SendError(value)) => assert_eq!(value, 1),
Ok(_) => panic!("expected SendError"),
}
}
#[test]
fn overflow_yields_lagged_on_recv() {
let (sender, mut receiver) = channel::<u32>(2);
sender.send(1).expect("send 1");
sender.send(2).expect("send 2");
sender.send(3).expect("send 3 (overflows 1)");
block_on(async move {
match receiver.recv().await {
Err(error::RecvError::Lagged(skipped)) => assert!(skipped >= 1),
other => panic!("expected Lagged, got {other:?}"),
}
});
}
}