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
171mod core;
172use self::core::Cell;
173use self::core::Header;
174
175mod error;
176pub use self::error::JoinError;
177
178mod harness;
179use self::harness::Harness;
180
181mod id;
182#[cfg_attr(not(tokio_unstable), allow(unreachable_pub, unused_imports))]
183pub use id::{id, try_id, Id};
184
185#[cfg(feature = "rt")]
186mod abort;
187mod join;
188
189#[cfg(feature = "rt")]
190pub use self::abort::AbortHandle;
191
192pub use self::join::JoinHandle;
193
194mod list;
195pub(crate) use self::list::{LocalOwnedTasks, OwnedTasks};
196
197mod raw;
198pub(crate) use self::raw::RawTask;
199
200mod state;
201use self::state::State;
202
203mod waker;
204
205cfg_taskdump! {
206 pub(crate) mod trace;
207}
208
209use crate::future::Future;
210use crate::util::linked_list;
211use crate::util::sharded_list;
212
213use std::marker::PhantomData;
214use std::ptr::NonNull;
215use std::{fmt, mem};
216
217/// An owned handle to the task, tracked by ref count.
218#[repr(transparent)]
219pub(crate) struct Task<S: 'static> {
220 raw: RawTask,
221 _p: PhantomData<S>,
222}
223
224unsafe impl<S> Send for Task<S> {}
225unsafe impl<S> Sync for Task<S> {}
226
227/// A task was notified.
228#[repr(transparent)]
229pub(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.
233unsafe impl<S: Schedule> Send for Notified<S> {}
234unsafe 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)]
239pub(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.
246pub(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.
252unsafe impl<S> Send for UnownedTask<S> {}
253unsafe impl<S> Sync for UnownedTask<S> {}
254
255/// Task result sent back.
256pub(crate) type Result<T> = std::result::Result<T, JoinError>;
257
258pub(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
281cfg_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
335impl<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
382impl<S: 'static> Notified<S> {
383 fn header(&self) -> &Header {
384 self.0.header()
385 }
386}
387
388impl<S: 'static> Notified<S> {
389 pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified<S> {
390 Notified(Task::new(ptr))
391 }
392}
393
394impl<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
402impl<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
411impl<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
420impl<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
463impl<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
473impl<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
483impl<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
489impl<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.
498unsafe 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.)
520unsafe 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