tracking_allocator/

allocator.rs

use std::alloc::{handle_alloc_error, GlobalAlloc, Layout, System};

use crate::token::try_with_suspended_allocation_group;
use crate::{get_global_tracker, AllocationGroupId};

/// Tracking allocator implementation.
///
/// This allocator must be installed via `#[global_allocator]` in order to take effect.  More
/// information on using this allocator can be found in the examples, or directly in the standard
/// library docs for [`GlobalAlloc`].
pub struct Allocator<A> {
    inner: A,
}

impl<A> Allocator<A> {
    /// Creates a new `Allocator` that wraps another allocator.
    #[must_use]
    pub const fn from_allocator(allocator: A) -> Self {
        Self { inner: allocator }
    }
}

impl Allocator<System> {
    /// Creates a new `Allocator` that wraps the system allocator.
    #[must_use]
    pub const fn system() -> Allocator<System> {
        Self::from_allocator(System)
    }
}

impl<A: GlobalAlloc> Allocator<A> {
    unsafe fn get_wrapped_allocation(
        &self,
        object_layout: Layout,
    ) -> (*mut usize, *mut u8, Layout) {
        // Allocate our wrapped layout and make sure the allocation succeeded.
        let (actual_layout, offset_to_object) = get_wrapped_layout(object_layout);
        let actual_ptr = self.inner.alloc(actual_layout);
        if actual_ptr.is_null() {
            handle_alloc_error(actual_layout);
        }

        // Zero out the group ID field to make sure it's in the `None` state.
        //
        // SAFETY: We know that `actual_ptr` is at least aligned enough for casting it to `*mut usize` as the layout for
        // the allocation backing this pointer ensures the first field in the layout is `usize.
        #[allow(clippy::cast_ptr_alignment)]
        let group_id_ptr = actual_ptr.cast::<usize>();
        group_id_ptr.write(0);

        // SAFETY: If the allocation succeeded and `actual_ptr` is valid, then it must be valid to advance by
        // `offset_to_object` as it would land within the allocation.
        let object_ptr = actual_ptr.wrapping_add(offset_to_object);

        (group_id_ptr, object_ptr, actual_layout)
    }
}

impl Default for Allocator<System> {
    fn default() -> Self {
        Self::from_allocator(System)
    }
}

unsafe impl<A: GlobalAlloc> GlobalAlloc for Allocator<A> {
    #[track_caller]
    unsafe fn alloc(&self, object_layout: Layout) -> *mut u8 {
        let (group_id_ptr, object_ptr, wrapped_layout) = self.get_wrapped_allocation(object_layout);
        let object_addr = object_ptr as usize;
        let object_size = object_layout.size();
        let wrapped_size = wrapped_layout.size();

        if let Some(tracker) = get_global_tracker() {
            try_with_suspended_allocation_group(
                #[inline(always)]
                |group_id| {
                    // We only set the group ID in the wrapper header if we're tracking an allocation, because when it
                    // comes back to us during deallocation, we want to skip doing any checks at all if it's already
                    // zero.
                    //
                    // If we never track the allocation, tracking the deallocation will only produce incorrect numbers,
                    // and that includes even if we just used the rule of "always attribute allocations to the root
                    // allocation group by default".
                    group_id_ptr.write(group_id.as_usize().get());
                    tracker.allocated(object_addr, object_size, wrapped_size, group_id);
                },
            );
        }

        object_ptr
    }

    #[track_caller]
    unsafe fn dealloc(&self, object_ptr: *mut u8, object_layout: Layout) {
        // Regenerate the wrapped layout so we know where we have to look, as the pointer we've given relates to the
        // requested layout, not the wrapped layout that was actually allocated.
        let (wrapped_layout, offset_to_object) = get_wrapped_layout(object_layout);

        // SAFETY: We only ever return pointers to the actual requested object layout, not our wrapped layout. Since
        // global allocators cannot be changed at runtime, we know that if we're here, then the given pointer, and the
        // allocation it refers to, was allocated by us. Thus, since we wrap _all_ allocations, we know that this object
        // pointer can be safely subtracted by `offset_to_object` to get back to the group ID field in our wrapper.
        let actual_ptr = object_ptr.wrapping_sub(offset_to_object);

        // SAFETY: We know that `actual_ptr` is at least aligned enough for casting it to `*mut usize` as the layout for
        // the allocation backing this pointer ensures the first field in the layout is `usize.
        #[allow(clippy::cast_ptr_alignment)]
        let raw_group_id = actual_ptr.cast::<usize>().read();

        // Deallocate before tracking, just to make sure we're reclaiming memory as soon as possible.
        self.inner.dealloc(actual_ptr, wrapped_layout);

        let object_addr = object_ptr as usize;
        let object_size = object_layout.size();
        let wrapped_size = wrapped_layout.size();

        if let Some(tracker) = get_global_tracker() {
            if let Some(source_group_id) = AllocationGroupId::from_raw(raw_group_id) {
                try_with_suspended_allocation_group(
                    #[inline(always)]
                    |current_group_id| {
                        tracker.deallocated(
                            object_addr,
                            object_size,
                            wrapped_size,
                            source_group_id,
                            current_group_id,
                        );
                    },
                );
            }
        }
    }
}

fn get_wrapped_layout(object_layout: Layout) -> (Layout, usize) {
    static HEADER_LAYOUT: Layout = Layout::new::<usize>();

    // We generate a new allocation layout that gives us a location to store the active allocation group ID ahead
    // of the requested allocation, which lets us always attempt to retrieve it on the deallocation path. We'll
    // always set this to zero, and conditionally update it to the actual allocation group ID if tracking is enabled.
    let (actual_layout, offset_to_object) = HEADER_LAYOUT
        .extend(object_layout)
        .expect("wrapping requested layout resulted in overflow");
    let actual_layout = actual_layout.pad_to_align();

    (actual_layout, offset_to_object)
}