Using arena for complex and long lived allocations?

Explain
1

So I am loading quite complex configuration where I don’t want to keep full configuration source loaded due to its size and lifetime (easily months or even years), so I need to copy some strings and construct some other non-u8 slices. I am using ArrayLists for creation these slices and their size is not easily determined in advance and later I do .toOwnedSlice.

So during init/deinit of this structure I do tens of different allocations of same lifetime, so If I want to use gpa I have to keep track of every single one. This seems like the perfect use case for arena. But I am reluctant to use it because it is a bump allocator and probably will leave some bytes in the middle due to the intermediate ArrayLists reallocations. And then there will also be some bytes left unused at the end because I believe arena preallocates some size in advance. So because I can’t do something like ArrayList.shrinkToLen on arena, leaving those unused bytes for months/years feels wasteful. But using gpa just feels like the wrong tool here. Any advice on how to apporach this?

2

you could use a gpa to construct the data and move it to an arena when done, piece by piece. But that might not be better if things are too convoluted.

The more extreme option is a custom allocator…

3

I have a few suggestions, maybe one of them will work.
ArrayList + arenas is also a problem that I am looking for better solutions to, for now these are the best I got:

  1. This Is probably the best one, but will need you to write more code. If you can count the items and then create an ArrayList with that capacity, then you can append assuming capacity. This will result in no wasted memory.
  2. Accept the wasted memory
  3. If your allocator has a large buffer, then it’s likely that you hit resize, which doesn’t allocate extra memory.
  4. I have a temporary allocator in my project, that’s like a stack allocator, you could use an allocator like that to do temporary calculations, and only save with the actual arena a copy of the final data you need
  5. This is unsatisfying, but you can use a GPA for the array lists and an arena for everything else, depending on how your project is structured.
4

It depends on how your code is structured, but my first attempt at this would probably be using the gpa for arraylists in the parsing code, but every declaration is immediately followed by a defer my_list.deinit(gpa). Then I would use an arena passed into the parse function to do a .my_list = try arena.dupe(My type, my_list.items), in the return value to move everything over to the arena.

In general, I don’t find dealing with a bunch of ArrayLists used in this fashion to be difficult. If you have a bunch which immediately defer deinit, one without would look suspicious. As long you run the code in debug where each list gets at least one item, and do a safe exit, then you’ll get a leak warning if your code is bad. It’s when you start having different code paths (including error paths) either need to free or not free some of the allocated values that things actually start to get complicated.

5

For this requirement, I once tried to modify SmpAllocator into a thread-unsafe version that tracks all slabs to allow one-time deinit.

/// single-owner-at-a-time allocator

const std = @import("std");
const slab_len: usize = @max(std.heap.page_size_max, 64 * 1024);
/// Because of storing free list pointers, the minimum size class is 3.
const min_class = std.math.log2(@sizeOf(usize));
const size_class_count = std.math.log2(slab_len) - min_class;
fn sizeClassIndex(len: usize, alignment: std.mem.Alignment) usize {
    return @max(@bitSizeOf(usize) - @clz(len - 1), @intFromEnum(alignment), min_class) - min_class;
}
fn slotSize(class: usize) usize {
    return @as(usize, 1) << @intCast(class + min_class);
}
const slab_alignment = @as(std.mem.Alignment, @enumFromInt(std.math.log2(slab_len)));
const max_manager_slabs = @divExact(slab_len - @sizeOf(?*anyopaque), @sizeOf([*]align(slab_alignment.toByteUnits()) u8));

