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
61use std::ptr::NonNull;
62
63/// Low-level interface to the underlying C library.
64pub mod sys {
65    #[allow(non_camel_case_types)]
66    mod c {
67        include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
68    }
69
70    pub use c::custom_labels_label_t as Label;
71    pub use c::custom_labels_labelset_t as Labelset;
72    pub use c::custom_labels_string_t as String;
73
74    impl<'a> From<&'a [u8]> for self::String {
75        fn from(value: &'a [u8]) -> Self {
76            Self {
77                len: value.len(),
78                buf: value.as_ptr(),
79            }
80        }
81    }
82
83    impl Clone for self::String {
84        fn clone(&self) -> Self {
85            unsafe {
86                let buf = libc::malloc(self.len);
87                if buf.is_null() {
88                    panic!("Out of memory");
89                }
90                libc::memcpy(buf, self.buf as *const _, self.len);
91                Self {
92                    len: self.len,
93                    buf: buf as *mut _,
94                }
95            }
96        }
97    }
98
99    impl Drop for self::String {
100        fn drop(&mut self) {
101            unsafe {
102                libc::free(self.buf as *mut _);
103            }
104        }
105    }
106
107    pub use c::custom_labels_delete as delete;
108    pub use c::custom_labels_get as get;
109    pub use c::custom_labels_set as set;
110
111    // these aren't used yet in the higher-level API,
112    // but use them here to prevent "unused" warnings.
113    pub use c::custom_labels_labelset_clone as labelset_clone;
114    pub use c::custom_labels_labelset_current as labelset_current;
115    pub use c::custom_labels_labelset_delete as labelset_delete;
116    pub use c::custom_labels_labelset_free as labelset_free;
117    pub use c::custom_labels_labelset_get as labelset_get;
118    pub use c::custom_labels_labelset_new as labelset_new;
119    pub use c::custom_labels_labelset_replace as labelset_replace;
120    pub use c::custom_labels_labelset_set as labelset_set;
121}
122
123/// Utilities for build scripts
124pub mod build {
125    /// Emit the instructions required for an
126    /// executable to expose custom labels data.
127    pub fn emit_build_instructions() {
128        let dlist_path = format!("{}/dlist", std::env::var("OUT_DIR").unwrap());
129        std::fs::write(&dlist_path, include_str!("../dlist")).unwrap();
130        println!("cargo:rustc-link-arg=-Wl,--dynamic-list={}", dlist_path);
131    }
132}
133
134/// A set of key-value labels that can be installed as the current label set.
135pub struct Labelset {
136    raw: NonNull<sys::Labelset>,
137}
138
139unsafe impl Send for Labelset {}
140
141impl Labelset {
142    /// Create a new label set.
143    pub fn new() -> Self {
144        Self::with_capacity(0)
145    }
146
147    /// Create a new label set with the specified capacity.
148    pub fn with_capacity(capacity: usize) -> Self {
149        let raw = unsafe { sys::labelset_new(capacity) };
150        let raw = NonNull::new(raw).expect("failed to allocate labelset");
151        Self { raw }
152    }
153
154    /// Create a new label set by cloning the current one, if it exists,
155    /// or creating a new one otherwise.
156    pub fn clone_from_current() -> Self {
157        Self::try_clone_from_current().unwrap_or_default()
158    }
159
160    /// Create a new label set by cloning the current one, if it exists,
161    pub fn try_clone_from_current() -> Option<Self> {
162        let raw = unsafe { sys::labelset_current() };
163        if raw.is_null() {
164            None
165        } else {
166            let raw = unsafe { sys::labelset_clone(raw) };
167            let raw = NonNull::new(raw).expect("failed to clone labelset");
168            Some(Self { raw })
169        }
170    }
171
172    /// Run a function with this set of labels applied.
173    pub fn enter<F, Ret>(&mut self, f: F) -> Ret
174    where
175        F: FnOnce() -> Ret,
176    {
177        struct Guard {
178            old: *mut sys::Labelset,
179        }
180
181        impl Drop for Guard {
182            fn drop(&mut self) {
183                unsafe { sys::labelset_replace(self.old) };
184            }
185        }
186
187        let old = unsafe { sys::labelset_replace(self.raw.as_ptr()) };
188        let _guard = Guard { old };
189        f()
190    }
191
192    /// Adds the specified key-value pair to the label set.
193    pub fn set<K, V>(&mut self, key: K, value: V)
194    where
195        K: AsRef<[u8]>,
196        V: AsRef<[u8]>,
197    {
198        unsafe {
199            sys::labelset_set(
200                self.raw.as_ptr(),
201                key.as_ref().into(),
202                value.as_ref().into(),
203            )
204        };
205    }
206}
207
208impl Default for Labelset {
209    fn default() -> Self {
210        Self::new()
211    }
212}
213
214impl Drop for Labelset {
215    fn drop(&mut self) {
216        unsafe { sys::labelset_free(self.raw.as_ptr()) }
217    }
218}
219
220impl Clone for Labelset {
221    fn clone(&self) -> Self {
222        let raw = unsafe { sys::labelset_clone(self.raw.as_ptr()) };
223        let raw = NonNull::new(raw).expect("failed to clone labelset");
224        Self { raw }
225    }
226}
227
228impl<K, V> Extend<(K, V)> for Labelset
229where
230    K: AsRef<[u8]>,
231    V: AsRef<[u8]>,
232{
233    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
234        for (k, v) in iter {
235            self.set(k, v);
236        }
237    }
238}
239
240/// Set the label for the specified key to the specified
241/// value while the given function is running.
242///
243/// All labels are thread-local: setting a label on one thread
244/// has no effect on its value on any other thread.
245pub fn with_label<K, V, F, Ret>(k: K, v: V, f: F) -> Ret
246where
247    K: AsRef<[u8]>,
248    V: AsRef<[u8]>,
249    F: FnOnce() -> Ret,
250{
251    unsafe {
252        if sys::labelset_current().is_null() {
253            let l = sys::labelset_new(0);
254            sys::labelset_replace(l);
255        }
256    }
257    struct Guard<'a> {
258        k: &'a [u8],
259        old_v: Option<sys::String>,
260    }
261
262    impl<'a> Drop for Guard<'a> {
263        fn drop(&mut self) {
264            if let Some(old_v) = std::mem::take(&mut self.old_v) {
265                let errno = unsafe { sys::set(self.k.into(), old_v) };
266                if errno != 0 {
267                    panic!("corruption in custom labels library: errno {errno}");
268                }
269            } else {
270                unsafe { sys::delete(self.k.into()) };
271            }
272        }
273    }
274
275    let old_v = unsafe { sys::get(k.as_ref().into()).as_ref() }.map(|lbl| lbl.value.clone());
276    let _g = Guard {
277        k: k.as_ref(),
278        old_v,
279    };
280
281    let errno = unsafe { sys::set(k.as_ref().into(), v.as_ref().into()) };
282    if errno != 0 {
283        panic!("corruption in custom labels library: errno {errno}")
284    }
285
286    f()
287}
288
289/// Set the labels for the specified keys to the specified
290/// values while the given function is running.
291///
292/// `i` is an iterator of key-value pairs.
293///
294/// The effect is the same as repeatedly nesting calls to the singular [`with_label`].
295pub fn with_labels<I, K, V, F, Ret>(i: I, f: F) -> Ret
296where
297    I: IntoIterator<Item = (K, V)>,
298    K: AsRef<[u8]>,
299    V: AsRef<[u8]>,
300    F: FnOnce() -> Ret,
301{
302    let mut i = i.into_iter();
303    if let Some((k, v)) = i.next() {
304        with_label(k, v, || with_labels(i, f))
305    } else {
306        f()
307    }
308}
309
310pub mod asynchronous {
311    use pin_project_lite::pin_project;
312    use std::future::Future;
313    use std::iter;
314    use std::pin::Pin;
315    use std::task::{Context, Poll};
316
317    use crate::Labelset;
318
319    pin_project! {
320        /// A [`Future`] with custom labels attached.
321        ///
322        /// This type is returned by the [`Label`] extension trait. See that
323        /// trait's documentation for details.
324        pub struct Labeled<Fut> {
325            #[pin]
326            inner: Fut,
327            labelset: Labelset,
328        }
329    }
330
331    impl<Fut, Ret> Future for Labeled<Fut>
332    where
333        Fut: Future<Output = Ret>,
334    {
335        type Output = Ret;
336
337        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
338            let p = self.project();
339            p.labelset.enter(|| p.inner.poll(cx))
340        }
341    }
342
343    /// Attaches custom labels to a [`Future`].
344    pub trait Label: Sized {
345        /// Attach the currently active labels to the future.
346        ///
347        /// This can be used to propagate the current labels when spawning
348        /// a new future.
349        fn with_current_labels(self) -> Labeled<Self>;
350
351        /// Attach a single label to the future.
352        ///
353        /// This is equivalent to calling [`with_labels`] with an iterator that
354        /// yields a single key–value pair.
355        fn with_label<K, V>(self, k: K, v: V) -> Labeled<Self>
356        where
357            K: AsRef<[u8]>,
358            V: AsRef<[u8]>;
359
360        /// Attaches the specified labels to the future.
361        ///
362        /// The labels will be installed in the current thread whenever the
363        /// future is polled, and removed when the poll completes.
364        ///
365        /// This is equivalent to calling [`with_labelset`] with a label
366        /// set constructed like so:
367        ///
368        /// ```rust
369        /// # use custom_labels::Labelset;
370        /// let mut labelset = Labelset::clone_from_current();
371        /// labelset.extend(i);
372        /// ```
373        fn with_labels<I, K, V>(self, i: I) -> Labeled<Self>
374        where
375            I: IntoIterator<Item = (K, V)>,
376            K: AsRef<[u8]>,
377            V: AsRef<[u8]>;
378
379        /// Attaches the specified labelset to the future.
380        ///
381        /// The labels in the set will be installed in the current thread
382        /// whenever the future is polled, and removed when the poll completes.
383        fn with_labelset(self, labelset: Labelset) -> Labeled<Self>;
384    }
385
386    impl<Fut: Future> Label for Fut {
387        fn with_current_labels(self) -> Labeled<Self> {
388            self.with_labels(iter::empty::<(&[u8], &[u8])>())
389        }
390
391        fn with_label<K, V>(self, k: K, v: V) -> Labeled<Self>
392        where
393            K: AsRef<[u8]>,
394            V: AsRef<[u8]>,
395        {
396            self.with_labels(iter::once((k, v)))
397        }
398
399        fn with_labels<I, K, V>(self, iter: I) -> Labeled<Self>
400        where
401            I: IntoIterator<Item = (K, V)>,
402            K: AsRef<[u8]>,
403            V: AsRef<[u8]>,
404        {
405            let mut labelset = Labelset::clone_from_current();
406            labelset.extend(iter);
407            Labeled {
408                inner: self,
409                labelset,
410            }
411        }
412
413        fn with_labelset(self, labelset: Labelset) -> Labeled<Self> {
414            Labeled {
415                inner: self,
416                labelset,
417            }
418        }
419    }
420}
custom_labels/lib.rs

custom_labels/
lib.rs
