use core::fmt;
use core::future::Future;
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 async fn send(&self, value: T) -> Result<(), SendError<T>> {
match self.reserve().await {
Ok(permit) => {
permit.send(value);
Ok(())
}
Err(_) => Err(SendError(value)),
}
}
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(self.send(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 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();
this.wait
.poll(&this.chan.limit.sem, this.n, cx.waker())
.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::test_support::Polled;
#[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_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");
}
}