const SlabManager = extern struct {
    next: ?*SlabManager,
    slabs: [max_manager_slabs][*]align(slab_alignment.toByteUnits()) u8,
};
const LargeObjectFooterNode = struct {
    prev: ?*LargeObjectFooterNode,
    next: ?*LargeObjectFooterNode,
    footer_offset: usize,
    alignment: std.mem.Alignment,
};
fn largeObjectActualAlignment(alignment: std.mem.Alignment) std.mem.Alignment {
    return @enumFromInt(@max(@intFromEnum(alignment), @intFromEnum(std.mem.Alignment.fromByteUnits(@alignOf(LargeObjectFooterNode)))));
}
pub fn ExclusiveRecyclingArena(comptime inline_slab_slot_count: comptime_int) type {
    return struct {
        pub const State = struct {
            next_addrs: [size_class_count]usize = @splat(0),
            frees: [size_class_count]usize = @splat(0),
            inline_slab_slots: [inline_slab_slot_count][*]align(slab_alignment.toByteUnits()) u8 = undefined,
            head_slab_manager: ?*SlabManager = null,
            current_slab_manager: ?*SlabManager = null,
            allocated_slab_count: usize = 0,
            used_slab_count: usize = 0,
            large_objects: ?*LargeObjectFooterNode = null,
            pub fn handle(self: *State, backing_allocator: std.mem.Allocator) Handle {
                return .{ .state = self, .backing_allocator = backing_allocator };
            }
        };
        pub const Handle = struct {
            state: *State,
            backing_allocator: std.mem.Allocator,
            pub fn allocator(self: *Handle) std.mem.Allocator {
                return .{
                    .ptr = self,
                    .vtable = &.{
                        .alloc = alloc,
                        .resize = resize,
                        .remap = remap,
                        .free = free,
                    },
                };
            }
            fn getOrAllocSlab(self: Handle) error{OutOfMemory}![*]align(slab_alignment.toByteUnits()) u8 {
                const idx = self.state.used_slab_count;
                if (idx < inline_slab_slot_count) {
                    if (idx < self.state.allocated_slab_count) {
                        self.state.used_slab_count += 1;
                        return self.state.inline_slab_slots[idx];
                    }
                    const new_slab_ptr = self.backing_allocator.rawAlloc(slab_len, slab_alignment, @returnAddress()) orelse return error.OutOfMemory;
                    errdefer comptime unreachable;
                    self.state.inline_slab_slots[idx] = @alignCast(new_slab_ptr);
                    self.state.allocated_slab_count += 1;
                    self.state.used_slab_count += 1;
                    return @alignCast(new_slab_ptr);
                }
                const managed_slab_idx = idx - inline_slab_slot_count;
                const manager_idx = managed_slab_idx / max_manager_slabs;
                const slot_in_manager = managed_slab_idx % max_manager_slabs;

                if (idx < self.state.allocated_slab_count) {
                    if (slot_in_manager == 0 and manager_idx > 0) {
                        self.state.current_slab_manager = self.state.current_slab_manager.?.next;
                    }
                    const slab = self.state.current_slab_manager.?.slabs[slot_in_manager];
                    self.state.used_slab_count += 1;
                    return slab;
                }

                const new_slab_ptr = self.backing_allocator.rawAlloc(slab_len, slab_alignment, @returnAddress()) orelse return error.OutOfMemory;
                errdefer self.backing_allocator.rawFree(new_slab_ptr[0..slab_len], slab_alignment, @returnAddress());
                if (slot_in_manager == 0) {
                    const manager_slice_ptr = self.backing_allocator.rawAlloc(slab_len, slab_alignment, @returnAddress()) orelse return error.OutOfMemory;
                    errdefer comptime unreachable;
                    comptime std.debug.assert(@sizeOf(SlabManager) == slab_len);
                    comptime std.debug.assert(@alignOf(SlabManager) <= slab_alignment.toByteUnits());
                    const new_manager: *SlabManager = @ptrCast(@alignCast(manager_slice_ptr));
                    new_manager.next = null;

                    if (self.state.current_slab_manager) |curr| {
                        curr.next = new_manager;
                    } else {
                        self.state.head_slab_manager = new_manager;
                    }
                    self.state.current_slab_manager = new_manager;
                }
                self.state.current_slab_manager.?.slabs[slot_in_manager] = @alignCast(new_slab_ptr);

                self.state.allocated_slab_count += 1;
                self.state.used_slab_count += 1;
                return @alignCast(new_slab_ptr);
            }
            fn alloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ra: usize) ?[*]u8 {
                const self: *const Handle = @ptrCast(@alignCast(ctx));
                const class = sizeClassIndex(len, alignment);

                if (class >= size_class_count) {
                    @branchHint(.unlikely);
                    const footer_offset = std.mem.alignForward(usize, len, @alignOf(LargeObjectFooterNode));
                    const total_len = footer_offset + @sizeOf(LargeObjectFooterNode);
                    const actual_alignment = largeObjectActualAlignment(alignment);
                    const raw_ptr = self.backing_allocator.rawAlloc(total_len, actual_alignment, ra) orelse return null;
                    const footer_ptr: *LargeObjectFooterNode = @ptrCast(@alignCast(raw_ptr + footer_offset));
                    footer_ptr.* = .{
                        .prev = null,
                        .next = self.state.large_objects,
                        .footer_offset = footer_offset,
                        .alignment = alignment,
                    };
                    if (self.state.large_objects) |head| head.prev = footer_ptr;
                    self.state.large_objects = footer_ptr;
                    return raw_ptr;
                }

                const slot_size = slotSize(class);

                const top_free_ptr = self.state.frees[class];
                if (top_free_ptr != 0) {
                    @branchHint(.likely);
                    const node: *usize = @ptrFromInt(top_free_ptr);
                    self.state.frees[class] = node.*;
                    return @ptrFromInt(top_free_ptr);
                }

                const next_addr = self.state.next_addrs[class];
                if ((next_addr % slab_len) != 0) {
                    @branchHint(.likely);
                    self.state.next_addrs[class] = next_addr + slot_size;
                    return @ptrFromInt(next_addr);
                }

                const slab_ptr = self.getOrAllocSlab() catch return null;

                self.state.next_addrs[class] = @intFromPtr(slab_ptr) + slot_size;
                return slab_ptr;
            }
            fn free(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ra: usize) void {
                const self: *const Handle = @ptrCast(@alignCast(ctx));
                const class = sizeClassIndex(memory.len, alignment);

                if (class >= size_class_count) {
                    @branchHint(.unlikely);
                    const footer_offset = std.mem.alignForward(usize, memory.len, @alignOf(LargeObjectFooterNode));
                    const footer_ptr: *LargeObjectFooterNode = @ptrCast(@alignCast(memory.ptr + footer_offset));
                    std.debug.assert(footer_ptr.footer_offset == footer_offset);
                    std.debug.assert(footer_ptr.alignment == alignment);
                    if (footer_ptr.prev) |p| p.next = footer_ptr.next else self.state.large_objects = footer_ptr.next;
                    if (footer_ptr.next) |n| n.prev = footer_ptr.prev;
                    const total_len = footer_offset + @sizeOf(LargeObjectFooterNode);
                    const actual_alignment = largeObjectActualAlignment(alignment);
                    self.backing_allocator.rawFree(memory.ptr[0..total_len], actual_alignment, ra);
                    return;
                }

                const node: *usize = @ptrCast(@alignCast(memory.ptr));
                node.* = self.state.frees[class];
                self.state.frees[class] = @intFromPtr(node);
            }
            fn resize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ra: usize) bool {
                const self: *const Handle = @ptrCast(@alignCast(ctx));
                const class = sizeClassIndex(memory.len, alignment);
                const new_class = sizeClassIndex(new_len, alignment);

                if (class < size_class_count) {
                    return class == new_class;
                }

                if (new_class < size_class_count) return false;

                const old_offset = std.mem.alignForward(usize, memory.len, @alignOf(LargeObjectFooterNode));
                const new_offset = std.mem.alignForward(usize, new_len, @alignOf(LargeObjectFooterNode));

                if (old_offset == new_offset) return true;

                const old_total = old_offset + @sizeOf(LargeObjectFooterNode);
                const new_total = new_offset + @sizeOf(LargeObjectFooterNode);
                const prev, const next = blk: {
                    const old_footer: *LargeObjectFooterNode = @ptrCast(@alignCast(memory.ptr + old_offset));
                    std.debug.assert(old_footer.footer_offset == old_offset);
                    std.debug.assert(old_footer.alignment == alignment);
                    break :blk .{
                        old_footer.prev,
                        old_footer.next,
                    };
                };

                const actual_alignment = largeObjectActualAlignment(alignment);
                if (self.backing_allocator.rawResize(memory.ptr[0..old_total], actual_alignment, new_total, ra)) {

                    const new_footer: *LargeObjectFooterNode = @ptrCast(@alignCast(memory.ptr + new_offset));
                    new_footer.* = .{
                        .prev = prev,
                        .next = next,
                        .footer_offset = new_offset,
                        .alignment = alignment,
                    };

                    if (prev) |p| p.next = new_footer else self.state.large_objects = new_footer;
                    if (next) |n| n.prev = new_footer;

                    return true;
                }
                return false;
            }
            fn remap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ra: usize) ?[*]u8 {
                const self: *const Handle = @ptrCast(@alignCast(ctx));
                const class = sizeClassIndex(memory.len, alignment);
                const new_class = sizeClassIndex(new_len, alignment);

                if (class < size_class_count) {
                    return if (class == new_class) memory.ptr else null;
                }

                if (new_class < size_class_count) return null;

                const old_offset = std.mem.alignForward(usize, memory.len, @alignOf(LargeObjectFooterNode));
                const new_offset = std.mem.alignForward(usize, new_len, @alignOf(LargeObjectFooterNode));

                const old_total = old_offset + @sizeOf(LargeObjectFooterNode);
                const new_total = new_offset + @sizeOf(LargeObjectFooterNode);
                const prev, const next = blk: {
                    const old_footer: *LargeObjectFooterNode = @ptrCast(@alignCast(memory.ptr + old_offset));
                    std.debug.assert(old_footer.footer_offset == old_offset);
                    std.debug.assert(old_footer.alignment == alignment);
                    break :blk .{
                        old_footer.prev,
                        old_footer.next,
                    };
                };

                const actual_alignment = largeObjectActualAlignment(alignment);
                if (self.backing_allocator.rawRemap(memory.ptr[0..old_total], actual_alignment, new_total, ra)) |new_ptr| {

                    const new_footer: *LargeObjectFooterNode = @ptrCast(@alignCast(new_ptr + new_offset));
                    new_footer.* = .{
                        .prev = prev,
                        .next = next,
                        .footer_offset = new_offset,
                        .alignment = alignment,
                    };

                    if (prev) |p| p.next = new_footer else self.state.large_objects = new_footer;
                    if (next) |n| n.prev = new_footer;

                    return new_ptr;
                }
                return null;
            }
            pub fn deinit(self: Handle) void {
                const inline_count = @min(self.state.allocated_slab_count, inline_slab_slot_count);
                for (0..inline_count) |i| {
                    self.backing_allocator.rawFree(self.state.inline_slab_slots[i][0..slab_len], slab_alignment, @returnAddress());
                }
                var curr_manager = self.state.head_slab_manager;
                var manager_idx: usize = 0;
                while (curr_manager) |manager| {
                    const next_manager = manager.next;

                    const start_idx = inline_slab_slot_count + manager_idx * max_manager_slabs;
                    const end_idx = @min(self.state.allocated_slab_count, start_idx + max_manager_slabs);
                    const slabs_in_manager = end_idx - start_idx;

                    for (0..slabs_in_manager) |i| {
                        self.backing_allocator.rawFree(manager.slabs[i][0..slab_len], slab_alignment, @returnAddress());
                    }

                    const manager_slice = @as([*]u8, @ptrCast(manager))[0..slab_len];
                    self.backing_allocator.rawFree(manager_slice, slab_alignment, @returnAddress());

                    curr_manager = next_manager;
                    manager_idx += 1;
                }
                self.clearLarge();
            }
            pub fn reset(self: Handle) void {
                self.state.next_addrs = @splat(0);
                self.state.frees = @splat(0);
                self.state.used_slab_count = 0;
                self.state.current_slab_manager = self.state.head_slab_manager;

                self.clearLarge();
                self.state.large_objects = null;
            }
            fn clearLarge(self: Handle) void {
                var curr_large = self.state.large_objects;
                while (curr_large) |node| {
                    curr_large = node.next;
                    const total_len = node.footer_offset + @sizeOf(LargeObjectFooterNode);
                    const original_addr = @as([*]u8, @ptrCast(node)) - node.footer_offset;
                    const raw_slice = original_addr[0..total_len];
                    const alignment = largeObjectActualAlignment(node.alignment);
                    self.backing_allocator.rawFree(raw_slice, alignment, @returnAddress());
                }
            }
        };
    };
}

