tokio/runtime/task/
atomic_notified.rs

1use crate::loom::sync::atomic::AtomicPtr;
2use crate::runtime::task::{Header, Notified, RawTask};
3
4use std::marker::PhantomData;
5use std::ptr;
6use std::ptr::NonNull;
7use std::sync::atomic::Ordering::SeqCst;
8
9/// An atomic cell which can contain a pointer to a [`Notified`] task.
10///
11/// This is similar to the `crate::util::AtomicCell` type, but specialized to
12/// hold a task pointer --- this type "remembers" the task's scheduler generic
13/// when a task is stored in the cell, so that the pointer can be turned back
14/// into a [`Notified`] task with the correct generic type when it is retrieved.
15pub(crate) struct AtomicNotified<S: 'static> {
16    task: AtomicPtr<Header>,
17    _scheduler: PhantomData<S>,
18}
19
20impl<S: 'static> AtomicNotified<S> {
21    pub(crate) fn empty() -> Self {
22        Self {
23            task: AtomicPtr::new(ptr::null_mut()),
24            _scheduler: PhantomData,
25        }
26    }
27
28    pub(crate) fn swap(&self, task: Option<Notified<S>>) -> Option<Notified<S>> {
29        let new = task
30            .map(|t| t.into_raw().header_ptr().as_ptr())
31            .unwrap_or_else(ptr::null_mut);
32        let old = self.task.swap(new, SeqCst);
33        NonNull::new(old).map(|ptr| unsafe {
34            // Safety: since we only allow tasks with the same scheduler type to
35            // be placed in this cell, we know that the pointed task's scheduler
36            // type matches the type parameter S.
37            Notified::from_raw(RawTask::from_raw(ptr))
38        })
39    }
40
41    pub(crate) fn take(&self) -> Option<Notified<S>> {
42        self.swap(None)
43    }
44
45    pub(crate) fn is_some(&self) -> bool {
46        !self.task.load(SeqCst).is_null()
47    }
48}
49
50unsafe impl<S: Send> Send for AtomicNotified<S> {}
51unsafe impl<S: Send> Sync for AtomicNotified<S> {}
52
53impl<S> Drop for AtomicNotified<S> {
54    fn drop(&mut self) {
55        // Ensure the task reference is dropped if this cell is dropped.
56        let _ = self.take();
57    }
58}
tokio/runtime/task/atomic_notified.rs

tokio/runtime/task/
atomic_notified.rs
