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        sync::Arc,
339        task::{Context, Poll, Wake, Waker},
340    };
341
342    use super::{ResourceGroupRegistry, ResourceGroupToken, Track};
343
344    struct NoopWaker;
345
346    impl Wake for NoopWaker {
347        fn wake(self: Arc<Self>) {}
348    }
349
350    /// Polls a future to completion on the current thread using a no-op `waker`.
351    fn poll_to_completion<F: Future>(future: F) -> F::Output {
352        let mut future = Box::pin(future);
353        let waker = Waker::from(Arc::new(NoopWaker));
354        let mut cx = Context::from_waker(&waker);
355        loop {
356            if let Poll::Ready(output) = future.as_mut().poll(&mut cx) {
357                return output;
358            }
359        }
360    }
361
362    /// A future that records, each time it is polled, whether the currently entered resource group is `expected`.
363    struct RecordCurrentGroup {
364        expected: ResourceGroupToken,
365        matched: Rc<Cell<bool>>,
366    }
367
368    impl Future for RecordCurrentGroup {
369        type Output = ();
370
371        fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
372            self.matched.set(ResourceGroupToken::current().ptr_eq(&self.expected));
373            Poll::Ready(())
374        }
375    }
376
377    // CPU attribution only happens on platforms where per-thread CPU time is available (Linux), so the helpers that
378    // read it back and burn CPU are gated to avoid dead-code warnings elsewhere.
379    #[cfg(target_os = "linux")]
380    fn cpu_time_nanos_for(registry: &ResourceGroupRegistry, target: &str) -> u64 {
381        use crate::resource_tracking::ResourceStatsSnapshot;
382
383        let mut cpu_time_nanos = 0;
384        registry.visit_resource_groups(|name, stats| {
385            if name == target {
386                cpu_time_nanos = stats.snapshot_delta(&ResourceStatsSnapshot::empty()).cpu_time_nanos;
387            }
388        });
389        cpu_time_nanos
390    }
391
392    #[cfg(target_os = "linux")]
393    fn burn_cpu() {
394        let mut sum = 0u64;
395        for i in 0..20_000_000u64 {
396            sum = sum.wrapping_add(i);
397        }
398        std::hint::black_box(sum);
399    }
400
401    #[test]
402    fn existing_group() {
403        let registry = ResourceGroupRegistry::new();
404        let token = registry.register_resource_group("test");
405        let token2 = registry.register_resource_group("test");
406        let token3 = registry.register_resource_group("test2");
407
408        assert!(token.ptr_eq(&token2));
409        assert!(!token.ptr_eq(&token3));
410    }
411
412    #[test]
413    fn visit_resource_groups() {
414        let registry = ResourceGroupRegistry::new();
415        let _token = registry.register_resource_group("my-group");
416
417        let mut visited = Vec::new();
418        registry.visit_resource_groups(|name, _stats| {
419            visited.push(name.to_string());
420        });
421
422        assert_eq!(visited.len(), 2);
423        assert_eq!(visited[0], "root");
424        assert_eq!(visited[1], "my-group");
425    }
426
427    #[test]
428    fn enter_swaps_current_group_and_restores_previous_on_drop() {
429        let registry = ResourceGroupRegistry::new();
430        let group = registry.register_resource_group("group-a");
431        let previous = ResourceGroupToken::current();
432
433        {
434            let _guard = group.enter();
435            assert!(
436                ResourceGroupToken::current().ptr_eq(&group),
437                "entering a group should make it the current group"
438            );
439        }
440
441        assert!(
442            ResourceGroupToken::current().ptr_eq(&previous),
443            "dropping the guard should restore the previously-entered group"
444        );
445    }
446
447    #[test]
448    fn nested_groups_restore_in_lifo_order() {
449        let registry = ResourceGroupRegistry::new();
450        let outer = registry.register_resource_group("outer");
451        let inner = registry.register_resource_group("inner");
452        let root = ResourceGroupToken::current();
453
454        let outer_guard = outer.enter();
455        assert!(ResourceGroupToken::current().ptr_eq(&outer));
456
457        {
458            let _inner_guard = inner.enter();
459            assert!(ResourceGroupToken::current().ptr_eq(&inner));
460        }
461
462        // Exiting the inner group restores the outer group, not the root.
463        assert!(ResourceGroupToken::current().ptr_eq(&outer));
464
465        drop(outer_guard);
466        assert!(ResourceGroupToken::current().ptr_eq(&root));
467    }
468
469    #[test]
470    fn tracked_future_enters_attached_group_during_poll() {
471        let registry = ResourceGroupRegistry::new();
472        let group = registry.register_resource_group("tracked");
473        let previous = ResourceGroupToken::current();
474
475        let matched = Rc::new(Cell::new(false));
476        let future = RecordCurrentGroup {
477            expected: group,
478            matched: Rc::clone(&matched),
479        }
480        .track_resources(group);
481
482        poll_to_completion(future);
483
484        assert!(
485            matched.get(),
486            "the attached group should be the current group while the future is polled"
487        );
488        assert!(
489            ResourceGroupToken::current().ptr_eq(&previous),
490            "the previous group should be restored once the poll returns"
491        );
492    }
493
494    #[test]
495    fn in_current_resource_group_captures_group_at_attach_time() {
496        let registry = ResourceGroupRegistry::new();
497        let group = registry.register_resource_group("captured");
498
499        let matched = Rc::new(Cell::new(false));
500        let future = {
501            // Attach while `group` is entered; the wrapper should remember it.
502            let _guard = group.enter();
503            RecordCurrentGroup {
504                expected: group,
505                matched: Rc::clone(&matched),
506            }
507            .in_current_resource_group()
508        };
509
510        // The guard has been dropped, so `group` is no longer current...
511        assert!(!ResourceGroupToken::current().ptr_eq(&group));
512
513        // ...yet polling the wrapper still re-enters the group captured at attach time.
514        poll_to_completion(future);
515        assert!(matched.get());
516    }
517
518    #[cfg(target_os = "linux")]
519    #[test]
520    fn cpu_time_is_attributed_to_the_entered_group() {
521        let registry = ResourceGroupRegistry::new();
522        let busy = registry.register_resource_group("busy");
523        let _idle = registry.register_resource_group("idle");
524
525        {
526            let _guard = busy.enter();
527            burn_cpu();
528        }
529
530        // CPU time consumed while `busy` was entered is attributed to it; `idle` was never entered.
531        assert!(
532            cpu_time_nanos_for(&registry, "busy") > 0,
533            "the entered group should accrue the CPU time spent inside the guard"
534        );
535        assert_eq!(
536            cpu_time_nanos_for(&registry, "idle"),
537            0,
538            "a group that was never entered should accrue no CPU time"
539        );
540    }
541
542    #[cfg(target_os = "linux")]
543    #[test]
544    fn tracked_future_attributes_poll_cpu_time_to_its_group() {
545        let registry = ResourceGroupRegistry::new();
546        let group = registry.register_resource_group("worker");
547
548        struct BurnCpu;
549
550        impl Future for BurnCpu {
551            type Output = ();
552
553            fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Self::Output> {
554                burn_cpu();
555                Poll::Ready(())
556            }
557        }
558
559        poll_to_completion(BurnCpu.track_resources(group));
560
561        assert!(
562            cpu_time_nanos_for(&registry, "worker") > 0,
563            "CPU time spent polling a tracked future should be attributed to its group"
564        );
565    }
566}
saluki_common/resource_tracking/groups.rs

saluki_common/resource_tracking/
groups.rs
