Rust嵌套线性分配器引发堆损坏的原因及解决方法

试验Rust的allocator_api特性时，实现了一个简单的线性分配器，单例测试正常，但嵌套分配器时抛出STATUS_HEAP_CORRUPTION错误。

测试代码：

let linear1 = LinearAllocator::with_capacity(32);
let linear2 = LinearAllocator::with_capacity_in(4, linear1.clone());

assert!(linear1.allocate(Layout::new::<u32>()).is_ok());
assert!(linear2.allocate(Layout::new::<u32>()).is_ok());

错误输出：

error: test failed, to rerun pass `-p allocators --bin allocators`

Caused by:
  process didn't exit successfully: `D:\Dev\allocators\target\debug\deps\allocators-a76735fe34fe79e2.exe test::linear_alloc_nested_with_linear_alloc --exact --nocapture` (exit code: 0xc0000374, STATUS_HEAP_CORRUPTION)

调试发现，deallocate时程序尝试从Global分配器而非内部LinearAllocator释放内存，且Drop trait中的deallocate调用是问题根源——移除该调用测试可通过，但需要保留释放逻辑以支持嵌套场景。

完整LinearAllocator代码：

#[derive(Clone, Debug)]
pub struct LinearAllocator<A: Allocator = Global> {
    size: usize,
    allocated: Arc<AtomicUsize>,
    data: *mut u8,
    layout: Layout,
    alloc: A,
}

impl Default for LinearAllocator<Global> {
    fn default() -> Self {
        Self {
            size: 0,
            allocated: Arc::new(AtomicUsize::new(0)),
            data: std::ptr::null_mut(),
            layout: Layout::new::<u8>(),
            alloc: Global,
        }
    }
}

impl LinearAllocator<Global> {
    pub fn with_capacity(size: usize) -> Self {
        Self::with_capacity_in(size, Global)
    }
}

impl<A: Allocator> LinearAllocator<A> {
    pub fn with_capacity_in(size: usize, alloc: A) -> Self {
        let layout = Layout::array::<u8>(size).unwrap();
        let data = unsafe { alloc.allocate(layout).unwrap().as_mut() as *mut [u8] as *mut u8 };

        Self {
            size,
            allocated: Arc::new(AtomicUsize::new(0)),
            data,
            layout,
            alloc,
        }
    }

    pub fn reset(&self) {
        self.allocated.store(0, Ordering::Relaxed);
    }
}

unsafe impl<A: Allocator + Debug> Allocator for LinearAllocator<A> {
    fn allocate(&self, layout: Layout) -> Result<std::ptr::NonNull<[u8]>, AllocError> {
        let size = layout.size();
        let align = layout.align();

        let remainder = size % align;
        let aligned_size = if remainder == 0 {
            size
        } else {
            size + align - remainder
        };

        let mut current = self.allocated.load(Ordering::Relaxed);
        loop {
            if self.size - current < aligned_size {
                return Err(AllocError);
            }

            let new = current + aligned_size;

            if let Err(actual) =
                self.allocated
                    .compare_exchange(current, new, Ordering::Relaxed, Ordering::Relaxed)
            {
                current = actual;
            } else {
                break;
            }
        }

        // Safety: We previously checked that there is enough space in data.
        let data = unsafe { self.data.add(current) };
        let ptr = std::ptr::slice_from_raw_parts_mut(data, aligned_size);

        std::ptr::NonNull::new(ptr).ok_or(AllocError)
    }

    unsafe fn deallocate(&self, _ptr: std::ptr::NonNull<u8>, _layout: Layout) {
        // Linear allocator does not deallocate.
        dbg!(&self);
    }
}

impl<A: Allocator> Drop for LinearAllocator<A> {
    fn drop(&mut self) {
        self.reset();

        if let Some(ptr) = std::ptr::NonNull::new(self.data) {
            unsafe { A::deallocate(&self.alloc, ptr, self.layout) };
        }
    }
}

1. 重复释放内存：自动派生的Clone会让多个LinearAllocator实例共享同一data指针和底层分配器。当第一个实例销毁时，Drop会释放data指向的内存；后续第二个实例销毁时，会再次尝试释放同一块已被释放的内存，直接触发堆损坏。
2. 分配器调用逻辑错误：Drop中使用A::deallocate(&self.alloc, ...)是静态调用方式，当A是LinearAllocator（嵌套场景）时，该调用无法正确触发底层分配器的动态分发逻辑，导致释放路径混乱。

