use crate::sync::Mutex;
use core::num::NonZeroUsize;
use futures::{Sink, Stream, stream::FusedStream};
use std::{
collections::VecDeque,
mem,
pin::Pin,
sync::Arc,
task::{Context, Poll, Waker},
};
use thiserror::Error;
#[derive(Debug, Error)]
#[error("channel closed")]
pub struct ChannelClosed;
#[derive(Debug, Error, PartialEq, Eq)]
pub enum TryRecvError {
#[error("channel empty")]
Empty,
#[error("channel closed")]
Disconnected,
}
#[derive(Debug)]
struct Shared<T> {
buffer: VecDeque<T>,
capacity: usize,
receiver_waker: Option<Waker>,
sender_count: usize,
receiver_dropped: bool,
}
#[derive(Debug)]
pub struct Sender<T> {
shared: Arc<Mutex<Shared<T>>>,
}
impl<T> Sender<T> {
pub fn is_closed(&self) -> bool {
let shared = self.shared.lock();
shared.receiver_dropped
}
pub fn send_lossy(&self, item: T) -> bool {
let mut shared = self.shared.lock();
if shared.receiver_dropped {
return false;
}
let old_item = if shared.buffer.len() >= shared.capacity {
shared.buffer.pop_front()
} else {
None
};
shared.buffer.push_back(item);
let waker = shared.receiver_waker.take();
drop(shared);
drop(old_item);
if let Some(w) = waker {
w.wake();
}
true
}
}
impl<T> Clone for Sender<T> {
fn clone(&self) -> Self {
let mut shared = self.shared.lock();
shared.sender_count += 1;
drop(shared);
Self {
shared: self.shared.clone(),
}
}
}
impl<T> Drop for Sender<T> {
fn drop(&mut self) {
let mut shared = self.shared.lock();
shared.sender_count -= 1;
let waker = if shared.sender_count == 0 {
shared.receiver_waker.take()
} else {
None
};
drop(shared);
if let Some(w) = waker {
w.wake();
}
}
}
impl<T> Sink<T> for Sender<T> {
type Error = ChannelClosed;
fn poll_ready(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
let shared = self.shared.lock();
if shared.receiver_dropped {
return Poll::Ready(Err(ChannelClosed));
}
Poll::Ready(Ok(()))
}
fn start_send(self: Pin<&mut Self>, item: T) -> Result<(), Self::Error> {
self.send_lossy(item).then_some(()).ok_or(ChannelClosed)
}
fn poll_flush(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
fn poll_close(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
Poll::Ready(Ok(()))
}
}
#[derive(Debug)]
pub struct Receiver<T> {
shared: Arc<Mutex<Shared<T>>>,
}
impl<T> Receiver<T> {
pub async fn recv(&mut self) -> Option<T> {
futures::future::poll_fn(|cx| Pin::new(&mut *self).poll_next(cx)).await
}
pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
let mut shared = self.shared.lock();
if let Some(item) = shared.buffer.pop_front() {
return Ok(item);
}
if shared.sender_count == 0 {
return Err(TryRecvError::Disconnected);
}
Err(TryRecvError::Empty)
}
}
impl<T> Stream for Receiver<T> {
type Item = T;
fn poll_next(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
let waker = cx.waker().clone();
let mut shared = self.shared.lock();
if let Some(item) = shared.buffer.pop_front() {
return Poll::Ready(Some(item));
}
if shared.sender_count == 0 {
return Poll::Ready(None);
}
let replaced = shared.receiver_waker.replace(waker);
drop(shared);
drop(replaced);
Poll::Pending
}
}
impl<T> FusedStream for Receiver<T> {
fn is_terminated(&self) -> bool {
let shared = self.shared.lock();
shared.sender_count == 0 && shared.buffer.is_empty()
}
}
impl<T> Drop for Receiver<T> {
fn drop(&mut self) {
let mut shared = self.shared.lock();
shared.receiver_dropped = true;
let buffer = mem::take(&mut shared.buffer);
let waker = shared.receiver_waker.take();
drop(shared);
drop(buffer);
drop(waker);
}
}
pub fn channel<T>(capacity: NonZeroUsize) -> (Sender<T>, Receiver<T>) {
let shared = Arc::new(Mutex::new(Shared {
buffer: VecDeque::with_capacity(capacity.get()),
capacity: capacity.get(),
receiver_waker: None,
sender_count: 1,
receiver_dropped: false,
}));
let sender = Sender {
shared: shared.clone(),
};
let receiver = Receiver { shared };
(sender, receiver)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::NZUsize;
use futures::{SinkExt, StreamExt, executor::block_on};
use std::{sync::Weak, task::Wake};
#[test]
fn test_basic_send_recv() {
block_on(async {
let (mut sender, mut receiver) = channel::<i32>(NZUsize!(10));
sender.send(1).await.unwrap();
sender.send(2).await.unwrap();
sender.send(3).await.unwrap();
assert_eq!(receiver.next().await, Some(1));
assert_eq!(receiver.next().await, Some(2));
assert_eq!(receiver.next().await, Some(3));
});
}
#[test]
fn test_overflow_drops_oldest() {
block_on(async {
let (mut sender, mut receiver) = channel::<i32>(NZUsize!(2));
sender.send(1).await.unwrap();
sender.send(2).await.unwrap();
sender.send(3).await.unwrap(); sender.send(4).await.unwrap();
assert_eq!(receiver.next().await, Some(3));
assert_eq!(receiver.next().await, Some(4));
});
}
#[test]
fn test_send_after_receiver_dropped() {
block_on(async {
let (mut sender, receiver) = channel::<i32>(NZUsize!(10));
drop(receiver);
let err = sender.send(1).await.unwrap_err();
assert!(matches!(err, ChannelClosed));
});
}
#[test]
fn test_recv_after_sender_dropped() {
block_on(async {
let (mut sender, mut receiver) = channel::<i32>(NZUsize!(10));
sender.send(1).await.unwrap();
sender.send(2).await.unwrap();
drop(sender);
assert_eq!(receiver.next().await, Some(1));
assert_eq!(receiver.next().await, Some(2));
assert_eq!(receiver.next().await, None);
});
}
#[test]
fn test_stream_collect() {
block_on(async {
let (mut sender, receiver) = channel::<i32>(NZUsize!(10));
sender.send(1).await.unwrap();
sender.send(2).await.unwrap();
sender.send(3).await.unwrap();
drop(sender);
let items: Vec<_> = receiver.collect().await;
assert_eq!(items, vec![1, 2, 3]);
});
}
#[test]
fn test_clone_sender() {
block_on(async {
let (mut sender1, mut receiver) = channel::<i32>(NZUsize!(10));
let mut sender2 = sender1.clone();
sender1.send(1).await.unwrap();
sender2.send(2).await.unwrap();
assert_eq!(receiver.next().await, Some(1));
assert_eq!(receiver.next().await, Some(2));
});
}
#[test]
fn test_sender_drop_with_clones() {
block_on(async {
let (sender1, mut receiver) = channel::<i32>(NZUsize!(10));
let mut sender2 = sender1.clone();
drop(sender1);
sender2.send(1).await.unwrap();
assert_eq!(receiver.next().await, Some(1));
drop(sender2);
assert_eq!(receiver.next().await, None);
});
}
#[test]
fn test_capacity_one() {
block_on(async {
let (mut sender, mut receiver) = channel::<i32>(NZUsize!(1));
sender.send(1).await.unwrap();
sender.send(2).await.unwrap();
assert_eq!(receiver.next().await, Some(2));
sender.send(1).await.unwrap();
sender.send(2).await.unwrap(); sender.send(3).await.unwrap();
assert_eq!(receiver.next().await, Some(3));
});
}
#[test]
fn test_send_all() {
block_on(async {
let (mut sender, receiver) = channel::<i32>(NZUsize!(10));
let items = futures::stream::iter(vec![1, 2, 3]);
sender.send_all(&mut items.map(Ok)).await.unwrap();
drop(sender);
let received: Vec<_> = receiver.collect().await;
assert_eq!(received, vec![1, 2, 3]);
});
}
#[test]
fn test_fused_stream() {
use futures::stream::FusedStream;
block_on(async {
let (mut sender, mut receiver) = channel::<i32>(NZUsize!(10));
assert!(!receiver.is_terminated());
sender.send(1).await.unwrap();
assert!(!receiver.is_terminated());
drop(sender);
assert!(!receiver.is_terminated());
assert_eq!(receiver.next().await, Some(1));
assert!(receiver.is_terminated());
assert_eq!(receiver.next().await, None);
assert!(receiver.is_terminated());
});
}
#[test]
fn test_is_closed() {
block_on(async {
let (sender, receiver) = channel::<i32>(NZUsize!(10));
assert!(!sender.is_closed());
drop(receiver);
assert!(sender.is_closed());
});
}
struct Probe(Option<Sender<Arc<Self>>>);
impl Probe {
fn new(sender: &Sender<Arc<Self>>) -> Arc<Self> {
Arc::new(Self(Some(sender.clone())))
}
}
#[expect(clippy::manual_noop_waker, reason = "a probe waker must own a sender")]
impl Wake for Probe {
fn wake(self: Arc<Self>) {}
}
impl Drop for Probe {
fn drop(&mut self) {
let sender = self.0.take().expect("a probe must still own its sender");
if sender.shared.try_lock().is_none() {
mem::forget(sender);
panic!("senders must be dropped outside the channel lock");
}
}
}
fn register(sender: &Sender<Arc<Probe>>, receiver: &mut Receiver<Arc<Probe>>) -> Weak<Probe> {
let probe = Probe::new(sender);
let weak = Arc::downgrade(&probe);
let waker = Waker::from(probe);
assert!(
receiver
.poll_next_unpin(&mut Context::from_waker(&waker))
.is_pending(),
"an empty open channel must be pending"
);
drop(waker);
assert!(
weak.upgrade().is_some(),
"a pending receiver must register its waker"
);
weak
}
#[test]
fn test_receiver_drop_releases_unread_items() {
let (sender, receiver) = channel(NZUsize!(1));
let item = Probe::new(&sender);
let weak = Arc::downgrade(&item);
assert!(sender.send_lossy(item));
drop(receiver);
assert!(
weak.upgrade().is_none(),
"receiver drop must release unread items"
);
}
#[test]
fn test_receiver_drop_releases_waker() {
let (sender, mut receiver) = channel(NZUsize!(1));
let shared = Arc::downgrade(&sender.shared);
let waker = register(&sender, &mut receiver);
drop(receiver);
assert!(
waker.upgrade().is_none(),
"receiver drop must release its waker"
);
drop(sender);
assert!(
shared.upgrade().is_none(),
"the channel must not outlive its handles"
);
}
#[test]
fn test_replaced_waker_dropped_outside_lock() {
let (sender, mut receiver) = channel(NZUsize!(1));
let first = register(&sender, &mut receiver);
let second = register(&sender, &mut receiver);
assert!(
first.upgrade().is_none(),
"a replaced waker must be released"
);
assert!(sender.send_lossy(Probe::new(&sender)));
assert!(
second.upgrade().is_none(),
"a send must consume the registered waker"
);
assert!(
receiver.try_recv().is_ok(),
"the channel must deliver after its waker is replaced"
);
}
}