#![allow(unsafe_code)]
use core::fmt;
use core::future::Future;
use core::marker::PhantomPinned;
use core::pin::Pin;
use core::task::{Context, Poll};
use super::chan::{Chan, Limit};
use super::error::{SendError, SendTimeoutError, TrySendError};
use super::permit::{OwnedPermit, Permit, PermitIterator};
use super::queue;
use super::receiver::Receiver;
use crate::sync::fair_semaphore::{FairSemaphore, FairWait};
use crate::sync::internal::{self, Arc};
use crate::sync::raw_semaphore::observed_closed;
pub(super) struct Bounded {
pub(super) sem: FairSemaphore,
pub(super) bound: usize,
}
impl Limit for Bounded {
type Queue<T> = queue::Bounded<T>;
fn queue<T>(&self) -> queue::Bounded<T> {
queue::Bounded::with_capacity(self.bound)
}
#[inline]
fn is_idle(&self) -> bool {
self.sem.available() == self.bound
}
#[inline]
fn release(&self, n: usize) {
self.sem.release(n);
}
fn close(&self) {
self.sem.close();
}
#[inline]
fn is_closed(&self) -> bool {
self.sem.is_closed()
}
}
#[track_caller]
pub fn channel<T>(buffer: usize) -> (Sender<T>, Receiver<T>) {
assert!(buffer > 0, "mpsc bounded channel requires buffer > 0");
let chan = Arc::new(Chan::new(Bounded {
sem: FairSemaphore::new(buffer),
bound: buffer,
}));
let tx = Sender {
chan: Arc::clone(&chan),
};
(tx, Receiver::new(chan))
}
pub struct Sender<T> {
pub(super) chan: Arc<Chan<T, Bounded>>,
}
pub struct WeakSender<T> {
chan: Arc<Chan<T, Bounded>>,
}
impl<T> Sender<T> {
pub fn send(&self, value: T) -> impl Future<Output = Result<(), SendError<T>>> + '_ {
self.sending(value)
}
fn sending(&self, value: T) -> Sending<'_, T> {
Sending {
chan: &self.chan,
wait: FairWait::new(),
value: Some(value),
_pinned: PhantomPinned,
}
}
pub async fn send_timeout(
&self,
value: T,
timeout: std::time::Duration,
) -> Result<(), SendTimeoutError<T>> {
match crate::time::timeout(timeout, self.reserve()).await {
Ok(Ok(permit)) => {
permit.send(value);
Ok(())
}
Ok(Err(_)) => Err(SendTimeoutError::Closed(value)),
Err(_) => Err(SendTimeoutError::Timeout(value)),
}
}
pub async fn closed(&self) {
self.chan.closed().await;
}
pub fn try_send(&self, message: T) -> Result<(), TrySendError<T>> {
match try_reserve_inner(&self.chan, 1) {
Ok(()) => {
self.chan.send(message);
Ok(())
}
Err(refused) => Err(refused.with_value(message)),
}
}
pub fn blocking_send(&self, value: T) -> Result<(), SendError<T>> {
internal::blocking_in_line(self.sending(value))
}
pub fn is_closed(&self) -> bool {
self.chan.is_closed()
}
pub async fn reserve(&self) -> Result<Permit<'_, T>, SendError<()>> {
reserve_inner(&self.chan, 1).await?;
Ok(Permit::new(&self.chan))
}
pub async fn reserve_many(&self, n: usize) -> Result<PermitIterator<'_, T>, SendError<()>> {
reserve_inner(&self.chan, n).await?;
Ok(PermitIterator::new(&self.chan, n))
}
pub async fn reserve_owned(self) -> Result<OwnedPermit<T>, SendError<()>> {
reserve_inner(&self.chan, 1).await?;
Ok(OwnedPermit::new(self))
}
pub fn try_reserve(&self) -> Result<Permit<'_, T>, TrySendError<()>> {
try_reserve_inner(&self.chan, 1)?;
Ok(Permit::new(&self.chan))
}
pub fn try_reserve_many(&self, n: usize) -> Result<PermitIterator<'_, T>, TrySendError<()>> {
if n > self.chan.limit.bound {
return Err(TrySendError::Full(()));
}
try_reserve_inner(&self.chan, n)?;
Ok(PermitIterator::new(&self.chan, n))
}
pub fn try_reserve_owned(self) -> Result<OwnedPermit<T>, TrySendError<Self>> {
match try_reserve_inner(&self.chan, 1) {
Ok(()) => Ok(OwnedPermit::new(self)),
Err(refused) => Err(refused.with_value(self)),
}
}
pub fn same_channel(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.chan, &other.chan)
}
pub fn capacity(&self) -> usize {
self.chan.limit.sem.available()
}
#[must_use = "Downgrade creates a WeakSender without destroying the original non-weak sender."]
pub fn downgrade(&self) -> WeakSender<T> {
self.chan.add_weak();
WeakSender {
chan: Arc::clone(&self.chan),
}
}
pub fn max_capacity(&self) -> usize {
self.chan.limit.bound
}
pub fn strong_count(&self) -> usize {
self.chan.strong_count()
}
pub fn weak_count(&self) -> usize {
self.chan.weak_count()
}
}
impl<T> Clone for Sender<T> {
fn clone(&self) -> Self {
self.chan.add_sender();
Sender {
chan: Arc::clone(&self.chan),
}
}
}
impl<T> Drop for Sender<T> {
fn drop(&mut self) {
self.chan.drop_sender();
}
}
impl<T> fmt::Debug for Sender<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Sender")
.field("closed", &self.is_closed())
.field("capacity", &self.capacity())
.field("max_capacity", &self.max_capacity())
.finish()
}
}
impl<T> WeakSender<T> {
pub fn upgrade(&self) -> Option<Sender<T>> {
self.chan.try_add_sender().then(|| Sender {
chan: Arc::clone(&self.chan),
})
}
pub fn strong_count(&self) -> usize {
self.chan.strong_count()
}
pub fn weak_count(&self) -> usize {
self.chan.weak_count()
}
}
impl<T> Clone for WeakSender<T> {
fn clone(&self) -> Self {
self.chan.add_weak();
WeakSender {
chan: Arc::clone(&self.chan),
}
}
}
impl<T> Drop for WeakSender<T> {
fn drop(&mut self) {
self.chan.drop_weak();
}
}
impl<T> fmt::Debug for WeakSender<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("WeakSender").finish()
}
}
fn try_reserve_inner<T>(chan: &Chan<T, Bounded>, n: usize) -> Result<(), TrySendError<()>> {
match chan.limit.sem.try_acquire(n) {
Ok(()) => Ok(()),
Err(seen) if observed_closed(seen) => Err(TrySendError::Closed(())),
Err(_) => Err(TrySendError::Full(())),
}
}
async fn reserve_inner<T>(chan: &Chan<T, Bounded>, n: usize) -> Result<(), SendError<()>> {
if n > chan.limit.bound {
return Err(SendError(()));
}
Reserve {
chan,
n,
wait: FairWait::new(),
taken: false,
}
.await
}
struct Sending<'a, T> {
chan: &'a Chan<T, Bounded>,
wait: FairWait,
value: Option<T>,
_pinned: PhantomPinned,
}
impl<T> Sending<'_, T> {
fn let_go(&mut self) {
self.wait.cancel(&self.chan.limit.sem, 1);
self.chan.wake_if_finished();
}
}
impl<T> Future for Sending<'_, T> {
type Output = Result<(), SendError<T>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = unsafe { self.get_unchecked_mut() };
if this.value.is_none() {
polled_after_completion();
}
let chan = this.chan;
let granted = core::task::ready!(this.wait.poll(&chan.limit.sem, 1, cx.waker(), |n| {
chan.give_back(n);
}));
let Some(value) = this.value.take() else {
polled_after_completion();
};
match granted {
Ok(()) => {
this.chan.send(value);
Poll::Ready(Ok(()))
}
Err(_) => {
this.let_go();
Poll::Ready(Err(SendError(value)))
}
}
}
}
impl<T> Drop for Sending<'_, T> {
fn drop(&mut self) {
if self.value.is_some() {
self.let_go();
}
}
}
impl<T> internal::Withdraw for Sending<'_, T> {
fn withdraw(self: Pin<&mut Self>) {
let this = unsafe { self.get_unchecked_mut() };
if this.value.is_some() {
this.let_go();
}
}
}
#[cold]
#[expect(
clippy::panic,
reason = "polled after it completed, which the `Future` contract rules out"
)]
fn polled_after_completion() -> ! {
panic!("`Sender::send` polled after it completed");
}
struct Reserve<'a, T> {
chan: &'a Chan<T, Bounded>,
n: usize,
wait: FairWait,
taken: bool,
}
impl<T> Future for Reserve<'_, T> {
type Output = Result<(), SendError<()>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
let chan = this.chan;
this.wait
.poll(&chan.limit.sem, this.n, cx.waker(), |n| chan.give_back(n))
.map(|taken| {
this.taken = taken.is_ok();
taken.map_err(|_| SendError(()))
})
}
}
impl<T> Drop for Reserve<'_, T> {
fn drop(&mut self) {
if !self.taken {
self.wait.cancel(&self.chan.limit.sem, self.n);
self.chan.wake_if_finished();
}
}
}
#[cfg(all(test, not(loom)))]
mod tests {
use super::*;
use crate::sync::mpsc::error::TryRecvError;
use crate::sync::test_support::{Polled, hooked, with_budget_spent};
#[test]
fn the_free_slots_plus_the_queued_values_are_the_buffer() {
let (tx, mut rx) = channel::<u32>(4);
let permit = tx.try_reserve().unwrap();
tx.try_send(1).unwrap();
assert_eq!(tx.capacity() + rx.len() + 1, 4);
permit.send(2);
assert_eq!(tx.capacity() + rx.len(), 4);
rx.try_recv().unwrap();
assert_eq!(tx.capacity() + rx.len(), 4);
}
#[test]
fn a_send_granted_as_it_queues_spends_the_budget() {
let (tx, rx) = channel::<u32>(1);
tx.try_send(0).unwrap();
let rx = std::rc::Rc::new(core::cell::RefCell::new(rx));
let frees = std::rc::Rc::clone(&rx);
let waker = hooked(
Some(Box::new(move || {
assert_eq!(frees.borrow_mut().try_recv(), Ok(0));
})),
None,
);
let mut send = Box::pin(tx.send(1));
let mut cx = Context::from_waker(&waker);
let polled = with_budget_spent(|| send.as_mut().poll(&mut cx));
assert!(
polled.is_pending(),
"a send completed past the spent budget"
);
assert!(matches!(send.as_mut().poll(&mut cx), Poll::Ready(Ok(()))));
assert_eq!(rx.borrow_mut().try_recv(), Ok(1));
assert_eq!(tx.capacity(), 1);
}
#[test]
fn a_send_refused_for_the_budget_gives_its_slot_back_before_it_wakes() {
let (tx, rx) = channel::<u32>(1);
let rx = std::rc::Rc::new(core::cell::RefCell::new(rx));
let receiver = Polled::new(core::future::pending::<()>());
let rx_waker = receiver.waker();
let closes = std::rc::Rc::clone(&rx);
let waker = hooked(
None,
Some(Box::new(move || {
let mut rx = closes.borrow_mut();
rx.close();
let looked =
crate::coop::outside(|| rx.poll_recv(&mut Context::from_waker(&rx_waker)));
assert!(matches!(looked, Poll::Ready(None)), "a slot is still out");
})),
);
let mut send = Box::pin(tx.send(1));
let mut cx = Context::from_waker(&waker);
let polled = with_budget_spent(|| send.as_mut().poll(&mut cx));
assert!(polled.is_pending(), "the send was refused for the budget");
assert_eq!(tx.capacity(), 1, "the refused send gave its slot back");
assert!(matches!(send.as_mut().poll(&mut cx), Poll::Ready(Err(_))));
}
#[test]
fn a_zero_buffer_is_refused() {
let refused = std::panic::catch_unwind(|| channel::<u8>(0));
assert!(refused.is_err());
}
#[test]
fn a_reserve_granted_then_dropped_after_the_close_wakes_the_receiver() {
let (tx, mut rx) = channel::<()>(1);
let held = tx.try_reserve().unwrap();
let mut reserve = Polled::new(tx.reserve());
reserve.pending();
drop(held);
assert_eq!(reserve.wakes(), 1, "the slot was granted to the waiter");
assert_eq!(tx.capacity(), 0);
rx.close();
let mut recv = Polled::new(rx.recv());
recv.pending();
drop(reserve);
assert_eq!(tx.capacity(), 1, "the granted slot came back");
assert_eq!(recv.wakes(), 1, "and the closed receiver was told");
assert_eq!(recv.ready(), None);
}
#[test]
fn a_failed_reserve_wakes_a_receiver_waiting_for_idle() {
let (tx, mut rx) = channel::<()>(1);
let permit = tx.try_reserve().unwrap();
let mut reserve = Polled::new(tx.reserve());
reserve.pending();
rx.close();
assert!(reserve.wakes() >= 1);
assert!(matches!(reserve.ready(), Err(SendError(()))));
let mut recv = Polled::new(rx.recv());
recv.pending();
drop(permit);
assert_eq!(recv.wakes(), 1, "the last permit back ends the wait");
assert_eq!(recv.ready(), None);
}
#[test]
fn dropping_the_receiver_fails_a_waiting_reserve_with_the_value_kept() {
let (tx, rx) = channel::<String>(1);
tx.try_send("held".to_string()).unwrap();
let mut send = Polled::new(tx.send("kept".to_string()));
send.pending();
drop(rx);
let SendError(value) = send.ready().unwrap_err();
assert_eq!(value, "kept");
}
fn panicking_sender() -> core::task::Waker {
hooked(None, Some(Box::new(|| panic!("a sender's waker panicked"))))
}
fn quiet() -> core::task::Waker {
hooked(None, None)
}
#[test]
fn a_send_refused_for_the_budget_whose_wake_panics_gives_its_slot_back() {
let (tx, _rx) = channel::<u32>(1);
let waker = panicking_sender();
let mut send = Box::pin(tx.send(1));
let mut cx = Context::from_waker(&waker);
let caught = with_budget_spent(|| {
std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| send.as_mut().poll(&mut cx)))
});
assert!(caught.is_err(), "the panic reaches the caller");
assert_eq!(tx.capacity(), 1, "the slot was lost to the panic");
}
#[test]
fn a_receive_whose_release_panics_keeps_its_value_for_the_next_one() {
let (tx, mut rx) = channel::<u32>(1);
tx.try_send(1).unwrap();
let waker = panicking_sender();
let mut send = Box::pin(tx.send(2));
assert!(
send.as_mut()
.poll(&mut Context::from_waker(&waker))
.is_pending()
);
let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| rx.try_recv()));
assert!(caught.is_err(), "the panic reaches the receiver");
assert_eq!(rx.len(), 1, "the stashed value counts");
assert_eq!(rx.try_recv(), Ok(1), "the value was lost to the panic");
assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
let quiet = quiet();
assert!(matches!(
send.as_mut().poll(&mut Context::from_waker(&quiet)),
Poll::Ready(Ok(()))
));
assert_eq!(rx.try_recv(), Ok(2));
assert_eq!(tx.capacity(), 1, "the stash made up an admission");
tx.try_send(3).unwrap();
assert!(matches!(tx.try_send(4), Err(TrySendError::Full(4))));
}
#[test]
fn a_polled_receive_whose_release_panics_keeps_its_value() {
let (tx, mut rx) = channel::<u32>(1);
tx.try_send(1).unwrap();
let waker = panicking_sender();
let mut send = Box::pin(tx.send(2));
assert!(
send.as_mut()
.poll(&mut Context::from_waker(&waker))
.is_pending()
);
let mut recv = Polled::new(rx.recv());
let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| recv.poll()));
assert!(caught.is_err(), "the panic reaches the poll");
drop(recv);
assert_eq!(Polled::new(rx.recv()).ready(), Some(1));
let quiet = quiet();
assert!(matches!(
send.as_mut().poll(&mut Context::from_waker(&quiet)),
Poll::Ready(Ok(()))
));
assert_eq!(Polled::new(rx.recv()).ready(), Some(2));
}
#[test]
fn a_batch_receive_takes_a_stashed_value_without_a_release() {
let (tx, mut rx) = channel::<u32>(2);
tx.try_send(1).unwrap();
tx.try_send(2).unwrap();
let waker = panicking_sender();
let mut send = Box::pin(tx.send(3));
assert!(
send.as_mut()
.poll(&mut Context::from_waker(&waker))
.is_pending()
);
let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| rx.try_recv()));
assert!(caught.is_err());
let mut buf = Vec::new();
let mut recv = Polled::new(rx.recv_many(&mut buf, 8));
assert_eq!(recv.ready(), 2);
drop(recv);
assert_eq!(buf, [1, 2]);
drop(send);
assert_eq!(tx.capacity(), 2, "every admission is back, none made up");
}
#[test]
fn a_channel_dropped_with_a_stashed_value_drops_it_once() {
use std::sync::Arc;
use std::sync::atomic::{AtomicUsize, Ordering};
struct Counted(Arc<AtomicUsize>);
impl Drop for Counted {
fn drop(&mut self) {
self.0.fetch_add(1, Ordering::SeqCst);
}
}
let drops = Arc::new(AtomicUsize::new(0));
let (tx, mut rx) = channel::<Counted>(1);
assert!(tx.try_send(Counted(Arc::clone(&drops))).is_ok());
let waker = panicking_sender();
let mut send = Box::pin(tx.send(Counted(Arc::clone(&drops))));
assert!(
send.as_mut()
.poll(&mut Context::from_waker(&waker))
.is_pending()
);
let caught = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| rx.try_recv()));
assert!(caught.is_err());
assert_eq!(drops.load(Ordering::SeqCst), 0, "the value is stashed");
drop(send);
drop(rx);
drop(tx);
assert_eq!(drops.load(Ordering::SeqCst), 2, "each dropped exactly once");
}
}