针对上述问题，修改代码如下：

1. 重写Clone实现，避免内存共享 #

将自动派生的Clone改为让每个克隆实例独立从底层分配器申请内存，避免多实例共享同一块内存：

impl<A: Allocator + Clone> Clone for LinearAllocator<A> {
    fn clone(&self) -> Self {
        Self::with_capacity_in(self.size, self.alloc.clone())
    }
}

2. 修正Drop中的释放逻辑 #

使用实例方法调用self.alloc.deallocate(...)替代静态调用，确保符合Allocator trait的语义，同时重置指针避免悬垂：

impl<A: Allocator> Drop for LinearAllocator<A> {
    fn drop(&mut self) {
        self.reset();

        if let Some(ptr) = std::ptr::NonNull::new(self.data) {
            unsafe { self.alloc.deallocate(ptr.cast(), self.layout) };
            self.data = std::ptr::null_mut();
        }
    }
}

3. 明确线性分配器的deallocate语义 #

保留空实现但补充注释，明确线性分配器不支持单独释放，仅在Drop时释放整块内存：

unsafe fn deallocate(&self, _ptr: std::ptr::NonNull<u8>, _layout: Layout) {
    // 线性分配器不支持单独释放内存，仅支持整体reset或Drop时释放整块内存
}

修改后的完整代码 #

use std::alloc::{Allocator, AllocError, Layout};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::Arc;

#[derive(Debug)]
pub struct LinearAllocator<A: Allocator = Global> {
    size: usize,
    allocated: Arc<AtomicUsize>,
    data: *mut u8,
    layout: Layout,
    alloc: A,
}

impl<A: Allocator + Clone> Clone for LinearAllocator<A> {
    fn clone(&self) -> Self {
        Self::with_capacity_in(self.size, self.alloc.clone())
    }
}

impl Default for LinearAllocator<Global> {
    fn default() -> Self {
        Self {
            size: 0,
            allocated: Arc::new(AtomicUsize::new(0)),
            data: std::ptr::null_mut(),
            layout: Layout::new::<u8>(),
            alloc: Global,
        }
    }
}

impl LinearAllocator<Global> {
    pub fn with_capacity(size: usize) -> Self {
        Self::with_capacity_in(size, Global)
    }
}

impl<A: Allocator> LinearAllocator<A> {
    pub fn with_capacity_in(size: usize, alloc: A) -> Self {
        let layout = Layout::array::<u8>(size).unwrap();
        let data = unsafe { alloc.allocate(layout).unwrap().as_mut() as *mut [u8] as *mut u8 };

        Self {
            size,
            allocated: Arc::new(AtomicUsize::new(0)),
            data,
            layout,
            alloc,
        }
    }

    pub fn reset(&self) {
        self.allocated.store(0, Ordering::Relaxed);
    }
}

unsafe impl<A: Allocator + Debug> Allocator for LinearAllocator<A> {
    fn allocate(&self, layout: Layout) -> Result<std::ptr::NonNull<[u8]>, AllocError> {
        let size = layout.size();
        let align = layout.align();

        let remainder = size % align;
        let aligned_size = if remainder == 0 {
            size
        } else {
            size + align - remainder
        };

        let mut current = self.allocated.load(Ordering::Relaxed);
        loop {
            if self.size - current < aligned_size {
                return Err(AllocError);
            }

            let new = current + aligned_size;

            if let Err(actual) =
                self.allocated
                    .compare_exchange(current, new, Ordering::Relaxed, Ordering::Relaxed)
            {
                current = actual;
            } else {
                break;
            }
        }

        // Safety: We previously checked that there is enough space in data.
        let data = unsafe { self.data.add(current) };
        let ptr = std::ptr::slice_from_raw_parts_mut(data, aligned_size);

        std::ptr::NonNull::new(ptr).ok_or(AllocError)
    }

    unsafe fn deallocate(&self, _ptr: std::ptr::NonNull<u8>, _layout: Layout) {
        // 线性分配器不支持单独释放内存，仅支持整体reset或Drop时释放整块内存
    }
}

impl<A: Allocator> Drop for LinearAllocator<A> {
    fn drop(&mut self) {
        self.reset();

        if let Some(ptr) = std::ptr::NonNull::new(self.data) {
            unsafe { self.alloc.deallocate(ptr.cast(), self.layout) };
            self.data = std::ptr::null_mut();
        }
    }
}

内容的提问来源于stack exchange，提问作者junglie85