async_backtrace/
tasks.rs

1use crate::Frame;
2use dashmap::DashSet as Set;
3use once_cell::sync::Lazy;
4use rustc_hash::FxHasher;
5use std::{hash::BuildHasherDefault, ops::Deref, ptr::NonNull};
6
7/// A top-level [framed](crate::framed) future.
8#[derive(Hash, Eq, PartialEq)]
9#[repr(transparent)]
10pub struct Task(NonNull<Frame>);
11
12unsafe impl Send for Task {}
13unsafe impl Sync for Task {}
14
15static TASK_SET: Lazy<Set<Task, BuildHasherDefault<FxHasher>>> = Lazy::new(Set::default);
16
17/// Register a given root frame as a task.
18///
19/// **SAFETY:** You vow to remove the given frame prior to it being dropped.
20pub(crate) unsafe fn register(root_frame: &Frame) {
21    let unique = TASK_SET.insert(Task(NonNull::from(root_frame)));
22    debug_assert!(unique);
23}
24
25/// De-register a given root frame as a task.
26pub(crate) fn deregister(root_frame: &Frame) {
27    TASK_SET.remove(&Task(NonNull::from(root_frame)));
28}
29
30/// An iterator over tasks.
31///
32/// **NOTE:** The creation and destruction of some or all tasks will be blocked
33/// for as long as the return value of this function is live.
34pub fn tasks() -> impl Iterator<Item = impl Deref<Target = Task>> {
35    TASK_SET.iter()
36}
37
38impl Task {
39    /// The location of this task.
40    pub fn location(&self) -> crate::Location {
41        // safety: we promise to not inspect the subframes without first locking
42        let frame = unsafe { self.0.as_ref() };
43        frame.location()
44    }
45
46    /// Pretty-prints this task as a tree.
47    ///
48    /// If `block_until_idle` is `true`, this routine will block until the task
49    /// is no longer being polled.  In this case, the caller should not hold any
50    /// locks which might be held by the task, otherwise deadlock may occur.
51    ///
52    /// If `block_until_idle` is `false`, and the task is being polled, the
53    /// output will not include the sub-frames, instead simply note that the
54    /// task is being polled.
55    pub fn pretty_tree(&self, block_until_idle: bool) -> String {
56        use crate::sync::TryLockError;
57
58        // safety: we promise to not inspect the subframes without first locking
59        let frame = unsafe { self.0.as_ref() };
60
61        let current_task: Option<NonNull<Frame>> =
62            Frame::with_active(|maybe_frame| maybe_frame.map(|frame| frame.root().into()));
63
64        let maybe_lock = &frame
65            .mutex()
66            // don't grab a lock if we're *in* the active task (it's already locked, then)
67            .filter(|_| Some(self.0) != current_task)
68            .map(|mutex| {
69                if block_until_idle {
70                    mutex.lock().map_err(TryLockError::from)
71                } else {
72                    mutex.try_lock()
73                }
74            });
75
76        let subframes_locked = match maybe_lock {
77            None | Some(Ok(..)) => true,
78            Some(Err(TryLockError::WouldBlock)) => false,
79            Some(Err(err @ TryLockError::Poisoned(..))) => panic!("{:?}", err),
80        };
81
82        let mut string = String::new();
83
84        unsafe {
85            frame.fmt(&mut string, subframes_locked).unwrap();
86        }
87
88        string
89    }
90}