custom_labels/lib.rs
1//! # Custom labels for profilers.
2//!
3//! ## Overview
4//!
5//! This library provides Rust bindings to [v1 of the Custom Labels ABI](../custom-labels-v1.md).
6//!
7//! It allows time ranges within a thread's execution to be annotated with labels (key/value pairs) in such a way
8//! that the labels are visible to a CPU profiler that may
9//! interrupt the running process at any time. The profiler can then report the labels
10//! that were active at any given time in the profiles it produces.
11//!
12//! The main interface is [`with_label`], which sets a label to a provided value,
13//! executes a function, and then resets the label to its old value (or removes it if there was no old value).
14//!
15//! For example, imagine a program that performs database queries on behalf of a user. It might have a function
16//! like the following:
17//!
18//! ```rust
19//! # struct DbResultSet;
20//! #
21//! # fn do_query(sql: &str) -> DbResultSet {
22//! # DbResultSet
23//! # }
24//! #
25//! fn query_for_user(username: &str, sql: &str) -> DbResultSet {
26//! custom_labels::with_label("username", username, || do_query(sql))
27//! }
28//! ```
29//!
30//! If two users named "Marta" and "Petros" repeatedly query the database, a profiler might produce a
31//! CPU profile like the following:
32//! ```text
33//! * all (13.4s)
34//! |
35//! +---- username: Marta (4.9s)
36//! | |
37//! | +---- query_for_user (4.9s)
38//! | |
39//! | + ---- custom_labels::with_label (4.9s)
40//! | |
41//! | + ---- do_query (4.9s)
42//! |
43//! +---- username: Petros (8.5s)
44//! |
45//! +---- query_for_user (8.5s)
46//! |
47//! + ---- custom_labels::with_label (8.5s)
48//! |
49//! + ---- do_query (8.5s)
50//! ```
51//!
52//! ## Profiler Support
53//!
54//! The following profilers can make use of the labels set with this library:
55//! * [Parca](parca.dev) (when using parca-agent v0.33.0 and later)
56//! * [Polar Signals Cloud](https://www.polarsignals.com) (when using parca-agent v0.33.0 and later).
57//!
58//! If you work on another profiler that also supports this format, [send us a PR](https://github.com/polarsignals/custom-labels)
59//! to update this list!
60
61/// Low-level interface to the underlying C library.
62pub mod sys {
63 #[allow(non_camel_case_types)]
64 mod c {
65 include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
66 }
67
68 pub use c::custom_labels_label_t as Label;
69 pub use c::custom_labels_string_t as String;
70
71 impl<'a> From<&'a [u8]> for self::String {
72 fn from(value: &'a [u8]) -> Self {
73 Self {
74 len: value.len(),
75 buf: value.as_ptr(),
76 }
77 }
78 }
79
80 impl Clone for self::String {
81 fn clone(&self) -> Self {
82 unsafe {
83 let buf = libc::malloc(self.len);
84 if buf.is_null() {
85 panic!("Out of memory");
86 }
87 libc::memcpy(buf, self.buf as *const _, self.len);
88 Self {
89 len: self.len,
90 buf: buf as *mut _,
91 }
92 }
93 }
94 }
95
96 impl Drop for self::String {
97 fn drop(&mut self) {
98 unsafe {
99 libc::free(self.buf as *mut _);
100 }
101 }
102 }
103
104 pub use c::custom_labels_delete as delete;
105 pub use c::custom_labels_get as get;
106 pub use c::custom_labels_set as set;
107
108 // these aren't used yet in the higher-level API,
109 // but use them here to prevent "unused" warnings.
110 pub use c::custom_labels_labelset_clone as labelset_clone;
111 pub use c::custom_labels_labelset_delete as labelset_delete;
112 pub use c::custom_labels_labelset_free as labelset_free;
113 pub use c::custom_labels_labelset_get as labelset_get;
114 pub use c::custom_labels_labelset_new as labelset_new;
115 pub use c::custom_labels_labelset_replace as labelset_replace;
116 pub use c::custom_labels_labelset_set as labelset_set;
117 pub use c::custom_labels_labelset_current as labelset_current;
118}
119
120/// Utilities for build scripts
121pub mod build {
122 /// Emit the instructions required for an
123 /// executable to expose custom labels data.
124 pub fn emit_build_instructions() {
125 let dlist_path = format!("{}/dlist", std::env::var("OUT_DIR").unwrap());
126 std::fs::write(&dlist_path, include_str!("../dlist")).unwrap();
127 println!("cargo:rustc-link-arg=-Wl,--dynamic-list={}", dlist_path);
128 }
129}
130
131/// Set the label for the specified key to the specified
132/// value while the given function is running.
133///
134/// All labels are thread-local: setting a label on one thread
135/// has no effect on its value on any other thread.
136pub fn with_label<K, V, F, Ret>(k: K, v: V, f: F) -> Ret
137where
138 K: AsRef<[u8]>,
139 V: AsRef<[u8]>,
140 F: FnOnce() -> Ret,
141{
142 unsafe {
143 if sys::labelset_current().is_null() {
144 let l = sys::labelset_new(0);
145 sys::labelset_replace(l);
146 }
147 }
148 struct Guard<'a> {
149 k: &'a [u8],
150 old_v: Option<sys::String>,
151 }
152
153 impl<'a> Drop for Guard<'a> {
154 fn drop(&mut self) {
155 if let Some(old_v) = std::mem::take(&mut self.old_v) {
156 let errno = unsafe { sys::set(self.k.into(), old_v) };
157 if errno != 0 {
158 panic!("corruption in custom labels library: errno {errno}");
159 }
160 } else {
161 unsafe { sys::delete(self.k.into()) };
162 }
163 }
164 }
165
166 let old_v = unsafe { sys::get(k.as_ref().into()).as_ref() }.map(|lbl| lbl.value.clone());
167 let _g = Guard {
168 k: k.as_ref(),
169 old_v,
170 };
171
172 let errno = unsafe { sys::set(k.as_ref().into(), v.as_ref().into()) };
173 if errno != 0 {
174 panic!("corruption in custom labels library: errno {errno}")
175 }
176
177 f()
178}
179
180/// Set the labels for the specified keys to the specified
181/// values while the given function is running.
182///
183/// `i` is an iterator of key-value pairs.
184///
185/// The effect is the same as repeatedly nesting calls to the singular [`with_label`].
186pub fn with_labels<I, K, V, F, Ret>(i: I, f: F) -> Ret
187where
188 I: IntoIterator<Item = (K, V)>,
189 K: AsRef<[u8]>,
190 V: AsRef<[u8]>,
191 F: FnOnce() -> Ret,
192{
193 let mut i = i.into_iter();
194 if let Some((k, v)) = i.next() {
195 with_label(k, v, || with_labels(i, f))
196 } else {
197 f()
198 }
199}