tokio_sync/task/
atomic_task.rs

1use loom::{
2    futures::task::{self, Task},
3    sync::atomic::AtomicUsize,
4    sync::CausalCell,
5};
6
7use std::fmt;
8use std::sync::atomic::Ordering::{AcqRel, Acquire, Release};
9
10/// A synchronization primitive for task notification.
11///
12/// `AtomicTask` will coordinate concurrent notifications with the consumer
13/// potentially "updating" the underlying task to notify. This is useful in
14/// scenarios where a computation completes in another thread and wants to
15/// notify the consumer, but the consumer is in the process of being migrated to
16/// a new logical task.
17///
18/// Consumers should call `register` before checking the result of a computation
19/// and producers should call `notify` after producing the computation (this
20/// differs from the usual `thread::park` pattern). It is also permitted for
21/// `notify` to be called **before** `register`. This results in a no-op.
22///
23/// A single `AtomicTask` may be reused for any number of calls to `register` or
24/// `notify`.
25///
26/// `AtomicTask` does not provide any memory ordering guarantees, as such the
27/// user should use caution and use other synchronization primitives to guard
28/// the result of the underlying computation.
29pub struct AtomicTask {
30    state: AtomicUsize,
31    task: CausalCell<Option<Task>>,
32}
33
34// `AtomicTask` is a multi-consumer, single-producer transfer cell. The cell
35// stores a `Task` value produced by calls to `register` and many threads can
36// race to take the task (to notify it) by calling `notify.
37//
38// If a new `Task` instance is produced by calling `register` before an existing
39// one is consumed, then the existing one is overwritten.
40//
41// While `AtomicTask` is single-producer, the implementation ensures memory
42// safety. In the event of concurrent calls to `register`, there will be a
43// single winner whose task will get stored in the cell. The losers will not
44// have their tasks notified. As such, callers should ensure to add
45// synchronization to calls to `register`.
46//
47// The implementation uses a single `AtomicUsize` value to coordinate access to
48// the `Task` cell. There are two bits that are operated on independently. These
49// are represented by `REGISTERING` and `NOTIFYING`.
50//
51// The `REGISTERING` bit is set when a producer enters the critical section. The
52// `NOTIFYING` bit is set when a consumer enters the critical section. Neither
53// bit being set is represented by `WAITING`.
54//
55// A thread obtains an exclusive lock on the task cell by transitioning the
56// state from `WAITING` to `REGISTERING` or `NOTIFYING`, depending on the
57// operation the thread wishes to perform. When this transition is made, it is
58// guaranteed that no other thread will access the task cell.
59//
60// # Registering
61//
62// On a call to `register`, an attempt to transition the state from WAITING to
63// REGISTERING is made. On success, the caller obtains a lock on the task cell.
64//
65// If the lock is obtained, then the thread sets the task cell to the task
66// provided as an argument. Then it attempts to transition the state back from
67// `REGISTERING` -> `WAITING`.
68//
69// If this transition is successful, then the registering process is complete
70// and the next call to `notify` will observe the task.
71//
72// If the transition fails, then there was a concurrent call to `notify` that
73// was unable to access the task cell (due to the registering thread holding the
74// lock). To handle this, the registering thread removes the task it just set
75// from the cell and calls `notify` on it. This call to notify represents the
76// attempt to notify by the other thread (that set the `NOTIFYING` bit). The
77// state is then transitioned from `REGISTERING | NOTIFYING` back to `WAITING`.
78// This transition must succeed because, at this point, the state cannot be
79// transitioned by another thread.
80//
81// # Notifying
82//
83// On a call to `notify`, an attempt to transition the state from `WAITING` to
84// `NOTIFYING` is made. On success, the caller obtains a lock on the task cell.
85//
86// If the lock is obtained, then the thread takes ownership of the current value
87// in teh task cell, and calls `notify` on it. The state is then transitioned
88// back to `WAITING`. This transition must succeed as, at this point, the state
89// cannot be transitioned by another thread.
90//
91// If the thread is unable to obtain the lock, the `NOTIFYING` bit is still.
92// This is because it has either been set by the current thread but the previous
93// value included the `REGISTERING` bit **or** a concurrent thread is in the
94// `NOTIFYING` critical section. Either way, no action must be taken.
95//
96// If the current thread is the only concurrent call to `notify` and another
97// thread is in the `register` critical section, when the other thread **exits**
98// the `register` critical section, it will observe the `NOTIFYING` bit and
99// handle the notify itself.
100//
101// If another thread is in the `notify` critical section, then it will handle
102// notifying the task.
103//
104// # A potential race (is safely handled).
105//
106// Imagine the following situation:
107//
108// * Thread A obtains the `notify` lock and notifies a task.
109//
110// * Before thread A releases the `notify` lock, the notified task is scheduled.
111//
112// * Thread B attempts to notify the task. In theory this should result in the
113//   task being notified, but it cannot because thread A still holds the notify
114//   lock.
115//
116// This case is handled by requiring users of `AtomicTask` to call `register`
117// **before** attempting to observe the application state change that resulted
118// in the task being notified. The notifiers also change the application state
119// before calling notify.
120//
121// Because of this, the task will do one of two things.
122//
123// 1) Observe the application state change that Thread B is notifying on. In
124//    this case, it is OK for Thread B's notification to be lost.
125//
126// 2) Call register before attempting to observe the application state. Since
127//    Thread A still holds the `notify` lock, the call to `register` will result
128//    in the task notifying itself and get scheduled again.
129
130/// Idle state
131const WAITING: usize = 0;
132
133/// A new task value is being registered with the `AtomicTask` cell.
134const REGISTERING: usize = 0b01;
135
136/// The task currently registered with the `AtomicTask` cell is being notified.
137const NOTIFYING: usize = 0b10;
138
139impl AtomicTask {
140    /// Create an `AtomicTask` initialized with the given `Task`
141    pub fn new() -> AtomicTask {
142        AtomicTask {
143            state: AtomicUsize::new(WAITING),
144            task: CausalCell::new(None),
145        }
146    }
147
148    /// Registers the current task to be notified on calls to `notify`.
149    ///
150    /// This is the same as calling `register_task` with `task::current()`.
151    pub fn register(&self) {
152        self.do_register(CurrentTask);
153    }
154
155    /// Registers the provided task to be notified on calls to `notify`.
156    ///
157    /// The new task will take place of any previous tasks that were registered
158    /// by previous calls to `register`. Any calls to `notify` that happen after
159    /// a call to `register` (as defined by the memory ordering rules), will
160    /// notify the `register` caller's task.
161    ///
162    /// It is safe to call `register` with multiple other threads concurrently
163    /// calling `notify`. This will result in the `register` caller's current
164    /// task being notified once.
165    ///
166    /// This function is safe to call concurrently, but this is generally a bad
167    /// idea. Concurrent calls to `register` will attempt to register different
168    /// tasks to be notified. One of the callers will win and have its task set,
169    /// but there is no guarantee as to which caller will succeed.
170    pub fn register_task(&self, task: Task) {
171        self.do_register(ExactTask(task));
172    }
173
174    fn do_register<R>(&self, reg: R)
175    where
176        R: Register,
177    {
178        debug!(" + register_task");
179        match self.state.compare_and_swap(WAITING, REGISTERING, Acquire) {
180            WAITING => {
181                unsafe {
182                    // Locked acquired, update the waker cell
183                    self.task.with_mut(|t| reg.register(&mut *t));
184
185                    // Release the lock. If the state transitioned to include
186                    // the `NOTIFYING` bit, this means that a notify has been
187                    // called concurrently, so we have to remove the task and
188                    // notify it.`
189                    //
190                    // Start by assuming that the state is `REGISTERING` as this
191                    // is what we jut set it to.
192                    let res = self
193                        .state
194                        .compare_exchange(REGISTERING, WAITING, AcqRel, Acquire);
195
196                    match res {
197                        Ok(_) => {}
198                        Err(actual) => {
199                            // This branch can only be reached if a
200                            // concurrent thread called `notify`. In this
201                            // case, `actual` **must** be `REGISTERING |
202                            // `NOTIFYING`.
203                            debug_assert_eq!(actual, REGISTERING | NOTIFYING);
204
205                            // Take the task to notify once the atomic operation has
206                            // completed.
207                            let notify = self.task.with_mut(|t| (*t).take()).unwrap();
208
209                            // Just swap, because no one could change state
210                            // while state == `Registering | `Waking`
211                            self.state.swap(WAITING, AcqRel);
212
213                            // The atomic swap was complete, now
214                            // notify the task and return.
215                            notify.notify();
216                        }
217                    }
218                }
219            }
220            NOTIFYING => {
221                // Currently in the process of notifying the task, i.e.,
222                // `notify` is currently being called on the old task handle.
223                // So, we call notify on the new task handle
224                reg.notify();
225            }
226            state => {
227                // In this case, a concurrent thread is holding the
228                // "registering" lock. This probably indicates a bug in the
229                // caller's code as racing to call `register` doesn't make much
230                // sense.
231                //
232                // We just want to maintain memory safety. It is ok to drop the
233                // call to `register`.
234                debug_assert!(state == REGISTERING || state == REGISTERING | NOTIFYING);
235            }
236        }
237    }
238
239    /// Notifies the task that last called `register`.
240    ///
241    /// If `register` has not been called yet, then this does nothing.
242    pub fn notify(&self) {
243        debug!(" + notify");
244        if let Some(task) = self.take_task() {
245            task.notify();
246        }
247    }
248
249    /// Attempts to take the `Task` value out of the `AtomicTask` with the
250    /// intention that the caller will notify the task later.
251    pub fn take_task(&self) -> Option<Task> {
252        debug!(" + take_task");
253        // AcqRel ordering is used in order to acquire the value of the `task`
254        // cell as well as to establish a `release` ordering with whatever
255        // memory the `AtomicTask` is associated with.
256        match self.state.fetch_or(NOTIFYING, AcqRel) {
257            WAITING => {
258                debug!(" + WAITING");
259                // The notifying lock has been acquired.
260                let task = unsafe { self.task.with_mut(|t| (*t).take()) };
261
262                // Release the lock
263                self.state.fetch_and(!NOTIFYING, Release);
264                debug!(" + Done taking");
265
266                task
267            }
268            state => {
269                debug!(" + state = {:?}", state);
270                // There is a concurrent thread currently updating the
271                // associated task.
272                //
273                // Nothing more to do as the `NOTIFYING` bit has been set. It
274                // doesn't matter if there are concurrent registering threads or
275                // not.
276                //
277                debug_assert!(
278                    state == REGISTERING || state == REGISTERING | NOTIFYING || state == NOTIFYING
279                );
280                None
281            }
282        }
283    }
284}
285
286impl Default for AtomicTask {
287    fn default() -> Self {
288        AtomicTask::new()
289    }
290}
291
292impl fmt::Debug for AtomicTask {
293    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
294        write!(fmt, "AtomicTask")
295    }
296}
297
298unsafe impl Send for AtomicTask {}
299unsafe impl Sync for AtomicTask {}
300
301trait Register {
302    fn register(self, slot: &mut Option<Task>);
303    fn notify(self);
304}
305
306struct CurrentTask;
307
308impl Register for CurrentTask {
309    fn register(self, slot: &mut Option<Task>) {
310        let should_update = (&*slot)
311            .as_ref()
312            .map(|prev| !prev.will_notify_current())
313            .unwrap_or(true);
314        if should_update {
315            *slot = Some(task::current());
316        }
317    }
318
319    fn notify(self) {
320        task::current().notify();
321    }
322}
323
324struct ExactTask(Task);
325
326impl Register for ExactTask {
327    fn register(self, slot: &mut Option<Task>) {
328        // When calling register_task with an exact task, it doesn't matter
329        // if the previous task would have notified current. We *always* want
330        // to save that exact task.
331        *slot = Some(self.0);
332    }
333
334    fn notify(self) {
335        self.0.notify();
336    }
337}