use grommet_core::waker_slot::WakerSlot;
use std::task::Waker;
#[cfg(loom)]
use loom::sync::atomic::{AtomicBool, Ordering};
#[cfg(not(loom))]
use std::sync::atomic::{AtomicBool, Ordering};
pub(crate) struct Doorbell {
owner: WakerSlot,
closed: AtomicBool,
}
impl Doorbell {
pub(crate) fn new() -> Self {
Self { owner: WakerSlot::new(), closed: AtomicBool::new(false) }
}
#[inline]
pub(crate) fn register(&self, waker: &Waker) {
self.owner.register(waker);
}
#[inline]
pub(crate) fn ring(&self) {
if self.is_closed() {
return;
}
self.owner.wake();
}
#[inline]
pub(crate) fn is_closed(&self) -> bool {
self.closed.load(Ordering::Acquire)
}
pub(crate) fn close(&self) {
self.closed.store(true, Ordering::Release);
drop(self.owner.take());
}
}
impl Default for Doorbell {
fn default() -> Self {
Self::new()
}
}
impl std::fmt::Debug for Doorbell {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Doorbell").field("closed", &self.is_closed()).finish_non_exhaustive()
}
}
#[cfg(all(test, not(loom)))]
mod tests {
use super::*;
use std::sync::Arc;
use std::sync::atomic::AtomicUsize;
use std::task::Wake;
struct Counter(AtomicUsize);
impl Counter {
fn waker() -> (Arc<Self>, Waker) {
let counter = Arc::new(Self(AtomicUsize::new(0)));
(counter.clone(), Waker::from(counter))
}
fn count(&self) -> usize {
self.0.load(Ordering::Relaxed)
}
}
impl Wake for Counter {
fn wake(self: Arc<Self>) {
self.0.fetch_add(1, Ordering::Relaxed);
}
fn wake_by_ref(self: &Arc<Self>) {
self.0.fetch_add(1, Ordering::Relaxed);
}
}
#[test]
fn a_ring_wakes_whoever_registered_last() {
let bell = Doorbell::new();
let (first, first_waker) = Counter::waker();
let (second, second_waker) = Counter::waker();
bell.register(&first_waker);
bell.register(&second_waker);
bell.ring();
assert_eq!(first.count(), 0, "a replaced owner must not be woken");
assert_eq!(second.count(), 1);
}
#[test]
fn a_ring_consumes_the_registration_so_repeats_coalesce() {
let bell = Doorbell::new();
let (counter, waker) = Counter::waker();
bell.register(&waker);
bell.ring();
bell.ring();
bell.ring();
assert_eq!(counter.count(), 1);
bell.register(&waker);
bell.ring();
assert_eq!(counter.count(), 2, "re-registering re-arms it");
}
#[test]
fn ringing_an_unregistered_doorbell_is_a_no_op() {
let bell = Doorbell::new();
bell.ring();
assert!(!bell.is_closed());
}
#[test]
fn closing_silences_later_rings_and_releases_the_waker() {
let bell = Doorbell::new();
let (counter, waker) = Counter::waker();
bell.register(&waker);
assert_eq!(Arc::strong_count(&counter), 3, "the Arc, the Waker, and the doorbell");
bell.close();
assert!(bell.is_closed());
assert_eq!(
Arc::strong_count(&counter),
2,
"closing must drop the registration, not merely ignore it"
);
bell.ring();
bell.ring();
assert_eq!(counter.count(), 0);
}
#[test]
fn a_registration_after_close_still_never_rings() {
let bell = Doorbell::new();
let (counter, waker) = Counter::waker();
bell.close();
bell.register(&waker);
bell.ring();
assert_eq!(counter.count(), 0);
}
#[test]
fn closing_twice_is_harmless() {
let bell = Doorbell::new();
bell.close();
bell.close();
assert!(bell.is_closed());
}
#[test]
fn the_debug_rendering_names_the_state_that_matters() {
let bell = Doorbell::default();
assert!(format!("{bell:?}").contains("closed: false"));
bell.close();
assert!(format!("{bell:?}").contains("closed: true"));
}
}
#[cfg(all(test, loom))]
mod loom_tests {
use super::*;
use loom::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::task::Wake;
struct Flag(AtomicBool);
impl Flag {
fn waker() -> (Arc<Self>, Waker) {
let flag = Arc::new(Self(AtomicBool::new(false)));
(flag.clone(), Waker::from(flag))
}
fn woken(&self) -> bool {
self.0.load(Ordering::Acquire)
}
}
impl Wake for Flag {
fn wake(self: Arc<Self>) {
self.0.store(true, Ordering::Release);
}
fn wake_by_ref(self: &Arc<Self>) {
self.0.store(true, Ordering::Release);
}
}
#[test]
fn loom_register_before_check_cannot_strand_a_ring() {
loom::model(|| {
let bell = Arc::new(Doorbell::new());
let work = Arc::new(AtomicBool::new(false));
let (flag, waker) = Flag::waker();
let producer = {
let bell = bell.clone();
let work = work.clone();
loom::thread::spawn(move || {
work.store(true, Ordering::Release);
bell.ring();
})
};
bell.register(&waker);
let observed = work.load(Ordering::Acquire);
producer.join().unwrap();
assert!(
observed || flag.woken(),
"the owner neither saw the work nor was woken for it"
);
});
}
#[test]
fn loom_a_ring_racing_close_never_delivers_afterwards() {
loom::model(|| {
let bell = Arc::new(Doorbell::new());
let (flag, waker) = Flag::waker();
bell.register(&waker);
let notifier = {
let bell = bell.clone();
loom::thread::spawn(move || bell.ring())
};
bell.close();
notifier.join().unwrap();
let after_close = flag.woken();
bell.register(&waker);
bell.ring();
assert_eq!(flag.woken(), after_close, "a ring after close delivered a wake");
});
}
#[test]
fn loom_concurrent_rings_still_schedule_the_owner() {
loom::model(|| {
let bell = Arc::new(Doorbell::new());
let (flag, waker) = Flag::waker();
bell.register(&waker);
let left = {
let bell = bell.clone();
loom::thread::spawn(move || bell.ring())
};
let right = {
let bell = bell.clone();
loom::thread::spawn(move || bell.ring())
};
left.join().unwrap();
right.join().unwrap();
assert!(flag.woken(), "two rings on an open doorbell woke nobody");
});
}
}