Zig 自定义分配器：从 Allocator 接口到内存池与组合

pub const Allocator = struct {
    ptr: *anyopaque,
    vtable: *const VTable,

    pub const VTable = struct {
        alloc: *const fn (ctx: *anyopaque, len: usize, alignment: Alignment, ret_addr: usize) ?[*]u8,
        resize: *const fn (ctx: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) bool,
        remap: *const fn (ctx: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) ?[*]u8,
        free: *const fn (ctx: *anyopaque, memory: []u8, alignment: Alignment, ret_addr: usize) void,
    };
};

const std = @import("std");

fn processRequest(gpa: std.mem.Allocator, payload: []const u8) ![]u8 {
    var arena = std.heap.ArenaAllocator.init(gpa);
    defer arena.deinit();                 // 一次性释放本次请求的全部分配
    const a = arena.allocator();
    const tokens = try splitTokens(a, payload);   // 无需逐个 free
    const tree = try parseIntoTree(a, tokens);
    return try render(a, tree);                   // 返回的切片仍属 arena
}

const std = @import("std");

fn handlePacket(bytes: []const u8) !void {
    var buf: [4096]u8 = undefined;               // 栈上 4 KB
    var fba = std.heap.FixedBufferAllocator.init(&buf);
    const a = fba.allocator();

    const header = try a.create(Header);
    const body = try a.alloc(u8, bytes.len);
    @memcpy(body, bytes);
    _ = header;                                  // 无需 free：随栈帧消失
}

var gpa = std.heap.GeneralPurposeAllocator(.{
    .safety = true,          // 越界/use-after-free 检测（Debug 默认开）
    .never_unmap = true,     // free 后保留页映射，更容易抓到悬空指针
    .retain_metadata = true, // 保留元数据，支持重复释放检测
    .verbose_log = true,     // 打印每次分配/释放的调用栈
}){};
defer _ = gpa.deinit();      // 返回 .leak / .ok，泄漏状态一目了然
const allocator = gpa.allocator();

const std = @import("std");

const BumpAllocator = struct {
    buffer: []u8,
    offset: usize = 0,

    pub fn init(buffer: []u8) BumpAllocator {
        return .{ .buffer = buffer };
    }

    pub fn allocator(self: *BumpAllocator) std.mem.Allocator {
        return .{ .ptr = self, .vtable = &.{
            .alloc = alloc, .resize = resize, .remap = remap, .free = free,
        } };
    }

    fn alloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
        _ = ret_addr;
        const self: *BumpAllocator = @ptrCast(@alignCast(ctx));
        const base = @intFromPtr(self.buffer.ptr);
        const aligned = std.mem.alignForward(usize, base + self.offset, alignment.toByteUnits());
        if (aligned + len > base + self.buffer.len) return null;   // 溢出
        self.offset = aligned + len - base;
        return @ptrFromInt(aligned);
    }

    fn resize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
        _ = .{ ctx, memory, alignment, new_len, ret_addr };
        return false;   // 不支持就地扩容
    }

    fn remap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
        const self: *BumpAllocator = @ptrCast(@alignCast(ctx));
        const base = @intFromPtr(self.buffer.ptr);
        const start = @intFromPtr(memory.ptr);
        // 若是最后一个分配，直接前移游标即可，无需拷贝
        if (start + memory.len == base + self.offset and start + new_len <= base + self.buffer.len) {
            self.offset = start + new_len - base;
            return memory.ptr;
        }
        return alloc(ctx, new_len, alignment, ret_addr);
    }

    fn free(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
        _ = .{ ctx, memory, alignment, ret_addr };   // 空操作：整体重置
    }

    pub fn reset(self: *BumpAllocator) void {
        self.offset = 0;
    }
};

const logged = std.heap.loggingAllocator(gpa_allocator);          // 记录每次分配/释放
var failing = std.testing.FailingAllocator.init(gpa_allocator, .{});  // 失败注入

test "cleanup on OOM" {
    var failing = std.testing.FailingAllocator.init(std.testing.allocator, .{});
    for (0..32) |i| {
        failing.fail_index = i;                    // 逐个注入失败点
        var list = std.ArrayList(u8).init(failing.allocator());
        defer list.deinit();
        list.appendSlice("data") catch continue;
    }
}

const Counting = struct {
    child: std.mem.Allocator,
    total_bytes: usize = 0,
    alloc_count: usize = 0,

    fn alloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
        const self: *Counting = @ptrCast(@alignCast(ctx));
        const r = self.child.rawAlloc(len, alignment, ret_addr);   // 转发到 raw 入口
        if (r != null) {
            self.total_bytes += len;
            self.alloc_count += 1;
        }
        return r;
    }
    // resize/remap/free 同样转发给 child.rawResize/rawRemap/rawFree
};

pub fn serve(conn: std.net.Stream, gpa: std.mem.Allocator) !void {
    var arena = std.heap.ArenaAllocator.init(gpa);
    defer arena.deinit();                                   // 请求级内存
    var pool = try std.heap.MemoryPool(Header).init(gpa);   // 定长对象池
    defer pool.deinit();

    const req = try readRequest(arena.allocator(), conn);   // 变长：用 arena
    const hdr = try pool.create();                          // 定长：用池
    defer pool.destroy(hdr);
    hdr.* = req.header;
    try handle(arena.allocator(), hdr, conn);
}