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.
31pub struct CallSite {
32    /// Unique identifier of this call site
33    id: CallSiteId,
34    /// The name of the call site
35    name: &'static str,
36    /// The name of the Rust module where the call site occurred
37    module_path: &'static str,
38    /// The name of the source code file where the call site occurred
39    file: &'static str,
40    /// The line number in the source code file where the call site occurred
41    line: u32,
42    /// Call sites are registered using an atomic, append only intrusive linked
43    /// list. If more than one call site are registered, this will be set to the
44    /// last registered call site.
45    next: AtomicPtr<CallSite>,
46}
47
48impl CallSite {
49    /// Create a new `CallSite` with the given metadata. This function is
50    /// private to this crate, and is only marked `pub` to be able to call it
51    /// from inside macros.
52    #[doc(hidden)]
53    pub fn new(name: &'static str, module_path: &'static str, file: &'static str, line: u32) -> CallSite {
54        let id = CallSiteId::new(NEXT_CALL_SITE_ID.fetch_add(1, Ordering::SeqCst));
55        let next = AtomicPtr::new(std::ptr::null_mut());
56        CallSite { id, name, module_path, file, line, next }
57    }
58
59    pub(crate) fn id(&self) -> CallSiteId {
60        self.id
61    }
62
63    /// Get the user-provided name for this call site
64    pub fn name(&self) -> &str {
65        self.name
66    }
67
68    /// Get the rust module path to the source code location of this call site
69    pub fn module_path(&self) -> &str {
70        self.module_path
71    }
72
73    /// Get the path to the file containing this call site
74    pub fn file(&self) -> &str {
75        self.file
76    }
77
78    /// Get the line of the source file containing this call site
79    pub fn line(&self) -> u32 {
80        self.line
81    }
82
83    /// Get the full name of this call site, containing both the name and the
84    /// module path.
85    pub fn full_name(&self) -> String {
86        let mut name = self.module_path.to_owned();
87        name += "::";
88
89        if self.name.contains(' ') {
90            name += "{";
91            name += self.name;
92            name += "}";
93        } else {
94            name += self.name;
95        }
96
97        return name;
98    }
99}
100
101/// Registry of CallSite, as the head pointer of an atomic, append-only linked
102/// list.
103struct Registry {
104    head: AtomicPtr<CallSite>,
105}
106
107impl Registry {
108    /// Register a new callsite within the list
109    fn register(&self, callsite: &'static CallSite) {
110        let mut head = self.head.load(Ordering::Acquire);
111
112        loop {
113            callsite.next.store(head, Ordering::Release);
114
115            assert_ne!(
116                callsite as *const _, head,
117                "Attempted to register a `Callsite` that already exists! \
118                This will cause an infinite loop when attempting to read from the \
119                callsite registry."
120            );
121
122            match self.head.compare_exchange(
123                head,
124                callsite as *const _ as *mut _,
125                Ordering::AcqRel,
126                Ordering::Acquire,
127            ) {
128                Ok(_) => {
129                    break;
130                }
131                Err(current) => head = current,
132            }
133        }
134    }
135
136    /// Execute the provided function on all elements of the list
137    fn for_each(&self, mut f: impl FnMut(&'static CallSite)) {
138        let mut head = self.head.load(Ordering::Acquire);
139
140        while let Some(registered) = unsafe { head.as_ref() } {
141            f(registered);
142            head = registered.next.load(Ordering::Acquire);
143        }
144    }
145}
146
147/// Register a call site. This function is a private function of this crate. It
148/// is only marked `pub` to be able to call it from inside macros.
149#[doc(hidden)]
150pub fn register_callsite(callsite: &'static CallSite) {
151    REGISTRY.register(callsite);
152}
153
154/// Execute the given function on all call sites we know about.
155///
156/// # Examples
157/// ```
158/// # use time_graph::traverse_registered_callsite;
159///
160/// traverse_registered_callsite(|callsite| {
161///     println!("got a callsite at {}:{}", callsite.file(), callsite.line());
162/// })
163/// ```
164pub fn traverse_registered_callsite(function: impl FnMut(&'static CallSite)) {
165    REGISTRY.for_each(function);
166}