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}