Skip to content

MemoryPool: Allocate only once for preheating & impl. proper resetting. #25429

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 2 commits into
base: master
Choose a base branch
from
+75 −36
Open

MemoryPool: Allocate only once for preheating & impl. proper resetting. #25429

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
617cd89
Allocate only once for preheating.
two-horned Oct 1, 2025
728b54c
Implement TODO: Resets for mempools.
two-horned Oct 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
111 changes: 75 additions & 36 deletions lib/std/heap/memory_pool.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,22 +39,24 @@ pub fn MemoryPoolExtra(comptime Item: type, comptime pool_options: Options) type
return struct {
const Pool = @This();

/// Size of the memory pool items. This is not necessarily the same
/// as `@sizeOf(Item)` as the pool also uses the items for internal means.
/// Size of the memory pool items. This may be larger than `@sizeOf(Item)`
/// as the pool also uses the items for internal means.
pub const item_size = @max(@sizeOf(Node), @sizeOf(Item));

// This needs to be kept in sync with Node.
const node_alignment: Alignment = .of(*anyopaque);

/// Alignment of the memory pool items. This is not necessarily the same
/// as `@alignOf(Item)` as the pool also uses the items for internal means.
/// Alignment of the memory pool items. This may be larger than `@alignOf(Item)`
/// as the pool also uses the items for internal means.
pub const item_alignment: Alignment = node_alignment.max(pool_options.alignment orelse .of(Item));

const Node = struct {
next: ?*align(item_alignment.toByteUnits()) @This(),
};
const NodePtr = *align(item_alignment.toByteUnits()) Node;
const ItemPtr = *align(item_alignment.toByteUnits()) Item;
const Node = struct { next: ?*align(unit_al_bytes) @This() };
const Byte = std.meta.Int(.unsigned, std.mem.byte_size_in_bits);
const Unit = [item_alignment.forward(item_size)]Byte;
const unit_al_bytes = item_alignment.toByteUnits();

const ItemPtr = *align(unit_al_bytes) Item;
const NodePtr = *align(unit_al_bytes) Node;

arena: std.heap.ArenaAllocator,
free_list: ?NodePtr = null,
Expand Down Expand Up @@ -84,18 +86,29 @@ pub fn MemoryPoolExtra(comptime Item: type, comptime pool_options: Options) type
/// This allows up to `size` active allocations before an
/// `OutOfMemory` error might happen when calling `create()`.
pub fn preheat(pool: *Pool, size: usize) MemoryPoolError!void {
var i: usize = 0;
while (i < size) : (i += 1) {
const raw_mem = try pool.allocNew();
const free_node = @as(NodePtr, @ptrCast(raw_mem));
free_node.* = Node{
.next = pool.free_list,
};
const raw_mem = try pool.allocNew(size);
const uni_slc = raw_mem[0..size];
for (uni_slc) |*unit| {
const free_node: NodePtr = @ptrCast(unit);
free_node.next = pool.free_list;
pool.free_list = free_node;
}
}

pub const ResetMode = std.heap.ArenaAllocator.ResetMode;
pub const ResetMode = union(enum) {
/// Releases all allocated memory in the arena.
free_all,
/// This will pre-heat the memory pool for future allocations by allocating a
/// large enough buffer to accomodate the highest amount of actively allocated items
/// in the past. Preheating will speed up the allocation process by invoking the
/// backing allocator less often than before. If `reset()` is used in a loop, this
/// means if the highest amount of actively allocated items is never being surpassed,
/// no memory allocations are performed anymore.
retain_capacity,
/// This is the same as `retain_capacity`, but the memory will be shrunk to
/// only hold at most this value of items.
retain_with_limit: usize,
};

/// Resets the memory pool and destroys all allocated items.
/// This can be used to batch-destroy all objects without invalidating the memory pool.
Expand All @@ -107,27 +120,34 @@ pub fn MemoryPoolExtra(comptime Item: type, comptime pool_options: Options) type
///
/// NOTE: If `mode` is `free_all`, the function will always return `true`.
pub fn reset(pool: *Pool, mode: ResetMode) bool {
// TODO: Potentially store all allocated objects in a list as well, allowing to
// just move them into the free list instead of actually releasing the memory.

const reset_successful = pool.arena.reset(mode);

const ArenaResetMode = std.heap.ArenaAllocator.ResetMode;
const arena_mode = switch (mode) {
.free_all => .free_all,
.retain_capacity => .retain_capacity,
.retain_with_limit => |limit| ArenaResetMode{ .retain_with_limit = limit * item_size },
};
pool.free_list = null;

return reset_successful;
if (!pool.arena.reset(arena_mode)) return false;
// When the backing arena allocator is being reset to
// a capacity greater than 0, then its internals consists
// of a *single* buffer node of said capacity. This means,
// we can safely pre-heat without causing additional allocations.
const arena_capacity = pool.arena.queryCapacity() / item_size;
if (arena_capacity != 0) pool.preheat(arena_capacity) catch unreachable;
return true;
}

/// Creates a new item and adds it to the memory pool.
pub fn create(pool: *Pool) !ItemPtr {
const node = if (pool.free_list) |item| blk: {
const node_ptr: NodePtr = if (pool.free_list) |item| blk: {
pool.free_list = item.next;
break :blk item;
} else if (pool_options.growable)
@as(NodePtr, @ptrCast(try pool.allocNew()))
@ptrCast(try pool.allocNew(1))
else
return error.OutOfMemory;

const ptr = @as(ItemPtr, @ptrCast(node));
const ptr: ItemPtr = @ptrCast(node_ptr);
ptr.* = undefined;
return ptr;
}
Expand All @@ -136,17 +156,14 @@ pub fn MemoryPoolExtra(comptime Item: type, comptime pool_options: Options) type
/// Only pass items to `ptr` that were previously created with `create()` of the same memory pool!
pub fn destroy(pool: *Pool, ptr: ItemPtr) void {
ptr.* = undefined;

const node = @as(NodePtr, @ptrCast(ptr));
node.* = Node{
.next = pool.free_list,
};
pool.free_list = node;
const node_ptr: NodePtr = @ptrCast(ptr);
node_ptr.next = pool.free_list;
pool.free_list = node_ptr;
}

fn allocNew(pool: *Pool) MemoryPoolError!*align(item_alignment.toByteUnits()) [item_size]u8 {
const mem = try pool.arena.allocator().alignedAlloc(u8, item_alignment, item_size);
return mem[0..item_size]; // coerce slice to array pointer
fn allocNew(pool: *Pool, n: usize) MemoryPoolError![*]align(unit_al_bytes) Unit {
const mem = try pool.arena.allocator().alignedAlloc(Unit, item_alignment, n);
return mem.ptr;
}
};
}
Expand Down Expand Up @@ -220,3 +237,25 @@ test "greater than pointer manual alignment" {
const foo: *align(16) Foo = try pool.create();
_ = foo;
}

test "reset" {
const pool_extra = MemoryPoolExtra(u32, .{ .growable = false });
var pool = try pool_extra.initPreheated(std.testing.allocator, 3);
defer pool.deinit();

try std.testing.expect(pool.create() != error.OutOfMemory);
try std.testing.expect(pool.create() != error.OutOfMemory);
try std.testing.expect(pool.create() != error.OutOfMemory);
try std.testing.expect(pool.create() == error.OutOfMemory);

try std.testing.expect(pool.reset(.{ .retain_with_limit = 2 }));

try std.testing.expect(pool.create() != error.OutOfMemory);
try std.testing.expect(pool.create() != error.OutOfMemory);
try std.testing.expect(pool.create() == error.OutOfMemory);

try std.testing.expect(pool.reset(.{ .retain_with_limit = 1 }));

try std.testing.expect(pool.create() != error.OutOfMemory);
try std.testing.expect(pool.create() == error.OutOfMemory);
}
Loading