use std::cell::{Cell, OnceCell, RefCell};
use std::future::{Future, poll_fn};
use std::sync::Arc;
use std::task::{Context, Poll, Wake, Waker};
use std::{fmt, marker::PhantomData, mem, pin::Pin, rc::Rc, thread, thread::ThreadId};
use crate::ctx::WaitersInner;
thread_local! {
static NOTIFICATION: RefCell<u64> = const { RefCell::new(0) };
}
#[derive(Clone)]
pub struct ReadyCallback(Rc<dyn Fn(usize)>);
impl ReadyCallback {
pub fn new(callback: impl Fn(usize) + 'static) -> Self {
Self(Rc::new(callback))
}
}
impl fmt::Debug for ReadyCallback {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ReadyCallback").finish_non_exhaustive()
}
}
#[derive(Clone, Debug, Default)]
pub struct ReadyNotifier(Rc<RefCell<Vec<Endpoint>>>);
impl ReadyNotifier {
pub fn new() -> Self {
Self::default()
}
pub fn notify(&self) {
with_notification(|epoch| {
let scopes = self.0.borrow();
for endpoint in &*scopes {
endpoint
.store
.scopes
.with(endpoint.scope, |s| s.direct_epoch.set(epoch));
}
for endpoint in &*scopes {
notify_scope(&endpoint.store, endpoint.scope, epoch, true, true);
}
});
}
}
fn with_notification(f: impl FnOnce(u64)) {
NOTIFICATION.with_borrow_mut(|epoch| {
*epoch = epoch
.checked_add(1)
.expect("readiness notification epoch exhausted");
f(*epoch);
});
}
fn notify_scope(store: &WaitersInner, id: usize, epoch: u64, direct: bool, callbacks: bool) {
store.scopes.with(id, |s| {
if !direct && (s.notify_epoch.get() == epoch || s.direct_epoch.get() == epoch) {
return;
}
s.notify_epoch.set(epoch);
s.invalidate(direct);
let _wake = NotifyObservers(&s.observers);
for &parent in s.parents.borrow().iter() {
notify_scope(store, parent, epoch, false, callbacks);
}
for parent in s.external_parents.borrow().iter() {
notify_scope(&parent.store, parent.scope, epoch, false, callbacks);
}
if callbacks {
for callback in s.callbacks.borrow().iter() {
(callback.0)(s.data.get());
}
}
});
}
struct NotifyObservers<'a>(&'a ReadyObservers);
impl Drop for NotifyObservers<'_> {
fn drop(&mut self) {
let skip = self.0.checker.get().map(|c| Arc::as_ptr(c).cast());
let capacity = self.0.slots.borrow().capacity();
for idx in 0..capacity {
let waker = self
.0
.slots
.borrow_mut()
.get_mut(idx)
.and_then(Option::take);
if let Some(waker) = waker
&& Some(waker.data()) != skip
{
waker.wake();
}
}
}
}
#[derive(Default)]
struct Wakers {
inline: [Option<Waker>; 8],
rest: Vec<Waker>,
}
impl Wakers {
fn push(&mut self, waker: Waker) {
if let Some(slot) = self.inline.iter_mut().find(|slot| slot.is_none()) {
*slot = Some(waker);
} else {
self.rest.push(waker);
}
}
}
impl Drop for Wakers {
fn drop(&mut self) {
for slot in &mut self.inline {
if let Some(waker) = slot.take() {
waker.wake();
}
}
for waker in self.rest.drain(..) {
waker.wake();
}
}
}
fn link<T: PartialEq>(links: &RefCell<Vec<T>>, item: T) -> bool {
let mut links = links.borrow_mut();
let new = !links.contains(&item);
if new {
links.push(item);
}
new
}
#[derive(Clone)]
pub struct ReadyLink {
store: Rc<WaitersInner>,
}
impl fmt::Debug for ReadyLink {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("ReadyLink")
.field("store", &Rc::as_ptr(&self.store))
.finish()
}
}
#[derive(Clone)]
struct Endpoint {
store: Rc<WaitersInner>,
scope: usize,
}
impl PartialEq for Endpoint {
fn eq(&self, other: &Self) -> bool {
self.scope == other.scope && Rc::ptr_eq(&self.store, &other.store)
}
}
impl fmt::Debug for Endpoint {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("Endpoint")
.field("store", &Rc::as_ptr(&self.store))
.field("scope", &self.scope)
.finish()
}
}
#[derive(Debug)]
pub(crate) struct ScopeStore {
slots: RefCell<Vec<Option<ScopeInner>>>,
}
impl Default for ScopeStore {
fn default() -> Self {
Self {
slots: RefCell::new(vec![Some(ScopeInner::default())]),
}
}
}
impl ScopeStore {
fn allocate(&self) -> usize {
let mut slots = self.slots.borrow_mut();
assert!(!slots.is_empty(), "readiness store has been retired");
let id = slots.len();
slots.push(Some(ScopeInner::default()));
id
}
fn with<R>(&self, id: usize, f: impl FnOnce(&ScopeInner) -> R) -> Option<R> {
self.slots.borrow().get(id).and_then(Option::as_ref).map(f)
}
fn release(&self, id: usize) {
if self.with(id, |s| {
s.owners.set(s.owners.get() - 1);
s.owners.get() == 0
}) != Some(true)
{
return;
}
if id == 0 {
let slots = self.slots.take();
for scope in slots.iter().flatten() {
scope.observers.close();
}
} else {
let scope = self.slots.borrow_mut()[id].take().unwrap();
scope.observers.close();
}
}
}
#[derive(Copy, Clone, Debug, Default, PartialEq, Eq)]
pub(crate) struct Generation {
value: u64,
direct: u64,
}
impl From<Generation> for () {
fn from(_: Generation) {}
}
#[derive(Clone, Debug, Default)]
pub(crate) struct ReadyScope(OnceCell<ScopeRef>);
#[derive(Debug)]
struct ScopeRef {
store: Rc<WaitersInner>,
id: usize,
}
impl Clone for ScopeRef {
fn clone(&self) -> Self {
self.store
.scopes
.with(self.id, |s| s.owners.set(s.owners.get() + 1))
.expect("readiness scope has been retired");
Self {
store: self.store.clone(),
id: self.id,
}
}
}
impl Drop for ScopeRef {
fn drop(&mut self) {
self.store.scopes.release(self.id);
}
}
impl ReadyScope {
pub(crate) fn new() -> Self {
Self::default()
}
pub(crate) fn root(store: Rc<WaitersInner>) -> Self {
Self(OnceCell::from(ScopeRef { store, id: 0 }))
}
fn bound(&self) -> &ScopeRef {
self.0
.get()
.expect("ReadyBinding must be registered by Service::bind before use")
}
fn with<R>(&self, f: impl FnOnce(&ScopeInner) -> R) -> R {
let scope = self.bound();
scope
.store
.scopes
.with(scope.id, f)
.expect("readiness scope has been retired")
}
fn endpoint(&self) -> Endpoint {
let scope = self.bound();
Endpoint {
store: scope.store.clone(),
scope: scope.id,
}
}
pub(crate) fn readiness_link(&self) -> ReadyLink {
assert_eq!(self.bound().id, 0);
ReadyLink {
store: self.bound().store.clone(),
}
}
pub(crate) fn subscribe_pipeline(&self, source_link: &ReadyLink) {
source_link.store.scopes.with(0, |source| {
let scope = self.bound();
if Rc::ptr_eq(&scope.store, &source_link.store) {
if scope.id != 0 {
link(&source.parents, scope.id);
}
} else {
link(&source.external_parents, self.endpoint());
}
});
}
pub(crate) fn generation(&self) -> Generation {
self.with(|s| s.generation.get())
}
pub(crate) fn is_current(&self, generation: Generation) -> bool {
self.generation() == generation
}
pub(crate) fn is_directly_invalidated(&self, generation: Generation) -> bool {
self.generation().direct != generation.direct
}
pub(crate) fn is_ready(&self) -> bool {
self.0.get().is_some_and(|scope| {
scope
.store
.scopes
.with(scope.id, |s| s.ready.get())
.unwrap_or(false)
})
}
pub(crate) fn reset_root(&self) {
let scope = self.bound();
with_notification(|epoch| {
scope
.store
.scopes
.with(scope.id, |s| s.direct_epoch.set(epoch));
notify_scope(&scope.store, scope.id, epoch, true, false);
});
}
pub(crate) fn set_ready(&self, generation: Generation) -> bool {
if self.is_current(generation) {
self.with(|s| s.ready.set(true));
true
} else {
false
}
}
pub(crate) fn subscribe(&self, source: &ReadyNotifier) {
let endpoint = self.endpoint();
source.0.borrow_mut().push(endpoint);
}
pub(crate) fn set_data(&self, data: usize) {
self.with(|s| s.data.set(data));
}
pub(crate) fn on_notify(&self, callback: &ReadyCallback) {
self.with(|s| {
let mut callbacks = s.callbacks.borrow_mut();
if !callbacks.iter().any(|c| Rc::ptr_eq(&c.0, &callback.0)) {
callbacks.push(callback.clone());
}
});
}
pub(crate) fn depend_on(&self, child: &Self) {
let parent = self.bound();
let child = child.0.get_or_init(|| {
let id = parent.store.scopes.allocate();
ScopeRef {
store: parent.store.clone(),
id,
}
});
assert!(
Rc::ptr_eq(&parent.store, &child.store),
"ReadyBinding cannot be shared across pipelines; use ReadyLink with a shared Pipeline"
);
if parent.id != child.id {
parent
.store
.scopes
.with(child.id, |s| link(&s.parents, parent.id));
}
}
pub(crate) fn assert_dependency(&self, child: &Self) {
assert!(
child.0.get().is_some_and(|child| {
let parent = self.bound();
Rc::ptr_eq(&parent.store, &child.store)
&& (parent.id == child.id
|| parent
.store
.scopes
.with(child.id, |s| s.parents.borrow().contains(&parent.id))
== Some(true))
}),
"ReadyBinding must be registered by Service::bind before use"
);
}
pub(crate) fn observer(&self) -> ReadyObserver {
ReadyObserver {
store: self.bound().store.clone(),
scope: self.bound().id,
idx: self.with(|s| s.observers.slots.borrow_mut().insert(None)),
}
}
pub(crate) fn checker_waker(&self) -> Waker {
let endpoint = self.endpoint();
self.with(|s| s.observers.checker_waker(endpoint))
}
pub(crate) async fn check<F, Fut, E>(&self, mut check: F) -> Result<(), E>
where
F: FnMut() -> Fut,
Fut: Future<Output = Result<(), E>>,
{
let mut observer = None;
loop {
if self.is_ready() {
return Ok(());
}
let observer = observer.get_or_insert_with(|| self.observer());
let owner = poll_fn(|cx| {
observer.register(cx.waker());
if self.is_ready() {
Poll::Ready(None)
} else if self.with(|s| s.checking.replace(true)) {
Poll::Pending
} else {
Poll::Ready(Some(CheckOwner(self)))
}
})
.await;
if let Some(owner) = owner {
let result = RunCheck::new(self, &mut check).await.map(|generation| {
self.set_ready(generation);
});
drop(owner);
result?;
}
}
}
pub(crate) fn wake_observers(&self) {
let mut wakers = Wakers::default();
self.with(|s| s.observers.take_wakers(&mut wakers, None));
}
}
#[derive(Debug)]
struct ScopeInner {
owners: Cell<usize>,
data: Cell<usize>,
generation: Cell<Generation>,
ready: Cell<bool>,
checking: Cell<bool>,
notify_epoch: Cell<u64>,
direct_epoch: Cell<u64>,
callbacks: RefCell<Vec<ReadyCallback>>,
parents: RefCell<Vec<usize>>,
external_parents: RefCell<Vec<Endpoint>>,
observers: ReadyObservers,
}
impl Default for ScopeInner {
fn default() -> Self {
Self {
owners: Cell::new(1),
data: Cell::new(0),
generation: Cell::default(),
ready: Cell::default(),
checking: Cell::default(),
notify_epoch: Cell::default(),
direct_epoch: Cell::default(),
callbacks: RefCell::default(),
parents: RefCell::default(),
external_parents: RefCell::default(),
observers: ReadyObservers::default(),
}
}
}
impl ScopeInner {
fn invalidate(&self, direct: bool) {
let mut generation = self.generation.get();
generation.value = generation
.value
.checked_add(1)
.expect("readiness generation exhausted");
if direct {
generation.direct = generation.value;
}
self.generation.set(generation);
self.ready.set(false);
}
}
#[derive(Debug, Default)]
struct ReadyObservers {
slots: RefCell<slab::Slab<Option<Waker>>>,
checker: OnceCell<Arc<CheckerWake>>,
_local: PhantomData<Rc<()>>,
}
impl ReadyObservers {
fn checker_waker(&self, endpoint: Endpoint) -> Waker {
let checker = self.checker.get_or_init(|| {
Arc::new(CheckerWake {
owner: thread::current().id(),
endpoint: mem::ManuallyDrop::new(endpoint),
closed: Cell::new(false),
})
});
Waker::from(checker.clone())
}
fn close(&self) {
if let Some(checker) = self.checker.get() {
checker.closed.set(true);
}
}
fn take_wakers(&self, wakers: &mut Wakers, skip: Option<*const ()>) {
for (_, waker) in self.slots.borrow_mut().iter_mut() {
if let Some(waker) = waker.take_if(|waker| Some(waker.data()) != skip) {
wakers.push(waker);
}
}
}
}
impl Drop for ReadyObservers {
fn drop(&mut self) {
self.close();
}
}
struct CheckerWake {
owner: ThreadId,
closed: Cell<bool>,
endpoint: mem::ManuallyDrop<Endpoint>,
}
unsafe impl Send for CheckerWake {}
unsafe impl Sync for CheckerWake {}
impl CheckerWake {
fn is_owner(&self) -> bool {
thread::current().id() == self.owner
}
}
impl fmt::Debug for CheckerWake {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
f.debug_struct("CheckerWake")
.field("owner", &self.owner)
.finish_non_exhaustive()
}
}
impl Drop for CheckerWake {
fn drop(&mut self) {
if self.is_owner() {
unsafe { mem::ManuallyDrop::drop(&mut self.endpoint) };
} else if !thread::panicking() {
panic!("ntex-service readiness observers can only be accessed on their owning thread");
}
}
}
impl Wake for CheckerWake {
fn wake(self: Arc<Self>) {
self.wake_by_ref();
}
fn wake_by_ref(self: &Arc<Self>) {
assert!(
self.is_owner(),
"ntex-service readiness observers can only be accessed on their owning thread"
);
if self.closed.get() {
return;
}
let this = Arc::as_ptr(self).cast::<()>();
let mut wakers = Wakers::default();
self.endpoint.store.scopes.with(self.endpoint.scope, |s| {
s.observers.take_wakers(&mut wakers, Some(this));
});
}
}
struct CheckOwner<'a>(&'a ReadyScope);
impl Drop for CheckOwner<'_> {
fn drop(&mut self) {
self.0.with(|s| s.checking.set(false));
self.0.wake_observers();
}
}
#[derive(Debug)]
pub(crate) struct ReadyObserver {
store: Rc<WaitersInner>,
scope: usize,
idx: usize,
}
impl ReadyObserver {
pub(crate) fn clear(&self) {
let waker = self.store.scopes.with(self.scope, |s| {
s.observers.slots.borrow_mut()[self.idx].take()
});
drop(waker);
}
pub(crate) fn register(&self, waker: &Waker) -> bool {
let Some(registered) = self.store.scopes.with(self.scope, |s| {
s.observers.slots.borrow()[self.idx]
.as_ref()
.is_some_and(|prev| prev.will_wake(waker))
}) else {
return false;
};
if registered {
return true;
}
let mut waker = Some(waker.clone());
let previous = self.store.scopes.with(self.scope, |s| {
std::mem::replace(&mut s.observers.slots.borrow_mut()[self.idx], waker.take())
});
let registered = previous.is_some();
drop(previous);
registered
}
}
impl Drop for ReadyObserver {
fn drop(&mut self) {
let waker = self.store.scopes.with(self.scope, |s| {
s.observers.slots.borrow_mut().remove(self.idx)
});
drop(waker);
}
}
pub(crate) struct RunCheck<'a, F, Fut> {
scope: &'a ReadyScope,
check: F,
fut: Option<Fut>,
observer: Option<ReadyObserver>,
generation: Generation,
}
impl<'a, F, Fut, E> RunCheck<'a, F, Fut>
where
F: FnMut() -> Fut,
Fut: Future<Output = Result<(), E>>,
{
pub(crate) fn new(scope: &'a ReadyScope, check: F) -> Self {
Self {
scope,
check,
fut: None,
observer: None,
generation: Generation::default(),
}
}
pub(crate) fn scope(&self) -> &'a ReadyScope {
self.scope
}
}
impl<F, Fut, E> Future for RunCheck<'_, F, Fut>
where
F: FnMut() -> Fut,
Fut: Future<Output = Result<(), E>>,
{
type Output = Result<Generation, E>;
fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
let this = unsafe { self.get_unchecked_mut() };
let mut restarted = false;
loop {
if this.fut.as_mut().is_none() {
this.generation = this.scope.generation();
this.fut = Some((this.check)());
}
if let Some(observer) = &this.observer {
observer.register(cx.waker());
}
let result = if this.scope.is_directly_invalidated(this.generation) {
None
} else {
let before_poll = this.scope.generation();
let fut = unsafe { Pin::new_unchecked(this.fut.as_mut().unwrap()) };
match fut.poll(cx) {
Poll::Ready(result) => Some(result),
Poll::Pending => {
let scope = this.scope;
this.observer
.get_or_insert_with(|| scope.observer())
.register(cx.waker());
if !this.scope.is_current(before_poll) {
cx.waker().wake_by_ref();
}
return Poll::Pending;
}
}
};
this.fut = None;
if let Some(result) = result {
this.observer = None;
result?;
if this.scope.is_current(this.generation) {
return Poll::Ready(Ok(this.generation));
}
}
if restarted {
cx.waker().wake_by_ref();
return Poll::Pending;
}
restarted = true;
}
}
}
#[cfg(test)]
#[path = "../tests/unit/readiness.rs"]
mod tests;