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    /// <!-- vale off -->
167    /// ```rust
168    /// use saluki_common::resource_tracking::{ResourceGroupRegistry, ResourceGroupToken, Track as _};
169    ///
170    /// # async fn doc() {
171    /// let future = async {
172    ///     // All allocations in this future will be attached to the resource group
173    ///     // represented by `token`...
174    /// };
175    ///
176    /// let token = ResourceGroupRegistry::global().register_resource_group("my-group");
177    /// future
178    ///     .track_resources(token)
179    ///     .await
180    /// # }
181    /// ```
182    /// <!-- vale on -->
183    fn track_resources(self, token: ResourceGroupToken) -> Tracked<Self> {
184        Tracked { token, inner: self }
185    }
186
187    /// Instruments this type by attaching the current resource group, returning a `Tracked` wrapper.
188    ///
189    /// The resource group will be entered every time the wrapped future is polled.
190    ///
191    /// This can be used to propagate the current resource group when spawning a new future.
192    ///
193    /// # Examples
194    ///
195    /// <!-- vale off -->
196    /// ```rust
197    /// use saluki_common::resource_tracking::{ResourceGroupRegistry, ResourceGroupToken, Track as _};
198    ///
199    /// # mod tokio {
200    /// #     pub(super) fn spawn(_: impl std::future::Future) {}
201    /// # }
202    /// # async fn doc() {
203    /// let token = ResourceGroupRegistry::global().register_resource_group("my-group");
204    /// let _enter = token.enter();
205    ///
206    /// // ...
207    ///
208    /// let future = async {
209    ///     // All allocations in this future will be attached to the resource group
210    ///     // represented by `token`...
211    /// };
212    /// tokio::spawn(future.in_current_resource_group());
213    /// # }
214    /// ```
215    /// <!-- vale on -->
216    fn in_current_resource_group(self) -> Tracked<Self> {
217        Tracked {
218            token: ResourceGroupToken::current(),
219            inner: self,
220        }
221    }
222}
223
224impl<T: Sized> Track for T {}
225
226/// A registry of resource groups and the statistics for each of them.
227///
228/// Resource groups are user-defined groups which can then be associated with memory and CPU usage in distinct code
229/// regions. This mechanism allows for granular resource accounting at the level which makes sense to the application,
230/// such as per-thread, per-async task, and so on.
231///
232/// # Token guard
233///
234/// When an resource group is registered, an `ResourceGroupToken` is returned. This token can be used to "enter" the
235/// group, which causes memory and CPU usage on the current thread to be attributed to that group. Entering the group
236/// returns a drop guard that restores the previously entered group when dropped.
237///
238/// This allows for arbitrarily nested resource groups.
239///
240/// Additionally, [`Tracked`] can be used to wrap a [`Future`], attaching it to a specific resource group token. This
241/// causes the future to track all memory and CPU usage during polls such that the usage is properly attributed to the
242/// resource group.
243///
244/// # Resources tracked
245///
246/// ## Memory usage
247///
248/// In order for memory usage to be tracked, [`TrackingAllocator`][super::TrackingAllocator] must be installed
249/// as the global allocator for the process.
250///
251/// ## CPU usage
252///
253/// CPU usage is automatically tracked if platform support is detected.
254///
255/// Currently, only Linux is supported for CPU usage tracking.
256pub struct ResourceGroupRegistry {
257    resource_groups: Mutex<HashMap<String, Box<ResourceStats>>>,
258}
259
260impl ResourceGroupRegistry {
261    fn new() -> Self {
262        in_root_resource_group(|| Self {
263            resource_groups: Mutex::new(HashMap::with_capacity(4)),
264        })
265    }
266
267    /// Gets a reference to the global resource group registry.
268    pub fn global() -> &'static Self {
269        REGISTRY.get_or_init(Self::new)
270    }
271
272    /// Returns `true` if `TrackingAllocator` is installed as the global allocator.
273    pub fn allocator_installed() -> bool {
274        // Essentially, when we load the group registry, and it gets created for the first time, it will specifically
275        // allocate its internal data structures while entered into the root resource group.
276        //
277        // This means that if the allocator is installed, we should always have some allocations in the root group by
278        // the time we call `ResourceStats::has_allocated`.
279        ROOT_GROUP.has_allocated()
280    }
281
282    /// Registers a new resource group with the given name.
283    ///
284    /// Returns an `ResourceGroupToken` that can be used to attribute CPU and memory usage to the
285    /// newly created resource group.
286    pub fn register_resource_group<S>(&self, name: S) -> ResourceGroupToken
287    where
288        S: AsRef<str>,
289    {
290        in_root_resource_group(|| {
291            let mut resource_groups = self.resource_groups.lock().unwrap();
292            match resource_groups.get(name.as_ref()) {
293                Some(stats) => ResourceGroupToken::new(NonNull::from(&**stats)),
294                None => {
295                    let resource_group_stats = Box::new(ResourceStats::new());
296                    let token = ResourceGroupToken::new(NonNull::from(&*resource_group_stats));
297
298                    resource_groups.insert(name.as_ref().to_string(), resource_group_stats);
299
300                    token
301                }
302            }
303        })
304    }
305
306    /// Visits all resource groups in the registry and calls the given closure with their names and statistics.
307    pub fn visit_resource_groups<F>(&self, mut f: F)
308    where
309        F: FnMut(&str, &ResourceStats),
310    {
311        in_root_resource_group(|| {
312            f("root", &ROOT_GROUP);
313
314            let resource_groups = self.resource_groups.lock().unwrap();
315            for (name, stats) in resource_groups.iter() {
316                f(name, stats);
317            }
318        });
319    }
320}
321
322fn in_root_resource_group<F, R>(f: F) -> R
323where
324    F: FnOnce() -> R,
325{
326    let token = ResourceGroupToken::root();
327    let _enter = token.enter();
328    f()
329}
330
331#[cfg(test)]
332mod tests {
333    use std::{
334        cell::Cell,
335        future::Future,
336        pin::Pin,
337        rc::Rc,
338        task::{Context, Poll, Waker},
339    };
340
341    use super::{ResourceGroupRegistry, ResourceGroupToken, Track};
342
343    /// Polls a future to completion on the current thread using a no-op `waker`.
344    fn poll_to_completion<F: Future>(future: F) -> F::Output {
345        let mut future = Box::pin(future);
346        let waker = Waker::noop();
347        let mut cx = Context::from_waker(waker);
348        loop {
349            if let Poll::Ready(output) = future.as_mut().poll(&mut cx) {
350                return output;
351            }
352        }
353    }
354
355    /// A future that records, each time it is polled, whether the currently entered resource group is `expected`.
356    struct RecordCurrentGroup {
357        expected: ResourceGroupToken,
358        matched: Rc<Cell<bool>>,
359    }
360
361    impl Future for RecordCurrentGroup {
362        type Output = ();
363
364        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
365            self.matched.set(ResourceGroupToken::current().ptr_eq(&self.expected));
366            Poll::Ready(())
367        }
368    }
369
370    // CPU attribution only happens on platforms where per-thread CPU time is available (Linux), so the helpers that
371    // read it back and burn CPU are gated to avoid dead-code warnings elsewhere.
372    #[cfg(target_os = "linux")]
373    fn cpu_time_nanos_for(registry: &ResourceGroupRegistry, target: &str) -> u64 {
374        use crate::resource_tracking::ResourceStatsSnapshot;
375
376        let mut cpu_time_nanos = 0;
377        registry.visit_resource_groups(|name, stats| {
378            if name == target {
379                cpu_time_nanos = stats.snapshot_delta(&ResourceStatsSnapshot::empty()).cpu_time_nanos;
380            }
381        });
382        cpu_time_nanos
383    }
384
385    #[cfg(target_os = "linux")]
386    fn burn_cpu() {
387        let mut sum = 0u64;
388        for i in 0..20_000_000u64 {
389            sum = sum.wrapping_add(i);
390        }
391        std::hint::black_box(sum);
392    }
393
394    #[test]
395    fn existing_group() {
396        let registry = ResourceGroupRegistry::new();
397        let token = registry.register_resource_group("test");
398        let token2 = registry.register_resource_group("test");
399        let token3 = registry.register_resource_group("test2");
400
401        assert!(token.ptr_eq(&token2));
402        assert!(!token.ptr_eq(&token3));
403    }
404
405    #[test]
406    fn visit_resource_groups() {
407        let registry = ResourceGroupRegistry::new();
408        let _token = registry.register_resource_group("my-group");
409
410        let mut visited = Vec::new();
411        registry.visit_resource_groups(|name, _stats| {
412            visited.push(name.to_string());
413        });
414
415        assert_eq!(visited.len(), 2);
416        assert_eq!(visited[0], "root");
417        assert_eq!(visited[1], "my-group");
418    }
419
420    #[test]
421    fn enter_swaps_current_group_and_restores_previous_on_drop() {
422        let registry = ResourceGroupRegistry::new();
423        let group = registry.register_resource_group("group-a");
424        let previous = ResourceGroupToken::current();
425
426        {
427            let _guard = group.enter();
428            assert!(
429                ResourceGroupToken::current().ptr_eq(&group),
430                "entering a group should make it the current group"
431            );
432        }
433
434        assert!(
435            ResourceGroupToken::current().ptr_eq(&previous),
436            "dropping the guard should restore the previously-entered group"
437        );
438    }
439
440    #[test]
441    fn nested_groups_restore_in_lifo_order() {
442        let registry = ResourceGroupRegistry::new();
443        let outer = registry.register_resource_group("outer");
444        let inner = registry.register_resource_group("inner");
445        let root = ResourceGroupToken::current();
446
447        let outer_guard = outer.enter();
448        assert!(ResourceGroupToken::current().ptr_eq(&outer));
449
450        {
451            let _inner_guard = inner.enter();
452            assert!(ResourceGroupToken::current().ptr_eq(&inner));
453        }
454
455        // Exiting the inner group restores the outer group, not the root.
456        assert!(ResourceGroupToken::current().ptr_eq(&outer));
457
458        drop(outer_guard);
459        assert!(ResourceGroupToken::current().ptr_eq(&root));
460    }
461
462    #[test]
463    fn tracked_future_enters_attached_group_during_poll() {
464        let registry = ResourceGroupRegistry::new();
465        let group = registry.register_resource_group("tracked");
466        let previous = ResourceGroupToken::current();
467
468        let matched = Rc::new(Cell::new(false));
469        let future = RecordCurrentGroup {
470            expected: group,
471            matched: Rc::clone(&matched),
472        }
473        .track_resources(group);
474
475        poll_to_completion(future);
476
477        assert!(
478            matched.get(),
479            "the attached group should be the current group while the future is polled"
480        );
481        assert!(
482            ResourceGroupToken::current().ptr_eq(&previous),
483            "the previous group should be restored once the poll returns"
484        );
485    }
486
487    #[test]
488    fn in_current_resource_group_captures_group_at_attach_time() {
489        let registry = ResourceGroupRegistry::new();
490        let group = registry.register_resource_group("captured");
491
492        let matched = Rc::new(Cell::new(false));
493        let future = {
494            // Attach while `group` is entered; the wrapper should remember it.
495            let _guard = group.enter();
496            RecordCurrentGroup {
497                expected: group,
498                matched: Rc::clone(&matched),
499            }
500            .in_current_resource_group()
501        };
502
503        // The guard has been dropped, so `group` is no longer current...
504        assert!(!ResourceGroupToken::current().ptr_eq(&group));
505
506        // ...yet polling the wrapper still re-enters the group captured at attach time.
507        poll_to_completion(future);
508        assert!(matched.get());
509    }
510
511    #[cfg(target_os = "linux")]
512    #[test]
513    fn cpu_time_is_attributed_to_the_entered_group() {
514        let registry = ResourceGroupRegistry::new();
515        let busy = registry.register_resource_group("busy");
516        let _idle = registry.register_resource_group("idle");
517
518        {
519            let _guard = busy.enter();
520            burn_cpu();
521        }
522
523        // CPU time consumed while `busy` was entered is attributed to it; `idle` was never entered.
524        assert!(
525            cpu_time_nanos_for(&registry, "busy") > 0,
526            "the entered group should accrue the CPU time spent inside the guard"
527        );
528        assert_eq!(
529            cpu_time_nanos_for(&registry, "idle"),
530            0,
531            "a group that was never entered should accrue no CPU time"
532        );
533    }
534
535    #[cfg(target_os = "linux")]
536    #[test]
537    fn tracked_future_attributes_poll_cpu_time_to_its_group() {
538        let registry = ResourceGroupRegistry::new();
539        let group = registry.register_resource_group("worker");
540
541        struct BurnCpu;
542
543        impl Future for BurnCpu {
544            type Output = ();
545
546            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
547                burn_cpu();
548                Poll::Ready(())
549            }
550        }
551
552        poll_to_completion(BurnCpu.track_resources(group));
553
554        assert!(
555            cpu_time_nanos_for(&registry, "worker") > 0,
556            "CPU time spent polling a tracked future should be attributed to its group"
557        );
558    }
559}
saluki_common/resource_tracking/groups.rs

saluki_common/resource_tracking/
groups.rs