test "mixed small and large with multiple resets" {
    var state: ExclusiveRecyclingArena(16).State = .{};
    var handle = state.handle(std.testing.allocator);
    defer handle.deinit();

    const gpa = handle.allocator();
    var rng = std.Random.DefaultPrng.init(0x12345678);
    const random = rng.random();

    var round: usize = 0;
    while (round < 10) : (round += 1) {
        handle.reset();

        var pending8: std.ArrayListUnmanaged([]align(8) u8) = .empty;
        defer {
            for (pending8.items) |s| gpa.free(s);
            pending8.deinit(gpa);
        }

        var pending32: std.ArrayListUnmanaged([]align(32) u8) = .empty;
        defer {
            for (pending32.items) |s| gpa.free(s);
            pending32.deinit(gpa);
        }

        var allocated: usize = 0;
        const target = random.intRangeAtMost(usize, 1024, 32 * 1024);

        while (allocated < target) {
            const size = random.intRangeAtMost(usize, 8, slab_len + 2048);

            if (random.boolean()) {
                const slice = try gpa.alignedAlloc(u8, .@"8", size);
                try std.testing.expect(std.mem.Alignment.@"8".check(@intFromPtr(slice.ptr)));
                @memset(slice, @truncate(round));

                if (random.boolean() and slice.len < slab_len) {
                    gpa.free(slice);
                } else {
                    try pending8.append(gpa, slice);
                    allocated += slice.len;
                }
            } else {
                const slice = try gpa.alignedAlloc(u8, .@"32", size);
                try std.testing.expect(std.mem.Alignment.@"32".check(@intFromPtr(slice.ptr)));
                @memset(slice, @truncate(round));

                if (random.boolean() and slice.len < slab_len) {
                    gpa.free(slice);
                } else {
                    try pending32.append(gpa, slice);
                    allocated += slice.len;
                }
            }
        }

        if (random.boolean() and pending8.items.len > 0) {
            const idx = random.intRangeLessThan(usize, 0, pending8.items.len);
            gpa.free(pending8.swapRemove(idx));
        }
        if (random.boolean() and pending32.items.len > 0) {
            const idx = random.intRangeLessThan(usize, 0, pending32.items.len);
            gpa.free(pending32.swapRemove(idx));
        }
    }
}

