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` represents a user-level unit of concurrency. Each task has an `id` that is unique within
259/// the execution, and a `state` reflecting whether the task is runnable (enabled) or not.
260#[derive(Debug)]
261pub struct Task {
262    pub(super) id: TaskId,
263    pub(super) parent_task_id: Option<TaskId>,
264    pub(super) state: TaskState,
265    pub(super) detached: bool,
266    park_state: ParkState,
267
268    pub(super) continuation: Rc<RefCell<PooledContinuation>>,
269    pub(super) yielder: *const Yielder<ContinuationInput, ContinuationOutput>,
270
271    pub clock: VectorClock,
272
273    waiter: Option<TaskId>,
274
275    waker: Waker,
276    // Remember whether the waker was invoked while we were running
277    woken: bool,
278
279    name: Option<String>,
280
281    local_storage: StorageMap,
282
283    // The `Span` which looks like this: step{task=task_id}, or, if step count recording is enabled, like this:
284    // step{task=task_id i=step_count}. Becomes the parent of the spans created by the `Task`.
285    pub step_span: Span,
286
287    // The current `Span` "stack" of the `Task`.
288    // `Span`s are stored such that the `Task`s current `Span` is at `span_stack[0]`, that `Span`s parent (if it exists)
289    // is at `span_stack[1]`, and so on, until `span_stack[span_stack.len()-1]`, which is the "outermost" (left-most when printed)
290    // `Span`. This means that `span_stack[span_stack.len()-1]` will usually be the `Span` saying `execution{i=X}`.
291    // We `pop` it empty when resuming a `Task`, and `push` + `exit` `tracing::Span::current()`
292    // until there is no entered `Span` when we switch out of the `Task`.
293    // There are two things to note:
294    // 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.
295    // 2: We have to store the stack of `Span`s in order to return to the correct `Span` once the `Entered<'_>` from an
296    //    `instrument`ed future is dropped.
297    pub(super) span_stack: Vec<Span>,
298
299    // Arbitrarily settable tag which is inherited from the parent.
300    #[allow(deprecated)]
301    tag: Option<Arc<dyn Tag>>,
302
303    /// If [`crate::CAPTURE_BACKTRACE`] is set then this will be populated on task block.
304    /// If the test then fails, then each task's backtrace will be printed.
305    pub backtrace: Option<Backtrace>,
306
307    /// The signature of a Task; this is an identifier that is *not* guaranteed to be unique but should be *mostly*
308    /// stable across iterations in a single Shuttle test. Tasks with the same signature are very likely to exhibit
309    /// similar behavior
310    pub signature: TaskSignature,
311}
312
313#[allow(deprecated)]
314impl Task {
315    /// Create a task from a continuation
316    #[allow(clippy::too_many_arguments)]
317    fn new(
318        f: Box<dyn FnOnce() + 'static>,
319        stack_size: usize,
320        id: TaskId,
321        name: Option<String>,
322        clock: VectorClock,
323        parent_span_id: Option<tracing::span::Id>,
324        schedule_len: usize,
325        tag: Option<Arc<dyn Tag>>,
326        parent_task_id: Option<TaskId>,
327        signature: TaskSignature,
328    ) -> Self {
329        #[cfg(all(any(test, feature = "vector-clocks"), not(feature = "bench-no-vector-clocks")))]
330        assert!(id.0 < clock.time.len());
331        let mut continuation = ContinuationPool::acquire(stack_size);
332        continuation.initialize(f);
333        let yielder = continuation.yielder;
334        let waker = make_waker(id);
335        let continuation = Rc::new(RefCell::new(continuation));
336
337        let step_span =
338            error_span!(parent: parent_span_id.clone(), "step", task = format!("{:?}", id), i = field::Empty);
339        // Note that this is slightly lazy — we are starting storing at the step_span, but could have gotten the
340        // full `Span` stack and stored that. It should be fine, but if any issues arise, then full storing should
341        // be tried.
342        let span_stack = vec![step_span.clone()];
343
344        let mut task = Self {
345            id,
346            parent_task_id,
347            state: TaskState::Runnable,
348            continuation,
349            yielder,
350            clock,
351            waiter: None,
352            waker,
353            woken: false,
354            detached: false,
355            park_state: ParkState::default(),
356            name,
357            step_span,
358            span_stack,
359            local_storage: StorageMap::new(),
360            tag: None,
361            backtrace: None,
362            signature,
363        };
364
365        if let Some(tag) = tag {
366            task.set_tag(tag);
367        }
368
369        // Note: the tests for the task signature in [`crate::tests::basic::task`] depend on tracing the task signature and creation point here
370        error_span!(parent: parent_span_id, "new_task", parent = ?parent_task_id, i = schedule_len).in_scope(
371            || event!(Level::DEBUG, task_id = ?task.id, signature = task.signature.signature_hash(), static_create_location = task.signature.static_create_location_hash(), "created task"),
372        );
373
374        task
375    }
376
377    #[allow(clippy::too_many_arguments)]
378    pub fn from_closure(
379        f: Box<dyn FnOnce() + 'static>,
380        stack_size: usize,
381        id: TaskId,
382        name: Option<String>,
383        clock: VectorClock,
384        parent_span_id: Option<tracing::span::Id>,
385        schedule_len: usize,
386        tag: Option<Arc<dyn Tag>>,
387        parent_task_id: Option<TaskId>,
388        signature: TaskSignature,
389    ) -> Self {
390        Self::new(
391            f,
392            stack_size,
393            id,
394            name,
395            clock,
396            parent_span_id,
397            schedule_len,
398            tag,
399            parent_task_id,
400            signature,
401        )
402    }
403
404    #[allow(clippy::too_many_arguments)]
405    pub fn from_future<F>(
406        future: F,
407        stack_size: usize,
408        id: TaskId,
409        name: Option<String>,
410        clock: VectorClock,
411        parent_span_id: Option<tracing::span::Id>,
412        schedule_len: usize,
413        tag: Option<Arc<dyn Tag>>,
414        parent_task_id: Option<TaskId>,
415        signature: TaskSignature,
416    ) -> Self
417    where
418        F: Future<Output = ()> + 'static,
419    {
420        let mut future = Box::pin(future);
421
422        Self::new(
423            Box::new(move || {
424                let waker = ExecutionState::with(|state| state.current_mut().waker());
425                let cx = &mut Context::from_waker(&waker);
426                while future.as_mut().poll(cx).is_pending() {
427                    ExecutionState::with(|state| state.current_mut().sleep_unless_woken());
428                    thread::switch();
429                }
430            }),
431            stack_size,
432            id,
433            name,
434            clock,
435            parent_span_id,
436            schedule_len,
437            tag,
438            parent_task_id,
439            signature,
440        )
441    }
442
443    /// Returns the identifier of this task.
444    pub fn id(&self) -> TaskId {
445        self.id
446    }
447
448    /// Returns the identifier of the task that spawned this task.
449    pub fn parent_task_id(&self) -> Option<TaskId> {
450        self.parent_task_id
451    }
452
453    pub fn runnable(&self) -> bool {
454        self.state == TaskState::Runnable
455    }
456
457    pub fn blocked(&self) -> bool {
458        matches!(self.state, TaskState::Blocked { .. })
459    }
460
461    pub fn can_spuriously_wakeup(&self) -> bool {
462        match self.state {
463            TaskState::Blocked { allow_spurious_wakeups } => allow_spurious_wakeups,
464            _ => false,
465        }
466    }
467
468    pub fn sleeping(&self) -> bool {
469        self.state == TaskState::Sleeping
470    }
471
472    pub fn finished(&self) -> bool {
473        self.state == TaskState::Finished
474    }
475
476    pub fn is_detached(&self) -> bool {
477        self.detached
478    }
479
480    pub fn detach(&mut self) {
481        self.detached = true;
482    }
483
484    /// Wake this task so the Wrapper future can observe the abort flag on its next poll.
485    pub fn abort(&mut self) {
486        if self.finished() {
487            return;
488        }
489        self.wake();
490    }
491
492    pub fn waker(&self) -> Waker {
493        self.waker.clone()
494    }
495
496    /// Block the current thread. If `allow_spurious_wakeups` is true, then the scheduler is
497    /// permitted to spuriously wake up the thread (though it will still not count as a live thread
498    /// for deadlock detection purposes for as long as it remains blocked).
499    pub fn block(&mut self, allow_spurious_wakeups: bool) {
500        self.backtrace = if backtrace_enabled() {
501            Some(Backtrace::force_capture())
502        } else {
503            None
504        };
505
506        assert!(self.state != TaskState::Finished);
507        self.state = TaskState::Blocked { allow_spurious_wakeups };
508    }
509
510    pub fn sleep(&mut self) {
511        self.backtrace = if backtrace_enabled() {
512            Some(Backtrace::force_capture())
513        } else {
514            None
515        };
516
517        assert!(self.state != TaskState::Finished);
518        self.state = TaskState::Sleeping;
519    }
520
521    pub fn unblock(&mut self) {
522        // Note we don't assert the task is blocked here. For example, a task invoking its own waker
523        // will not be blocked when this is called.
524        assert!(self.state != TaskState::Finished);
525        self.state = TaskState::Runnable;
526
527        // When a task gets unblocked, it's definitely no longer blocked in a call to `park`. This
528        // is necessary to do here because a parked task could be spuriously woken up outside of the
529        // `unpark` path. If it later becomes blocked by something else, we don't want a later
530        // `unpark` to be able to unblock the task.
531        self.park_state.blocked_in_park = false;
532    }
533
534    pub fn finish(&mut self) {
535        assert!(self.state != TaskState::Finished);
536        self.state = TaskState::Finished;
537    }
538
539    /// Potentially put this task to sleep after it was polled by the executor, unless someone has
540    /// called its waker first.
541    ///
542    /// A synchronous Task should never call this, because we want threads to be enabled-by-default
543    /// to avoid bugs where Shuttle incorrectly omits a potential execution.
544    pub fn sleep_unless_woken(&mut self) {
545        let was_woken = std::mem::replace(&mut self.woken, false);
546        if !was_woken {
547            self.sleep();
548        }
549    }
550
551    /// Remember that our waker has been called, and so we should not block the next time the
552    /// executor tries to put us to sleep.
553    pub(super) fn wake(&mut self) {
554        self.woken = true;
555        if self.state == TaskState::Sleeping {
556            self.unblock();
557        }
558    }
559
560    /// Register a waiter for this thread to terminate. Returns a boolean indicating whether the
561    /// waiter should block or not. If false, this task has already finished, and so the waiter need
562    /// not block.
563    pub fn set_waiter(&mut self, waiter: TaskId) -> bool {
564        assert!(
565            self.waiter.is_none() || self.waiter == Some(waiter),
566            "Task cannot have more than one waiter"
567        );
568        if self.finished() {
569            false
570        } else {
571            self.waiter = Some(waiter);
572            true
573        }
574    }
575
576    pub fn take_waiter(&mut self) -> Option<TaskId> {
577        self.waiter.take()
578    }
579
580    pub fn name(&self) -> Option<String> {
581        self.name.clone()
582    }
583
584    /// Retrieve a reference to the given thread-local storage slot.
585    ///
586    /// Returns Some(Err(_)) if the slot has already been destructed. Returns None if the slot has
587    /// not yet been initialized.
588    pub fn local<T: 'static>(&self, key: &'static LocalKey<T>) -> Option<Result<&T, AlreadyDestructedError>> {
589        self.local_storage.get(key.into())
590    }
591
592    /// Initialize the given thread-local storage slot with a new value.
593    ///
594    /// Panics if the slot has already been initialized.
595    pub fn init_local<T: 'static>(&mut self, key: &'static LocalKey<T>, value: T) {
596        self.local_storage.init(key.into(), value)
597    }
598
599    /// Return ownership of the next still-initialized thread-local storage slot, to be used when
600    /// running thread-local storage destructors.
601    ///
602    /// TLS destructors are a little tricky:
603    /// 1. Their code can perform synchronization operations (and so require Shuttle to call back
604    ///    into ExecutionState), so we can't drop them from within an ExecutionState borrow. Instead
605    ///    we move the contents of a slot to the caller to be dropped outside the borrow.
606    /// 2. It's valid for destructors to read other TLS slots, although destructor order is
607    ///    undefined. This also means it's valid for a destructor to *initialize* another TLS slot.
608    ///    To make this work, we run the destructors incrementally, so one destructor can initialize
609    ///    another slot that just gets added via `init_local` like normal, and then will be
610    ///    available to be popped on a future call to `pop_local`. To prevent an infinite loop, we
611    ///    forbid *reinitializing* a TLS slot whose destructor has already run, or is currently
612    ///    being run.
613    pub fn pop_local(&mut self) -> Option<Box<dyn Any>> {
614        self.local_storage.pop()
615    }
616
617    /// Park the task if its park token is unavailable. If the task blocks, then it will be woken up
618    /// when the token becomes available or spuriously without consuming the token (see the
619    /// documentation for [`std::thread::park`], which says that "it may also return spuriously,
620    /// without consuming the token"). Returns true if the execution should switch to a different
621    /// task (e.g., if the token was unavailable).
622    pub fn park(&mut self) -> bool {
623        assert!(
624            !self.park_state.blocked_in_park,
625            "task cannot park while already parked"
626        );
627        assert!(!self.blocked(), "task cannot park while blocked by something else");
628
629        if self.park_state.token_available {
630            self.park_state.token_available = false;
631            false
632        } else {
633            self.park_state.blocked_in_park = true;
634            self.block(true);
635            true
636        }
637    }
638
639    /// Make the task's park token available, and unblock the task if it was parked.
640    pub fn unpark(&mut self) {
641        if self.park_state.blocked_in_park {
642            assert!(
643                self.blocked() && self.can_spuriously_wakeup(),
644                "parked tasks should be blocked"
645            );
646            assert!(
647                !self.park_state.token_available,
648                "token shouldn't be available for parked task"
649            );
650
651            self.unblock();
652        } else {
653            // If the thread isn't currently blocked in `park`, then make the token available. If
654            // the token already is available, then this does nothing.
655            self.park_state.token_available = true;
656        }
657    }
658
659    pub fn get_tag(&self) -> Option<Arc<dyn Tag>> {
660        self.tag.clone()
661    }
662
663    /// Sets the `tag` field of the current task.
664    /// Returns the `tag` which was there previously.
665    pub fn set_tag(&mut self, tag: Arc<dyn Tag>) -> Option<Arc<dyn Tag>> {
666        TASK_ID_TO_TAGS.with(|cell| cell.borrow_mut().insert(self.id(), tag.clone()));
667        self.tag.replace(tag)
668    }
669
670    pub fn format_for_deadlock(&self) -> String {
671        use crate::backtrace_enabled;
672        format!(
673            "{} (task {:?}{}{}){}",
674            self.name().unwrap_or_else(|| "<unknown>".to_string()),
675            self.id(),
676            if self.detached { ", detached" } else { "" },
677            if self.sleeping() { ", pending future" } else { "" },
678            if backtrace_enabled() {
679                format!("\nBacktrace:\n{:#?}\n", self.backtrace)
680            } else {
681                "".into()
682            }
683        )
684    }
685}
686
687#[derive(PartialEq, Eq, Clone, Copy, Debug)]
688pub enum TaskState {
689    /// Available to be scheduled
690    Runnable,
691    /// Blocked in a synchronization operation
692    Blocked { allow_spurious_wakeups: bool },
693    /// A `Future` that returned `Pending` is waiting to be woken up
694    Sleeping,
695    /// Task has finished
696    Finished,
697}
698
699#[derive(PartialEq, Eq, Clone, Copy, Debug, Default)]
700pub struct ParkState {
701    /// Whether the task's park token is currently available. If it's available, then the next time
702    /// the task calls `park`, the token will be atomically consumed and the task will continue
703    /// executing. If it's not available, then the task will block until either another task makes
704    /// it available with `unpark`, or a spurious wakeup occurs.
705    token_available: bool,
706
707    /// Whether the task is currently blocked in a call to `park`.
708    /// Invariant: `!(token_available && blocked_in_park)`. If the token is available, then the task
709    /// shouldn't be blocked in a call to `park`---the task should either have been woken up when
710    /// the token became available, or never have blocked in the first place if the token was
711    /// available before the call to `park`.
712    blocked_in_park: bool,
713}
714
715/// A `TaskId` is a unique identifier for a task. `TaskId`s are never reused within a single
716/// execution.
717#[derive(PartialEq, Eq, Hash, Clone, Copy, PartialOrd, Ord)]
718pub struct TaskId(pub(super) usize);
719
720impl Debug for TaskId {
721    // If the `TaskName` label is set, use that when generating the Debug string
722    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
723        if let Some(name) = get_name_for_task(*self) {
724            f.write_str(&format!("{:?}({})", name, self.0))
725        } else {
726            f.debug_tuple("TaskId").field(&self.0).finish()
727        }
728    }
729}
730
731impl std::fmt::Display for TaskId {
732    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
733        std::fmt::Display::fmt(&self.0, f)
734    }
735}
736
737impl From<usize> for TaskId {
738    fn from(id: usize) -> Self {
739        TaskId(id)
740    }
741}
742
743impl From<TaskId> for usize {
744    fn from(tid: TaskId) -> usize {
745        tid.0
746    }
747}
748
749/// A `TaskSet` is a set of `TaskId`s but implemented efficiently as a BitVec
750#[derive(PartialEq, Eq)]
751pub struct TaskSet {
752    tasks: BitVec,
753}
754
755impl TaskSet {
756    pub const fn new() -> Self {
757        Self { tasks: BitVec::EMPTY }
758    }
759
760    pub fn contains(&self, tid: TaskId) -> bool {
761        // Return false if tid is outside the TaskSet
762        (tid.0 < self.tasks.len()) && self.tasks[tid.0]
763    }
764
765    pub fn is_empty(&self) -> bool {
766        self.tasks.iter().all(|b| !*b)
767    }
768
769    /// Add a task to the set. If the set did not have this value present, `true` is returned. If
770    /// the set did have this value present, `false` is returned.
771    pub fn insert(&mut self, tid: TaskId) -> bool {
772        if tid.0 >= self.tasks.len() {
773            self.tasks.resize(DEFAULT_INLINE_TASKS.max(1 + tid.0), false);
774        }
775        !std::mem::replace(&mut *self.tasks.get_mut(tid.0).unwrap(), true)
776    }
777
778    /// Removes a value from the set. Returns whether the value was present in the set.
779    pub fn remove(&mut self, tid: TaskId) -> bool {
780        if tid.0 >= self.tasks.len() {
781            return false;
782        }
783        std::mem::replace(&mut self.tasks.get_mut(tid.0).unwrap(), false)
784    }
785
786    pub fn iter(&self) -> impl Iterator<Item = TaskId> + '_ {
787        self.tasks
788            .iter()
789            .enumerate()
790            .filter(|(_, b)| **b)
791            .map(|(i, _)| TaskId(i))
792    }
793}
794
795impl Debug for TaskSet {
796    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
797        write!(f, "TaskSet {{ ")?;
798        for (i, t) in self.iter().enumerate() {
799            if i > 0 {
800                write!(f, ", ")?;
801            }
802            write!(f, "{t:?}")?;
803        }
804        write!(f, " }}")
805    }
806}
807
808impl<T: 'static> From<&'static LocalKey<T>> for StorageKey {
809    fn from(key: &'static LocalKey<T>) -> Self {
810        Self(key as *const _ as usize, 0x1)
811    }
812}