1 | //! The task module. |
2 | //! |
3 | //! The task module contains the code that manages spawned tasks and provides a |
4 | //! safe API for the rest of the runtime to use. Each task in a runtime is |
5 | //! stored in an OwnedTasks or LocalOwnedTasks object. |
6 | //! |
7 | //! # Task reference types |
8 | //! |
9 | //! A task is usually referenced by multiple handles, and there are several |
10 | //! types of handles. |
11 | //! |
12 | //! * OwnedTask - tasks stored in an OwnedTasks or LocalOwnedTasks are of this |
13 | //! reference type. |
14 | //! |
15 | //! * JoinHandle - each task has a JoinHandle that allows access to the output |
16 | //! of the task. |
17 | //! |
18 | //! * Waker - every waker for a task has this reference type. There can be any |
19 | //! number of waker references. |
20 | //! |
21 | //! * Notified - tracks whether the task is notified. |
22 | //! |
23 | //! * Unowned - this task reference type is used for tasks not stored in any |
24 | //! runtime. Mainly used for blocking tasks, but also in tests. |
25 | //! |
26 | //! The task uses a reference count to keep track of how many active references |
27 | //! exist. The Unowned reference type takes up two ref-counts. All other |
28 | //! reference types take up a single ref-count. |
29 | //! |
30 | //! Besides the waker type, each task has at most one of each reference type. |
31 | //! |
32 | //! # State |
33 | //! |
34 | //! The task stores its state in an atomic usize with various bitfields for the |
35 | //! necessary information. The state has the following bitfields: |
36 | //! |
37 | //! * RUNNING - Tracks whether the task is currently being polled or cancelled. |
38 | //! This bit functions as a lock around the task. |
39 | //! |
40 | //! * COMPLETE - Is one once the future has fully completed and has been |
41 | //! dropped. Never unset once set. Never set together with RUNNING. |
42 | //! |
43 | //! * NOTIFIED - Tracks whether a Notified object currently exists. |
44 | //! |
45 | //! * CANCELLED - Is set to one for tasks that should be cancelled as soon as |
46 | //! possible. May take any value for completed tasks. |
47 | //! |
48 | //! * JOIN_INTEREST - Is set to one if there exists a JoinHandle. |
49 | //! |
50 | //! * JOIN_WAKER - Acts as an access control bit for the join handle waker. The |
51 | //! protocol for its usage is described below. |
52 | //! |
53 | //! The rest of the bits are used for the ref-count. |
54 | //! |
55 | //! # Fields in the task |
56 | //! |
57 | //! The task has various fields. This section describes how and when it is safe |
58 | //! to access a field. |
59 | //! |
60 | //! * The state field is accessed with atomic instructions. |
61 | //! |
62 | //! * The OwnedTask reference has exclusive access to the `owned` field. |
63 | //! |
64 | //! * The Notified reference has exclusive access to the `queue_next` field. |
65 | //! |
66 | //! * The `owner_id` field can be set as part of construction of the task, but |
67 | //! is otherwise immutable and anyone can access the field immutably without |
68 | //! synchronization. |
69 | //! |
70 | //! * If COMPLETE is one, then the JoinHandle has exclusive access to the |
71 | //! stage field. If COMPLETE is zero, then the RUNNING bitfield functions as |
72 | //! a lock for the stage field, and it can be accessed only by the thread |
73 | //! that set RUNNING to one. |
74 | //! |
75 | //! * The waker field may be concurrently accessed by different threads: in one |
76 | //! thread the runtime may complete a task and *read* the waker field to |
77 | //! invoke the waker, and in another thread the task's JoinHandle may be |
78 | //! polled, and if the task hasn't yet completed, the JoinHandle may *write* |
79 | //! a waker to the waker field. The JOIN_WAKER bit ensures safe access by |
80 | //! multiple threads to the waker field using the following rules: |
81 | //! |
82 | //! 1. JOIN_WAKER is initialized to zero. |
83 | //! |
84 | //! 2. If JOIN_WAKER is zero, then the JoinHandle has exclusive (mutable) |
85 | //! access to the waker field. |
86 | //! |
87 | //! 3. If JOIN_WAKER is one, then the JoinHandle has shared (read-only) |
88 | //! access to the waker field. |
89 | //! |
90 | //! 4. If JOIN_WAKER is one and COMPLETE is one, then the runtime has shared |
91 | //! (read-only) access to the waker field. |
92 | //! |
93 | //! 5. If the JoinHandle needs to write to the waker field, then the |
94 | //! JoinHandle needs to (i) successfully set JOIN_WAKER to zero if it is |
95 | //! not already zero to gain exclusive access to the waker field per rule |
96 | //! 2, (ii) write a waker, and (iii) successfully set JOIN_WAKER to one. |
97 | //! |
98 | //! 6. The JoinHandle can change JOIN_WAKER only if COMPLETE is zero (i.e. |
99 | //! the task hasn't yet completed). |
100 | //! |
101 | //! Rule 6 implies that the steps (i) or (iii) of rule 5 may fail due to a |
102 | //! race. If step (i) fails, then the attempt to write a waker is aborted. If |
103 | //! step (iii) fails because COMPLETE is set to one by another thread after |
104 | //! step (i), then the waker field is cleared. Once COMPLETE is one (i.e. |
105 | //! task has completed), the JoinHandle will not modify JOIN_WAKER. After the |
106 | //! runtime sets COMPLETE to one, it invokes the waker if there is one. |
107 | //! |
108 | //! All other fields are immutable and can be accessed immutably without |
109 | //! synchronization by anyone. |
110 | //! |
111 | //! # Safety |
112 | //! |
113 | //! This section goes through various situations and explains why the API is |
114 | //! safe in that situation. |
115 | //! |
116 | //! ## Polling or dropping the future |
117 | //! |
118 | //! Any mutable access to the future happens after obtaining a lock by modifying |
119 | //! the RUNNING field, so exclusive access is ensured. |
120 | //! |
121 | //! When the task completes, exclusive access to the output is transferred to |
122 | //! the JoinHandle. If the JoinHandle is already dropped when the transition to |
123 | //! complete happens, the thread performing that transition retains exclusive |
124 | //! access to the output and should immediately drop it. |
125 | //! |
126 | //! ## Non-Send futures |
127 | //! |
128 | //! If a future is not Send, then it is bound to a LocalOwnedTasks. The future |
129 | //! will only ever be polled or dropped given a LocalNotified or inside a call |
130 | //! to LocalOwnedTasks::shutdown_all. In either case, it is guaranteed that the |
131 | //! future is on the right thread. |
132 | //! |
133 | //! If the task is never removed from the LocalOwnedTasks, then it is leaked, so |
134 | //! there is no risk that the task is dropped on some other thread when the last |
135 | //! ref-count drops. |
136 | //! |
137 | //! ## Non-Send output |
138 | //! |
139 | //! When a task completes, the output is placed in the stage of the task. Then, |
140 | //! a transition that sets COMPLETE to true is performed, and the value of |
141 | //! JOIN_INTEREST when this transition happens is read. |
142 | //! |
143 | //! If JOIN_INTEREST is zero when the transition to COMPLETE happens, then the |
144 | //! output is immediately dropped. |
145 | //! |
146 | //! If JOIN_INTEREST is one when the transition to COMPLETE happens, then the |
147 | //! JoinHandle is responsible for cleaning up the output. If the output is not |
148 | //! Send, then this happens: |
149 | //! |
150 | //! 1. The output is created on the thread that the future was polled on. Since |
151 | //! only non-Send futures can have non-Send output, the future was polled on |
152 | //! the thread that the future was spawned from. |
153 | //! 2. Since `JoinHandle<Output>` is not Send if Output is not Send, the |
154 | //! JoinHandle is also on the thread that the future was spawned from. |
155 | //! 3. Thus, the JoinHandle will not move the output across threads when it |
156 | //! takes or drops the output. |
157 | //! |
158 | //! ## Recursive poll/shutdown |
159 | //! |
160 | //! Calling poll from inside a shutdown call or vice-versa is not prevented by |
161 | //! the API exposed by the task module, so this has to be safe. In either case, |
162 | //! the lock in the RUNNING bitfield makes the inner call return immediately. If |
163 | //! the inner call is a `shutdown` call, then the CANCELLED bit is set, and the |
164 | //! poll call will notice it when the poll finishes, and the task is cancelled |
165 | //! at that point. |
166 | |
167 | // Some task infrastructure is here to support `JoinSet`, which is currently |
168 | // unstable. This should be removed once `JoinSet` is stabilized. |
169 | #![cfg_attr (not(tokio_unstable), allow(dead_code))] |
170 | |
171 | mod core; |
172 | use self::core::Cell; |
173 | use self::core::Header; |
174 | |
175 | mod error; |
176 | pub use self::error::JoinError; |
177 | |
178 | mod harness; |
179 | use self::harness::Harness; |
180 | |
181 | mod id; |
182 | #[cfg_attr (not(tokio_unstable), allow(unreachable_pub, unused_imports))] |
183 | pub use id::{id, try_id, Id}; |
184 | |
185 | #[cfg (feature = "rt" )] |
186 | mod abort; |
187 | mod join; |
188 | |
189 | #[cfg (feature = "rt" )] |
190 | pub use self::abort::AbortHandle; |
191 | |
192 | pub use self::join::JoinHandle; |
193 | |
194 | mod list; |
195 | pub(crate) use self::list::{LocalOwnedTasks, OwnedTasks}; |
196 | |
197 | mod raw; |
198 | pub(crate) use self::raw::RawTask; |
199 | |
200 | mod state; |
201 | use self::state::State; |
202 | |
203 | mod waker; |
204 | |
205 | cfg_taskdump! { |
206 | pub(crate) mod trace; |
207 | } |
208 | |
209 | use crate::future::Future; |
210 | use crate::util::linked_list; |
211 | use crate::util::sharded_list; |
212 | |
213 | use std::marker::PhantomData; |
214 | use std::ptr::NonNull; |
215 | use std::{fmt, mem}; |
216 | |
217 | /// An owned handle to the task, tracked by ref count. |
218 | #[repr (transparent)] |
219 | pub(crate) struct Task<S: 'static> { |
220 | raw: RawTask, |
221 | _p: PhantomData<S>, |
222 | } |
223 | |
224 | unsafe impl<S> Send for Task<S> {} |
225 | unsafe impl<S> Sync for Task<S> {} |
226 | |
227 | /// A task was notified. |
228 | #[repr (transparent)] |
229 | pub(crate) struct Notified<S: 'static>(Task<S>); |
230 | |
231 | // safety: This type cannot be used to touch the task without first verifying |
232 | // that the value is on a thread where it is safe to poll the task. |
233 | unsafe impl<S: Schedule> Send for Notified<S> {} |
234 | unsafe impl<S: Schedule> Sync for Notified<S> {} |
235 | |
236 | /// A non-Send variant of Notified with the invariant that it is on a thread |
237 | /// where it is safe to poll it. |
238 | #[repr (transparent)] |
239 | pub(crate) struct LocalNotified<S: 'static> { |
240 | task: Task<S>, |
241 | _not_send: PhantomData<*const ()>, |
242 | } |
243 | |
244 | /// A task that is not owned by any OwnedTasks. Used for blocking tasks. |
245 | /// This type holds two ref-counts. |
246 | pub(crate) struct UnownedTask<S: 'static> { |
247 | raw: RawTask, |
248 | _p: PhantomData<S>, |
249 | } |
250 | |
251 | // safety: This type can only be created given a Send task. |
252 | unsafe impl<S> Send for UnownedTask<S> {} |
253 | unsafe impl<S> Sync for UnownedTask<S> {} |
254 | |
255 | /// Task result sent back. |
256 | pub(crate) type Result<T> = std::result::Result<T, JoinError>; |
257 | |
258 | pub(crate) trait Schedule: Sync + Sized + 'static { |
259 | /// The task has completed work and is ready to be released. The scheduler |
260 | /// should release it immediately and return it. The task module will batch |
261 | /// the ref-dec with setting other options. |
262 | /// |
263 | /// If the scheduler has already released the task, then None is returned. |
264 | fn release(&self, task: &Task<Self>) -> Option<Task<Self>>; |
265 | |
266 | /// Schedule the task |
267 | fn schedule(&self, task: Notified<Self>); |
268 | |
269 | /// Schedule the task to run in the near future, yielding the thread to |
270 | /// other tasks. |
271 | fn yield_now(&self, task: Notified<Self>) { |
272 | self.schedule(task); |
273 | } |
274 | |
275 | /// Polling the task resulted in a panic. Should the runtime shutdown? |
276 | fn unhandled_panic(&self) { |
277 | // By default, do nothing. This maintains the 1.0 behavior. |
278 | } |
279 | } |
280 | |
281 | cfg_rt! { |
282 | /// This is the constructor for a new task. Three references to the task are |
283 | /// created. The first task reference is usually put into an OwnedTasks |
284 | /// immediately. The Notified is sent to the scheduler as an ordinary |
285 | /// notification. |
286 | fn new_task<T, S>( |
287 | task: T, |
288 | scheduler: S, |
289 | id: Id, |
290 | ) -> (Task<S>, Notified<S>, JoinHandle<T::Output>) |
291 | where |
292 | S: Schedule, |
293 | T: Future + 'static, |
294 | T::Output: 'static, |
295 | { |
296 | let raw = RawTask::new::<T, S>(task, scheduler, id); |
297 | let task = Task { |
298 | raw, |
299 | _p: PhantomData, |
300 | }; |
301 | let notified = Notified(Task { |
302 | raw, |
303 | _p: PhantomData, |
304 | }); |
305 | let join = JoinHandle::new(raw); |
306 | |
307 | (task, notified, join) |
308 | } |
309 | |
310 | /// Creates a new task with an associated join handle. This method is used |
311 | /// only when the task is not going to be stored in an `OwnedTasks` list. |
312 | /// |
313 | /// Currently only blocking tasks use this method. |
314 | pub(crate) fn unowned<T, S>(task: T, scheduler: S, id: Id) -> (UnownedTask<S>, JoinHandle<T::Output>) |
315 | where |
316 | S: Schedule, |
317 | T: Send + Future + 'static, |
318 | T::Output: Send + 'static, |
319 | { |
320 | let (task, notified, join) = new_task(task, scheduler, id); |
321 | |
322 | // This transfers the ref-count of task and notified into an UnownedTask. |
323 | // This is valid because an UnownedTask holds two ref-counts. |
324 | let unowned = UnownedTask { |
325 | raw: task.raw, |
326 | _p: PhantomData, |
327 | }; |
328 | std::mem::forget(task); |
329 | std::mem::forget(notified); |
330 | |
331 | (unowned, join) |
332 | } |
333 | } |
334 | |
335 | impl<S: 'static> Task<S> { |
336 | unsafe fn new(raw: RawTask) -> Task<S> { |
337 | Task { |
338 | raw, |
339 | _p: PhantomData, |
340 | } |
341 | } |
342 | |
343 | unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> { |
344 | Task::new(RawTask::from_raw(ptr)) |
345 | } |
346 | |
347 | #[cfg (all( |
348 | tokio_unstable, |
349 | tokio_taskdump, |
350 | feature = "rt" , |
351 | target_os = "linux" , |
352 | any(target_arch = "aarch64" , target_arch = "x86" , target_arch = "x86_64" ) |
353 | ))] |
354 | pub(super) fn as_raw(&self) -> RawTask { |
355 | self.raw |
356 | } |
357 | |
358 | fn header(&self) -> &Header { |
359 | self.raw.header() |
360 | } |
361 | |
362 | fn header_ptr(&self) -> NonNull<Header> { |
363 | self.raw.header_ptr() |
364 | } |
365 | |
366 | cfg_taskdump! { |
367 | /// Notify the task for task dumping. |
368 | /// |
369 | /// Returns `None` if the task has already been notified. |
370 | pub(super) fn notify_for_tracing(&self) -> Option<Notified<S>> { |
371 | if self.as_raw().state().transition_to_notified_for_tracing() { |
372 | // SAFETY: `transition_to_notified_for_tracing` increments the |
373 | // refcount. |
374 | Some(unsafe { Notified(Task::new(self.raw)) }) |
375 | } else { |
376 | None |
377 | } |
378 | } |
379 | } |
380 | } |
381 | |
382 | impl<S: 'static> Notified<S> { |
383 | fn header(&self) -> &Header { |
384 | self.0.header() |
385 | } |
386 | } |
387 | |
388 | impl<S: 'static> Notified<S> { |
389 | pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified<S> { |
390 | Notified(Task::new(ptr)) |
391 | } |
392 | } |
393 | |
394 | impl<S: 'static> Notified<S> { |
395 | pub(crate) fn into_raw(self) -> RawTask { |
396 | let raw = self.0.raw; |
397 | mem::forget(self); |
398 | raw |
399 | } |
400 | } |
401 | |
402 | impl<S: Schedule> Task<S> { |
403 | /// Preemptively cancels the task as part of the shutdown process. |
404 | pub(crate) fn shutdown(self) { |
405 | let raw = self.raw; |
406 | mem::forget(self); |
407 | raw.shutdown(); |
408 | } |
409 | } |
410 | |
411 | impl<S: Schedule> LocalNotified<S> { |
412 | /// Runs the task. |
413 | pub(crate) fn run(self) { |
414 | let raw = self.task.raw; |
415 | mem::forget(self); |
416 | raw.poll(); |
417 | } |
418 | } |
419 | |
420 | impl<S: Schedule> UnownedTask<S> { |
421 | // Used in test of the inject queue. |
422 | #[cfg (test)] |
423 | #[cfg_attr (target_family = "wasm" , allow(dead_code))] |
424 | pub(super) fn into_notified(self) -> Notified<S> { |
425 | Notified(self.into_task()) |
426 | } |
427 | |
428 | fn into_task(self) -> Task<S> { |
429 | // Convert into a task. |
430 | let task = Task { |
431 | raw: self.raw, |
432 | _p: PhantomData, |
433 | }; |
434 | mem::forget(self); |
435 | |
436 | // Drop a ref-count since an UnownedTask holds two. |
437 | task.header().state.ref_dec(); |
438 | |
439 | task |
440 | } |
441 | |
442 | pub(crate) fn run(self) { |
443 | let raw = self.raw; |
444 | mem::forget(self); |
445 | |
446 | // Transfer one ref-count to a Task object. |
447 | let task = Task::<S> { |
448 | raw, |
449 | _p: PhantomData, |
450 | }; |
451 | |
452 | // Use the other ref-count to poll the task. |
453 | raw.poll(); |
454 | // Decrement our extra ref-count |
455 | drop(task); |
456 | } |
457 | |
458 | pub(crate) fn shutdown(self) { |
459 | self.into_task().shutdown(); |
460 | } |
461 | } |
462 | |
463 | impl<S: 'static> Drop for Task<S> { |
464 | fn drop(&mut self) { |
465 | // Decrement the ref count |
466 | if self.header().state.ref_dec() { |
467 | // Deallocate if this is the final ref count |
468 | self.raw.dealloc(); |
469 | } |
470 | } |
471 | } |
472 | |
473 | impl<S: 'static> Drop for UnownedTask<S> { |
474 | fn drop(&mut self) { |
475 | // Decrement the ref count |
476 | if self.raw.header().state.ref_dec_twice() { |
477 | // Deallocate if this is the final ref count |
478 | self.raw.dealloc(); |
479 | } |
480 | } |
481 | } |
482 | |
483 | impl<S> fmt::Debug for Task<S> { |
484 | fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { |
485 | write!(fmt, "Task({:p})" , self.header()) |
486 | } |
487 | } |
488 | |
489 | impl<S> fmt::Debug for Notified<S> { |
490 | fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { |
491 | write!(fmt, "task::Notified({:p})" , self.0.header()) |
492 | } |
493 | } |
494 | |
495 | /// # Safety |
496 | /// |
497 | /// Tasks are pinned. |
498 | unsafe impl<S> linked_list::Link for Task<S> { |
499 | type Handle = Task<S>; |
500 | type Target = Header; |
501 | |
502 | fn as_raw(handle: &Task<S>) -> NonNull<Header> { |
503 | handle.raw.header_ptr() |
504 | } |
505 | |
506 | unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> { |
507 | Task::from_raw(ptr) |
508 | } |
509 | |
510 | unsafe fn pointers(target: NonNull<Header>) -> NonNull<linked_list::Pointers<Header>> { |
511 | self::core::Trailer::addr_of_owned(Header::get_trailer(target)) |
512 | } |
513 | } |
514 | |
515 | /// # Safety |
516 | /// |
517 | /// The id of a task is never changed after creation of the task, so the return value of |
518 | /// `get_shard_id` will not change. (The cast may throw away the upper 32 bits of the task id, but |
519 | /// the shard id still won't change from call to call.) |
520 | unsafe impl<S> sharded_list::ShardedListItem for Task<S> { |
521 | unsafe fn get_shard_id(target: NonNull<Self::Target>) -> usize { |
522 | // SAFETY: The caller guarantees that `target` points at a valid task. |
523 | let task_id = unsafe { Header::get_id(target) }; |
524 | task_id.0 as usize |
525 | } |
526 | } |
527 | |