authorgravatar for justus@klausecker.deJustus Klausecker <justus@klausecker.de> 2026-02-15 17:05:00+01:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2026-02-25 19:12:35+01:00
loga3a9dc111da7a8455f9a476782f291f349d0ed8d
tree32e005d5fdc079a979068d61041cc77a85369951
parent2f8e66080580c1ef3814cc4e2f27ae418d8d6e0c

std.heap.ArenaAllocator: make it threadsafe

Modifies the `Allocator` implementation provided by `ArenaAllocator` to be threadsafe using only atomics and no synchronization primitives locked behind an `Io` implementation. At its core this is a lock-free singly linked list which uses CAS loops to exchange the head node. A nice property of `ArenaAllocator` is that the only functions that can ever remove nodes from its linked list are `reset` and `deinit`, both of which are not part of the `Allocator` interface and thus aren't threadsafe, so node-related ABA problems are impossible. There *are* some trade-offs: end index tracking is now per node instead of per allocator instance. It's not possible to publish a head node and its end index at the same time if the latter isn't part of the former. Another compromise had to be made in regards to resizing existing nodes. Annoyingly, `rawResize` of an arbitrary thread-safe child allocator can of course never be guaranteed to be an atomic operation, so only one `alloc` call can ever resize at the same time, other threads have to consider any resizes they attempt during that time failed. This causes slightly less optimal behavior than what could be achieved with a mutex. The LSB of `Node.size` is used to signal that a node is being resized. This means that all nodes have to have an even size. Calls to `alloc` have to allocate new nodes optimistically as they can only know whether any CAS on a head node will succeed after attempting it, and to attempt the CAS they of course already need to know the address of the freshly allocated node they are trying to make the new head. The simplest solution to this would be to just free the new node again if a CAS fails, however this can be expensive and would mean that in practice arenas could only really be used with a GPA as their child allocator. To work around this, this implementation keeps its own free list of nodes which didn't make their CAS to be reused by a later `alloc` invocation. To keep things simple and avoid ABA problems the free list is only ever be accessed beyond its head by 'stealing' the head node (and thus the entire list) with an atomic swap. This makes iteration and removal trivial since there's only ever one thread doing it at a time which also owns all nodes it's holding. When the thread is done it can just push its list onto the free list again. This implementation offers comparable performance to the previous one when only being accessed by a single thread and a slight speedup compared to the previous implementation wrapped into a `ThreadSafeAllocator` up to ~7 threads performing operations on it concurrently. (measured on a base model MacBook Pro M1)

7 files changed, 650 insertions(+), 322 deletions(-)

