use alloc::boxed::Box;
use core::future::Future;
use core::pin::Pin;
use core::task::{Context, Poll};
use event_listener::{Event, EventListener};
use crate::sync::loom_atomic::{AtomicBool, Ordering};
#[derive(Debug, Default)]
pub struct Notify {
event: Event,
permit: AtomicBool,
}
impl Notify {
#[cfg(not(loom))]
#[must_use]
pub const fn const_new() -> Self {
Self {
event: Event::new(),
permit: AtomicBool::new(false),
}
}
#[cfg(loom)]
#[must_use]
pub fn const_new() -> Self {
Self {
event: Event::new(),
permit: AtomicBool::new(false),
}
}
#[must_use]
pub fn new() -> Self {
Self::const_new()
}
pub fn notify_one(&self) {
if self
.permit
.compare_exchange(false, true, Ordering::AcqRel, Ordering::Acquire)
.is_ok()
{
self.event.notify(1);
}
}
pub fn notify_waiters(&self) {
self.event.notify(usize::MAX);
}
#[must_use = "futures are lazy and do nothing unless awaited"]
pub fn notified(&self) -> Notified<'_> {
Notified {
notify: self,
listener: None,
}
}
}
pub struct Notified<'lifetime> {
notify: &'lifetime Notify,
listener: Option<Pin<Box<EventListener>>>,
}
impl Notified<'_> {
pub fn enable(self: Pin<&mut Self>) -> bool {
let this = self.get_mut();
this.ensure_listener();
if this.notify.permit.swap(false, Ordering::AcqRel) {
this.listener = None;
true
} else {
false
}
}
fn ensure_listener(&mut self) {
if self.listener.is_none() {
self.listener = Some(Box::pin(self.notify.event.listen()));
}
}
}
impl Future for Notified<'_> {
type Output = ();
fn poll(self: Pin<&mut Self>, context: &mut Context<'_>) -> Poll<Self::Output> {
let this = self.get_mut();
if this.notify.permit.swap(false, Ordering::AcqRel) {
this.listener = None;
return Poll::Ready(());
}
this.ensure_listener();
if this.notify.permit.swap(false, Ordering::AcqRel) {
this.listener = None;
return Poll::Ready(());
}
let Some(listener) = this.listener.as_mut() else {
return Poll::Pending;
};
match listener.as_mut().poll(context) {
Poll::Ready(()) => {
this.listener = None;
let _ = this.notify.permit.swap(false, Ordering::AcqRel);
Poll::Ready(())
}
Poll::Pending => Poll::Pending,
}
}
}
#[cfg(all(test, not(loom)))]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod tests {
use super::*;
use futures::executor::block_on;
use std::sync::Arc;
#[test]
fn notify_one_then_notified_returns_immediately() {
let notify = Notify::new();
notify.notify_one();
block_on(notify.notified());
}
#[test]
fn second_notified_parks_until_next_notify() {
let notify = Arc::new(Notify::new());
notify.notify_one();
block_on(notify.notified());
let started = std::sync::atomic::AtomicBool::new(false);
let started_ref = &started;
let notify_ref = notify.clone();
std::thread::scope(|scope| {
scope.spawn(move || {
while !started_ref.load(Ordering::Acquire) {
std::hint::spin_loop();
}
notify_ref.notify_one();
});
started.store(true, Ordering::Release);
block_on(notify.notified());
});
}
#[test]
fn notify_waiters_does_not_save_permit() {
let notify = Notify::new();
notify.notify_waiters();
notify.notify_one();
block_on(notify.notified());
}
#[test]
fn enable_returns_true_when_permit_already_pending() {
let notify = Notify::new();
notify.notify_one();
let waiter = notify.notified();
futures::pin_mut!(waiter);
assert!(waiter.as_mut().enable());
}
#[test]
fn enable_returns_false_when_no_permit() {
let notify = Notify::new();
let waiter = notify.notified();
futures::pin_mut!(waiter);
assert!(!waiter.as_mut().enable());
}
}
#[cfg(all(test, loom))]
#[allow(clippy::unwrap_used, clippy::expect_used)]
mod loom_tests {
use super::Notify;
use loom::sync::Arc;
fn check(model: impl Fn() + Sync + Send + 'static) {
let mut builder = loom::model::Builder::new();
builder.preemption_bound = Some(3);
builder.check(model);
}
#[test]
fn notify_one_then_notified_returns_immediately() {
check(|| {
let notify = Arc::new(Notify::new());
notify.notify_one();
loom::future::block_on(notify.notified());
});
}
#[test]
fn notified_then_notify_one_resolves() {
check(|| {
let notify = Arc::new(Notify::new());
let notify_producer = notify.clone();
let producer = loom::thread::spawn(move || {
notify_producer.notify_one();
});
loom::future::block_on(notify.notified());
producer.join().expect("producer thread panicked");
});
}
#[test]
fn concurrent_notify_one_and_notified() {
check(|| {
let notify = Arc::new(Notify::new());
let notify_producer = notify.clone();
let producer = loom::thread::spawn(move || {
notify_producer.notify_one();
});
loom::future::block_on(notify.notified());
producer.join().expect("producer thread panicked");
});
}
#[test]
fn multiple_notify_one_does_not_stack() {
check(|| {
let notify = Arc::new(Notify::new());
notify.notify_one();
notify.notify_one();
loom::future::block_on(notify.notified());
let notify_for_producer = notify.clone();
let producer = loom::thread::spawn(move || {
notify_for_producer.notify_one(); });
loom::future::block_on(notify.notified());
producer.join().expect("producer thread panicked");
});
}
}