use std::{
any::Any,
fmt,
future::Future,
mem,
ops::{Deref, DerefMut},
panic::{self, AssertUnwindSafe},
pin::Pin,
sync::{
Arc, Mutex, MutexGuard, OnceLock, PoisonError,
atomic::{AtomicUsize, Ordering, fence},
},
task::{Context, Poll, Waker},
};
pub struct Event {
shared: OnceLock<Arc<Shared>>,
}
impl Event {
pub const fn new() -> Self {
Self {
shared: OnceLock::new(),
}
}
pub fn listen(&self) -> EventListener {
self.add_listener(Order::Fence)
}
pub fn notify(&self, n: usize) -> usize {
self.send(n, Notification::Counting, Order::Fence)
}
pub fn notify_additional(&self, n: usize) -> usize {
self.send(n, Notification::Additional, Order::Fence)
}
#[cfg(test)]
pub(crate) fn slots(&self) -> Option<usize> {
self.shared.get().map(|shared| shared.lock().slots.len())
}
#[cfg(any(feature = "broadcast", feature = "lock", feature = "mpmc"))]
pub(crate) fn listen_unfenced(&self) -> EventListener {
self.add_listener(Order::SeqCst)
}
#[cfg(any(feature = "broadcast", feature = "lock", feature = "mpmc"))]
pub(crate) fn notify_unfenced(&self, n: usize) -> usize {
self.send(n, Notification::Counting, Order::SeqCst)
}
#[cfg(any(feature = "lock", feature = "mpmc"))]
pub(crate) fn notify_additional_unfenced(&self, n: usize) -> usize {
self.send(n, Notification::Additional, Order::SeqCst)
}
fn shared(&self) -> &Arc<Shared> {
self.shared.get_or_init(Arc::default)
}
fn add_listener(&self, order: Order) -> EventListener {
let shared = self.shared();
let key = shared.lock_writing(order.write()).insert();
order.fence();
EventListener {
shared: shared.clone(),
key: Some(key),
}
}
#[inline]
fn send(&self, n: usize, notification: Notification, order: Order) -> usize {
let Some(shared) = self.shared.get() else {
return self.send_to_new(n, notification, order);
};
let mut reaches = shared.would_reach(n, notification, order.look());
if !reaches && order.fence() {
reaches = shared.would_reach(n, notification, order.look());
}
if !reaches {
return 0;
}
shared.send(n, notification)
}
#[cold]
#[inline(never)]
fn send_to_new(&self, n: usize, notification: Notification, order: Order) -> usize {
self.shared();
self.send(n, notification, order)
}
}
impl Default for Event {
fn default() -> Self {
Self::new()
}
}
impl fmt::Debug for Event {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let (listeners, notified) = match self.shared.get() {
Some(shared) => {
let list = shared.lock();
(list.len, list.notified)
}
None => (0, 0),
};
f.debug_struct("Event")
.field("listeners", &listeners)
.field("notified", ¬ified)
.finish()
}
}
pub struct EventListener {
shared: Arc<Shared>,
key: Option<usize>,
}
impl Future for EventListener {
type Output = ();
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
let this = self.get_mut();
let Some(key) = this.key else {
return Poll::Ready(());
};
let mut clone = None;
let mut polled = this.shared.lock().poll(key, cx.waker(), &mut clone);
if let Polled::WantsClone = polled {
clone = Some(cx.waker().clone());
polled = this.shared.lock().poll(key, cx.waker(), &mut clone);
}
match polled {
Polled::Notified => {
this.key = None;
Poll::Ready(())
}
Polled::Waiting(replaced) => {
drop(replaced);
Poll::Pending
}
Polled::WantsClone => {
unreachable!("a second look, with a clone in hand, never asks for one")
}
}
}
}
impl Drop for EventListener {
fn drop(&mut self) {
let Some(key) = self.key.take() else {
return;
};
let mut wakers = Wakers::default();
let removed = {
let mut list = self.shared.lock();
let removed = list.remove(key);
if let State::Notified(notification) = removed {
list.notify(1, notification, &mut wakers);
}
removed
};
wakers.wake();
drop(removed);
}
}
impl fmt::Debug for EventListener {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let notified = self
.key
.is_none_or(|key| matches!(self.shared.lock().entry(key).state, State::Notified(_)));
f.debug_struct("EventListener")
.field("notified", ¬ified)
.finish()
}
}
struct Shared {
list: Mutex<List>,
notified: AtomicUsize,
}
impl Shared {
#[inline(never)]
fn send(&self, n: usize, notification: Notification) -> usize {
let mut wakers = Wakers::default();
let notified = self.lock().notify(n, notification, &mut wakers);
wakers.wake();
notified
}
fn lock(&self) -> Locked<'_> {
self.lock_writing(Ordering::Release)
}
fn lock_writing(&self, write: Ordering) -> Locked<'_> {
Locked {
list: lock(&self.list),
word: &self.notified,
write,
}
}
fn would_reach(&self, n: usize, notification: Notification, look: Ordering) -> bool {
let notified = self.notified.load(look);
match notification {
Notification::Counting => notified < n,
Notification::Additional => n > 0 && notified != NONE_WAITING,
}
}
}
impl Default for Shared {
fn default() -> Self {
Self {
list: Mutex::default(),
notified: AtomicUsize::new(NONE_WAITING),
}
}
}
struct Locked<'a> {
list: MutexGuard<'a, List>,
word: &'a AtomicUsize,
write: Ordering,
}
impl Deref for Locked<'_> {
type Target = List;
fn deref(&self) -> &List {
&self.list
}
}
impl DerefMut for Locked<'_> {
fn deref_mut(&mut self) -> &mut List {
&mut self.list
}
}
impl Drop for Locked<'_> {
fn drop(&mut self) {
let notified = if self.list.notified < self.list.len {
self.list.notified
} else {
NONE_WAITING
};
self.word.store(notified, self.write);
}
}
#[derive(Clone, Copy)]
enum Order {
Fence,
#[cfg(any(feature = "broadcast", feature = "lock", feature = "mpmc"))]
SeqCst,
}
impl Order {
fn write(self) -> Ordering {
match self {
Order::Fence => Ordering::Release,
#[cfg(any(feature = "broadcast", feature = "lock", feature = "mpmc"))]
Order::SeqCst => Ordering::SeqCst,
}
}
fn look(self) -> Ordering {
match self {
Order::Fence => Ordering::Acquire,
#[cfg(any(feature = "broadcast", feature = "lock", feature = "mpmc"))]
Order::SeqCst => Ordering::SeqCst,
}
}
fn fence(self) -> bool {
match self {
Order::Fence => {
fence(Ordering::SeqCst);
true
}
#[cfg(any(feature = "broadcast", feature = "lock", feature = "mpmc"))]
Order::SeqCst => false,
}
}
}
const NONE_WAITING: usize = usize::MAX;
#[derive(Default)]
struct List {
slots: Vec<Slot>,
vacant: Option<usize>,
tail: Option<usize>,
first_waiting: Option<usize>,
len: usize,
notified: usize,
}
impl List {
fn insert(&mut self) -> usize {
let entry = Entry {
prev: self.tail,
next: None,
state: State::Waiting(None),
};
let key = match self.vacant {
Some(key) => {
let Slot::Vacant(next) = mem::replace(&mut self.slots[key], Slot::Occupied(entry))
else {
unreachable!("the chain of vacant slots holds vacant slots only");
};
self.vacant = next;
key
}
None => {
self.slots.push(Slot::Occupied(entry));
self.slots.len() - 1
}
};
if let Some(tail) = self.tail {
self.entry_mut(tail).next = Some(key);
}
self.tail = Some(key);
if self.first_waiting.is_none() {
self.first_waiting = Some(key);
}
self.len += 1;
key
}
fn poll(&mut self, key: usize, waker: &Waker, clone: &mut Option<Waker>) -> Polled {
let State::Waiting(stored) = &mut self.entry_mut(key).state else {
self.remove(key);
return Polled::Notified;
};
if stored
.as_ref()
.is_some_and(|stored| stored.will_wake(waker))
{
return Polled::Waiting(None);
}
match clone.take() {
Some(clone) => Polled::Waiting(stored.replace(clone)),
None => Polled::WantsClone,
}
}
fn remove(&mut self, key: usize) -> State {
let Slot::Occupied(entry) = mem::replace(&mut self.slots[key], Slot::Vacant(self.vacant))
else {
unreachable!("a listener's key names a slot holding its entry");
};
self.vacant = Some(key);
if let Some(prev) = entry.prev {
self.entry_mut(prev).next = entry.next;
}
match entry.next {
Some(next) => self.entry_mut(next).prev = entry.prev,
None => self.tail = entry.prev,
}
if self.first_waiting == Some(key) {
self.first_waiting = entry.next;
}
self.len -= 1;
if let State::Notified(_) = entry.state {
self.notified -= 1;
}
entry.state
}
fn notify(&mut self, n: usize, notification: Notification, wakers: &mut Wakers) -> usize {
let wanted = match notification {
Notification::Counting => n.saturating_sub(self.notified),
Notification::Additional => n,
};
let mut reached = 0;
while reached < wanted {
let Some(key) = self.first_waiting else {
break;
};
let entry = self.entry_mut(key);
let next = entry.next;
let State::Waiting(waker) =
mem::replace(&mut entry.state, State::Notified(notification))
else {
unreachable!("every entry from the first one waiting onwards is waiting");
};
self.first_waiting = next;
if let Some(waker) = waker {
wakers.push(waker);
}
reached += 1;
}
self.notified += reached;
reached
}
fn entry(&self, key: usize) -> &Entry {
match &self.slots[key] {
Slot::Occupied(entry) => entry,
Slot::Vacant(_) => unreachable!("a listener's key names a slot holding its entry"),
}
}
fn entry_mut(&mut self, key: usize) -> &mut Entry {
match &mut self.slots[key] {
Slot::Occupied(entry) => entry,
Slot::Vacant(_) => unreachable!("a listener's key names a slot holding its entry"),
}
}
}
enum Slot {
Occupied(Entry),
Vacant(Option<usize>),
}
struct Entry {
prev: Option<usize>,
next: Option<usize>,
state: State,
}
enum State {
Waiting(Option<Waker>),
Notified(Notification),
}
#[derive(Clone, Copy)]
enum Notification {
Counting,
Additional,
}
enum Polled {
Notified,
Waiting(Option<Waker>),
WantsClone,
}
#[derive(Default)]
struct Wakers {
first: Option<Waker>,
rest: Vec<Waker>,
}
impl Wakers {
fn push(&mut self, waker: Waker) {
match self.first {
None => self.first = Some(waker),
Some(_) => self.rest.push(waker),
}
}
fn wake(self) {
let mut first_panic = None;
for waker in self.first.into_iter().chain(self.rest) {
if let Err(panic) = panic::catch_unwind(AssertUnwindSafe(|| waker.wake())) {
match first_panic {
None => first_panic = Some(panic),
Some(_) => dispose(panic),
}
}
}
if let Some(panic) = first_panic {
panic::resume_unwind(panic);
}
}
}
fn dispose(payload: Box<dyn Any + Send>) {
let Err(payload) = panic::catch_unwind(AssertUnwindSafe(move || drop(payload))) else {
return;
};
if let Err(payload) = panic::catch_unwind(AssertUnwindSafe(move || drop(payload))) {
mem::forget(payload);
}
}
fn lock<T>(mutex: &Mutex<T>) -> MutexGuard<'_, T> {
mutex.lock().unwrap_or_else(PoisonError::into_inner)
}