test "alignment is respected for small objects" {
    var state: ExclusiveRecyclingArena(0).State = .{};
    var handle = state.handle(std.testing.allocator);
    defer handle.deinit();

    const gpa = handle.allocator();

    const alignments = [_]std.mem.Alignment{ .@"1", .@"2", .@"4", .@"8", .@"16", .@"32", .@"64" };

    inline for (alignments) |alignment| {
        const size = 48;
        const slice = try gpa.alignedAlloc(u8, alignment, size);
        try std.testing.expect(alignment.check(@intFromPtr(slice.ptr)));
        try std.testing.expect(slice.len == size);
        gpa.free(slice);
    }
}

test "large object with weak alignment" {
    var state: ExclusiveRecyclingArena(0).State = .{};
    var handle = state.handle(std.testing.allocator);
    defer handle.deinit();

    const gpa = handle.allocator();

    const weak_align: std.mem.Alignment = .@"1";
    const big_len = slab_len + 512;

    const ptr = try gpa.alignedAlloc(u8, weak_align, big_len);
    defer gpa.free(ptr);

    try std.testing.expect(weak_align.check(@intFromPtr(ptr.ptr)));
    try std.testing.expect(std.mem.Alignment.fromByteUnits(@alignOf(LargeObjectFooterNode)).check(@intFromPtr(ptr.ptr)));
}

