1 use std::cell::UnsafeCell;
2 use std::sync::atomic::AtomicUsize;
3 use std::sync::atomic::Ordering::SeqCst;
4 
5 /// A "lock" around data `D`, which employs a *helping* strategy.
6 ///
7 /// Used to ensure that concurrent `unpark` invocations lead to (1) `poll` being
8 /// invoked on only a single thread at a time (2) `poll` being invoked at least
9 /// once after each `unpark` (unless the future has completed).
10 pub(crate) struct UnparkMutex<D> {
11     // The state of task execution (state machine described below)
12     status: AtomicUsize,
13 
14     // The actual task data, accessible only in the POLLING state
15     inner: UnsafeCell<Option<D>>,
16 }
17 
18 // `UnparkMutex<D>` functions in many ways like a `Mutex<D>`, except that on
19 // acquisition failure, the current lock holder performs the desired work --
20 // re-polling.
21 //
22 // As such, these impls mirror those for `Mutex<D>`. In particular, a reference
23 // to `UnparkMutex` can be used to gain `&mut` access to the inner data, which
24 // must therefore be `Send`.
25 unsafe impl<D: Send> Send for UnparkMutex<D> {}
26 unsafe impl<D: Send> Sync for UnparkMutex<D> {}
27 
28 // There are four possible task states, listed below with their possible
29 // transitions:
30 
31 // The task is blocked, waiting on an event
32 const WAITING: usize = 0; // --> POLLING
33 
34 // The task is actively being polled by a thread; arrival of additional events
35 // of interest should move it to the REPOLL state
36 const POLLING: usize = 1; // --> WAITING, REPOLL, or COMPLETE
37 
38 // The task is actively being polled, but will need to be re-polled upon
39 // completion to ensure that all events were observed.
40 const REPOLL: usize = 2; // --> POLLING
41 
42 // The task has finished executing (either successfully or with an error/panic)
43 const COMPLETE: usize = 3; // No transitions out
44 
45 impl<D> UnparkMutex<D> {
new() -> Self46     pub(crate) fn new() -> Self {
47         Self { status: AtomicUsize::new(WAITING), inner: UnsafeCell::new(None) }
48     }
49 
50     /// Attempt to "notify" the mutex that a poll should occur.
51     ///
52     /// An `Ok` result indicates that the `POLLING` state has been entered, and
53     /// the caller can proceed to poll the future. An `Err` result indicates
54     /// that polling is not necessary (because the task is finished or the
55     /// polling has been delegated).
notify(&self) -> Result<D, ()>56     pub(crate) fn notify(&self) -> Result<D, ()> {
57         let mut status = self.status.load(SeqCst);
58         loop {
59             match status {
60                 // The task is idle, so try to run it immediately.
61                 WAITING => {
62                     match self.status.compare_exchange(WAITING, POLLING, SeqCst, SeqCst) {
63                         Ok(_) => {
64                             let data = unsafe {
65                                 // SAFETY: we've ensured mutual exclusion via
66                                 // the status protocol; we are the only thread
67                                 // that has transitioned to the POLLING state,
68                                 // and we won't transition back to QUEUED until
69                                 // the lock is "released" by this thread. See
70                                 // the protocol diagram above.
71                                 (*self.inner.get()).take().unwrap()
72                             };
73                             return Ok(data);
74                         }
75                         Err(cur) => status = cur,
76                     }
77                 }
78 
79                 // The task is being polled, so we need to record that it should
80                 // be *repolled* when complete.
81                 POLLING => match self.status.compare_exchange(POLLING, REPOLL, SeqCst, SeqCst) {
82                     Ok(_) => return Err(()),
83                     Err(cur) => status = cur,
84                 },
85 
86                 // The task is already scheduled for polling, or is complete, so
87                 // we've got nothing to do.
88                 _ => return Err(()),
89             }
90         }
91     }
92 
93     /// Alert the mutex that polling is about to begin, clearing any accumulated
94     /// re-poll requests.
95     ///
96     /// # Safety
97     ///
98     /// Callable only from the `POLLING`/`REPOLL` states, i.e. between
99     /// successful calls to `notify` and `wait`/`complete`.
start_poll(&self)100     pub(crate) unsafe fn start_poll(&self) {
101         self.status.store(POLLING, SeqCst);
102     }
103 
104     /// Alert the mutex that polling completed with `Pending`.
105     ///
106     /// # Safety
107     ///
108     /// Callable only from the `POLLING`/`REPOLL` states, i.e. between
109     /// successful calls to `notify` and `wait`/`complete`.
wait(&self, data: D) -> Result<(), D>110     pub(crate) unsafe fn wait(&self, data: D) -> Result<(), D> {
111         *self.inner.get() = Some(data);
112 
113         match self.status.compare_exchange(POLLING, WAITING, SeqCst, SeqCst) {
114             // no unparks came in while we were running
115             Ok(_) => Ok(()),
116 
117             // guaranteed to be in REPOLL state; just clobber the
118             // state and run again.
119             Err(status) => {
120                 assert_eq!(status, REPOLL);
121                 self.status.store(POLLING, SeqCst);
122                 Err((*self.inner.get()).take().unwrap())
123             }
124         }
125     }
126 
127     /// Alert the mutex that the task has completed execution and should not be
128     /// notified again.
129     ///
130     /// # Safety
131     ///
132     /// Callable only from the `POLLING`/`REPOLL` states, i.e. between
133     /// successful calls to `notify` and `wait`/`complete`.
complete(&self)134     pub(crate) unsafe fn complete(&self) {
135         self.status.store(COMPLETE, SeqCst);
136     }
137 }
138