resource_accounting/
allocator.rs

1//! Global allocator implementation that allows tracking allocations on a per-group basis.
2
3// TODO: The current design does not allow for deregistering groups, which is currently fine and
4// likely will be for a while, but would be a limitation in a world where we dynamically launched
5// data pipelines and wanted to clean up removed components, and so on.
6
7use std::alloc::{GlobalAlloc, Layout};
8
9use crate::{groups::CURRENT_GROUP, ResourceStats};
10
11const STATS_LAYOUT: Layout = Layout::new::<*const ResourceStats>();
12
13/// A global allocator that tracks allocations on a per-group basis.
14///
15/// This allocator provides the ability to track the allocations/deallocations, both in bytes and objects, for
16/// different, user-defined resource groups.
17///
18/// # Resource groups
19///
20/// Allocation (and deallocations) are tracked by **resource group**. When this allocator is used, every allocation is
21/// associated with an resource group. Resource groups are user-defined, except for the default "root" resource
22/// group which acts as a catch-all when a user-defined group isn't currently entered.
23///
24/// # Token guard
25///
26/// When an resource group is registered, an `ResourceGroupToken` is returned. This token can be used to "enter" the
27/// group, which attribute all allocations on the current thread to that group. Entering the group returns a drop guard
28/// that restores the previously entered allocation when it's dropped.
29///
30/// This allows for arbitrarily nested resource groups.
31///
32/// ## Changes to memory layout
33///
34/// In order to associate an allocation with the current resource group, a small trailer is added to the requested
35/// allocation layout, in the form of a pointer to the statistics for the resource group. This allows updating the
36/// statistics directly when an allocation is deallocated, without having to externally keep track of what group a given
37/// allocation belongs to. These statistics are updated directly when the allocation is initially made, and when it's
38/// deallocated.
39///
40/// This means that all requested allocations end up being one machine word larger: 4 bytes on 32-bit systems, and 8
41/// bytes on 64-bit systems.
42pub struct TrackingAllocator<A> {
43    allocator: A,
44}
45
46impl<A> TrackingAllocator<A> {
47    /// Creates a new `TrackingAllocator` that wraps another allocator.
48    ///
49    /// The wrapped allocator is used to actually allocate and deallocate memory, while this allocator is responsible
50    /// purely for tracking the allocations and deallocations themselves.
51    pub const fn new(allocator: A) -> Self {
52        Self { allocator }
53    }
54}
55
56unsafe impl<A> GlobalAlloc for TrackingAllocator<A>
57where
58    A: GlobalAlloc,
59{
60    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
61        // Adjust the requested layout to fit our trailer and then try and allocate it.
62        let (layout, trailer_start) = get_layout_with_group_trailer(layout);
63        let layout_size = layout.size();
64        let ptr = self.allocator.alloc(layout);
65        if ptr.is_null() {
66            return ptr;
67        }
68
69        // Store the pointer to the current resource group in the trailer, and also update the statistics.
70        let trailer_ptr = ptr.add(trailer_start) as *mut *mut ResourceStats;
71        CURRENT_GROUP.with(|current_group| {
72            let group_ptr = current_group.borrow();
73            group_ptr.as_ref().track_allocation(layout_size);
74
75            trailer_ptr.write(group_ptr.as_ptr());
76        });
77
78        ptr
79    }
80
81    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
82        // Read the pointer to the owning resource group from the trailer and update the statistics.
83        let (layout, trailer_start) = get_layout_with_group_trailer(layout);
84        let trailer_ptr = ptr.add(trailer_start) as *mut *mut ResourceStats;
85        let group = (*trailer_ptr).as_ref().unwrap();
86        group.track_deallocation(layout.size());
87
88        // Deallocate the memory.
89        self.allocator.dealloc(ptr, layout);
90    }
91}
92
93fn get_layout_with_group_trailer(layout: Layout) -> (Layout, usize) {
94    let (new_layout, trailer_start) = layout.extend(STATS_LAYOUT).unwrap();
95    (new_layout.pad_to_align(), trailer_start)
96}
resource_accounting/allocator.rs

resource_accounting/
allocator.rs
