saluki_common/resource_tracking/
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::{
8    alloc::{GlobalAlloc, Layout},
9    ptr,
10};
11
12use super::{groups::CURRENT_GROUP, stats::ResourceStats};
13
14const STATS_LAYOUT: Layout = Layout::new::<*const ResourceStats>();
15
16/// A global allocator that tracks allocations on a per-group basis.
17///
18/// This allocator provides the ability to track the allocations/deallocations, both in bytes and objects, for
19/// different, user-defined resource groups.
20///
21/// # Resource groups
22///
23/// Allocation (and deallocations) are tracked by **resource group**. When this allocator is used, every allocation is
24/// associated with an resource group. Resource groups are user-defined, except for the default "root" resource
25/// group which acts as a catch-all when a user-defined group isn't currently entered.
26///
27/// # Token guard
28///
29/// When an resource group is registered, an `ResourceGroupToken` is returned. This token can be used to "enter" the
30/// group, which attribute all allocations on the current thread to that group. Entering the group returns a drop guard
31/// that restores the previously entered allocation when it's dropped.
32///
33/// This allows for arbitrarily nested resource groups.
34///
35/// ## Changes to memory layout
36///
37/// In order to associate an allocation with the current resource group, a small trailer is added to the requested
38/// allocation layout, in the form of a pointer to the statistics for the resource group. This allows updating the
39/// statistics directly when an allocation is deallocated, without having to externally keep track of what group a given
40/// allocation belongs to. These statistics are updated directly when the allocation is initially made, and when it's
41/// deallocated.
42///
43/// This means that all requested allocations end up being one machine word larger: 4 bytes on 32-bit systems, and 8
44/// bytes on 64-bit systems.
45pub struct TrackingAllocator<A> {
46    allocator: A,
47}
48
49impl<A> TrackingAllocator<A> {
50    /// Creates a new `TrackingAllocator` that wraps another allocator.
51    ///
52    /// The wrapped allocator is used to actually allocate and deallocate memory, while this allocator is responsible
53    /// purely for tracking the allocations and deallocations themselves.
54    pub const fn new(allocator: A) -> Self {
55        Self { allocator }
56    }
57}
58
59unsafe impl<A> GlobalAlloc for TrackingAllocator<A>
60where
61    A: GlobalAlloc,
62{
63    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
64        // Adjust the requested layout to fit our trailer and then try and allocate it.
65        let (layout, trailer_start) = get_layout_with_group_trailer(layout);
66        let layout_size = layout.size();
67        let ptr = self.allocator.alloc(layout);
68        if ptr.is_null() {
69            return ptr;
70        }
71
72        // Store the pointer to the current resource group in the trailer, and also update the statistics.
73        let trailer_ptr = trailer_ptr_from(ptr, trailer_start);
74        CURRENT_GROUP.with(|current_group| {
75            let group_ptr = current_group.borrow();
76            group_ptr.as_ref().track_allocation(layout_size);
77
78            trailer_ptr.write(group_ptr.as_ptr());
79        });
80
81        ptr
82    }
83
84    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
85        // Read the pointer to the owning resource group from the trailer and update the statistics.
86        let (layout, trailer_start) = get_layout_with_group_trailer(layout);
87        let trailer_ptr = trailer_ptr_from(ptr, trailer_start);
88        let group = (*trailer_ptr).as_ref().unwrap();
89        group.track_deallocation(layout.size());
90
91        // Deallocate the memory.
92        self.allocator.dealloc(ptr, layout);
93    }
94}
95
96/// Returns a pointer to the group trailer of the allocation based at `ptr`.
97///
98/// # Provenance
99///
100/// The trailer sits past the end of the layout the caller asked for, and the caller's pointer does not carry provenance
101/// that reaches it. `<*mut u8>::add` requires its result to stay within the allocated object, and the object the
102/// compiler believes it is offsetting within is exactly `layout.size()` bytes: the allocation shim generated for a
103/// global allocator advertises the requested size to the optimizer, not whatever larger block the allocator actually
104/// carved out. Reaching the trailer with ordinary pointer arithmetic is therefore out of bounds, and the resulting
105/// offset folds to a poison value wherever the optimizer can see both the allocation and the offset -- which it can
106/// whenever the shim is inlined into the allocating function. A poisoned pointer then fails checks it should pass,
107/// including the alignment check inserted under debug assertions.
108///
109/// Synthesizing the pointer from the address instead keeps the trailer reachable. `alloc` exposes the provenance of the
110/// whole padded block, so a pointer recovered from an address within that block carries provenance for all of it,
111/// trailer included, on both the allocation and deallocation paths.
112fn trailer_ptr_from(ptr: *mut u8, trailer_start: usize) -> *mut *mut ResourceStats {
113    ptr::with_exposed_provenance_mut(ptr.expose_provenance().wrapping_add(trailer_start))
114}
115
116fn get_layout_with_group_trailer(layout: Layout) -> (Layout, usize) {
117    let (new_layout, trailer_start) = layout.extend(STATS_LAYOUT).unwrap();
118    (new_layout.pad_to_align(), trailer_start)
119}
saluki_common/resource_tracking/allocator.rs

saluki_common/resource_tracking/
allocator.rs
