Skip to main content

shuttle_engine/runtime/task/
mod.rs

1use crate::backtrace_enabled;
2use crate::current::get_name_for_task;
3use crate::runtime::execution::{ExecutionState, TASK_ID_TO_TAGS};
4use crate::runtime::storage::{AlreadyDestructedError, StorageKey, StorageMap};
5use crate::runtime::task::clock::VectorClock;
6use crate::runtime::task::labels::Labels;
7use crate::runtime::thread;
8use crate::runtime::thread::continuation::{
9    ContinuationInput, ContinuationOutput, ContinuationPool, PooledContinuation,
10};
11use crate::sync_types::{ResourceSignature, ResourceType};
12use crate::thread_support::LocalKey;
13use bitvec::prelude::*;
14use corosensei::Yielder;
15use std::any::Any;
16use std::backtrace::Backtrace;
17use std::cell::RefCell;
18use std::collections::HashMap;
19use std::fmt::Debug;
20use std::future::Future;
21use std::hash::{DefaultHasher, Hash, Hasher};
22use std::panic::Location;
23use std::rc::Rc;
24use std::sync::Arc;
25use std::task::{Context, Waker};
26use tracing::{error_span, event, field, Level, Span};
27
28pub mod clock;
29pub mod labels;
30pub mod waker;
31use waker::make_waker;
32
33// A note on terminology: we have competing notions of threads floating around. Here's the
34// convention for disambiguating them:
35// * A "thread" is a user-level unit of concurrency. User code creates threads, passes data
36//   between them, etc.
37// * A "future" is another user-level unit of concurrency, corresponding directly to Rust's notion
38//   in std::future::Future. A future has a single method `poll` that can be used to resume
39//   executing its computation. Both futures and threads are implemented in Task,
40//   which wraps a continuation that is resumed when the task is scheduled.
41// * A "task" is the Shuttle executor's reflection of a user-level unit of concurrency. Each task
42//   has a corresponding continuation, which is the user-level code it runs, as well as a state like
43//   "blocked", "runnable", etc. Scheduling algorithms take as input the state of all tasks
44//   and decide which task should execute next. A context switch is when one task stops executing
45//   and another begins.
46// * A "continuation" is a low-level implementation of green threading for concurrency. Each
47//   Task contains a corresponding continuation. When the Shuttle executor context switches to a
48//   Task, the executor resumes that task's continuation until it yields, which happens when its
49//   thread decides it might want to context switch (e.g., because it's blocked on a lock).
50
51pub const DEFAULT_INLINE_TASKS: usize = 16;
52
53/// A reserved label that is used to assign readable names to tasks for debugging.
54///
55/// To make debugging easier, if a task is assigned a `TaskName(s)` Label,
56/// Shuttle will display the String `s` in addition to the `TaskId` in debug output.
57#[derive(Clone, PartialEq, Eq)]
58pub struct TaskName(String);
59
60impl From<String> for TaskName {
61    fn from(s: String) -> Self {
62        Self(s)
63    }
64}
65
66impl From<&str> for TaskName {
67    fn from(s: &str) -> Self {
68        Self(String::from(s))
69    }
70}
71
72impl std::fmt::Debug for TaskName {
73    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
74        write!(f, "{}", self.0)
75    }
76}
77
78impl From<TaskName> for String {
79    fn from(task_name: TaskName) -> Self {
80        task_name.0
81    }
82}
83
84impl<'a> From<&'a TaskName> for &'a String {
85    fn from(task_name: &'a TaskName) -> Self {
86        &task_name.0
87    }
88}
89
90/// A special label that can be used to set labels for a task when it is spawned.
91///
92/// By default, when a task or thread T is spawned, it inherits all labels from its parent.
93/// It's often useful to modify or add new Labels to T.  One approach is to put label changes
94/// at the beginning of the closure that is passed to `spawn`, but this approach has the drawback
95/// that the changes are applied only when T is first selected for execution, and the closure
96/// is invoked.  To overcome this drawback, we introduce the `ChildLabelFn` label.  If a parent
97/// task or thread has a `ChildLabelFn` set when it spawns a new child task or thread, the
98/// child's label set at spawn time will be modified by applying the function inside the `ChildLabelFn`.
99///
100/// # Example
101/// The following example shows how a `ChildLabelFn` can be used to set up names for the next child(ren)
102/// that will be spawned by a parent task (see `shuttle/tests/basic/labels.rs` for runnable versions).
103/// ```ignore
104/// # use shuttle_engine::current::{me, set_label_for_task, get_name_for_task, ChildLabelFn, TaskName};
105/// # use std::sync::Arc;
106/// // In the parent, set up a `ChildLabelFn` that assigns a name to the child task
107/// shuttle::check_dfs(|| {
108///   set_label_for_task(me(), ChildLabelFn(Arc::new(|_task_id, labels| { labels.insert(TaskName::from("ChildTask")); })));
109///   shuttle::thread::spawn(|| {
110///     assert_eq!(get_name_for_task(me()).unwrap(), TaskName::from("ChildTask")); // child task already has the name
111///     // ... rest of child
112///   }).join().unwrap();
113/// }, None);
114/// ```
115#[derive(Clone)]
116#[allow(clippy::type_complexity)]
117pub struct ChildLabelFn(pub Arc<dyn Fn(TaskId, &mut Labels) + 'static>);
118
119impl Debug for ChildLabelFn {
120    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
121        write!(f, "ChildLabelFn")
122    }
123}
124
125/// A `Tag` is an optional piece of metadata associated with a task (a thread or spawned future) to
126/// aid debugging.
127///
128/// It is automatically implemented for types which implement `Taggable` and are `Any`.
129///
130/// When set, the tag will be included in the [Debug] representation of [TaskId]s, which can help
131/// identify tasks in failing Shuttle tests. A task's [Tag] can be set with the
132/// [set_tag_for_current_task](crate::current::set_tag_for_current_task) function. Newly spawned
133/// threads and futures inherit the tag of their parent at spawn time.
134#[deprecated]
135#[allow(deprecated)]
136pub trait Tag: Taggable {
137    /// Return the tag as `Any`, typically so that it can be downcast to a known concrete type
138    fn as_any(&self) -> &dyn Any;
139}
140/// `Taggable` is a marker trait which types implementing `Tag` have to implement.
141/// It exists since we both want to provide a blanket implementation of `as_any`, and have users
142/// opt in to a type being able to be used as a tag. If we did not have this trait, then `Tag`
143/// would be automatically implemented for most types (as most types are `Debug + Any`), which
144/// opens up for accidentally using a type which was not intended to be used as a tag as a tag.
145#[deprecated]
146pub trait Taggable: Debug {}
147
148#[allow(deprecated)]
149impl<T> Tag for T
150where
151    T: Taggable + Any,
152{
153    fn as_any(&self) -> &dyn Any {
154        self
155    }
156}
157
158/// A task signature is an identifier that is intended to be *mostly* stable across executions
159/// and allow for categorization of tasks according to how they were created. It provides two
160/// levels of granularity: static (compile-time) spawn location and dynamic (run-time) context where
161/// that spawn location was reached. The static spawn location and signature are each represented
162/// by a u64 so that the details of how they are computed can be non-breaking changes in the future.
163/// Hashes are all pre-computed for fast checking of equality of signatures at runtime.
164#[derive(Debug, Clone)]
165pub struct TaskSignature {
166    /// The task creation stack is a tuple of (create location, number of tasks created at that location in the parent)
167    task_creation_stack: Vec<(&'static Location<'static>, u32)>,
168    spawn_call_site_hash: u64,
169    parent_signature_hash: u64,
170    signature_hash: u64,
171    child_counters: HashMap<&'static Location<'static>, u32>,
172}
173
174impl TaskSignature {
175    pub fn new_parentless(spawn_call_site: &'static Location<'static>) -> TaskSignature {
176        let mut hasher = DefaultHasher::new();
177        let task_creation_stack = vec![(spawn_call_site, 0)];
178        task_creation_stack.hash(&mut hasher);
179        let signature_hash = hasher.finish();
180        spawn_call_site.hash(&mut hasher);
181
182        Self {
183            task_creation_stack,
184            spawn_call_site_hash: hasher.finish(),
185            parent_signature_hash: 0,
186            signature_hash,
187            child_counters: HashMap::new(),
188        }
189    }
190
191    pub fn new_child(&mut self, spawn_call_site: &'static Location<'static>) -> Self {
192        let mut hasher = DefaultHasher::new();
193        let counter = self
194            .child_counters
195            .entry(spawn_call_site)
196            .and_modify(|c| *c += 1)
197            .or_insert(1);
198        let mut task_creation_stack = self.task_creation_stack.clone();
199        task_creation_stack.push((spawn_call_site, *counter));
200
201        spawn_call_site.hash(&mut hasher);
202        let spawn_call_site_hash = hasher.finish();
203
204        task_creation_stack.hash(&mut hasher);
205
206        Self {
207            task_creation_stack,
208            parent_signature_hash: self.signature_hash,
209            spawn_call_site_hash,
210            signature_hash: hasher.finish(),
211            child_counters: HashMap::new(),
212        }
213    }
214
215    #[track_caller]
216    pub fn new_resource(&mut self, resource_type: ResourceType) -> ResourceSignature {
217        let static_create_location = Location::caller();
218        let counter = self
219            .child_counters
220            .entry(static_create_location)
221            .and_modify(|c| *c += 1)
222            .or_insert(1);
223
224        ResourceSignature::new(resource_type, static_create_location, self.signature_hash, *counter)
225    }
226
227    /// Hash of the static location within the source code where the task was spawned
228    pub fn static_create_location_hash(&self) -> u64 {
229        self.spawn_call_site_hash
230    }
231
232    /// Combined signature of the static location and dynamic context
233    /// context where the task was spawned.
234    pub fn signature_hash(&self) -> u64 {
235        self.signature_hash
236    }
237
238    /// Signature hash of the parent of this task
239    pub fn parent_signature_hash(&self) -> u64 {
240        self.parent_signature_hash
241    }
242}
243
244impl Hash for TaskSignature {
245    fn hash<H: Hasher>(&self, state: &mut H) {
246        self.task_creation_stack.hash(state);
247    }
248}
249
250impl PartialEq for TaskSignature {
251    fn eq(&self, other: &Self) -> bool {
252        self.signature_hash == other.signature_hash
253    }
254}
255
256impl Eq for TaskSignature {}
257
258/// A task's default `tracing` dispatcher, parked while the task is switched out (see
259/// `Execution::park_task_default`).
260#[derive(Debug)]
261pub(crate) struct ParkedDefault(Option<tracing::dispatcher::DefaultGuard>);
262
263impl ParkedDefault {
264    pub(crate) fn new(guard: tracing::dispatcher::DefaultGuard) -> Self {
265        Self(Some(guard))
266    }
267
268    /// Reinstates the task's default dispatcher.
269    pub(crate) fn reinstate(mut self) {
270        drop(self.0.take());
271    }
272}
273
274impl Drop for ParkedDefault {
275    // Only reached when the task is torn down while switched out. Reinstating the task's default then
276    // lets the guards on its stack restore their priors in order as the stack unwinds, which leaves
277    // the execution's default in place. A stack that is leaked instead of unwound (see
278    // `Continuation::drop`) would leave the task's default installed for good, so in that case keep
279    // the execution's default, and leak the guard along with the stack.
280    fn drop(&mut self) {
281        if let Some(guard) = self.0.take() {
282            if std::thread::panicking() || ExecutionState::execution_stopped() {
283                std::mem::forget(guard);
284            }
285        }
286    }
287}
288
289/// A `Task` represents a user-level unit of concurrency. Each task has an `id` that is unique within
290/// the execution, and a `state` reflecting whether the task is runnable (enabled) or not.
291#[derive(Debug)]
292pub struct Task {
293    pub(super) id: TaskId,
294    pub(super) parent_task_id: Option<TaskId>,
295    pub(super) state: TaskState,
296    pub(super) detached: bool,
297    park_state: ParkState,
298
299    // The task's default `tracing` dispatcher while the task is switched out. Declared before
300    // `continuation`, so that it is dropped before the task's stack is (see `ParkedDefault`).
301    pub(super) parked_default: Option<ParkedDefault>,
302
303    pub(super) continuation: Rc<RefCell<PooledContinuation>>,
304    pub(super) yielder: *const Yielder<ContinuationInput, ContinuationOutput>,
305
306    pub clock: VectorClock,
307
308    waiter: Option<TaskId>,
309
310    waker: Waker,
311    // Remember whether the waker was invoked while we were running
312    woken: bool,
313
314    name: Option<String>,
315
316    local_storage: StorageMap,
317
318    // The `Span` which looks like this: step{task=task_id}, or, if step count recording is enabled, like this:
319    // step{task=task_id i=step_count}. Becomes the parent of the spans created by the `Task`.
320    pub step_span: Span,
321
322    // The current `Span` "stack" of the `Task`.
323    // `Span`s are stored such that the `Task`s current `Span` is at `span_stack[0]`, that `Span`s parent (if it exists)
324    // is at `span_stack[1]`, and so on, until `span_stack[span_stack.len()-1]`, which is the "outermost" (left-most when printed)
325    // `Span`. This means that `span_stack[span_stack.len()-1]` will usually be the `Span` saying `execution{i=X}`.
326    // We `pop` it empty when resuming a `Task`, and `push` + `exit` `tracing::Span::current()`
327    // until there is no entered `Span` when we switch out of the `Task`.
328    // There are two things to note:
329    // 1: We have to own the `Span`s (versus storing `Id`s) for the `Span` to not get dropped while the task is switched out.
330    // 2: We have to store the stack of `Span`s in order to return to the correct `Span` once the `Entered<'_>` from an
331    //    `instrument`ed future is dropped.
332    pub(super) span_stack: Vec<Span>,
333
334    // Arbitrarily settable tag which is inherited from the parent.
335    #[allow(deprecated)]
336    tag: Option<Arc<dyn Tag>>,
337
338    /// If [`crate::CAPTURE_BACKTRACE`] is set then this will be populated on task block.
339    /// If the test then fails, then each task's backtrace will be printed.
340    pub backtrace: Option<Backtrace>,
341
342    /// The signature of a Task; this is an identifier that is *not* guaranteed to be unique but should be *mostly*
343    /// stable across iterations in a single Shuttle test. Tasks with the same signature are very likely to exhibit
344    /// similar behavior
345    pub signature: TaskSignature,
346}
347
348#[allow(deprecated)]
349impl Task {
350    /// Create a task from a continuation
351    #[allow(clippy::too_many_arguments)]
352    fn new(
353        f: Box<dyn FnOnce() + 'static>,
354        stack_size: usize,
355        id: TaskId,
356        name: Option<String>,
357        clock: VectorClock,
358        parent_span_id: Option<tracing::span::Id>,
359        schedule_len: usize,
360        tag: Option<Arc<dyn Tag>>,
361        parent_task_id: Option<TaskId>,
362        signature: TaskSignature,
363    ) -> Self {
364        #[cfg(all(any(test, feature = "vector-clocks"), not(feature = "bench-no-vector-clocks")))]
365        assert!(id.0 < clock.time.len());
366        let mut continuation = ContinuationPool::acquire(stack_size);
367        continuation.initialize(f);
368        let yielder = continuation.yielder;
369        let waker = make_waker(id);
370        let continuation = Rc::new(RefCell::new(continuation));
371
372        let step_span =
373            error_span!(parent: parent_span_id.clone(), "step", task = format!("{:?}", id), i = field::Empty);
374        // Note that this is slightly lazy — we are starting storing at the step_span, but could have gotten the
375        // full `Span` stack and stored that. It should be fine, but if any issues arise, then full storing should
376        // be tried.
377        let span_stack = vec![step_span.clone()];
378
379        let mut task = Self {
380            id,
381            parent_task_id,
382            state: TaskState::Runnable,
383            parked_default: None,
384            continuation,
385            yielder,
386            clock,
387            waiter: None,
388            waker,
389            woken: false,
390            detached: false,
391            park_state: ParkState::default(),
392            name,
393            step_span,
394            span_stack,
395            local_storage: StorageMap::new(),
396            tag: None,
397            backtrace: None,
398            signature,
399        };
400
401        if let Some(tag) = tag {
402            task.set_tag(tag);
403        }
404
405        // Note: the tests for the task signature in [`crate::tests::basic::task`] depend on tracing the task signature and creation point here
406        error_span!(parent: parent_span_id, "new_task", parent = ?parent_task_id, i = schedule_len).in_scope(
407            || event!(Level::DEBUG, task_id = ?task.id, signature = task.signature.signature_hash(), static_create_location = task.signature.static_create_location_hash(), "created task"),
408        );
409
410        task
411    }
412
413    #[allow(clippy::too_many_arguments)]
414    pub fn from_closure(
415        f: Box<dyn FnOnce() + 'static>,
416        stack_size: usize,
417        id: TaskId,
418        name: Option<String>,
419        clock: VectorClock,
420        parent_span_id: Option<tracing::span::Id>,
421        schedule_len: usize,
422        tag: Option<Arc<dyn Tag>>,
423        parent_task_id: Option<TaskId>,
424        signature: TaskSignature,
425    ) -> Self {
426        Self::new(
427            f,
428            stack_size,
429            id,
430            name,
431            clock,
432            parent_span_id,
433            schedule_len,
434            tag,
435            parent_task_id,
436            signature,
437        )
438    }
439
440    #[allow(clippy::too_many_arguments)]
441    pub fn from_future<F>(
442        future: F,
443        stack_size: usize,
444        id: TaskId,
445        name: Option<String>,
446        clock: VectorClock,
447        parent_span_id: Option<tracing::span::Id>,
448        schedule_len: usize,
449        tag: Option<Arc<dyn Tag>>,
450        parent_task_id: Option<TaskId>,
451        signature: TaskSignature,
452    ) -> Self
453    where
454        F: Future<Output = ()> + 'static,
455    {
456        let mut future = Box::pin(future);
457
458        Self::new(
459            Box::new(move || {
460                let waker = ExecutionState::with(|state| state.current_mut().waker());
461                let cx = &mut Context::from_waker(&waker);
462                while future.as_mut().poll(cx).is_pending() {
463                    ExecutionState::with(|state| state.current_mut().sleep_unless_woken());
464                    thread::switch();
465                }
466            }),
467            stack_size,
468            id,
469            name,
470            clock,
471            parent_span_id,
472            schedule_len,
473            tag,
474            parent_task_id,
475            signature,
476        )
477    }
478
479    /// Returns the identifier of this task.
480    pub fn id(&self) -> TaskId {
481        self.id
482    }
483
484    /// Returns the identifier of the task that spawned this task.
485    pub fn parent_task_id(&self) -> Option<TaskId> {
486        self.parent_task_id
487    }
488
489    pub fn runnable(&self) -> bool {
490        self.state == TaskState::Runnable
491    }
492
493    pub fn blocked(&self) -> bool {
494        matches!(self.state, TaskState::Blocked { .. })
495    }
496
497    pub fn can_spuriously_wakeup(&self) -> bool {
498        match self.state {
499            TaskState::Blocked { allow_spurious_wakeups } => allow_spurious_wakeups,
500            _ => false,
501        }
502    }
503
504    pub fn sleeping(&self) -> bool {
505        self.state == TaskState::Sleeping
506    }
507
508    pub fn finished(&self) -> bool {
509        self.state == TaskState::Finished
510    }
511
512    pub fn is_detached(&self) -> bool {
513        self.detached
514    }
515
516    pub fn detach(&mut self) {
517        self.detached = true;
518    }
519
520    /// Wake this task so the Wrapper future can observe the abort flag on its next poll.
521    pub fn abort(&mut self) {
522        if self.finished() {
523            return;
524        }
525        self.wake();
526    }
527
528    pub fn waker(&self) -> Waker {
529        self.waker.clone()
530    }
531
532    /// Block the current thread. If `allow_spurious_wakeups` is true, then the scheduler is
533    /// permitted to spuriously wake up the thread (though it will still not count as a live thread
534    /// for deadlock detection purposes for as long as it remains blocked).
535    pub fn block(&mut self, allow_spurious_wakeups: bool) {
536        self.backtrace = if backtrace_enabled() {
537            Some(Backtrace::force_capture())
538        } else {
539            None
540        };
541
542        assert!(self.state != TaskState::Finished);
543        self.state = TaskState::Blocked { allow_spurious_wakeups };
544    }
545
546    pub fn sleep(&mut self) {
547        assert!(self.state != TaskState::Finished);
548        self.state = TaskState::Sleeping;
549    }
550
551    pub fn unblock(&mut self) {
552        // Note we don't assert the task is blocked here. For example, a task invoking its own waker
553        // will not be blocked when this is called.
554        assert!(self.state != TaskState::Finished);
555        self.state = TaskState::Runnable;
556
557        // When a task gets unblocked, it's definitely no longer blocked in a call to `park`. This
558        // is necessary to do here because a parked task could be spuriously woken up outside of the
559        // `unpark` path. If it later becomes blocked by something else, we don't want a later
560        // `unpark` to be able to unblock the task.
561        self.park_state.blocked_in_park = false;
562    }
563
564    pub fn finish(&mut self) {
565        assert!(self.state != TaskState::Finished);
566        self.state = TaskState::Finished;
567    }
568
569    /// Potentially put this task to sleep after it was polled by the executor, unless someone has
570    /// called its waker first.
571    ///
572    /// A synchronous Task should never call this, because we want threads to be enabled-by-default
573    /// to avoid bugs where Shuttle incorrectly omits a potential execution.
574    pub fn sleep_unless_woken(&mut self) {
575        let was_woken = std::mem::replace(&mut self.woken, false);
576        if !was_woken {
577            self.sleep();
578        }
579    }
580
581    /// Remember that our waker has been called, and so we should not block the next time the
582    /// executor tries to put us to sleep.
583    pub(super) fn wake(&mut self) {
584        self.woken = true;
585        if self.state == TaskState::Sleeping {
586            self.unblock();
587        }
588    }
589
590    /// Register a waiter for this thread to terminate. Returns a boolean indicating whether the
591    /// waiter should block or not. If false, this task has already finished, and so the waiter need
592    /// not block.
593    pub fn set_waiter(&mut self, waiter: TaskId) -> bool {
594        assert!(
595            self.waiter.is_none() || self.waiter == Some(waiter),
596            "Task cannot have more than one waiter"
597        );
598        if self.finished() {
599            false
600        } else {
601            self.waiter = Some(waiter);
602            true
603        }
604    }
605
606    pub fn take_waiter(&mut self) -> Option<TaskId> {
607        self.waiter.take()
608    }
609
610    pub fn name(&self) -> Option<String> {
611        self.name.clone()
612    }
613
614    /// Retrieve a reference to the given thread-local storage slot.
615    ///
616    /// Returns Some(Err(_)) if the slot has already been destructed. Returns None if the slot has
617    /// not yet been initialized.
618    pub fn local<T: 'static>(&self, key: &'static LocalKey<T>) -> Option<Result<&T, AlreadyDestructedError>> {
619        self.local_storage.get(key.into())
620    }
621
622    /// Initialize the given thread-local storage slot with a new value.
623    ///
624    /// Panics if the slot has already been initialized.
625    pub fn init_local<T: 'static>(&mut self, key: &'static LocalKey<T>, value: T) {
626        self.local_storage.init(key.into(), value)
627    }
628
629    /// Return ownership of the next still-initialized thread-local storage slot, to be used when
630    /// running thread-local storage destructors.
631    ///
632    /// TLS destructors are a little tricky:
633    /// 1. Their code can perform synchronization operations (and so require Shuttle to call back
634    ///    into ExecutionState), so we can't drop them from within an ExecutionState borrow. Instead
635    ///    we move the contents of a slot to the caller to be dropped outside the borrow.
636    /// 2. It's valid for destructors to read other TLS slots, although destructor order is
637    ///    undefined. This also means it's valid for a destructor to *initialize* another TLS slot.
638    ///    To make this work, we run the destructors incrementally, so one destructor can initialize
639    ///    another slot that just gets added via `init_local` like normal, and then will be
640    ///    available to be popped on a future call to `pop_local`. To prevent an infinite loop, we
641    ///    forbid *reinitializing* a TLS slot whose destructor has already run, or is currently
642    ///    being run.
643    pub fn pop_local(&mut self) -> Option<Box<dyn Any>> {
644        self.local_storage.pop()
645    }
646
647    /// Park the task if its park token is unavailable. If the task blocks, then it will be woken up
648    /// when the token becomes available or spuriously without consuming the token (see the
649    /// documentation for [`std::thread::park`], which says that "it may also return spuriously,
650    /// without consuming the token"). Returns true if the execution should switch to a different
651    /// task (e.g., if the token was unavailable).
652    pub fn park(&mut self) -> bool {
653        assert!(
654            !self.park_state.blocked_in_park,
655            "task cannot park while already parked"
656        );
657        assert!(!self.blocked(), "task cannot park while blocked by something else");
658
659        if self.park_state.token_available {
660            self.park_state.token_available = false;
661            false
662        } else {
663            self.park_state.blocked_in_park = true;
664            self.block(true);
665            true
666        }
667    }
668
669    /// Make the task's park token available, and unblock the task if it was parked.
670    pub fn unpark(&mut self) {
671        if self.park_state.blocked_in_park {
672            assert!(
673                self.blocked() && self.can_spuriously_wakeup(),
674                "parked tasks should be blocked"
675            );
676            assert!(
677                !self.park_state.token_available,
678                "token shouldn't be available for parked task"
679            );
680
681            self.unblock();
682        } else {
683            // If the thread isn't currently blocked in `park`, then make the token available. If
684            // the token already is available, then this does nothing.
685            self.park_state.token_available = true;
686        }
687    }
688
689    pub fn get_tag(&self) -> Option<Arc<dyn Tag>> {
690        self.tag.clone()
691    }
692
693    /// Sets the `tag` field of the current task.
694    /// Returns the `tag` which was there previously.
695    pub fn set_tag(&mut self, tag: Arc<dyn Tag>) -> Option<Arc<dyn Tag>> {
696        TASK_ID_TO_TAGS.with(|cell| cell.borrow_mut().insert(self.id(), tag.clone()));
697        self.tag.replace(tag)
698    }
699
700    pub fn format_for_deadlock(&self) -> String {
701        use crate::backtrace_enabled;
702        format!(
703            "{} (task {:?}{}{}){}",
704            self.name().unwrap_or_else(|| "<unknown>".to_string()),
705            self.id(),
706            if self.detached { ", detached" } else { "" },
707            if self.sleeping() { ", pending future" } else { "" },
708            match (backtrace_enabled(), &self.backtrace) {
709                // `Display` prints the numbered `N: function` / `at file:line:col` layout that panics
710                // use under `RUST_BACKTRACE=1`, and ends each frame with a newline.
711                (true, Some(backtrace)) => format!("\nBacktrace:\n{backtrace}"),
712                (true, None) => "\nBacktrace: <not captured>\n".into(),
713                (false, _) => "".into(),
714            }
715        )
716    }
717}
718
719#[derive(PartialEq, Eq, Clone, Copy, Debug)]
720pub enum TaskState {
721    /// Available to be scheduled
722    Runnable,
723    /// Blocked in a synchronization operation
724    Blocked { allow_spurious_wakeups: bool },
725    /// A `Future` that returned `Pending` is waiting to be woken up
726    Sleeping,
727    /// Task has finished
728    Finished,
729}
730
731#[derive(PartialEq, Eq, Clone, Copy, Debug, Default)]
732pub struct ParkState {
733    /// Whether the task's park token is currently available. If it's available, then the next time
734    /// the task calls `park`, the token will be atomically consumed and the task will continue
735    /// executing. If it's not available, then the task will block until either another task makes
736    /// it available with `unpark`, or a spurious wakeup occurs.
737    token_available: bool,
738
739    /// Whether the task is currently blocked in a call to `park`.
740    /// Invariant: `!(token_available && blocked_in_park)`. If the token is available, then the task
741    /// shouldn't be blocked in a call to `park`---the task should either have been woken up when
742    /// the token became available, or never have blocked in the first place if the token was
743    /// available before the call to `park`.
744    blocked_in_park: bool,
745}
746
747/// A `TaskId` is a unique identifier for a task. `TaskId`s are never reused within a single
748/// execution.
749#[derive(PartialEq, Eq, Hash, Clone, Copy, PartialOrd, Ord)]
750pub struct TaskId(pub(super) usize);
751
752impl Debug for TaskId {
753    // If the `TaskName` label is set, use that when generating the Debug string
754    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
755        if let Some(name) = get_name_for_task(*self) {
756            f.write_str(&format!("{:?}({})", name, self.0))
757        } else {
758            f.debug_tuple("TaskId").field(&self.0).finish()
759        }
760    }
761}
762
763impl std::fmt::Display for TaskId {
764    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
765        std::fmt::Display::fmt(&self.0, f)
766    }
767}
768
769impl From<usize> for TaskId {
770    fn from(id: usize) -> Self {
771        TaskId(id)
772    }
773}
774
775impl From<TaskId> for usize {
776    fn from(tid: TaskId) -> usize {
777        tid.0
778    }
779}
780
781/// A `TaskSet` is a set of `TaskId`s but implemented efficiently as a BitVec
782#[derive(PartialEq, Eq)]
783pub struct TaskSet {
784    tasks: BitVec,
785}
786
787impl TaskSet {
788    pub const fn new() -> Self {
789        Self { tasks: BitVec::EMPTY }
790    }
791
792    pub fn contains(&self, tid: TaskId) -> bool {
793        // Return false if tid is outside the TaskSet
794        (tid.0 < self.tasks.len()) && self.tasks[tid.0]
795    }
796
797    pub fn is_empty(&self) -> bool {
798        self.tasks.iter().all(|b| !*b)
799    }
800
801    /// Add a task to the set. If the set did not have this value present, `true` is returned. If
802    /// the set did have this value present, `false` is returned.
803    pub fn insert(&mut self, tid: TaskId) -> bool {
804        if tid.0 >= self.tasks.len() {
805            self.tasks.resize(DEFAULT_INLINE_TASKS.max(1 + tid.0), false);
806        }
807        !std::mem::replace(&mut *self.tasks.get_mut(tid.0).unwrap(), true)
808    }
809
810    /// Removes a value from the set. Returns whether the value was present in the set.
811    pub fn remove(&mut self, tid: TaskId) -> bool {
812        if tid.0 >= self.tasks.len() {
813            return false;
814        }
815        std::mem::replace(&mut self.tasks.get_mut(tid.0).unwrap(), false)
816    }
817
818    pub fn iter(&self) -> impl Iterator<Item = TaskId> + '_ {
819        self.tasks
820            .iter()
821            .enumerate()
822            .filter(|(_, b)| **b)
823            .map(|(i, _)| TaskId(i))
824    }
825}
826
827impl Debug for TaskSet {
828    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
829        write!(f, "TaskSet {{ ")?;
830        for (i, t) in self.iter().enumerate() {
831            if i > 0 {
832                write!(f, ", ")?;
833            }
834            write!(f, "{t:?}")?;
835        }
836        write!(f, " }}")
837    }
838}
839
840impl<T: 'static> From<&'static LocalKey<T>> for StorageKey {
841    fn from(key: &'static LocalKey<T>) -> Self {
842        Self(key as *const _ as usize, 0x1)
843    }
844}