time_graph/
callsite.rs

1use std::num::NonZeroU64;
2use std::sync::atomic::{Ordering, AtomicU64, AtomicPtr};
3
4use once_cell::sync::Lazy;
5
6/// Store the id to be assigned to the next call site created.
7static NEXT_CALL_SITE_ID: AtomicU64 = AtomicU64::new(1);
8/// Store the global registry of call sites
9static REGISTRY: Lazy<Registry> = Lazy::new(|| {
10    Registry {
11        head: AtomicPtr::new(std::ptr::null_mut()),
12    }
13});
14
15/// Unique identifier of a [`CallSite`], attributed the first time the call site
16/// is entered.
17#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
18pub struct CallSiteId(NonZeroU64);
19
20impl CallSiteId {
21    pub(crate) fn new(value: u64) -> CallSiteId {
22        CallSiteId(std::num::NonZeroU64::new(value).expect("got a zero value for span id"))
23    }
24}
25
26/// A [`CallSite`] identify uniquely a location in the source code, and record
27/// multiple attributes associated with this location.
28///
29/// The only way to create a [`CallSite`] is with the [`macro@callsite`] macro,
30/// which also takes care of registering the call site globally.
31#[derive(Debug)]
32pub struct CallSite {
33    /// Unique identifier of this call site
34    id: CallSiteId,
35    /// Tracking the current call depth for recursive calls
36    pub(crate) depth: &'static std::thread::LocalKey<std::cell::Cell<usize>>,
37    /// The name of the call site
38    name: &'static str,
39    /// The name of the Rust module where the call site occurred
40    module_path: &'static str,
41    /// The name of the source code file where the call site occurred
42    file: &'static str,
43    /// The line number in the source code file where the call site occurred
44    line: u32,
45    /// Call sites are registered using an atomic, append only intrusive linked
46    /// list. If more than one call site are registered, this will be set to the
47    /// last registered call site.
48    next: AtomicPtr<CallSite>,
49}
50
51impl CallSite {
52    /// Create a new `CallSite` with the given metadata. This function is
53    /// private to this crate, and is only marked `pub` to be able to call it
54    /// from inside macros.
55    #[doc(hidden)]
56    pub fn new(
57        depth: &'static std::thread::LocalKey<std::cell::Cell<usize>>,
58        name: &'static str,
59        module_path: &'static str,
60        file: &'static str,
61        line: u32
62    ) -> CallSite {
63        let id = CallSiteId::new(NEXT_CALL_SITE_ID.fetch_add(1, Ordering::SeqCst));
64        let next = AtomicPtr::new(std::ptr::null_mut());
65        CallSite { id, depth, name, module_path, file, line, next }
66    }
67
68    pub(crate) fn id(&self) -> CallSiteId {
69        self.id
70    }
71
72    /// Get the user-provided name for this call site
73    pub fn name(&self) -> &str {
74        self.name
75    }
76
77    /// Get the rust module path to the source code location of this call site
78    pub fn module_path(&self) -> &str {
79        self.module_path
80    }
81
82    /// Get the path to the file containing this call site
83    pub fn file(&self) -> &str {
84        self.file
85    }
86
87    /// Get the line of the source file containing this call site
88    pub fn line(&self) -> u32 {
89        self.line
90    }
91
92    /// Get the full name of this call site, containing both the name and the
93    /// module path.
94    pub fn full_name(&self) -> String {
95        let mut name = self.module_path.to_owned();
96        name += "::";
97
98        if self.name.contains(' ') {
99            name += "{";
100            name += self.name;
101            name += "}";
102        } else {
103            name += self.name;
104        }
105
106        return name;
107    }
108}
109
110/// Registry of CallSite, as the head pointer of an atomic, append-only linked
111/// list.
112struct Registry {
113    head: AtomicPtr<CallSite>,
114}
115
116impl Registry {
117    /// Register a new callsite within the list
118    fn register(&self, callsite: &'static CallSite) {
119        let mut head = self.head.load(Ordering::Acquire);
120
121        loop {
122            callsite.next.store(head, Ordering::Release);
123
124            assert_ne!(
125                callsite as *const _, head,
126                "Attempted to register a `Callsite` that already exists! \
127                This will cause an infinite loop when attempting to read from the \
128                callsite registry."
129            );
130
131            match self.head.compare_exchange(
132                head,
133                callsite as *const _ as *mut _,
134                Ordering::AcqRel,
135                Ordering::Acquire,
136            ) {
137                Ok(_) => {
138                    break;
139                }
140                Err(current) => head = current,
141            }
142        }
143    }
144
145    /// Execute the provided function on all elements of the list
146    fn for_each(&self, mut f: impl FnMut(&'static CallSite)) {
147        let mut head = self.head.load(Ordering::Acquire);
148
149        while let Some(registered) = unsafe { head.as_ref() } {
150            f(registered);
151            head = registered.next.load(Ordering::Acquire);
152        }
153    }
154}
155
156/// Register a call site. This function is a private function of this crate. It
157/// is only marked `pub` to be able to call it from inside macros.
158#[doc(hidden)]
159pub fn register_callsite(callsite: &'static CallSite) {
160    REGISTRY.register(callsite);
161}
162
163/// Execute the given function on all call sites we know about.
164///
165/// # Examples
166/// ```
167/// # use time_graph::traverse_registered_callsite;
168///
169/// traverse_registered_callsite(|callsite| {
170///     println!("got a callsite at {}:{}", callsite.file(), callsite.line());
171/// })
172/// ```
173pub fn traverse_registered_callsite(function: impl FnMut(&'static CallSite)) {
174    REGISTRY.for_each(function);
175}
time_graph/callsite.rs

time_graph/
callsite.rs