CMakeLists.txt+1-1
...@@ -263,7 +263,7 @@ set(ZIG_STAGE2_SOURCES...@@ -263,7 +263,7 @@ set(ZIG_STAGE2_SOURCES
263 lib/std/hash/wyhash.zig263 lib/std/hash/wyhash.zig
264 lib/std/hash_map.zig264 lib/std/hash_map.zig
265 lib/std/heap.zig265 lib/std/heap.zig
266 lib/std/heap/arena_allocator.zig266 lib/std/heap/ArenaAllocator.zig
267 lib/std/json.zig267 lib/std/json.zig
268 lib/std/leb128.zig268 lib/std/leb128.zig
269 lib/std/log.zig269 lib/std/log.zig
lib/compiler/build_runner.zig+4-8
...@@ -38,13 +38,9 @@ pub fn main(init: process.Init.Minimal) !void {...@@ -38,13 +38,9 @@ pub fn main(init: process.Init.Minimal) !void {
38 const io = threaded.io();38 const io = threaded.io();
3939
40 // ...but we'll back our arena by `std.heap.page_allocator` for efficiency.40 // ...but we'll back our arena by `std.heap.page_allocator` for efficiency.
41 var single_threaded_arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator);41 var arena_instance: std.heap.ArenaAllocator = .init(std.heap.page_allocator);
42 defer single_threaded_arena.deinit();42 defer arena_instance.deinit();
43 var thread_safe_arena: std.heap.ThreadSafeAllocator = .{43 const arena = arena_instance.allocator();
44 .child_allocator = single_threaded_arena.allocator(),
45 .io = io,
46 };
47 const arena = thread_safe_arena.allocator();
4844
49 const args = try init.args.toSlice(arena);45 const args = try init.args.toSlice(arena);
5046
...@@ -86,7 +82,7 @@ pub fn main(init: process.Init.Minimal) !void {...@@ -86,7 +82,7 @@ pub fn main(init: process.Init.Minimal) !void {
86 .io = io,82 .io = io,
87 .gpa = gpa,83 .gpa = gpa,
88 .manifest_dir = try local_cache_directory.handle.createDirPathOpen(io, "h", .{}),84 .manifest_dir = try local_cache_directory.handle.createDirPathOpen(io, "h", .{}),
89 .cwd = try process.currentPathAlloc(io, single_threaded_arena.allocator()),85 .cwd = try process.currentPathAlloc(io, arena),
90 },86 },
91 .zig_exe = zig_exe,87 .zig_exe = zig_exe,
92 .environ_map = try init.environ.createMap(arena),88 .environ_map = try init.environ.createMap(arena),
lib/std/debug.zig+1-5
...@@ -1346,12 +1346,8 @@ pub fn getDebugInfoAllocator() Allocator {...@@ -1346,12 +1346,8 @@ pub fn getDebugInfoAllocator() Allocator {
1346 // Otherwise, use a global arena backed by the page allocator1346 // Otherwise, use a global arena backed by the page allocator
1347 const S = struct {1347 const S = struct {
1348 var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator);1348 var arena: std.heap.ArenaAllocator = .init(std.heap.page_allocator);
1349 var ts_arena: std.heap.ThreadSafeAllocator = .{
1350 .child_allocator = arena.allocator(),
1351 .io = std.Options.debug_io,
1352 };
1353 };1349 };
1354 return S.ts_arena.allocator();1350 return S.arena.allocator();
1355}1351}
13561352
1357/// Whether or not the current target can print useful debug information when a segfault occurs.1353/// Whether or not the current target can print useful debug information when a segfault occurs.
lib/std/heap.zig+1-1
...@@ -9,7 +9,7 @@ const Allocator = std.mem.Allocator;...@@ -9,7 +9,7 @@ const Allocator = std.mem.Allocator;
9const windows = std.os.windows;9const windows = std.os.windows;
10const Alignment = std.mem.Alignment;10const Alignment = std.mem.Alignment;
1111
12pub const ArenaAllocator = @import("heap/arena_allocator.zig").ArenaAllocator;12pub const ArenaAllocator = @import("heap/ArenaAllocator.zig");
13pub const SmpAllocator = @import("heap/SmpAllocator.zig");13pub const SmpAllocator = @import("heap/SmpAllocator.zig");
14pub const FixedBufferAllocator = @import("heap/FixedBufferAllocator.zig");14pub const FixedBufferAllocator = @import("heap/FixedBufferAllocator.zig");
15pub const PageAllocator = @import("heap/PageAllocator.zig");15pub const PageAllocator = @import("heap/PageAllocator.zig");
lib/std/heap/ArenaAllocator.zig created+642
...@@ -0,0 +1,642 @@
1//! This allocator takes an existing allocator, wraps it, and provides an interface where
2//! you can allocate and then free it all together. Calls to free an individual item only
3//! free the item if it was the most recent allocation, otherwise calls to free do
4//! nothing.
5//!
6//! The `Allocator` implementation provided is threadsafe, given that `child_allocator`
7//! is threadsafe as well.
8const ArenaAllocator = @This();
9
10child_allocator: Allocator,
11state: State,
12
13/// Inner state of ArenaAllocator. Can be stored rather than the entire ArenaAllocator
14/// as a memory-saving optimization.
15///
16/// Default initialization of this struct is deprecated; use `init` instead.
17pub const State = struct {
18 used_list: ?*Node = null,
19 free_list: ?*Node = null,
20
21 pub const init: State = .{
22 .used_list = null,
23 .free_list = null,
24 };
25
26 pub fn promote(state: State, child_allocator: Allocator) ArenaAllocator {
27 return .{
28 .child_allocator = child_allocator,
29 .state = state,
30 };
31 }
32};
33
34pub fn allocator(arena: *ArenaAllocator) Allocator {
35 return .{
36 .ptr = arena,
37 .vtable = &.{
38 .alloc = alloc,
39 .resize = resize,
40 .remap = remap,
41 .free = free,
42 },
43 };
44}
45
46pub fn init(child_allocator: Allocator) ArenaAllocator {
47 return State.init.promote(child_allocator);
48}
49
50/// Not threadsafe.
51pub fn deinit(arena: ArenaAllocator) void {
52 // NOTE: When changing this, make sure `reset()` is adjusted accordingly!
53
54 for ([_]?*Node{ arena.state.used_list, arena.state.free_list }) |first_node| {
55 var it = first_node;
56 while (it) |node| {
57 // this has to occur before the free because the free frees node
58 it = node.next;
59 arena.child_allocator.rawFree(node.allocatedSliceUnsafe(), .of(Node), @returnAddress());
60 }
61 }
62}
63
64/// Queries the current memory use of this arena.
65/// This will **not** include the storage required for internal keeping.
66///
67/// Not threadsafe.
68pub fn queryCapacity(arena: ArenaAllocator) usize {
69 var capacity: usize = 0;
70 for ([_]?*Node{ arena.state.used_list, arena.state.free_list }) |first_node| {
71 capacity += countListCapacity(first_node);
72 }
73 return capacity;
74}
75fn countListCapacity(first_node: ?*Node) usize {
76 var capacity: usize = 0;
77 var it = first_node;
78 while (it) |node| : (it = node.next) {
79 // Compute the actually allocated size excluding the
80 // linked list node.
81 capacity += node.size - @sizeOf(Node);
82 }
83 return capacity;
84}
85
86pub const ResetMode = union(enum) {
87 /// Releases all allocated memory in the arena.
88 free_all,
89 /// This will pre-heat the arena for future allocations by allocating a
90 /// large enough buffer for all previously done allocations.
91 /// Preheating will speed up the allocation process by invoking the backing allocator
92 /// less often than before. If `reset()` is used in a loop, this means that after the
93 /// biggest operation, no memory allocations are performed anymore.
94 retain_capacity,
95 /// This is the same as `retain_capacity`, but the memory will be shrunk to
96 /// this value if it exceeds the limit.
97 retain_with_limit: usize,
98};
99/// Resets the arena allocator and frees all allocated memory.
100///
101/// `mode` defines how the currently allocated memory is handled.
102/// See the variant documentation for `ResetMode` for the effects of each mode.
103///
104/// The function will return whether the reset operation was successful or not.
105/// If the reallocation failed `false` is returned. The arena will still be fully
106/// functional in that case, all memory is released. Future allocations just might
107/// be slower.
108///
109/// Not threadsafe.
110///
111/// NOTE: If `mode` is `free_all`, the function will always return `true`.
112pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
113 // Some words on the implementation:
114 // The reset function can be implemented with two basic approaches:
115 // - Counting how much bytes were allocated since the last reset, and storing that
116 // information in State. This will make reset fast and alloc only a teeny tiny bit
117 // slower.
118 // - Counting how much bytes were allocated by iterating the chunk linked list. This
119 // will make reset slower, but alloc() keeps the same speed when reset() as if reset()
120 // would not exist.
121 //
122 // The second variant was chosen for implementation, as with more and more calls to reset(),
123 // the function will get faster and faster. At one point, the complexity of the function
124 // will drop to amortized O(1), as we're only ever having a single chunk that will not be
125 // reallocated, and we're not even touching the backing allocator anymore.
126 //
127 // Thus, only the first hand full of calls to reset() will actually need to iterate the linked
128 // list, all future calls are just taking the first node, and only resetting the `end_index`
129 // value.
130
131 const limit: ?usize = switch (mode) {
132 .retain_capacity => null,
133 .retain_with_limit => |limit| limit,
134 .free_all => 0,
135 };
136 if (limit == 0) {
137 // just reset when we don't have anything to reallocate
138 arena.deinit();
139 arena.state = .init;
140 return true;
141 }
142
143 const used_capacity = countListCapacity(arena.state.used_list);
144 const free_capacity = countListCapacity(arena.state.free_list);
145
146 const new_used_capacity = if (limit) |lim| @min(lim, used_capacity) else used_capacity;
147 const new_free_capacity = if (limit) |lim| @min(lim - new_used_capacity, free_capacity) else free_capacity;
148
149 var ok = true;
150
151 for (
152 [_]*?*Node{ &arena.state.used_list, &arena.state.free_list },
153 [_]usize{ new_used_capacity, new_free_capacity },
154 ) |first_node_ptr, new_capacity| {
155 // Free all nodes except for the last one
156 var it = first_node_ptr.*;
157 const node: *Node = while (it) |node| {
158 // this has to occur before the free because the free frees node
159 it = node.next;
160 if (it == null) break node;
161 arena.child_allocator.rawFree(node.allocatedSliceUnsafe(), .of(Node), @returnAddress());
162 } else {
163 continue;
164 };
165 const allocated_slice = node.allocatedSliceUnsafe();
166
167 if (new_capacity == 0) {
168 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());
169 first_node_ptr.* = null;
170 continue;
171 }
172
173 node.end_index = 0;
174 first_node_ptr.* = node;
175
176 const adjusted_capacity: usize = mem.alignForward(usize, new_capacity, 2);
177
178 if (allocated_slice.len - @sizeOf(Node) == adjusted_capacity) {
179 // perfect, no need to invoke the child_allocator
180 continue;
181 }
182
183 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), adjusted_capacity, @returnAddress())) {
184 // successful resize
185 node.size = adjusted_capacity;
186 } else {
187 // manual realloc
188 const new_ptr = arena.child_allocator.rawAlloc(adjusted_capacity, .of(Node), @returnAddress()) orelse {
189 // we failed to preheat the arena properly, signal this to the user.
190 ok = false;
191 continue;
192 };
193 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());
194 const new_first_node: *Node = @ptrCast(@alignCast(new_ptr));
195 new_first_node.* = .{
196 .size = adjusted_capacity,
197 .end_index = 0,
198 .next = null,
199 };
200 first_node_ptr.* = new_first_node;
201 }
202 }
203
204 return ok;
205}
206
207/// Concurrent accesses to node pointers generally have to have acquire/release
208/// semantics to guarantee that newly allocated notes are in a valid state when
209/// being inserted into a list. Exceptions are possible, e.g. a CAS loop that
210/// never accesses the node returned on failure can use monotonic semantics on
211/// failure, but must still use release semantics on success to protect the node
212/// it's trying to push.
213const Node = struct {
214 /// Only meant to be accessed indirectly via the methods supplied by this type,
215 /// except if the node is owned by the thread accessing it.
216 /// Must always be an even number to accomodate `resize_bit`.
217 size: usize,
218 /// Concurrent accesses to `end_index` can be monotonic since it is only ever
219 /// incremented in `alloc` and `resize` after being compared to `size`.
220 /// Since `size` can only grow and never shrink, memory access depending on
221 /// `end_index` can never be OOB.
222 end_index: usize,
223 /// This field should only be accessed if the node is owned by the thread
224 /// accessing it.
225 next: ?*Node,
226
227 const resize_bit: usize = 1;
228
229 fn loadEndIndex(node: *Node) usize {
230 return @atomicLoad(usize, &node.end_index, .monotonic);
231 }
232
233 /// Returns `null` on success and previous value on failure.
234 fn trySetEndIndex(node: *Node, from: usize, to: usize) ?usize {
235 assert(from != to); // check this before attempting to set `end_index`!
236 return @cmpxchgWeak(usize, &node.end_index, from, to, .monotonic, .monotonic);
237 }
238
239 fn loadBuf(node: *Node) []u8 {
240 // monotonic is fine since `size` can only ever grow, so the buffer returned
241 // by this function is always valid memory.
242 const size = @atomicLoad(usize, &node.size, .monotonic);
243 return @as([*]u8, @ptrCast(node))[0 .. size & ~resize_bit][@sizeOf(Node)..];
244 }
245
246 /// Returns allocated slice or `null` if node is already (being) resized.
247 fn beginResize(node: *Node) ?[]u8 {
248 const size = @atomicRmw(usize, &node.size, .Or, resize_bit, .acquire); // syncs with release in `endResize`
249 if (size & resize_bit != 0) return null;
250 return @as([*]u8, @ptrCast(node))[0..size];
251 }
252
253 fn endResize(node: *Node, size: usize) void {
254 assert(size & resize_bit == 0);
255 return @atomicStore(usize, &node.size, size, .release); // syncs with acquire in `beginResize`
256 }
257
258 /// Not threadsafe.
259 fn allocatedSliceUnsafe(node: *Node) []u8 {
260 return @as([*]u8, @ptrCast(node))[0 .. node.size & ~resize_bit];
261 }
262};
263
264fn loadFirstNode(arena: *ArenaAllocator) ?*Node {
265 return @atomicLoad(?*Node, &arena.state.used_list, .acquire); // syncs with release in successful `tryPushNode`
266}
267
268const PushResult = union(enum) {
269 success,
270 failure: ?*Node,
271};
272fn tryPushNode(arena: *ArenaAllocator, node: *Node) PushResult {
273 assert(node != node.next);
274 if (@cmpxchgStrong( // strong because retrying means discarding a fitting node -> expensive
275 ?*Node,
276 &arena.state.used_list,
277 node.next,
278 node,
279 .release, // syncs with acquire in failure path or `loadFirstNode`
280 .acquire, // syncs with release in success path
281 )) |old_node| {
282 return .{ .failure = old_node };
283 } else {
284 return .success;
285 }
286}
287
288fn stealFreeList(arena: *ArenaAllocator) ?*Node {
289 // syncs with acq_rel in other `stealFreeList` calls or release in `pushFreeList`
290 return @atomicRmw(?*Node, &arena.state.free_list, .Xchg, null, .acq_rel);
291}
292
293fn pushFreeList(arena: *ArenaAllocator, first: *Node, last: *Node) void {
294 assert(first != last.next);
295 while (@cmpxchgWeak(
296 ?*Node,
297 &arena.state.free_list,
298 last.next,
299 first,
300 .release, // syncs with acquire part of acq_rel in `stealFreeList`
301 .monotonic, // we never access any fields of `old_free_list`, we only care about the pointer
302 )) |old_free_list| {
303 last.next = old_free_list;
304 }
305}
306
307fn alignedIndex(buf_ptr: [*]u8, end_index: usize, alignment: Alignment) usize {
308 return end_index +
309 mem.alignPointerOffset(buf_ptr + end_index, alignment.toByteUnits()).?;
310}
311
312fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u8 {
313 const arena: *ArenaAllocator = @ptrCast(@alignCast(ctx));
314 _ = ret_addr;
315
316 assert(n > 0);
317
318 var cur_first_node = arena.loadFirstNode();
319
320 var cur_new_node: ?*Node = null;
321 defer if (cur_new_node) |node| {
322 node.next = null; // optimize for empty free list
323 arena.pushFreeList(node, node);
324 };
325
326 retry: while (true) {
327 const first_node: ?*Node, const prev_size: usize = first_node: {
328 const node = cur_first_node orelse break :first_node .{ null, 0 };
329 var end_index = node.loadEndIndex();
330 while (true) {
331 const buf = node.loadBuf();
332 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
333
334 if (aligned_index + n > buf.len) {
335 break :first_node .{ node, buf.len };
336 }
337
338 end_index = node.trySetEndIndex(end_index, aligned_index + n) orelse {
339 return buf[aligned_index..][0..n].ptr;
340 };
341 }
342 };
343
344 resize: {
345 // Before attempting to get our hands on a new node, we try to resize
346 // the one we're currently holding. This is an exclusive operation;
347 // if another thread is already in this section we can never resize.
348
349 const node = first_node orelse break :resize;
350 const allocated_slice = node.beginResize() orelse break :resize;
351 var size = allocated_slice.len;
352 defer node.endResize(size);
353
354 const buf = allocated_slice[@sizeOf(Node)..];
355 const end_index = node.loadEndIndex();
356 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
357 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
358
359 if (new_size <= allocated_slice.len) {
360 // a `resize` or `free` call managed to sneak in and we need to
361 // guarantee that `size` is only ever increased; retry!
362 continue :retry;
363 }
364
365 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
366 size = new_size;
367
368 if (@cmpxchgStrong( // strong because a spurious failure could result in suboptimal usage of this node
369 usize,
370 &node.end_index,
371 end_index,
372 aligned_index + n,
373 .monotonic,
374 .monotonic,
375 ) == null) {
376 const new_buf = allocated_slice.ptr[0..new_size][@sizeOf(Node)..];
377 return new_buf[aligned_index..][0..n].ptr;
378 }
379 }
380 }
381
382 // We need a new node! First, we search `free_list` for one that's big
383 // enough, if we don't find one there we fall back to allocating a new
384 // node with `child_allocator` (if we haven't already done that!).
385
386 from_free_list: {
387 // We 'steal' the entire free list to operate on it without other
388 // threads getting up into our business.
389 // This is a rather pragmatic approach, but since the free list isn't
390 // used very frequently it's fine performance-wise, even under load.
391 // Also this avoids the ABA problem; stealing the list with an atomic
392 // swap doesn't introduce any potentially stale `next` pointers.
393
394 const free_list = arena.stealFreeList();
395 var first_free: ?*Node = free_list;
396 var last_free: ?*Node = free_list;
397 defer {
398 // Push remaining stolen free list back onto `arena.state.free_list`.
399 if (first_free) |first| {
400 const last = last_free.?;
401 assert(last.next == null); // optimize for no new nodes added during steal
402 arena.pushFreeList(first, last);
403 }
404 }
405
406 var best_fit_prev: ?*Node = null;
407 var best_fit: ?*Node = null;
408 var best_fit_diff: usize = std.math.maxInt(usize);
409
410 var it_prev: ?*Node = null;
411 var it = free_list;
412 const candidate: ?*Node, const prev: ?*Node = find: while (it) |node| : ({
413 it_prev = it;
414 it = node.next;
415 }) {
416 last_free = node;
417 assert(node.size & Node.resize_bit == 0);
418 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
419 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
420 if (buf.len < aligned_index + n) {
421 const diff = aligned_index + n - buf.len;
422 if (diff <= best_fit_diff) {
423 best_fit_prev = it_prev;
424 best_fit = node;
425 best_fit_diff = diff;
426 }
427 continue :find;
428 }
429 break :find .{ node, it_prev };
430 } else {
431 // Ideally we want to use all nodes in `free_list` eventually,
432 // so even if none fit we'll try to resize the one that was the
433 // closest to being large enough.
434 if (best_fit) |node| {
435 const allocated_slice = node.allocatedSliceUnsafe();
436 const buf = allocated_slice[@sizeOf(Node)..];
437 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
438 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
439
440 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
441 node.size = new_size;
442 break :find .{ node, best_fit_prev };
443 }
444 }
445 break :from_free_list;
446 };
447
448 it = last_free;
449 while (it) |node| : (it = node.next) {
450 last_free = node;
451 }
452
453 const node = candidate orelse break :from_free_list;
454
455 const old_next = node.next;
456
457 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
458 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
459
460 node.end_index = aligned_index + n;
461 node.next = first_node;
462
463 switch (arena.tryPushNode(node)) {
464 .success => {
465 // finish removing node from free list
466 if (prev) |p| p.next = old_next;
467 if (node == first_free) first_free = old_next;
468 if (node == last_free) last_free = prev;
469 return buf[aligned_index..][0..n].ptr;
470 },
471 .failure => |old_first_node| {
472 cur_first_node = old_first_node;
473 // restore free list to as we found it
474 node.next = old_next;
475 continue :retry;
476 },
477 }
478 }
479
480 const new_node: *Node = new_node: {
481 if (cur_new_node) |new_node| {
482 break :new_node new_node;
483 } else {
484 @branchHint(.cold);
485 }
486
487 const size: usize = size: {
488 const min_size = @sizeOf(Node) + alignment.toByteUnits() + n;
489 const big_enough_size = prev_size + min_size + 16;
490 break :size mem.alignForward(usize, big_enough_size + big_enough_size / 2, 2);
491 };
492 assert(size & Node.resize_bit == 0);
493 const ptr = arena.child_allocator.rawAlloc(size, .of(Node), @returnAddress()) orelse
494 return null;
495 const new_node: *Node = @ptrCast(@alignCast(ptr));
496 new_node.* = .{
497 .size = size,
498 .end_index = undefined, // set below
499 .next = undefined, // set below
500 };
501 cur_new_node = new_node;
502 break :new_node new_node;
503 };
504
505 const buf = new_node.allocatedSliceUnsafe()[@sizeOf(Node)..];
506 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
507 assert(new_node.size >= @sizeOf(Node) + aligned_index + n);
508
509 new_node.end_index = aligned_index + n;
510 new_node.next = first_node;
511
512 switch (arena.tryPushNode(new_node)) {
513 .success => {
514 cur_new_node = null;
515 return buf[aligned_index..][0..n].ptr;
516 },
517 .failure => |old_first_node| {
518 cur_first_node = old_first_node;
519 },
520 }
521 }
522}
523
524fn resize(ctx: *anyopaque, buf: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) bool {
525 const arena: *ArenaAllocator = @ptrCast(@alignCast(ctx));
526 _ = alignment;
527 _ = ret_addr;
528
529 assert(buf.len > 0);
530 assert(new_len > 0);
531 if (buf.len == new_len) return true;
532
533 const node = arena.loadFirstNode().?;
534 const cur_buf_ptr = @as([*]u8, @ptrCast(node)) + @sizeOf(Node);
535
536 var cur_end_index = node.loadEndIndex();
537 while (true) {
538 if (cur_buf_ptr + cur_end_index != buf.ptr + buf.len) {
539 // It's not the most recent allocation, so it cannot be expanded,
540 // but it's fine if they want to make it smaller.
541 return new_len <= buf.len;
542 }
543
544 const new_end_index: usize = new_end_index: {
545 if (buf.len >= new_len) {
546 break :new_end_index cur_end_index - (buf.len - new_len);
547 }
548 const cur_buf_len: usize = node.loadBuf().len;
549 // Saturating arithmetic because `end_index` and `size` are not
550 // guaranteed to be in sync.
551 if (cur_buf_len -| cur_end_index >= new_len - buf.len) {
552 break :new_end_index cur_end_index + (new_len - buf.len);
553 }
554 return false;
555 };
556
557 cur_end_index = node.trySetEndIndex(cur_end_index, new_end_index) orelse {
558 return true;
559 };
560 }
561}
562
563fn remap(
564 context: *anyopaque,
565 memory: []u8,
566 alignment: Alignment,
567 new_len: usize,
568 return_address: usize,
569) ?[*]u8 {
570 return if (resize(context, memory, alignment, new_len, return_address)) memory.ptr else null;
571}
572
573fn free(ctx: *anyopaque, buf: []u8, alignment: Alignment, ret_addr: usize) void {
574 const arena: *ArenaAllocator = @ptrCast(@alignCast(ctx));
575 _ = alignment;
576 _ = ret_addr;
577
578 assert(buf.len > 0);
579
580 const node = arena.loadFirstNode().?;
581 const cur_buf_ptr: [*]u8 = @as([*]u8, @ptrCast(node)) + @sizeOf(Node);
582
583 var cur_end_index = node.loadEndIndex();
584 while (true) {
585 if (cur_buf_ptr + cur_end_index != buf.ptr + buf.len) {
586 // Not the most recent allocation; we cannot free it.
587 return;
588 }
589 const new_end_index = cur_end_index - buf.len;
590 cur_end_index = node.trySetEndIndex(cur_end_index, new_end_index) orelse {
591 return;
592 };
593 }
594}
595
596const std = @import("std");
597const assert = std.debug.assert;
598const mem = std.mem;
599const Allocator = std.mem.Allocator;
600const Alignment = std.mem.Alignment;
601
602test "reset with preheating" {
603 var arena_allocator = ArenaAllocator.init(std.testing.allocator);
604 defer arena_allocator.deinit();
605 // provides some variance in the allocated data
606 var rng_src = std.Random.DefaultPrng.init(std.testing.random_seed);
607 const random = rng_src.random();
608 var rounds: usize = 25;
609 while (rounds > 0) {
610 rounds -= 1;
611 _ = arena_allocator.reset(.retain_capacity);
612 var alloced_bytes: usize = 0;
613 const total_size: usize = random.intRangeAtMost(usize, 256, 16384);
614 while (alloced_bytes < total_size) {
615 const size = random.intRangeAtMost(usize, 16, 256);
616 const alignment: Alignment = .@"32";
617 const slice = try arena_allocator.allocator().alignedAlloc(u8, alignment, size);
618 try std.testing.expect(alignment.check(@intFromPtr(slice.ptr)));
619 try std.testing.expectEqual(size, slice.len);
620 alloced_bytes += slice.len;
621 }
622 }
623}
624
625test "reset while retaining a buffer" {
626 var arena_allocator = ArenaAllocator.init(std.testing.allocator);
627 defer arena_allocator.deinit();
628 const a = arena_allocator.allocator();
629
630 // Create two internal buffers
631 _ = try a.alloc(u8, 1);
632 _ = try a.alloc(u8, 1000);
633
634 try std.testing.expect(arena_allocator.state.used_list != null);
635
636 // Check that we have at least two buffers
637 try std.testing.expect(arena_allocator.state.used_list.?.next != null);
638
639 // This retains the first allocated buffer
640 try std.testing.expect(arena_allocator.reset(.{ .retain_with_limit = 1 }));
641 try std.testing.expect(arena_allocator.state.used_list.?.next == null);
642}
lib/std/heap/arena_allocator.zig deleted-306
...@@ -1,306 +0,0 @@
1const std = @import("../std.zig");
2const assert = std.debug.assert;
3const mem = std.mem;
4const Allocator = std.mem.Allocator;
5const Alignment = std.mem.Alignment;
6
7/// This allocator takes an existing allocator, wraps it, and provides an interface where
8/// you can allocate and then free it all together. Calls to free an individual item only
9/// free the item if it was the most recent allocation, otherwise calls to free do
10/// nothing.
11pub const ArenaAllocator = struct {
12 child_allocator: Allocator,
13 state: State,
14
15 /// Inner state of ArenaAllocator. Can be stored rather than the entire ArenaAllocator
16 /// as a memory-saving optimization.
17 pub const State = struct {
18 buffer_list: std.SinglyLinkedList = .{},
19 end_index: usize = 0,
20
21 pub fn promote(self: State, child_allocator: Allocator) ArenaAllocator {
22 return .{
23 .child_allocator = child_allocator,
24 .state = self,
25 };
26 }
27 };
28
29 pub fn allocator(self: *ArenaAllocator) Allocator {
30 return .{
31 .ptr = self,
32 .vtable = &.{
33 .alloc = alloc,
34 .resize = resize,
35 .remap = remap,
36 .free = free,
37 },
38 };
39 }
40
41 const BufNode = struct {
42 data: usize,
43 node: std.SinglyLinkedList.Node = .{},
44 };
45 const BufNode_alignment: Alignment = .of(BufNode);
46
47 pub fn init(child_allocator: Allocator) ArenaAllocator {
48 return (State{}).promote(child_allocator);
49 }
50
51 pub fn deinit(self: ArenaAllocator) void {
52 // NOTE: When changing this, make sure `reset()` is adjusted accordingly!
53
54 var it = self.state.buffer_list.first;
55 while (it) |node| {
56 // this has to occur before the free because the free frees node
57 const next_it = node.next;
58 const buf_node: *BufNode = @fieldParentPtr("node", node);
59 const alloc_buf = @as([*]u8, @ptrCast(buf_node))[0..buf_node.data];
60 self.child_allocator.rawFree(alloc_buf, BufNode_alignment, @returnAddress());
61 it = next_it;
62 }
63 }
64
65 pub const ResetMode = union(enum) {
66 /// Releases all allocated memory in the arena.
67 free_all,
68 /// This will pre-heat the arena for future allocations by allocating a
69 /// large enough buffer for all previously done allocations.
70 /// Preheating will speed up the allocation process by invoking the backing allocator
71 /// less often than before. If `reset()` is used in a loop, this means that after the
72 /// biggest operation, no memory allocations are performed anymore.
73 retain_capacity,
74 /// This is the same as `retain_capacity`, but the memory will be shrunk to
75 /// this value if it exceeds the limit.
76 retain_with_limit: usize,
77 };
78 /// Queries the current memory use of this arena.
79 /// This will **not** include the storage required for internal keeping.
80 pub fn queryCapacity(self: ArenaAllocator) usize {
81 var size: usize = 0;
82 var it = self.state.buffer_list.first;
83 while (it) |node| : (it = node.next) {
84 // Compute the actually allocated size excluding the
85 // linked list node.
86 const buf_node: *BufNode = @fieldParentPtr("node", node);
87 size += buf_node.data - @sizeOf(BufNode);
88 }
89 return size;
90 }
91 /// Resets the arena allocator and frees all allocated memory.
92 ///
93 /// `mode` defines how the currently allocated memory is handled.
94 /// See the variant documentation for `ResetMode` for the effects of each mode.
95 ///
96 /// The function will return whether the reset operation was successful or not.
97 /// If the reallocation failed `false` is returned. The arena will still be fully
98 /// functional in that case, all memory is released. Future allocations just might
99 /// be slower.
100 ///
101 /// NOTE: If `mode` is `free_all`, the function will always return `true`.
102 pub fn reset(self: *ArenaAllocator, mode: ResetMode) bool {
103 // Some words on the implementation:
104 // The reset function can be implemented with two basic approaches:
105 // - Counting how much bytes were allocated since the last reset, and storing that
106 // information in State. This will make reset fast and alloc only a teeny tiny bit
107 // slower.
108 // - Counting how much bytes were allocated by iterating the chunk linked list. This
109 // will make reset slower, but alloc() keeps the same speed when reset() as if reset()
110 // would not exist.
111 //
112 // The second variant was chosen for implementation, as with more and more calls to reset(),
113 // the function will get faster and faster. At one point, the complexity of the function
114 // will drop to amortized O(1), as we're only ever having a single chunk that will not be
115 // reallocated, and we're not even touching the backing allocator anymore.
116 //
117 // Thus, only the first hand full of calls to reset() will actually need to iterate the linked
118 // list, all future calls are just taking the first node, and only resetting the `end_index`
119 // value.
120 const requested_capacity = switch (mode) {
121 .retain_capacity => self.queryCapacity(),
122 .retain_with_limit => |limit| @min(limit, self.queryCapacity()),
123 .free_all => 0,
124 };
125 if (requested_capacity == 0) {
126 // just reset when we don't have anything to reallocate
127 self.deinit();
128 self.state = State{};
129 return true;
130 }
131 const total_size = requested_capacity + @sizeOf(BufNode);
132 // Free all nodes except for the last one
133 var it = self.state.buffer_list.first;
134 const maybe_first_node = while (it) |node| {
135 // this has to occur before the free because the free frees node
136 const next_it = node.next;
137 if (next_it == null)
138 break node;
139 const buf_node: *BufNode = @fieldParentPtr("node", node);
140 const alloc_buf = @as([*]u8, @ptrCast(buf_node))[0..buf_node.data];
141 self.child_allocator.rawFree(alloc_buf, BufNode_alignment, @returnAddress());
142 it = next_it;
143 } else null;
144 std.debug.assert(maybe_first_node == null or maybe_first_node.?.next == null);
145 // reset the state before we try resizing the buffers, so we definitely have reset the arena to 0.
146 self.state.end_index = 0;
147 if (maybe_first_node) |first_node| {
148 self.state.buffer_list.first = first_node;
149 // perfect, no need to invoke the child_allocator
150 const first_buf_node: *BufNode = @fieldParentPtr("node", first_node);
151 if (first_buf_node.data == total_size)
152 return true;
153 const first_alloc_buf = @as([*]u8, @ptrCast(first_buf_node))[0..first_buf_node.data];
154 if (self.child_allocator.rawResize(first_alloc_buf, BufNode_alignment, total_size, @returnAddress())) {
155 // successful resize
156 first_buf_node.data = total_size;
157 } else {
158 // manual realloc
159 const new_ptr = self.child_allocator.rawAlloc(total_size, BufNode_alignment, @returnAddress()) orelse {
160 // we failed to preheat the arena properly, signal this to the user.
161 return false;
162 };
163 self.child_allocator.rawFree(first_alloc_buf, BufNode_alignment, @returnAddress());
164 const buf_node: *BufNode = @ptrCast(@alignCast(new_ptr));
165 buf_node.* = .{ .data = total_size };
166 self.state.buffer_list.first = &buf_node.node;
167 }
168 }
169 return true;
170 }
171
172 fn createNode(self: *ArenaAllocator, prev_len: usize, minimum_size: usize) ?*BufNode {
173 const actual_min_size = minimum_size + (@sizeOf(BufNode) + 16);
174 const big_enough_len = prev_len + actual_min_size;
175 const len = big_enough_len + big_enough_len / 2;
176 const ptr = self.child_allocator.rawAlloc(len, BufNode_alignment, @returnAddress()) orelse
177 return null;
178 const buf_node: *BufNode = @ptrCast(@alignCast(ptr));
179 buf_node.* = .{ .data = len };
180 self.state.buffer_list.prepend(&buf_node.node);
181 self.state.end_index = 0;
182 return buf_node;
183 }
184
185 fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ra: usize) ?[*]u8 {
186 const self: *ArenaAllocator = @ptrCast(@alignCast(ctx));
187 _ = ra;
188
189 const ptr_align = alignment.toByteUnits();
190 var cur_node: *BufNode = if (self.state.buffer_list.first) |first_node|
191 @fieldParentPtr("node", first_node)
192 else
193 (self.createNode(0, n + ptr_align) orelse return null);
194 while (true) {
195 const cur_alloc_buf = @as([*]u8, @ptrCast(cur_node))[0..cur_node.data];
196 const cur_buf = cur_alloc_buf[@sizeOf(BufNode)..];
197 const addr = @intFromPtr(cur_buf.ptr) + self.state.end_index;
198 const adjusted_addr = mem.alignForward(usize, addr, ptr_align);
199 const adjusted_index = self.state.end_index + (adjusted_addr - addr);
200 const new_end_index = adjusted_index + n;
201
202 if (new_end_index <= cur_buf.len) {
203 const result = cur_buf[adjusted_index..new_end_index];
204 self.state.end_index = new_end_index;
205 return result.ptr;
206 }
207
208 const bigger_buf_size = @sizeOf(BufNode) + new_end_index;
209 if (self.child_allocator.rawResize(cur_alloc_buf, BufNode_alignment, bigger_buf_size, @returnAddress())) {
210 cur_node.data = bigger_buf_size;
211 } else {
212 // Allocate a new node if that's not possible
213 cur_node = self.createNode(cur_buf.len, n + ptr_align) orelse return null;
214 }
215 }
216 }
217
218 fn resize(ctx: *anyopaque, buf: []u8, alignment: Alignment, new_len: usize, ret_addr: usize) bool {
219 const self: *ArenaAllocator = @ptrCast(@alignCast(ctx));
220 _ = alignment;
221 _ = ret_addr;
222
223 const cur_node = self.state.buffer_list.first orelse return false;
224 const cur_buf_node: *BufNode = @fieldParentPtr("node", cur_node);
225 const cur_buf = @as([*]u8, @ptrCast(cur_buf_node))[@sizeOf(BufNode)..cur_buf_node.data];
226 if (@intFromPtr(cur_buf.ptr) + self.state.end_index != @intFromPtr(buf.ptr) + buf.len) {
227 // It's not the most recent allocation, so it cannot be expanded,
228 // but it's fine if they want to make it smaller.
229 return new_len <= buf.len;
230 }
231
232 if (buf.len >= new_len) {
233 self.state.end_index -= buf.len - new_len;
234 return true;
235 } else if (cur_buf.len - self.state.end_index >= new_len - buf.len) {
236 self.state.end_index += new_len - buf.len;
237 return true;
238 } else {
239 return false;
240 }
241 }
242
243 fn remap(
244 context: *anyopaque,
245 memory: []u8,
246 alignment: Alignment,
247 new_len: usize,
248 return_address: usize,
249 ) ?[*]u8 {
250 return if (resize(context, memory, alignment, new_len, return_address)) memory.ptr else null;
251 }
252
253 fn free(ctx: *anyopaque, buf: []u8, alignment: Alignment, ret_addr: usize) void {
254 _ = alignment;
255 _ = ret_addr;
256
257 const self: *ArenaAllocator = @ptrCast(@alignCast(ctx));
258
259 const cur_node = self.state.buffer_list.first orelse return;
260 const cur_buf_node: *BufNode = @fieldParentPtr("node", cur_node);
261 const cur_buf = @as([*]u8, @ptrCast(cur_buf_node))[@sizeOf(BufNode)..cur_buf_node.data];
262
263 if (@intFromPtr(cur_buf.ptr) + self.state.end_index == @intFromPtr(buf.ptr) + buf.len) {
264 self.state.end_index -= buf.len;
265 }
266 }
267};
268
269test "reset with preheating" {
270 var arena_allocator = ArenaAllocator.init(std.testing.allocator);
271 defer arena_allocator.deinit();
272 // provides some variance in the allocated data
273 var rng_src = std.Random.DefaultPrng.init(std.testing.random_seed);
274 const random = rng_src.random();
275 var rounds: usize = 25;
276 while (rounds > 0) {
277 rounds -= 1;
278 _ = arena_allocator.reset(.retain_capacity);
279 var alloced_bytes: usize = 0;
280 const total_size: usize = random.intRangeAtMost(usize, 256, 16384);
281 while (alloced_bytes < total_size) {
282 const size = random.intRangeAtMost(usize, 16, 256);
283 const alignment: Alignment = .@"32";
284 const slice = try arena_allocator.allocator().alignedAlloc(u8, alignment, size);
285 try std.testing.expect(alignment.check(@intFromPtr(slice.ptr)));
286 try std.testing.expectEqual(size, slice.len);
287 alloced_bytes += slice.len;
288 }
289 }
290}
291
292test "reset while retaining a buffer" {
293 var arena_allocator = ArenaAllocator.init(std.testing.allocator);
294 defer arena_allocator.deinit();
295 const a = arena_allocator.allocator();
296
297 // Create two internal buffers
298 _ = try a.alloc(u8, 1);
299 _ = try a.alloc(u8, 1000);
300
301 // Check that we have at least two buffers
302 try std.testing.expect(arena_allocator.state.buffer_list.first.?.next != null);
303
304 // This retains the first allocated buffer
305 try std.testing.expect(arena_allocator.reset(.{ .retain_with_limit = 1 }));
306}
lib/std/process.zig+1-1
...@@ -31,7 +31,7 @@ pub const Init = struct {...@@ -31,7 +31,7 @@ pub const Init = struct {
31 /// `Init` is a superset of `Minimal`; the latter is included here.31 /// `Init` is a superset of `Minimal`; the latter is included here.
32 minimal: Minimal,32 minimal: Minimal,
33 /// Permanent storage for the entire process, cleaned automatically on33 /// Permanent storage for the entire process, cleaned automatically on
34 /// exit. Not threadsafe.34 /// exit. Threadsafe.
35 arena: *std.heap.ArenaAllocator,35 arena: *std.heap.ArenaAllocator,
36 /// A default-selected general purpose allocator for temporary heap36 /// A default-selected general purpose allocator for temporary heap
37 /// allocations. Debug mode will set up leak checking if possible.37 /// allocations. Debug mode will set up leak checking if possible.