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}