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))]
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;
211
212use std::marker::PhantomData;
213use std::ptr::NonNull;
214use std::{fmt, mem};
215
216/// An owned handle to the task, tracked by ref count.
217#[repr(transparent)]
218pub(crate) struct Task<S: 'static> {
219 raw: RawTask,
220 _p: PhantomData<S>,
221}
222
223unsafe impl<S> Send for Task<S> {}
224unsafe impl<S> Sync for Task<S> {}
225
226/// A task was notified.
227#[repr(transparent)]
228pub(crate) struct Notified<S: 'static>(Task<S>);
229
230// safety: This type cannot be used to touch the task without first verifying
231// that the value is on a thread where it is safe to poll the task.
232unsafe impl<S: Schedule> Send for Notified<S> {}
233unsafe impl<S: Schedule> Sync for Notified<S> {}
234
235/// A non-Send variant of Notified with the invariant that it is on a thread
236/// where it is safe to poll it.
237#[repr(transparent)]
238pub(crate) struct LocalNotified<S: 'static> {
239 task: Task<S>,
240 _not_send: PhantomData<*const ()>,
241}
242
243/// A task that is not owned by any OwnedTasks. Used for blocking tasks.
244/// This type holds two ref-counts.
245pub(crate) struct UnownedTask<S: 'static> {
246 raw: RawTask,
247 _p: PhantomData<S>,
248}
249
250// safety: This type can only be created given a Send task.
251unsafe impl<S> Send for UnownedTask<S> {}
252unsafe impl<S> Sync for UnownedTask<S> {}
253
254/// Task result sent back.
255pub(crate) type Result<T> = std::result::Result<T, JoinError>;
256
257pub(crate) trait Schedule: Sync + Sized + 'static {
258 /// The task has completed work and is ready to be released. The scheduler
259 /// should release it immediately and return it. The task module will batch
260 /// the ref-dec with setting other options.
261 ///
262 /// If the scheduler has already released the task, then None is returned.
263 fn release(&self, task: &Task<Self>) -> Option<Task<Self>>;
264
265 /// Schedule the task
266 fn schedule(&self, task: Notified<Self>);
267
268 /// Schedule the task to run in the near future, yielding the thread to
269 /// other tasks.
270 fn yield_now(&self, task: Notified<Self>) {
271 self.schedule(task);
272 }
273
274 /// Polling the task resulted in a panic. Should the runtime shutdown?
275 fn unhandled_panic(&self) {
276 // By default, do nothing. This maintains the 1.0 behavior.
277 }
278}
279
280cfg_rt! {
281 /// This is the constructor for a new task. Three references to the task are
282 /// created. The first task reference is usually put into an OwnedTasks
283 /// immediately. The Notified is sent to the scheduler as an ordinary
284 /// notification.
285 fn new_task<T, S>(
286 task: T,
287 scheduler: S,
288 id: Id,
289 ) -> (Task<S>, Notified<S>, JoinHandle<T::Output>)
290 where
291 S: Schedule,
292 T: Future + 'static,
293 T::Output: 'static,
294 {
295 let raw = RawTask::new::<T, S>(task, scheduler, id);
296 let task = Task {
297 raw,
298 _p: PhantomData,
299 };
300 let notified = Notified(Task {
301 raw,
302 _p: PhantomData,
303 });
304 let join = JoinHandle::new(raw);
305
306 (task, notified, join)
307 }
308
309 /// Creates a new task with an associated join handle. This method is used
310 /// only when the task is not going to be stored in an `OwnedTasks` list.
311 ///
312 /// Currently only blocking tasks use this method.
313 pub(crate) fn unowned<T, S>(task: T, scheduler: S, id: Id) -> (UnownedTask<S>, JoinHandle<T::Output>)
314 where
315 S: Schedule,
316 T: Send + Future + 'static,
317 T::Output: Send + 'static,
318 {
319 let (task, notified, join) = new_task(task, scheduler, id);
320
321 // This transfers the ref-count of task and notified into an UnownedTask.
322 // This is valid because an UnownedTask holds two ref-counts.
323 let unowned = UnownedTask {
324 raw: task.raw,
325 _p: PhantomData,
326 };
327 std::mem::forget(task);
328 std::mem::forget(notified);
329
330 (unowned, join)
331 }
332}
333
334impl<S: 'static> Task<S> {
335 unsafe fn new(raw: RawTask) -> Task<S> {
336 Task {
337 raw,
338 _p: PhantomData,
339 }
340 }
341
342 unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
343 Task::new(RawTask::from_raw(ptr))
344 }
345
346 #[cfg(all(
347 tokio_unstable,
348 tokio_taskdump,
349 feature = "rt",
350 target_os = "linux",
351 any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
352 ))]
353 pub(super) fn as_raw(&self) -> RawTask {
354 self.raw
355 }
356
357 fn header(&self) -> &Header {
358 self.raw.header()
359 }
360
361 fn header_ptr(&self) -> NonNull<Header> {
362 self.raw.header_ptr()
363 }
364}
365
366impl<S: 'static> Notified<S> {
367 fn header(&self) -> &Header {
368 self.0.header()
369 }
370}
371
372impl<S: 'static> Notified<S> {
373 pub(crate) unsafe fn from_raw(ptr: RawTask) -> Notified<S> {
374 Notified(Task::new(raw:ptr))
375 }
376}
377
378impl<S: 'static> Notified<S> {
379 pub(crate) fn into_raw(self) -> RawTask {
380 let raw: RawTask = self.0.raw;
381 mem::forget(self);
382 raw
383 }
384}
385
386impl<S: Schedule> Task<S> {
387 /// Preemptively cancels the task as part of the shutdown process.
388 pub(crate) fn shutdown(self) {
389 let raw: RawTask = self.raw;
390 mem::forget(self);
391 raw.shutdown();
392 }
393}
394
395impl<S: Schedule> LocalNotified<S> {
396 /// Runs the task.
397 pub(crate) fn run(self) {
398 let raw: RawTask = self.task.raw;
399 mem::forget(self);
400 raw.poll();
401 }
402}
403
404impl<S: Schedule> UnownedTask<S> {
405 // Used in test of the inject queue.
406 #[cfg(test)]
407 #[cfg_attr(tokio_wasm, allow(dead_code))]
408 pub(super) fn into_notified(self) -> Notified<S> {
409 Notified(self.into_task())
410 }
411
412 fn into_task(self) -> Task<S> {
413 // Convert into a task.
414 let task = Task {
415 raw: self.raw,
416 _p: PhantomData,
417 };
418 mem::forget(self);
419
420 // Drop a ref-count since an UnownedTask holds two.
421 task.header().state.ref_dec();
422
423 task
424 }
425
426 pub(crate) fn run(self) {
427 let raw = self.raw;
428 mem::forget(self);
429
430 // Transfer one ref-count to a Task object.
431 let task = Task::<S> {
432 raw,
433 _p: PhantomData,
434 };
435
436 // Use the other ref-count to poll the task.
437 raw.poll();
438 // Decrement our extra ref-count
439 drop(task);
440 }
441
442 pub(crate) fn shutdown(self) {
443 self.into_task().shutdown()
444 }
445}
446
447impl<S: 'static> Drop for Task<S> {
448 fn drop(&mut self) {
449 // Decrement the ref count
450 if self.header().state.ref_dec() {
451 // Deallocate if this is the final ref count
452 self.raw.dealloc();
453 }
454 }
455}
456
457impl<S: 'static> Drop for UnownedTask<S> {
458 fn drop(&mut self) {
459 // Decrement the ref count
460 if self.raw.header().state.ref_dec_twice() {
461 // Deallocate if this is the final ref count
462 self.raw.dealloc();
463 }
464 }
465}
466
467impl<S> fmt::Debug for Task<S> {
468 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
469 write!(fmt, "Task({:p})", self.header())
470 }
471}
472
473impl<S> fmt::Debug for Notified<S> {
474 fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
475 write!(fmt, "task::Notified({:p})", self.0.header())
476 }
477}
478
479/// # Safety
480///
481/// Tasks are pinned.
482unsafe impl<S> linked_list::Link for Task<S> {
483 type Handle = Task<S>;
484 type Target = Header;
485
486 fn as_raw(handle: &Task<S>) -> NonNull<Header> {
487 handle.raw.header_ptr()
488 }
489
490 unsafe fn from_raw(ptr: NonNull<Header>) -> Task<S> {
491 Task::from_raw(ptr)
492 }
493
494 unsafe fn pointers(target: NonNull<Header>) -> NonNull<linked_list::Pointers<Header>> {
495 self::core::Trailer::addr_of_owned(me:Header::get_trailer(me:target))
496 }
497}
498