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}