自定义Rust内存分配器追踪异常：原子计数与输出记录不匹配

为追踪进程内的内存分配情况，基于Rust实现了自定义GlobalAlloc分配器，通过原子计数器统计分配和释放次数，同时将操作记录写入指定缓冲区。但运行后发现原子计数器显示8次分配、1次释放，而缓冲区仅输出5次分配记录，二者数据不匹配。

实现代码 #

#![feature(allocator_api)]

use std::alloc::{Layout, System};
use std::io::Write;
use std::mem::MaybeUninit;
use std::sync::{
    atomic::{AtomicU32, Ordering},
    Mutex,
};

static ALLOC_COUNTER: AtomicU32 = AtomicU32::new(0);
static DEALLOC_COUNTER: AtomicU32 = AtomicU32::new(0);

struct TrackingAllocator<W: Write>(MaybeUninit<Mutex<W>>);

unsafe impl<W: Write> std::alloc::GlobalAlloc for TrackingAllocator<W> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let mut guard = self.0.assume_init_ref().lock().unwrap();
        writeln!(guard, "alloc: {layout:?}").unwrap();
        ALLOC_COUNTER.fetch_add(1, Ordering::SeqCst);
        std::alloc::System.alloc(layout)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let mut guard = self.0.assume_init_ref().lock().unwrap();
        writeln!(guard, "dealloc: {ptr:?} {layout:?}").unwrap();
        DEALLOC_COUNTER.fetch_add(1, Ordering::SeqCst);
        std::alloc::System.dealloc(ptr, layout)
    }
}

static SYSTEM: System = System;

// 使用原生系统分配器存储追踪日志的缓冲区
#[global_allocator]
static mut GLOBAL: TrackingAllocator<Vec<u8, System>> = TrackingAllocator(MaybeUninit::uninit());

fn main() {
    // 初始化分配器的日志目标
    unsafe {
        GLOBAL
            .0
            .write(Mutex::new(Vec::<u8, System>::new_in(SYSTEM)));
    }

    // 执行一些测试分配操作
    let x = vec![1, 2, 3];
    let _y = x
        .into_iter()
        .map(Vec::<u8>::with_capacity)
        .collect::<Vec<_>>();

    unsafe {
        GLOBAL.0.assume_init_ref().lock().unwrap().flush().unwrap();

        let guard = GLOBAL.0.assume_init_ref().lock().unwrap().clone();
        let s = std::str::from_utf8(&guard).unwrap();
        println!("{s}");
        println!(
            "{} | {}",
            ALLOC_COUNTER.load(Ordering::SeqCst),
            DEALLOC_COUNTER.load(Ordering::SeqCst)
        );

        GLOBAL.0.assume_init_drop();
    }
}

运行输出 #

Compiling test-tracing-allocator v0.1.0 (/home/ec2-user/testing-tracing-allocator)
    Finished dev [unoptimized + debuginfo] target(s) in 0.22s
     Running `target/debug/test-tracing-allocator`
alloc: Layout { size: 24, align: 8 (1 << 3) }
alloc: Layout { size: 72, align: 8 (1 << 3) }
alloc: Layout { size: 1, align: 1 (1 << 0) }
alloc: Layout { size: 2, align: 1 (1 << 0) }
alloc: Layout { size: 3, align: 1 (1 << 0) }
dealloc: 0x55e1c23bac10 Layout { size: 24, align: 8 (1 << 3) }

8 | 1

问题根源 #

1. 早期分配未被记录：Rust程序进入main前，运行时会执行初始化操作（如全局数据结构设置、静态变量初始化），这些操作会触发内存分配，但此时全局分配器内部的Mutex和日志缓冲区还未初始化（初始化逻辑在main中执行），调用assume_init_ref()会触发未定义行为，导致这部分分配的日志无法写入，但计数器仍正常累加。
2. 未定义行为导致日志丢失：分配器内部状态未初始化时调用其方法，程序行为不可预测，部分分配操作的日志写入逻辑可能被跳过，而计数器更新逻辑仍执行，最终造成统计数据与日志不一致。

修复方案 #

使用std::sync::Once确保分配器内部状态在第一次分配操作发生前完成初始化，覆盖main之前的运行时分配，消除未定义行为风险。

修改后的代码 #

#![feature(allocator_api)]

use std::alloc::{Layout, System};
use std::io::Write;
use std::sync::{
    atomic::{AtomicU32, Ordering},
    Mutex, MutexGuard, Once,
};

static ALLOC_COUNTER: AtomicU32 = AtomicU32::new(0);
static DEALLOC_COUNTER: AtomicU32 = AtomicU32::new(0);

struct TrackingAllocator<W: Write + 'static>(Once<Mutex<W>>);

impl<W: Write + 'static> TrackingAllocator<W> {
    // 创建未初始化的追踪分配器
    const fn uninit() -> Self {
        TrackingAllocator(Once::new())
    }

    // 初始化或获取日志写入器的锁
    fn get_writer(&self, init: impl FnOnce() -> W) -> MutexGuard<'_, W> {
        self.0.call_once(|| Mutex::new(init()));
        self.0.get_mut().unwrap().lock().unwrap()
    }
}

unsafe impl<W: Write + 'static> std::alloc::GlobalAlloc for TrackingAllocator<W> {
    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
        let mut guard = self.get_writer(|| Vec::new_in(System));
        writeln!(guard, "alloc: {layout:?}").unwrap();
        ALLOC_COUNTER.fetch_add(1, Ordering::SeqCst);
        System.alloc(layout)
    }

    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
        let mut guard = self.get_writer(|| Vec::new_in(System));
        writeln!(guard, "dealloc: {ptr:?} {layout:?}").unwrap();
        DEALLOC_COUNTER.fetch_add(1, Ordering::SeqCst);
        System.dealloc(ptr, layout)
    }
}

// 全局分配器静态初始化
#[global_allocator]
static GLOBAL: TrackingAllocator<Vec<u8, System>> = TrackingAllocator::uninit();

fn main() {
    // 执行测试分配操作
    let x = vec![1, 2, 3];
    let _y = x
        .into_iter()
        .map(Vec::<u8>::with_capacity)
        .collect::<Vec<_>>();

    // 打印日志与统计数据
    let guard = GLOBAL.get_writer(|| Vec::new_in(System));
    let s = std::str::from_utf8(&guard).unwrap();
    println!("{s}");
    println!(
        "{} | {}",
        ALLOC_COUNTER.load(Ordering::SeqCst),
        DEALLOC_COUNTER.load(Ordering::SeqCst)
    );
}

修改说明 #

• 用Once替代MaybeUninit，确保内部Mutex<W>在第一次alloc/dealloc调用时完成初始化，覆盖main之前的运行时分配。
• 新增get_writer方法统一处理初始化与锁的获取，避免重复代码，消除未定义行为。
• 移除手动的assume_init相关操作，确保代码安全。

修改后所有分配操作（包括main之前的早期分配）都会被正确记录，计数器与日志数量将保持一致。

内容的提问来源于stack exchange，提问作者Jonathan Woollett-light