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;
62use std::{fmt, slice};
63
64/// Low-level interface to the underlying C library.
65pub mod sys {
66    #[allow(non_camel_case_types)]
67    mod c {
68        include!(concat!(env!("OUT_DIR"), "/bindings.rs"));
69    }
70
71    pub use c::custom_labels_label_t as Label;
72    pub use c::custom_labels_labelset_t as Labelset;
73    pub use c::custom_labels_string_t as String;
74
75    impl<'a> From<&'a [u8]> for self::String {
76        fn from(value: &'a [u8]) -> Self {
77            Self {
78                len: value.len(),
79                buf: value.as_ptr(),
80            }
81        }
82    }
83
84    impl self::String {
85        pub fn to_owned(&self) -> OwnedString {
86            unsafe {
87                let buf = libc::malloc(self.len);
88                if buf.is_null() {
89                    panic!("Out of memory");
90                }
91                libc::memcpy(buf, self.buf as *const _, self.len);
92                OwnedString(Self {
93                    len: self.len,
94                    buf: buf as *mut _,
95                })
96            }
97        }
98    }
99
100    pub struct OwnedString(self::String);
101
102    impl OwnedString {
103        /// Creates a new empty owned string.
104        pub fn new() -> Self {
105            OwnedString(self::String {
106                len: 0,
107                buf: std::ptr::null(),
108            })
109        }
110    }
111
112    impl std::ops::Deref for OwnedString {
113        type Target = self::String;
114
115        fn deref(&self) -> &Self::Target {
116            &self.0
117        }
118    }
119
120    impl std::ops::DerefMut for OwnedString {
121        fn deref_mut(&mut self) -> &mut Self::Target {
122            &mut self.0
123        }
124    }
125
126    impl Drop for OwnedString {
127        fn drop(&mut self) {
128            unsafe {
129                libc::free(self.0.buf as *mut _);
130            }
131        }
132    }
133
134    pub use c::custom_labels_delete as delete;
135    pub use c::custom_labels_get as get;
136    pub use c::custom_labels_set as set;
137
138    pub use c::custom_labels_labelset_clone as labelset_clone;
139    pub use c::custom_labels_labelset_current as labelset_current;
140    pub use c::custom_labels_labelset_debug_string as labelset_debug_string;
141    pub use c::custom_labels_labelset_delete as labelset_delete;
142    pub use c::custom_labels_labelset_free as labelset_free;
143    pub use c::custom_labels_labelset_get as labelset_get;
144    pub use c::custom_labels_labelset_new as labelset_new;
145    pub use c::custom_labels_labelset_replace as labelset_replace;
146    pub use c::custom_labels_labelset_set as labelset_set;
147}
148
149/// Utilities for build scripts
150pub mod build {
151    /// Emit the instructions required for an
152    /// executable to expose custom labels data.
153    pub fn emit_build_instructions() {
154        let dlist_path = format!("{}/dlist", std::env::var("OUT_DIR").unwrap());
155        std::fs::write(&dlist_path, include_str!("../dlist")).unwrap();
156        println!("cargo:rustc-link-arg=-Wl,--dynamic-list={}", dlist_path);
157    }
158}
159
160/// A set of key-value labels that can be installed as the current label set.
161pub struct Labelset {
162    raw: NonNull<sys::Labelset>,
163}
164
165unsafe impl Send for Labelset {}
166
167impl Labelset {
168    /// Create a new label set.
169    pub fn new() -> Self {
170        Self::with_capacity(0)
171    }
172
173    /// Create a new label set with the specified capacity.
174    pub fn with_capacity(capacity: usize) -> Self {
175        let raw = unsafe { sys::labelset_new(capacity) };
176        let raw = NonNull::new(raw).expect("failed to allocate labelset");
177        Self { raw }
178    }
179
180    /// Create a new label set by cloning the current one, if it exists,
181    /// or creating a new one otherwise.
182    pub fn clone_from_current() -> Self {
183        Self::try_clone_from_current().unwrap_or_default()
184    }
185
186    /// Create a new label set by cloning the current one, if it exists,
187    pub fn try_clone_from_current() -> Option<Self> {
188        let raw = unsafe { sys::labelset_current() };
189        if raw.is_null() {
190            None
191        } else {
192            let raw = unsafe { sys::labelset_clone(raw) };
193            let raw = NonNull::new(raw).expect("failed to clone labelset");
194            Some(Self { raw })
195        }
196    }
197
198    /// Run a function with this set of labels applied.
199    pub fn enter<F, Ret>(&mut self, f: F) -> Ret
200    where
201        F: FnOnce() -> Ret,
202    {
203        struct Guard {
204            old: *mut sys::Labelset,
205        }
206
207        impl Drop for Guard {
208            fn drop(&mut self) {
209                unsafe { sys::labelset_replace(self.old) };
210            }
211        }
212
213        let old = unsafe { sys::labelset_replace(self.raw.as_ptr()) };
214        let _guard = Guard { old };
215        f()
216    }
217
218    /// Adds the specified key-value pair to the label set.
219    pub fn set<K, V>(&mut self, key: K, value: V)
220    where
221        K: AsRef<[u8]>,
222        V: AsRef<[u8]>,
223    {
224        let errno = unsafe {
225            sys::labelset_set(
226                self.raw.as_ptr(),
227                key.as_ref().into(),
228                value.as_ref().into(),
229            )
230        };
231        if errno != 0 {
232            panic!("out of memory");
233        }
234    }
235
236    /// Deletes the specified label, if it exists, from the label set.
237    pub fn delete<K>(&mut self, key: K)
238    where
239        K: AsRef<[u8]>,
240    {
241        unsafe { sys::labelset_delete(self.raw.as_ptr(), key.as_ref().into()) }
242    }
243
244    /// Gets the label corresponding to a key on the given label set,
245    /// or `None` if no such label exists.
246    pub fn get<K>(&self, key: K) -> Option<&[u8]>
247    where
248        K: AsRef<[u8]>,
249    {
250        unsafe {
251            sys::labelset_get(self.raw.as_ptr(), key.as_ref().into())
252                .as_ref()
253                .map(|lbl| slice::from_raw_parts(lbl.value.buf, lbl.value.len))
254        }
255    }
256}
257
258impl Default for Labelset {
259    fn default() -> Self {
260        Self::new()
261    }
262}
263
264impl Drop for Labelset {
265    fn drop(&mut self) {
266        unsafe { sys::labelset_free(self.raw.as_ptr()) }
267    }
268}
269
270impl Clone for Labelset {
271    fn clone(&self) -> Self {
272        let raw = unsafe { sys::labelset_clone(self.raw.as_ptr()) };
273        let raw = NonNull::new(raw).expect("failed to clone labelset");
274        Self { raw }
275    }
276}
277
278impl<K, V> Extend<(K, V)> for Labelset
279where
280    K: AsRef<[u8]>,
281    V: AsRef<[u8]>,
282{
283    fn extend<T: IntoIterator<Item = (K, V)>>(&mut self, iter: T) {
284        for (k, v) in iter {
285            self.set(k, v);
286        }
287    }
288}
289
290impl fmt::Debug for Labelset {
291    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
292        debug_labelset(f, self.raw.as_ptr())
293    }
294}
295
296fn debug_labelset(f: &mut fmt::Formatter<'_>, labelset: *const sys::Labelset) -> fmt::Result {
297    let mut cstr = sys::OwnedString::new();
298    let errno = unsafe { sys::labelset_debug_string(labelset, &mut *cstr) };
299    if errno != 0 {
300        panic!("out of memory");
301    }
302    let bytes = unsafe { slice::from_raw_parts(cstr.buf, cstr.len) };
303    let str = String::from_utf8_lossy(bytes);
304    f.write_str(&str)
305}
306
307/// The active label set for the current thread.
308pub const CURRENT_LABELSET: CurrentLabelset = CurrentLabelset { _priv: () };
309
310/// The type of [`CURRENT_LABELSET`].
311pub struct CurrentLabelset {
312    _priv: (),
313}
314
315impl CurrentLabelset {
316    /// Adds the specified key-value pair to the current label set.
317    ///
318    /// # Panics
319    ///
320    /// Panics if there is no current label set.
321    pub fn set<K, V>(&self, key: K, value: V)
322    where
323        K: AsRef<[u8]>,
324        V: AsRef<[u8]>,
325    {
326        if unsafe { sys::labelset_current() }.is_null() {
327            panic!("no current label set");
328        }
329        let errno = unsafe { sys::set(key.as_ref().into(), value.as_ref().into()) };
330        if errno != 0 {
331            panic!("out of memory");
332        }
333    }
334
335    /// Deletes the specified label, if it exists, from the current label set.
336    pub fn delete<K>(&self, key: K)
337    where
338        K: AsRef<[u8]>,
339    {
340        unsafe { sys::delete(key.as_ref().into()) }
341    }
342
343    /// Gets the label corresponding to a key on the current label set,
344    /// or `None` if no such label exists.
345    pub fn get<K>(&self, key: K) -> Option<Vec<u8>>
346    where
347        K: AsRef<[u8]>,
348    {
349        unsafe {
350            sys::get(key.as_ref().into()).as_ref().map(|lbl| {
351                let v = slice::from_raw_parts(lbl.value.buf, lbl.value.len);
352                v.to_vec()
353            })
354        }
355    }
356}
357
358impl fmt::Debug for CurrentLabelset {
359    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
360        let current = unsafe { sys::labelset_current() };
361        if current.is_null() {
362            panic!("no current labelset");
363        }
364        debug_labelset(f, current)
365    }
366}
367
368/// Set the label for the specified key to the specified
369/// value while the given function is running.
370///
371/// All labels are thread-local: setting a label on one thread
372/// has no effect on its value on any other thread.
373pub fn with_label<K, V, F, Ret>(k: K, v: V, f: F) -> Ret
374where
375    K: AsRef<[u8]>,
376    V: AsRef<[u8]>,
377    F: FnOnce() -> Ret,
378{
379    unsafe {
380        if sys::labelset_current().is_null() {
381            let l = sys::labelset_new(0);
382            sys::labelset_replace(l);
383        }
384    }
385    struct Guard<'a> {
386        k: &'a [u8],
387        old_v: Option<sys::OwnedString>,
388    }
389
390    impl<'a> Drop for Guard<'a> {
391        fn drop(&mut self) {
392            if let Some(old_v) = std::mem::take(&mut self.old_v) {
393                let errno = unsafe { sys::set(self.k.into(), *old_v) };
394                if errno != 0 {
395                    panic!("corruption in custom labels library: errno {errno}");
396                }
397            } else {
398                unsafe { sys::delete(self.k.into()) };
399            }
400        }
401    }
402
403    let old_v = unsafe { sys::get(k.as_ref().into()).as_ref() }.map(|lbl| lbl.value.to_owned());
404    let _g = Guard {
405        k: k.as_ref(),
406        old_v,
407    };
408
409    let errno = unsafe { sys::set(k.as_ref().into(), v.as_ref().into()) };
410    if errno != 0 {
411        panic!("corruption in custom labels library: errno {errno}")
412    }
413
414    f()
415}
416
417/// Set the labels for the specified keys to the specified
418/// values while the given function is running.
419///
420/// `i` is an iterator of key-value pairs.
421///
422/// The effect is the same as repeatedly nesting calls to the singular [`with_label`].
423pub fn with_labels<I, K, V, F, Ret>(i: I, f: F) -> Ret
424where
425    I: IntoIterator<Item = (K, V)>,
426    K: AsRef<[u8]>,
427    V: AsRef<[u8]>,
428    F: FnOnce() -> Ret,
429{
430    let mut i = i.into_iter();
431    if let Some((k, v)) = i.next() {
432        with_label(k, v, || with_labels(i, f))
433    } else {
434        f()
435    }
436}
437
438pub mod asynchronous {
439    use pin_project_lite::pin_project;
440    use std::future::Future;
441    use std::iter;
442    use std::pin::Pin;
443    use std::task::{Context, Poll};
444
445    use crate::Labelset;
446
447    pin_project! {
448        /// A [`Future`] with custom labels attached.
449        ///
450        /// This type is returned by the [`Label`] extension trait. See that
451        /// trait's documentation for details.
452        pub struct Labeled<Fut> {
453            #[pin]
454            inner: Fut,
455            labelset: Labelset,
456        }
457    }
458
459    impl<Fut, Ret> Future for Labeled<Fut>
460    where
461        Fut: Future<Output = Ret>,
462    {
463        type Output = Ret;
464
465        fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
466            let p = self.project();
467            p.labelset.enter(|| p.inner.poll(cx))
468        }
469    }
470
471    /// Attaches custom labels to a [`Future`].
472    pub trait Label: Sized {
473        /// Attach the currently active labels to the future.
474        ///
475        /// This can be used to propagate the current labels when spawning
476        /// a new future.
477        fn with_current_labels(self) -> Labeled<Self>;
478
479        /// Attach a single label to the future.
480        ///
481        /// This is equivalent to calling [`with_labels`] with an iterator that
482        /// yields a single key–value pair.
483        fn with_label<K, V>(self, k: K, v: V) -> Labeled<Self>
484        where
485            K: AsRef<[u8]>,
486            V: AsRef<[u8]>;
487
488        /// Attaches the specified labels to the future.
489        ///
490        /// The labels will be installed in the current thread whenever the
491        /// future is polled, and removed when the poll completes.
492        ///
493        /// This is equivalent to calling [`with_labelset`] with a label
494        /// set constructed like so:
495        ///
496        /// ```rust
497        /// # use custom_labels::Labelset;
498        /// let mut labelset = Labelset::clone_from_current();
499        /// labelset.extend(i);
500        /// ```
501        fn with_labels<I, K, V>(self, i: I) -> Labeled<Self>
502        where
503            I: IntoIterator<Item = (K, V)>,
504            K: AsRef<[u8]>,
505            V: AsRef<[u8]>;
506
507        /// Attaches the specified labelset to the future.
508        ///
509        /// The labels in the set will be installed in the current thread
510        /// whenever the future is polled, and removed when the poll completes.
511        fn with_labelset(self, labelset: Labelset) -> Labeled<Self>;
512    }
513
514    impl<Fut: Future> Label for Fut {
515        fn with_current_labels(self) -> Labeled<Self> {
516            self.with_labels(iter::empty::<(&[u8], &[u8])>())
517        }
518
519        fn with_label<K, V>(self, k: K, v: V) -> Labeled<Self>
520        where
521            K: AsRef<[u8]>,
522            V: AsRef<[u8]>,
523        {
524            self.with_labels(iter::once((k, v)))
525        }
526
527        fn with_labels<I, K, V>(self, iter: I) -> Labeled<Self>
528        where
529            I: IntoIterator<Item = (K, V)>,
530            K: AsRef<[u8]>,
531            V: AsRef<[u8]>,
532        {
533            let mut labelset = Labelset::clone_from_current();
534            labelset.extend(iter);
535            Labeled {
536                inner: self,
537                labelset,
538            }
539        }
540
541        fn with_labelset(self, labelset: Labelset) -> Labeled<Self> {
542            Labeled {
543                inner: self,
544                labelset,
545            }
546        }
547    }
548}
custom_labels/lib.rs

custom_labels/
lib.rs