test "resize small object" {
    var state: ExclusiveRecyclingArena(0).State = .{};
    var handle = state.handle(std.testing.allocator);
    defer handle.deinit();

    const gpa = handle.allocator();

    var buf = try gpa.alloc(u8, 32);
    try std.testing.expect(gpa.resize(buf, 24));
    buf = buf.ptr[0..24];

    try std.testing.expect(!gpa.resize(buf, 40));
}

test "resize large object" {
    var state: ExclusiveRecyclingArena(0).State = .{};
    var handle = state.handle(std.testing.allocator);
    defer handle.deinit();

    const gpa = handle.allocator();

    const orig_len = slab_len + 256;
    var buf = try gpa.alloc(u8, orig_len);
    @memset(buf, 0xCD);

    const new_len = orig_len - 64;
    if (gpa.resize(buf, new_len)) {
        buf = buf.ptr[0..new_len];
        try std.testing.expect(buf[0] == 0xCD);
    }

    if (gpa.resize(buf, orig_len + 128)) {
        buf = buf.ptr[0 .. orig_len + 128];
    }
    gpa.free(buf);
}

test "remap large object may move" {
    var state: ExclusiveRecyclingArena(0).State = .{};
    var handle = state.handle(std.testing.allocator);
    defer handle.deinit();

    const gpa = handle.allocator();

    const orig_len = slab_len + 128;
    const buf = try gpa.alloc(u8, orig_len);
    @memset(buf, 0xEF);

    if (gpa.remap(buf, orig_len + 512)) |new_buf| {
        try std.testing.expect(new_buf[0] == 0xEF);
        gpa.free(new_buf);
    } else {
        gpa.free(buf);
    }
}