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

custom_labels/
lib.rs
