#![allow(unsafe_code)]
use core::fmt;
use kovan::{
Atomic, CachePadded, RetiredNode, Shared as Ptr, alloc_node, drop_node, pin, retire_node,
};
use super::Notify;
use super::internal::{Arc, AtomicBool, AtomicU64, AtomicUsize, Ordering};
pub use self::error::{RecvError, SendError, TryRecvError};
pub mod error {
use core::fmt;
use std::error::Error;
#[derive(PartialEq, Eq, Clone, Copy)]
pub struct SendError<T>(pub T);
impl<T> fmt::Debug for SendError<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("SendError").finish_non_exhaustive()
}
}
impl<T> fmt::Display for SendError<T> {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(fmt, "channel closed")
}
}
impl<T> Error for SendError<T> {}
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum RecvError {
Closed,
Lagged(u64),
}
impl fmt::Display for RecvError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
RecvError::Closed => write!(fmt, "channel closed"),
RecvError::Lagged(missed) => write!(fmt, "channel lagged by {missed}"),
}
}
}
impl Error for RecvError {}
#[derive(Debug, PartialEq, Eq, Clone)]
pub enum TryRecvError {
Empty,
Closed,
Lagged(u64),
}
impl fmt::Display for TryRecvError {
fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
TryRecvError::Empty => write!(fmt, "channel empty"),
TryRecvError::Closed => write!(fmt, "channel closed"),
TryRecvError::Lagged(missed) => write!(fmt, "channel lagged by {missed}"),
}
}
}
impl Error for TryRecvError {}
}
#[repr(C)]
struct Node<T> {
retired: RetiredNode,
pos: u64,
value: T,
}
struct Cursor {
next: AtomicU64,
}
struct Shared<T: Send + Sync + 'static> {
slots: Box<[Atomic<Node<T>>]>,
mask: u64,
tail: CachePadded<AtomicU64>,
closed: AtomicBool,
senders: AtomicUsize,
weak_senders: AtomicUsize,
cursors: kovan_map::HashMap<u64, Arc<Cursor>>,
receivers: AtomicUsize,
notify_rx: Notify,
notify_tx: Notify,
}
unsafe impl<T: Send + Sync + 'static> Send for Shared<T> {}
unsafe impl<T: Send + Sync + 'static> Sync for Shared<T> {}
impl<T: Send + Sync + 'static> Shared<T> {
fn capacity(&self) -> u64 {
self.mask + 1
}
fn subscribe(shared: &Arc<Self>) -> Receiver<T> {
static IDS: core::sync::atomic::AtomicU64 = core::sync::atomic::AtomicU64::new(0);
let id = IDS.fetch_add(1, core::sync::atomic::Ordering::Relaxed);
let cursor = Arc::new(Cursor {
next: AtomicU64::new(shared.tail.load(Ordering::SeqCst)),
});
shared.receivers.fetch_add(1, Ordering::SeqCst);
shared.cursors.insert(id, Arc::clone(&cursor));
Receiver {
shared: Arc::clone(shared),
id,
cursor,
}
}
}
impl<T: Send + Sync + 'static> Drop for Shared<T> {
fn drop(&mut self) {
let guard = pin();
for slot in self.slots.iter() {
let node = slot.load(Ordering::Acquire, &guard);
if !node.is_null() {
unsafe { drop_node(node.as_raw()) };
}
}
}
}
pub fn channel<T: Clone + Send + Sync + 'static>(capacity: usize) -> (Sender<T>, Receiver<T>) {
let tx = Sender::new(capacity);
let rx = tx.subscribe();
(tx, rx)
}
pub struct Sender<T: Send + Sync + 'static> {
shared: Arc<Shared<T>>,
}
pub struct Receiver<T: Send + Sync + 'static> {
shared: Arc<Shared<T>>,
id: u64,
cursor: Arc<Cursor>,
}
pub struct WeakSender<T: Send + Sync + 'static> {
shared: Arc<Shared<T>>,
}
impl<T: Send + Sync + 'static> Sender<T> {
pub fn new(capacity: usize) -> Self {
assert!(
capacity > 0,
"a broadcast channel's capacity cannot be zero"
);
assert!(
capacity <= usize::MAX >> 1,
"a broadcast channel's capacity cannot exceed usize::MAX / 2"
);
let capacity = capacity.next_power_of_two();
Self {
shared: Arc::new(Shared {
slots: (0..capacity).map(|_| Atomic::null()).collect(),
mask: capacity as u64 - 1,
tail: CachePadded::new(AtomicU64::new(0)),
closed: AtomicBool::new(false),
senders: AtomicUsize::new(1),
weak_senders: AtomicUsize::new(0),
cursors: kovan_map::HashMap::with_capacity(64),
receivers: AtomicUsize::new(0),
notify_rx: Notify::new(),
notify_tx: Notify::new(),
}),
}
}
pub fn send(&self, value: T) -> Result<usize, SendError<T>> {
let shared = &*self.shared;
let receivers = shared.receivers.load(Ordering::SeqCst);
if receivers == 0 {
return Err(SendError(value));
}
let pos = shared.tail.fetch_add(1, Ordering::SeqCst);
let slot = &shared.slots[(pos & shared.mask) as usize];
let node = alloc_node(Node {
retired: RetiredNode::new(),
pos,
value,
});
let guard = pin();
let mut current = slot.load(Ordering::Acquire, &guard);
loop {
if unsafe { current.as_ref() }.is_some_and(|old| old.pos > pos) {
unsafe { drop_node(node) };
break;
}
let new = unsafe { Ptr::from_raw(node) };
match slot.compare_exchange(current, new, Ordering::AcqRel, Ordering::Acquire, &guard) {
Ok(old) => {
if !old.is_null() {
unsafe { retire_node(old.as_raw()) };
}
break;
}
Err(now) => current = now,
}
}
drop(guard);
shared.notify_rx.notify_waiters();
Ok(receivers)
}
pub fn subscribe(&self) -> Receiver<T> {
Shared::subscribe(&self.shared)
}
pub fn receiver_count(&self) -> usize {
self.shared.receivers.load(Ordering::SeqCst)
}
#[must_use = "a WeakSender that is not kept does nothing"]
pub fn downgrade(&self) -> WeakSender<T> {
self.shared.weak_senders.fetch_add(1, Ordering::Relaxed);
WeakSender {
shared: Arc::clone(&self.shared),
}
}
pub fn strong_count(&self) -> usize {
self.shared.senders.load(Ordering::Acquire)
}
pub fn weak_count(&self) -> usize {
self.shared.weak_senders.load(Ordering::Relaxed)
}
pub fn len(&self) -> usize {
let shared = &*self.shared;
let tail = shared.tail.load(Ordering::SeqCst);
let behind = shared
.cursors
.iter()
.map(|(_, cursor)| tail.saturating_sub(cursor.next.load(Ordering::Acquire)))
.max()
.unwrap_or(0);
behind.min(shared.capacity()) as usize
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn same_channel(&self, other: &Self) -> bool {
Arc::ptr_eq(&self.shared, &other.shared)
}
pub async fn closed(&self) {
while self.receiver_count() > 0 {
let notified = self.shared.notify_tx.notified();
if self.receiver_count() == 0 {
return;
}
notified.await;
}
}
}
impl<T: Send + Sync + 'static> Clone for Sender<T> {
fn clone(&self) -> Self {
self.shared.senders.fetch_add(1, Ordering::Relaxed);
Self {
shared: Arc::clone(&self.shared),
}
}
}
impl<T: Send + Sync + 'static> Drop for Sender<T> {
fn drop(&mut self) {
if self.shared.senders.fetch_sub(1, Ordering::AcqRel) == 1 {
self.shared.closed.store(true, Ordering::SeqCst);
self.shared.notify_rx.notify_waiters();
}
}
}
impl<T: Send + Sync + 'static> fmt::Debug for Sender<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Sender")
.field("receivers", &self.receiver_count())
.finish_non_exhaustive()
}
}
impl<T: Send + Sync + 'static> WeakSender<T> {
pub fn upgrade(&self) -> Option<Sender<T>> {
let senders = &self.shared.senders;
let mut count = senders.load(Ordering::Acquire);
while count > 0 {
match senders.compare_exchange_weak(
count,
count + 1,
Ordering::AcqRel,
Ordering::Acquire,
) {
Ok(_) => {
return Some(Sender {
shared: Arc::clone(&self.shared),
});
}
Err(now) => count = now,
}
}
None
}
pub fn strong_count(&self) -> usize {
self.shared.senders.load(Ordering::Acquire)
}
pub fn weak_count(&self) -> usize {
self.shared.weak_senders.load(Ordering::Relaxed)
}
}
impl<T: Send + Sync + 'static> Clone for WeakSender<T> {
fn clone(&self) -> Self {
self.shared.weak_senders.fetch_add(1, Ordering::Relaxed);
Self {
shared: Arc::clone(&self.shared),
}
}
}
impl<T: Send + Sync + 'static> Drop for WeakSender<T> {
fn drop(&mut self) {
self.shared.weak_senders.fetch_sub(1, Ordering::Relaxed);
}
}
impl<T: Send + Sync + 'static> fmt::Debug for WeakSender<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("WeakSender").finish_non_exhaustive()
}
}
impl<T: Send + Sync + 'static> Receiver<T> {
fn next(&self) -> u64 {
self.cursor.next.load(Ordering::Relaxed)
}
fn advance(&self, next: u64) {
self.cursor.next.store(next, Ordering::Release);
}
pub fn len(&self) -> usize {
let tail = self.shared.tail.load(Ordering::SeqCst);
tail.saturating_sub(self.next()) as usize
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn resubscribe(&self) -> Self {
Shared::subscribe(&self.shared)
}
}
impl<T: Clone + Send + Sync + 'static> Receiver<T> {
pub fn try_recv(&mut self) -> Result<T, TryRecvError> {
self.try_recv_if(|| true).map_err(|look| match look {
Look::Lagged(missed) => TryRecvError::Lagged(missed),
Look::Closed => TryRecvError::Closed,
Look::Empty | Look::Refused => TryRecvError::Empty,
})
}
fn try_recv_if(&mut self, admit: impl FnOnce() -> bool) -> Result<T, Look> {
let shared = &*self.shared;
let next = self.next();
let guard = pin();
let slot = &shared.slots[(next & shared.mask) as usize];
match unsafe { slot.load(Ordering::Acquire, &guard).as_ref() } {
Some(node) if node.pos == next => {
if !admit() {
return Err(Look::Refused);
}
let value = node.value.clone();
drop(guard);
self.advance(next + 1);
return Ok(value);
}
Some(node) if node.pos > next => {
drop(guard);
if !admit() {
return Err(Look::Refused);
}
let oldest = shared
.tail
.load(Ordering::SeqCst)
.saturating_sub(shared.capacity());
self.advance(oldest);
return Err(Look::Lagged(oldest - next));
}
_ => {}
}
drop(guard);
if shared.closed.load(Ordering::SeqCst) && shared.tail.load(Ordering::SeqCst) <= next {
return Err(if admit() { Look::Closed } else { Look::Refused });
}
Err(Look::Empty)
}
pub async fn recv(&mut self) -> Result<T, RecvError> {
let shared = Arc::clone(&self.shared);
loop {
let notified = shared.notify_rx.notified();
match self.try_recv_if(crate::coop::try_spend) {
Ok(value) => return Ok(value),
Err(Look::Lagged(missed)) => return Err(RecvError::Lagged(missed)),
Err(Look::Closed) => return Err(RecvError::Closed),
Err(Look::Empty) => notified.await,
Err(Look::Refused) => crate::coop::yield_spent().await,
}
}
}
pub fn blocking_recv(&mut self) -> Result<T, RecvError> {
super::internal::blocking(self.recv())
}
}
enum Look {
Empty,
Lagged(u64),
Closed,
Refused,
}
impl<T: Send + Sync + 'static> Drop for Receiver<T> {
fn drop(&mut self) {
self.shared.cursors.remove(&self.id);
if self.shared.receivers.fetch_sub(1, Ordering::SeqCst) == 1 {
self.shared.notify_tx.notify_waiters();
}
}
}
impl<T: Send + Sync + 'static> fmt::Debug for Receiver<T> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Receiver")
.field("next", &self.next())
.finish_non_exhaustive()
}
}