saluki_common/resource_tracking/
groups.rs

1use std::{
2    cell::RefCell,
3    collections::HashMap,
4    future::Future,
5    marker::PhantomData,
6    pin::Pin,
7    ptr::NonNull,
8    sync::{Mutex, OnceLock},
9    task::{Context, Poll},
10};
11
12use pin_project::pin_project;
13
14use super::stats::{thread_cpu_time_nanos, ResourceStats};
15
16static REGISTRY: OnceLock<ResourceGroupRegistry> = OnceLock::new();
17static ROOT_GROUP: ResourceStats = ResourceStats::new();
18
19thread_local! {
20    pub(super) static CURRENT_GROUP: RefCell<NonNull<ResourceStats>> = RefCell::new(NonNull::from(&ROOT_GROUP));
21}
22
23/// A token associated with a specific resource group.
24///
25/// Used to attribute allocations and deallocations to a specific group with a scope guard [`ResourceTrackingGuard`], or
26/// through helpers provided by the [`Track`] trait.
27#[derive(Clone, Copy)]
28pub struct ResourceGroupToken {
29    group_ptr: NonNull<ResourceStats>,
30}
31
32impl ResourceGroupToken {
33    fn new(group_ptr: NonNull<ResourceStats>) -> Self {
34        Self { group_ptr }
35    }
36
37    /// Returns an `ResourceGroupToken` for the current resource group.
38    pub fn current() -> Self {
39        CURRENT_GROUP.with(|current_group| {
40            let group_ptr = current_group.borrow();
41            Self::new(*group_ptr)
42        })
43    }
44
45    #[cfg(test)]
46    fn ptr_eq(&self, other: &Self) -> bool {
47        self.group_ptr == other.group_ptr
48    }
49
50    /// Returns the token for the root resource group.
51    pub fn root() -> Self {
52        Self::new(NonNull::from(&ROOT_GROUP))
53    }
54
55    /// Enters this resource group, returning a guard that will exit the resource group when dropped.
56    pub fn enter(&self) -> ResourceTrackingGuard<'_> {
57        // Track our starting point for this thread's CPU usage.
58        let thread_cpu_usage_start = thread_cpu_time_nanos().unwrap_or(0);
59
60        // Swap the current group to the one we're tracking.
61        CURRENT_GROUP.with(|current_group| {
62            let mut group_ptr = current_group.borrow_mut();
63            let previous_group_ptr = *group_ptr;
64            *group_ptr = self.group_ptr;
65
66            ResourceTrackingGuard {
67                previous_group_ptr,
68                thread_cpu_usage_start,
69                _token: PhantomData,
70            }
71        })
72    }
73}
74
75// SAFETY: There's nothing inherently thread-specific about the token.
76unsafe impl Send for ResourceGroupToken {}
77
78// SAFETY: There's nothing unsafe about sharing the token between threads, as it's safe to enter the same token on
79// multiple threads at the same time, and the token itself has no internal state or interior mutability.
80unsafe impl Sync for ResourceGroupToken {}
81
82/// A guard representing an resource group which has been entered.
83///
84/// When the guard is dropped, the resource group will be exited and the previously entered resource group will be
85/// restored.
86///
87/// This is returned by the [`ResourceGroupToken::enter`] method.
88pub struct ResourceTrackingGuard<'a> {
89    previous_group_ptr: NonNull<ResourceStats>,
90    thread_cpu_usage_start: u64,
91    _token: PhantomData<&'a ResourceGroupToken>,
92}
93
94impl Drop for ResourceTrackingGuard<'_> {
95    fn drop(&mut self) {
96        // Grab our current total CPU usage for the thread, and calculate the delta.
97        let thread_cpu_usage_end = thread_cpu_time_nanos().unwrap_or(0);
98        let cpu_usage_delta = thread_cpu_usage_end.saturating_sub(self.thread_cpu_usage_start);
99
100        // Reset the current group to the one that existed before we entered.
101        CURRENT_GROUP.with(|current_group| {
102            let mut group_ptr = current_group.borrow_mut();
103
104            // Now track the delta in CPU usage, if available, before resetting the group.
105            if cpu_usage_delta != 0 {
106                // SAFETY: We only construct the pointer to `ResourceStats` from a leaked heap allocation, and we never
107                // deallocate it, so it's always non-null/aligned/valid-for-`T`, etc.
108                unsafe { group_ptr.as_ref().track_cpu_time(cpu_usage_delta) }
109            }
110
111            *group_ptr = self.previous_group_ptr;
112        });
113    }
114}
115
116/// An object wrapper that tracks allocations and attributes them to a specific group.
117///
118/// Provides methods and implementations to help ensure that operations against/using the wrapped object have all
119/// allocations properly tracked and attributed to a given group.
120///
121/// Implements [`Future`] when the wrapped object itself implements [`Future`].
122//
123// TODO: A more complete example of this sort of thing is `tracing::Instrumented`, where they also have some fancy code
124// to trace execution even in the drop logic of the wrapped future. I'm not sure we need that here, because we don't
125// care about what components an object is deallocated in, and I don't think we expect to have any futures where the
126// drop logic actually _allocates_, and certainly not in a way where we want to attribute it to that future's attached
127// component... but for posterity, I'm mentioning it here since we _might_ consider doing it. Might.
128#[pin_project]
129#[must_use = "futures do nothing unless you `.await` or poll them"]
130pub struct Tracked<Inner> {
131    token: ResourceGroupToken,
132
133    #[pin]
134    inner: Inner,
135}
136
137impl<Inner> Tracked<Inner> {
138    /// Consumes this object and returns the inner object and tracking token.
139    pub fn into_parts(self) -> (ResourceGroupToken, Inner) {
140        (self.token, self.inner)
141    }
142}
143
144impl<Inner> Future for Tracked<Inner>
145where
146    Inner: Future,
147{
148    type Output = Inner::Output;
149
150    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
151        let this = self.project();
152        let _enter = this.token.enter();
153
154        this.inner.poll(cx)
155    }
156}
157
158/// Attaches resource groups to a [`Future`].
159pub trait Track: Sized {
160    /// Instruments this type by attaching the given resource group token, returning a `Tracked` wrapper.
161    ///
162    /// The resource group will be entered every time the wrapped future is polled.
163    ///
164    /// # Examples
165    ///
166    /// ```rust
167    /// use saluki_common::resource_tracking::{ResourceGroupRegistry, ResourceGroupToken, Track as _};
168    ///
169    /// # async fn doc() {
170    /// let future = async {
171    ///     // All allocations in this future will be attached to the resource group
172    ///     // represented by `token`...
173    /// };
174    ///
175    /// let token = ResourceGroupRegistry::global().register_resource_group("my-group");
176    /// future
177    ///     .track_resources(token)
178    ///     .await
179    /// # }
180    fn track_resources(self, token: ResourceGroupToken) -> Tracked<Self> {
181        Tracked { token, inner: self }
182    }
183
184    /// Instruments this type by attaching the current resource group, returning a `Tracked` wrapper.
185    ///
186    /// The resource group will be entered every time the wrapped future is polled.
187    ///
188    /// This can be used to propagate the current resource group when spawning a new future.
189    ///
190    /// # Examples
191    ///
192    /// ```rust
193    /// use saluki_common::resource_tracking::{ResourceGroupRegistry, ResourceGroupToken, Track as _};
194    ///
195    /// # mod tokio {
196    /// #     pub(super) fn spawn(_: impl std::future::Future) {}
197    /// # }
198    /// # async fn doc() {
199    /// let token = ResourceGroupRegistry::global().register_resource_group("my-group");
200    /// let _enter = token.enter();
201    ///
202    /// // ...
203    ///
204    /// let future = async {
205    ///     // All allocations in this future will be attached to the resource group
206    ///     // represented by `token`...
207    /// };
208    /// tokio::spawn(future.in_current_resource_group());
209    /// # }
210    /// ```
211    fn in_current_resource_group(self) -> Tracked<Self> {
212        Tracked {
213            token: ResourceGroupToken::current(),
214            inner: self,
215        }
216    }
217}
218
219impl<T: Sized> Track for T {}
220
221/// A registry of resource groups and the statistics for each of them.
222///
223/// Resource groups are user-defined groups which can then be associated with memory and CPU usage in distinct code
224/// regions. This mechanism allows for granular resource accounting at the level which makes sense to the application,
225/// such as per-thread, per-async task, and so on.
226///
227/// # Token guard
228///
229/// When an resource group is registered, an `ResourceGroupToken` is returned. This token can be used to "enter" the
230/// group, which causes memory and CPU usage on the current thread to be attributed to that group. Entering the group
231/// returns a drop guard that restores the previously entered group when dropped.
232///
233/// This allows for arbitrarily nested resource groups.
234///
235/// Additionally, [`Tracked`] can be used to wrap a [`Future`], attaching it to a specific resource group token. This
236/// causes the future to track all memory and CPU usage during polls such that the usage is properly attributed to the
237/// resource group.
238///
239/// # Resources tracked
240///
241/// ## Memory usage
242///
243/// In order for memory usage to be tracked, [`TrackingAllocator`][super::TrackingAllocator] must be installed
244/// as the global allocator for the process.
245///
246/// ## CPU usage
247///
248/// CPU usage is automatically tracked if platform support is detected.
249///
250/// Currently, only Linux is supported for CPU usage tracking.
251pub struct ResourceGroupRegistry {
252    resource_groups: Mutex<HashMap<String, Box<ResourceStats>>>,
253}
254
255impl ResourceGroupRegistry {
256    fn new() -> Self {
257        in_root_resource_group(|| Self {
258            resource_groups: Mutex::new(HashMap::with_capacity(4)),
259        })
260    }
261
262    /// Gets a reference to the global resource group registry.
263    pub fn global() -> &'static Self {
264        REGISTRY.get_or_init(Self::new)
265    }
266
267    /// Returns `true` if `TrackingAllocator` is installed as the global allocator.
268    pub fn allocator_installed() -> bool {
269        // Essentially, when we load the group registry, and it gets created for the first time, it will specifically
270        // allocate its internal data structures while entered into the root resource group.
271        //
272        // This means that if the allocator is installed, we should always have some allocations in the root group by
273        // the time we call `ResourceStats::has_allocated`.
274        ROOT_GROUP.has_allocated()
275    }
276
277    /// Registers a new resource group with the given name.
278    ///
279    /// Returns an `ResourceGroupToken` that can be used to attribute CPU and memory usage to the
280    /// newly created resource group.
281    pub fn register_resource_group<S>(&self, name: S) -> ResourceGroupToken
282    where
283        S: AsRef<str>,
284    {
285        in_root_resource_group(|| {
286            let mut resource_groups = self.resource_groups.lock().unwrap();
287            match resource_groups.get(name.as_ref()) {
288                Some(stats) => ResourceGroupToken::new(NonNull::from(&**stats)),
289                None => {
290                    let resource_group_stats = Box::new(ResourceStats::new());
291                    let token = ResourceGroupToken::new(NonNull::from(&*resource_group_stats));
292
293                    resource_groups.insert(name.as_ref().to_string(), resource_group_stats);
294
295                    token
296                }
297            }
298        })
299    }
300
301    /// Visits all resource groups in the registry and calls the given closure with their names and statistics.
302    pub fn visit_resource_groups<F>(&self, mut f: F)
303    where
304        F: FnMut(&str, &ResourceStats),
305    {
306        in_root_resource_group(|| {
307            f("root", &ROOT_GROUP);
308
309            let resource_groups = self.resource_groups.lock().unwrap();
310            for (name, stats) in resource_groups.iter() {
311                f(name, stats);
312            }
313        });
314    }
315}
316
317fn in_root_resource_group<F, R>(f: F) -> R
318where
319    F: FnOnce() -> R,
320{
321    let token = ResourceGroupToken::root();
322    let _enter = token.enter();
323    f()
324}
325
326#[cfg(test)]
327mod tests {
328    use std::{
329        cell::Cell,
330        future::Future,
331        pin::Pin,
332        rc::Rc,
333        sync::Arc,
334        task::{Context, Poll, Wake, Waker},
335    };
336
337    use super::{ResourceGroupRegistry, ResourceGroupToken, Track};
338
339    struct NoopWaker;
340
341    impl Wake for NoopWaker {
342        fn wake(self: Arc<Self>) {}
343    }
344
345    /// Polls a future to completion on the current thread using a no-op `waker`.
346    fn poll_to_completion<F: Future>(future: F) -> F::Output {
347        let mut future = Box::pin(future);
348        let waker = Waker::from(Arc::new(NoopWaker));
349        let mut cx = Context::from_waker(&waker);
350        loop {
351            if let Poll::Ready(output) = future.as_mut().poll(&mut cx) {
352                return output;
353            }
354        }
355    }
356
357    /// A future that records, each time it is polled, whether the currently entered resource group is `expected`.
358    struct RecordCurrentGroup {
359        expected: ResourceGroupToken,
360        matched: Rc<Cell<bool>>,
361    }
362
363    impl Future for RecordCurrentGroup {
364        type Output = ();
365
366        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
367            self.matched.set(ResourceGroupToken::current().ptr_eq(&self.expected));
368            Poll::Ready(())
369        }
370    }
371
372    // CPU attribution only happens on platforms where per-thread CPU time is available (Linux), so the helpers that
373    // read it back and burn CPU are gated to avoid dead-code warnings elsewhere.
374    #[cfg(target_os = "linux")]
375    fn cpu_time_nanos_for(registry: &ResourceGroupRegistry, target: &str) -> u64 {
376        use crate::resource_tracking::ResourceStatsSnapshot;
377
378        let mut cpu_time_nanos = 0;
379        registry.visit_resource_groups(|name, stats| {
380            if name == target {
381                cpu_time_nanos = stats.snapshot_delta(&ResourceStatsSnapshot::empty()).cpu_time_nanos;
382            }
383        });
384        cpu_time_nanos
385    }
386
387    #[cfg(target_os = "linux")]
388    fn burn_cpu() {
389        let mut sum = 0u64;
390        for i in 0..20_000_000u64 {
391            sum = sum.wrapping_add(i);
392        }
393        std::hint::black_box(sum);
394    }
395
396    #[test]
397    fn existing_group() {
398        let registry = ResourceGroupRegistry::new();
399        let token = registry.register_resource_group("test");
400        let token2 = registry.register_resource_group("test");
401        let token3 = registry.register_resource_group("test2");
402
403        assert!(token.ptr_eq(&token2));
404        assert!(!token.ptr_eq(&token3));
405    }
406
407    #[test]
408    fn visit_resource_groups() {
409        let registry = ResourceGroupRegistry::new();
410        let _token = registry.register_resource_group("my-group");
411
412        let mut visited = Vec::new();
413        registry.visit_resource_groups(|name, _stats| {
414            visited.push(name.to_string());
415        });
416
417        assert_eq!(visited.len(), 2);
418        assert_eq!(visited[0], "root");
419        assert_eq!(visited[1], "my-group");
420    }
421
422    #[test]
423    fn enter_swaps_current_group_and_restores_previous_on_drop() {
424        let registry = ResourceGroupRegistry::new();
425        let group = registry.register_resource_group("group-a");
426        let previous = ResourceGroupToken::current();
427
428        {
429            let _guard = group.enter();
430            assert!(
431                ResourceGroupToken::current().ptr_eq(&group),
432                "entering a group should make it the current group"
433            );
434        }
435
436        assert!(
437            ResourceGroupToken::current().ptr_eq(&previous),
438            "dropping the guard should restore the previously-entered group"
439        );
440    }
441
442    #[test]
443    fn nested_groups_restore_in_lifo_order() {
444        let registry = ResourceGroupRegistry::new();
445        let outer = registry.register_resource_group("outer");
446        let inner = registry.register_resource_group("inner");
447        let root = ResourceGroupToken::current();
448
449        let outer_guard = outer.enter();
450        assert!(ResourceGroupToken::current().ptr_eq(&outer));
451
452        {
453            let _inner_guard = inner.enter();
454            assert!(ResourceGroupToken::current().ptr_eq(&inner));
455        }
456
457        // Exiting the inner group restores the outer group, not the root.
458        assert!(ResourceGroupToken::current().ptr_eq(&outer));
459
460        drop(outer_guard);
461        assert!(ResourceGroupToken::current().ptr_eq(&root));
462    }
463
464    #[test]
465    fn tracked_future_enters_attached_group_during_poll() {
466        let registry = ResourceGroupRegistry::new();
467        let group = registry.register_resource_group("tracked");
468        let previous = ResourceGroupToken::current();
469
470        let matched = Rc::new(Cell::new(false));
471        let future = RecordCurrentGroup {
472            expected: group,
473            matched: Rc::clone(&matched),
474        }
475        .track_resources(group);
476
477        poll_to_completion(future);
478
479        assert!(
480            matched.get(),
481            "the attached group should be the current group while the future is polled"
482        );
483        assert!(
484            ResourceGroupToken::current().ptr_eq(&previous),
485            "the previous group should be restored once the poll returns"
486        );
487    }
488
489    #[test]
490    fn in_current_resource_group_captures_group_at_attach_time() {
491        let registry = ResourceGroupRegistry::new();
492        let group = registry.register_resource_group("captured");
493
494        let matched = Rc::new(Cell::new(false));
495        let future = {
496            // Attach while `group` is entered; the wrapper should remember it.
497            let _guard = group.enter();
498            RecordCurrentGroup {
499                expected: group,
500                matched: Rc::clone(&matched),
501            }
502            .in_current_resource_group()
503        };
504
505        // The guard has been dropped, so `group` is no longer current...
506        assert!(!ResourceGroupToken::current().ptr_eq(&group));
507
508        // ...yet polling the wrapper still re-enters the group captured at attach time.
509        poll_to_completion(future);
510        assert!(matched.get());
511    }
512
513    #[cfg(target_os = "linux")]
514    #[test]
515    fn cpu_time_is_attributed_to_the_entered_group() {
516        let registry = ResourceGroupRegistry::new();
517        let busy = registry.register_resource_group("busy");
518        let _idle = registry.register_resource_group("idle");
519
520        {
521            let _guard = busy.enter();
522            burn_cpu();
523        }
524
525        // CPU time consumed while `busy` was entered is attributed to it; `idle` was never entered.
526        assert!(
527            cpu_time_nanos_for(&registry, "busy") > 0,
528            "the entered group should accrue the CPU time spent inside the guard"
529        );
530        assert_eq!(
531            cpu_time_nanos_for(&registry, "idle"),
532            0,
533            "a group that was never entered should accrue no CPU time"
534        );
535    }
536
537    #[cfg(target_os = "linux")]
538    #[test]
539    fn tracked_future_attributes_poll_cpu_time_to_its_group() {
540        let registry = ResourceGroupRegistry::new();
541        let group = registry.register_resource_group("worker");
542
543        struct BurnCpu;
544
545        impl Future for BurnCpu {
546            type Output = ();
547
548            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
549                burn_cpu();
550                Poll::Ready(())
551            }
552        }
553
554        poll_to_completion(BurnCpu.track_resources(group));
555
556        assert!(
557            cpu_time_nanos_for(&registry, "worker") > 0,
558            "CPU time spent polling a tracked future should be attributed to its group"
559        );
560    }
561}
saluki_common/resource_tracking/groups.rs

saluki_common/resource_tracking/
groups.rs
