use std::{
collections::VecDeque,
error, fmt,
future::Future,
mem,
num::NonZeroUsize,
pin::Pin,
sync::{Arc, Mutex, MutexGuard, PoisonError},
task::{Context, Poll},
};
use futures_core::{ready, stream::Stream};
use crate::{Event, EventListener};
pub fn channel<T>(cap: NonZeroUsize) -> (Sender<T>, Receiver<T>) {
let channel = Arc::new(Channel {
inner: Mutex::new(Inner {
queue: VecDeque::with_capacity(cap.get()),
capacity: cap,
overflow: false,
await_active: true,
receiver_count: 1,
inactive_receiver_count: 0,
sender_count: 1,
head_pos: 0,
is_closed: false,
}),
send_ops: Event::new(),
recv_ops: Event::new(),
});
let s = Sender {
channel: channel.clone(),
};
let r = Receiver {
channel,
pos: 0,
listener: None,
};
(s, r)
}
#[derive(Debug)]
pub struct Sender<T> {
channel: Arc<Channel<T>>,
}
impl<T> Sender<T> {
pub fn capacity(&self) -> NonZeroUsize {
lock(&self.channel.inner).capacity
}
pub fn set_capacity(&mut self, new_cap: NonZeroUsize) {
self.channel.set_capacity(new_cap);
}
pub fn overflow(&self) -> bool {
lock(&self.channel.inner).overflow
}
pub fn set_overflow(&mut self, overflow: bool) {
self.channel.set_overflow(overflow);
}
pub fn await_active(&self) -> bool {
lock(&self.channel.inner).await_active
}
pub fn set_await_active(&mut self, await_active: bool) {
self.channel.set_await_active(await_active);
}
pub fn close(&self) -> bool {
self.channel.close()
}
pub fn is_closed(&self) -> bool {
lock(&self.channel.inner).is_closed
}
pub fn is_empty(&self) -> bool {
lock(&self.channel.inner).queue.is_empty()
}
pub fn is_full(&self) -> bool {
let inner = lock(&self.channel.inner);
inner.queue.len() == inner.capacity.get()
}
pub fn len(&self) -> usize {
lock(&self.channel.inner).queue.len()
}
pub fn receiver_count(&self) -> usize {
lock(&self.channel.inner).receiver_count
}
pub fn inactive_receiver_count(&self) -> usize {
lock(&self.channel.inner).inactive_receiver_count
}
pub fn sender_count(&self) -> usize {
lock(&self.channel.inner).sender_count
}
pub fn new_receiver(&self) -> Receiver<T> {
self.channel.new_receiver()
}
}
impl<T: Clone> Sender<T> {
pub fn broadcast(&self, msg: T) -> Send<'_, T> {
Send {
sender: self,
listener: None,
msg: Some(msg),
}
}
pub fn try_broadcast(&self, msg: T) -> Result<Option<T>, TrySendError<T>> {
let mut ret = None;
let mut inner = lock(&self.channel.inner);
if inner.is_closed {
return Err(TrySendError::Closed(msg));
} else if inner.receiver_count == 0 {
assert!(inner.inactive_receiver_count != 0);
return Err(TrySendError::Inactive(msg));
} else if inner.queue.len() == inner.capacity.get() {
if inner.overflow {
ret = inner.queue.pop_front().map(|(m, _)| m);
} else {
return Err(TrySendError::Full(msg));
}
}
let receiver_count = inner.receiver_count;
inner.queue.push_back((msg, receiver_count));
if ret.is_some() {
inner.head_pos += 1;
}
drop(inner);
self.channel.recv_ops.notify_unfenced(usize::MAX);
Ok(ret)
}
}
impl<T> Drop for Sender<T> {
fn drop(&mut self) {
let closed = {
let mut inner = lock(&self.channel.inner);
inner.sender_count -= 1;
inner.sender_count == 0 && inner.close()
};
if closed {
self.channel.notify_closed();
}
}
}
impl<T> Clone for Sender<T> {
fn clone(&self) -> Self {
lock(&self.channel.inner).sender_count += 1;
Sender {
channel: self.channel.clone(),
}
}
}
#[derive(Debug)]
pub struct Receiver<T> {
channel: Arc<Channel<T>>,
pos: u64,
listener: Option<EventListener>,
}
impl<T> Receiver<T> {
pub fn capacity(&self) -> NonZeroUsize {
lock(&self.channel.inner).capacity
}
pub fn set_capacity(&mut self, new_cap: NonZeroUsize) {
self.channel.set_capacity(new_cap);
}
pub fn overflow(&self) -> bool {
lock(&self.channel.inner).overflow
}
pub fn set_overflow(&mut self, overflow: bool) {
self.channel.set_overflow(overflow);
}
pub fn await_active(&self) -> bool {
lock(&self.channel.inner).await_active
}
pub fn set_await_active(&mut self, await_active: bool) {
self.channel.set_await_active(await_active);
}
pub fn close(&self) -> bool {
self.channel.close()
}
pub fn is_closed(&self) -> bool {
lock(&self.channel.inner).is_closed
}
pub fn is_empty(&self) -> bool {
lock(&self.channel.inner).queue.is_empty()
}
pub fn is_full(&self) -> bool {
let inner = lock(&self.channel.inner);
inner.queue.len() == inner.capacity.get()
}
pub fn len(&self) -> usize {
lock(&self.channel.inner).queue.len()
}
pub fn receiver_count(&self) -> usize {
lock(&self.channel.inner).receiver_count
}
pub fn inactive_receiver_count(&self) -> usize {
lock(&self.channel.inner).inactive_receiver_count
}
pub fn sender_count(&self) -> usize {
lock(&self.channel.inner).sender_count
}
pub fn deactivate(self) -> InactiveReceiver<T> {
lock(&self.channel.inner).inactive_receiver_count += 1;
InactiveReceiver {
channel: self.channel.clone(),
}
}
}
impl<T: Clone> Receiver<T> {
pub fn recv(&mut self) -> Recv<'_, T> {
Recv {
receiver: self,
listener: None,
}
}
pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
let mut inner = lock(&self.channel.inner);
let overflow = inner.overflow;
let received = inner.try_recv_at(&mut self.pos);
let popped = matches!(received, Ok(Ok(_)));
let received = received.map(|cow| cow.unwrap_or_else(T::clone));
drop(inner);
if popped && !overflow {
self.channel.send_ops.notify_unfenced(1);
}
received
}
pub fn new_sender(&self) -> Sender<T> {
lock(&self.channel.inner).sender_count += 1;
Sender {
channel: self.channel.clone(),
}
}
pub fn new_receiver(&self) -> Self {
self.channel.new_receiver()
}
pub fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<T, RecvError>>> {
loop {
if let Some(listener) = self.listener.as_mut() {
ready!(Pin::new(listener).poll(cx));
self.listener = None;
}
loop {
match self.try_recv() {
Ok(msg) => {
self.listener = None;
return Poll::Ready(Some(Ok(msg)));
}
Err(TryRecvError::Closed) => {
self.listener = None;
return Poll::Ready(None);
}
Err(TryRecvError::Overflowed(n)) => {
self.listener = None;
return Poll::Ready(Some(Err(RecvError::Overflowed(n))));
}
Err(TryRecvError::Empty) => {}
}
match self.listener.as_mut() {
None => {
self.listener = Some(self.channel.recv_ops.listen_unfenced());
}
Some(_) => {
break;
}
}
}
}
}
}
impl<T> Drop for Receiver<T> {
fn drop(&mut self) {
let mut last = Vec::new();
let (closed, senders_fail, room) = {
let mut inner = lock(&self.channel.inner);
loop {
match inner.try_recv_at(&mut self.pos) {
Ok(Ok(msg)) => last.push(msg),
Ok(Err(_)) | Err(TryRecvError::Overflowed(_)) => continue,
Err(TryRecvError::Closed | TryRecvError::Empty) => break,
}
}
inner.receiver_count -= 1;
let closed = inner.close_channel();
(
closed,
!inner.is_closed && inner.receiver_count == 0 && !inner.await_active,
!last.is_empty() && !inner.overflow,
)
};
if closed {
self.channel.notify_closed();
} else if senders_fail {
self.channel.send_ops.notify_unfenced(usize::MAX);
} else if room {
self.channel.send_ops.notify_unfenced(1);
}
drop(last);
}
}
impl<T> Clone for Receiver<T> {
fn clone(&self) -> Self {
let mut inner = lock(&self.channel.inner);
inner.receiver_count += 1;
let n = self.pos.saturating_sub(inner.head_pos) as usize;
for (_elt, waiters) in inner.queue.iter_mut().skip(n) {
*waiters += 1;
}
Receiver {
channel: self.channel.clone(),
pos: self.pos,
listener: None,
}
}
}
impl<T: Clone> Stream for Receiver<T> {
type Item = T;
fn poll_next(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
loop {
match ready!(self.poll_recv(cx)) {
Some(Ok(val)) => return Poll::Ready(Some(val)),
Some(Err(RecvError::Overflowed(_))) => continue,
None | Some(Err(RecvError::Closed)) => return Poll::Ready(None),
}
}
}
}
impl<T: Clone> futures_core::stream::FusedStream for Receiver<T> {
fn is_terminated(&self) -> bool {
let inner = lock(&self.channel.inner);
inner.is_closed && self.pos.max(inner.head_pos) == inner.head_pos + inner.queue.len() as u64
}
}
#[derive(Debug)]
pub struct InactiveReceiver<T> {
channel: Arc<Channel<T>>,
}
impl<T> InactiveReceiver<T> {
pub fn activate(self) -> Receiver<T> {
self.activate_cloned()
}
pub fn activate_cloned(&self) -> Receiver<T> {
self.channel.new_receiver()
}
pub fn capacity(&self) -> NonZeroUsize {
lock(&self.channel.inner).capacity
}
pub fn set_capacity(&mut self, new_cap: NonZeroUsize) {
self.channel.set_capacity(new_cap);
}
pub fn overflow(&self) -> bool {
lock(&self.channel.inner).overflow
}
pub fn set_overflow(&mut self, overflow: bool) {
self.channel.set_overflow(overflow);
}
pub fn await_active(&self) -> bool {
lock(&self.channel.inner).await_active
}
pub fn set_await_active(&mut self, await_active: bool) {
self.channel.set_await_active(await_active);
}
pub fn close(&self) -> bool {
self.channel.close()
}
pub fn is_closed(&self) -> bool {
lock(&self.channel.inner).is_closed
}
pub fn is_empty(&self) -> bool {
lock(&self.channel.inner).queue.is_empty()
}
pub fn is_full(&self) -> bool {
let inner = lock(&self.channel.inner);
inner.queue.len() == inner.capacity.get()
}
pub fn len(&self) -> usize {
lock(&self.channel.inner).queue.len()
}
pub fn receiver_count(&self) -> usize {
lock(&self.channel.inner).receiver_count
}
pub fn inactive_receiver_count(&self) -> usize {
lock(&self.channel.inner).inactive_receiver_count
}
pub fn sender_count(&self) -> usize {
lock(&self.channel.inner).sender_count
}
}
impl<T> Clone for InactiveReceiver<T> {
fn clone(&self) -> Self {
lock(&self.channel.inner).inactive_receiver_count += 1;
InactiveReceiver {
channel: self.channel.clone(),
}
}
}
impl<T> Drop for InactiveReceiver<T> {
fn drop(&mut self) {
let closed = {
let mut inner = lock(&self.channel.inner);
inner.inactive_receiver_count -= 1;
inner.close_channel()
};
if closed {
self.channel.notify_closed();
}
}
}
#[derive(Debug)]
#[must_use = "futures do nothing unless .awaited"]
pub struct Send<'a, T> {
sender: &'a Sender<T>,
listener: Option<EventListener>,
msg: Option<T>,
}
impl<T> Unpin for Send<'_, T> {}
impl<T: Clone> Future for Send<'_, T> {
type Output = Result<Option<T>, SendError<T>>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
loop {
let msg = this.msg.take().unwrap();
let channel = &this.sender.channel;
match this.sender.try_broadcast(msg) {
Ok(msg) => {
this.listener = None;
let room = {
let inner = lock(&channel.inner);
inner.queue.len() < inner.capacity.get()
};
if room {
channel.send_ops.notify_unfenced(1);
}
return Poll::Ready(Ok(msg));
}
Err(TrySendError::Closed(msg)) => {
this.listener = None;
return Poll::Ready(Err(SendError(msg)));
}
Err(TrySendError::Full(m)) => this.msg = Some(m),
Err(TrySendError::Inactive(m)) if lock(&channel.inner).await_active => {
this.msg = Some(m)
}
Err(TrySendError::Inactive(m)) => {
this.listener = None;
return Poll::Ready(Err(SendError(m)));
}
}
match &mut this.listener {
None => {
this.listener = Some(channel.send_ops.listen_unfenced());
}
Some(listener) => {
ready!(Pin::new(listener).poll(cx));
this.listener = None;
}
}
}
}
}
#[derive(Debug)]
#[must_use = "futures do nothing unless .awaited"]
pub struct Recv<'a, T> {
receiver: &'a mut Receiver<T>,
listener: Option<EventListener>,
}
impl<T: Clone> Future for Recv<'_, T> {
type Output = Result<T, RecvError>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
loop {
match this.receiver.try_recv() {
Ok(msg) => return Poll::Ready(Ok(msg)),
Err(TryRecvError::Closed) => return Poll::Ready(Err(RecvError::Closed)),
Err(TryRecvError::Overflowed(n)) => {
return Poll::Ready(Err(RecvError::Overflowed(n)));
}
Err(TryRecvError::Empty) => {}
}
match &mut this.listener {
None => {
this.listener = Some(this.receiver.channel.recv_ops.listen_unfenced());
}
Some(listener) => {
ready!(Pin::new(listener).poll(cx));
this.listener = None;
}
}
}
}
}
#[derive(PartialEq, Eq, Clone, Copy)]
pub struct SendError<T>(pub T);
impl<T> SendError<T> {
pub fn into_inner(self) -> T {
self.0
}
}
impl<T> error::Error for SendError<T> {}
impl<T> fmt::Debug for SendError<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "SendError(..)")
}
}
impl<T> fmt::Display for SendError<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "sending into a closed channel")
}
}
#[derive(PartialEq, Eq, Clone, Copy)]
pub enum TrySendError<T> {
Full(T),
Closed(T),
Inactive(T),
}
impl<T> TrySendError<T> {
pub fn into_inner(self) -> T {
match self {
TrySendError::Full(t) => t,
TrySendError::Closed(t) => t,
TrySendError::Inactive(t) => t,
}
}
pub fn is_full(&self) -> bool {
match self {
TrySendError::Full(_) => true,
TrySendError::Closed(_) | TrySendError::Inactive(_) => false,
}
}
pub fn is_closed(&self) -> bool {
match self {
TrySendError::Full(_) | TrySendError::Inactive(_) => false,
TrySendError::Closed(_) => true,
}
}
pub fn is_disconnected(&self) -> bool {
match self {
TrySendError::Full(_) | TrySendError::Closed(_) => false,
TrySendError::Inactive(_) => true,
}
}
}
impl<T> error::Error for TrySendError<T> {}
impl<T> fmt::Debug for TrySendError<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match *self {
TrySendError::Full(..) => write!(f, "Full(..)"),
TrySendError::Closed(..) => write!(f, "Closed(..)"),
TrySendError::Inactive(..) => write!(f, "Inactive(..)"),
}
}
}
impl<T> fmt::Display for TrySendError<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match *self {
TrySendError::Full(..) => write!(f, "sending into a full channel"),
TrySendError::Closed(..) => write!(f, "sending into a closed channel"),
TrySendError::Inactive(..) => write!(f, "sending into the void (no active receivers)"),
}
}
}
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum RecvError {
Overflowed(u64),
Closed,
}
impl error::Error for RecvError {}
impl fmt::Display for RecvError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Overflowed(n) => write!(f, "receiving skipped {} messages", n),
Self::Closed => write!(f, "receiving from an empty and closed channel"),
}
}
}
#[derive(PartialEq, Eq, Clone, Copy, Debug)]
pub enum TryRecvError {
Overflowed(u64),
Empty,
Closed,
}
impl TryRecvError {
pub fn is_empty(&self) -> bool {
match self {
TryRecvError::Empty => true,
TryRecvError::Closed => false,
TryRecvError::Overflowed(_) => false,
}
}
pub fn is_closed(&self) -> bool {
match self {
TryRecvError::Empty => false,
TryRecvError::Closed => true,
TryRecvError::Overflowed(_) => false,
}
}
pub fn is_overflowed(&self) -> bool {
match self {
TryRecvError::Empty => false,
TryRecvError::Closed => false,
TryRecvError::Overflowed(_) => true,
}
}
}
impl error::Error for TryRecvError {}
impl fmt::Display for TryRecvError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match *self {
TryRecvError::Empty => write!(f, "receiving from an empty channel"),
TryRecvError::Closed => write!(f, "receiving from an empty and closed channel"),
TryRecvError::Overflowed(n) => {
write!(f, "receiving operation observed {} lost messages", n)
}
}
}
}
#[derive(Debug)]
struct Channel<T> {
inner: Mutex<Inner<T>>,
send_ops: Event,
recv_ops: Event,
}
impl<T> Channel<T> {
fn new_receiver(self: &Arc<Self>) -> Receiver<T> {
let (pos, first) = {
let mut inner = lock(&self.inner);
inner.receiver_count += 1;
(
inner.head_pos + inner.queue.len() as u64,
inner.receiver_count == 1,
)
};
let receiver = Receiver {
channel: self.clone(),
pos,
listener: None,
};
if first {
self.send_ops.notify_unfenced(1);
}
receiver
}
fn close(&self) -> bool {
let closed = lock(&self.inner).close();
if closed {
self.notify_closed();
}
closed
}
fn notify_closed(&self) {
self.send_ops.notify_unfenced(usize::MAX);
self.recv_ops.notify_unfenced(usize::MAX);
}
fn set_capacity(&self, new_cap: NonZeroUsize) {
let (grown, dropped) = {
let mut inner = lock(&self.inner);
let old_cap = inner.capacity;
(new_cap > old_cap, inner.set_capacity(new_cap))
};
if grown {
self.send_ops.notify_unfenced(1);
}
drop(dropped);
}
fn set_overflow(&self, overflow: bool) {
let was = mem::replace(&mut lock(&self.inner).overflow, overflow);
if overflow && !was {
self.send_ops.notify_unfenced(usize::MAX);
}
}
fn set_await_active(&self, await_active: bool) {
let was = mem::replace(&mut lock(&self.inner).await_active, await_active);
if was && !await_active {
self.send_ops.notify_unfenced(usize::MAX);
}
}
}
#[derive(Debug)]
struct Inner<T> {
queue: VecDeque<(T, usize)>,
capacity: NonZeroUsize,
receiver_count: usize,
inactive_receiver_count: usize,
sender_count: usize,
head_pos: u64,
overflow: bool,
await_active: bool,
is_closed: bool,
}
impl<T> Inner<T> {
fn try_recv_at(&mut self, pos: &mut u64) -> Result<Result<T, &T>, TryRecvError> {
let i = match pos.checked_sub(self.head_pos) {
Some(i) => i
.try_into()
.expect("Head position more than usize::MAX behind a receiver"),
None => {
let count = self.head_pos - *pos;
*pos = self.head_pos;
return Err(TryRecvError::Overflowed(count));
}
};
let last_waiter;
if let Some((_elt, waiters)) = self.queue.get_mut(i) {
*pos += 1;
*waiters -= 1;
last_waiter = *waiters == 0;
} else {
debug_assert_eq!(i, self.queue.len());
if self.is_closed {
return Err(TryRecvError::Closed);
} else {
return Err(TryRecvError::Empty);
}
}
if last_waiter {
assert_eq!(i, 0);
let elt = self.queue.pop_front().unwrap().0;
self.head_pos += 1;
Ok(Ok(elt))
} else {
Ok(Err(&self.queue[i].0))
}
}
fn close(&mut self) -> bool {
if self.is_closed {
return false;
}
self.is_closed = true;
true
}
fn set_capacity(&mut self, new_cap: NonZeroUsize) -> Vec<T> {
self.capacity = new_cap;
let new_cap = new_cap.get();
if new_cap > self.queue.capacity() {
let diff = new_cap - self.queue.capacity();
self.queue.reserve(diff);
}
if new_cap < self.queue.len() {
let diff = self.queue.len() - new_cap;
self.head_pos += diff as u64;
self.queue.drain(0..diff).map(|(msg, _)| msg).collect()
} else {
Vec::new()
}
}
fn close_channel(&mut self) -> bool {
self.receiver_count == 0 && self.inactive_receiver_count == 0 && self.close()
}
}
fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(PoisonError::into_inner)
}