authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2026-02-26 20:42:30+01:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2026-02-26 20:42:30+01:00
log7bc6546fdfadf3e3c91ee2be1387913599a63e4a
tree35566c638e59179349f37b2ef2c30620a657a446
parent56253d9e31c0576f024d95929a8fe26428b35176
parentde4112395714b43890aa1e9d1e7dd952e8947288

Merge pull request 'std.heap.ArenaAllocator: Get rid of cmpxchg loop in hot path' (#31343) from justusk/zig:lock-free-arena-overshoot into master

Reviewed-on: https://codeberg.org/ziglang/zig/pulls/31343 Reviewed-by: Andrew Kelley <andrew@ziglang.org>

1 files changed, 119 insertions(+), 87 deletions(-)

lib/std/heap/ArenaAllocator.zig+119-87
......@@ -78,7 +78,7 @@ fn countListCapacity(first_node: ?*Node) usize {
7878 while (it) |node| : (it = node.next) {
7979 // Compute the actually allocated size excluding the
8080 // linked list node.
81 capacity += node.size - @sizeOf(Node);
81 capacity += node.size.toInt() - @sizeOf(Node);
8282 }
8383 return capacity;
8484}
......@@ -164,7 +164,10 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
164164 };
165165 const allocated_slice = node.allocatedSliceUnsafe();
166166
167 if (new_capacity == 0) {
167 // Align backwards to always stay below limit.
168 const new_size = mem.alignBackward(usize, @sizeOf(Node) + new_capacity, 2);
169
170 if (new_size == @sizeOf(Node)) {
168171 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());
169172 first_node_ptr.* = null;
170173 continue;
......@@ -173,19 +176,17 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
173176 node.end_index = 0;
174177 first_node_ptr.* = node;
175178
176 const adjusted_capacity: usize = mem.alignForward(usize, new_capacity, 2);
177
178 if (allocated_slice.len - @sizeOf(Node) == adjusted_capacity) {
179 if (allocated_slice.len == new_size) {
179180 // perfect, no need to invoke the child_allocator
180181 continue;
181182 }
182183
183 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), adjusted_capacity, @returnAddress())) {
184 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
184185 // successful resize
185 node.size = adjusted_capacity;
186 node.size = .fromInt(new_size);
186187 } else {
187188 // manual realloc
188 const new_ptr = arena.child_allocator.rawAlloc(adjusted_capacity, .of(Node), @returnAddress()) orelse {
189 const new_ptr = arena.child_allocator.rawAlloc(new_size, .of(Node), @returnAddress()) orelse {
189190 // we failed to preheat the arena properly, signal this to the user.
190191 ok = false;
191192 continue;
......@@ -193,7 +194,7 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
193194 arena.child_allocator.rawFree(allocated_slice, .of(Node), @returnAddress());
194195 const new_first_node: *Node = @ptrCast(@alignCast(new_ptr));
195196 new_first_node.* = .{
196 .size = adjusted_capacity,
197 .size = .fromInt(new_size),
197198 .end_index = 0,
198199 .next = null,
199200 };
......@@ -213,51 +214,60 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
213214const Node = struct {
214215 /// Only meant to be accessed indirectly via the methods supplied by this type,
215216 /// 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`.
217 /// Must always be an even number to accomodate `resize` bit.
218 size: Size,
219 /// Concurrent accesses to `end_index` can be monotonic as long as its value
220 /// is compared to a version of `size` before using it to access memory.
220221 /// Since `size` can only grow and never shrink, memory access depending on
221 /// `end_index` can never be OOB.
222 /// any `end_index` <= any `size` can never be OOB.
222223 end_index: usize,
223224 /// This field should only be accessed if the node is owned by the thread
224225 /// accessing it.
225226 next: ?*Node,
226227
227 const resize_bit: usize = 1;
228 const Size = packed struct(usize) {
229 resizing: bool,
230 _: @Int(.unsigned, @bitSizeOf(usize) - 1) = 0,
228231
229 fn loadEndIndex(node: *Node) usize {
230 return @atomicLoad(usize, &node.end_index, .monotonic);
231 }
232 fn fromInt(int: usize) Size {
233 assert(int >= @sizeOf(Node));
234 const size: Size = @bitCast(int);
235 assert(!size.resizing);
236 return size;
237 }
232238
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 }
239 fn toInt(size: Size) usize {
240 var int = size;
241 int.resizing = false;
242 return @bitCast(int);
243 }
244
245 comptime {
246 assert(Size{ .resizing = true } == @as(Size, @bitCast(@as(usize, 1))));
247 }
248 };
238249
239250 fn loadBuf(node: *Node) []u8 {
240251 // monotonic is fine since `size` can only ever grow, so the buffer returned
241252 // 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)..];
253 const size = @atomicLoad(Size, &node.size, .monotonic);
254 return @as([*]u8, @ptrCast(node))[0..size.toInt()][@sizeOf(Node)..];
244255 }
245256
246257 /// Returns allocated slice or `null` if node is already (being) resized.
247258 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];
259 const size = @atomicRmw(Size, &node.size, .Or, .{ .resizing = true }, .acquire); // syncs with release in `endResize`
260 if (size.resizing) return null;
261 return @as([*]u8, @ptrCast(node))[0..size.toInt()];
251262 }
252263
253264 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`
265 return @atomicStore(Size, &node.size, .fromInt(size), .release); // syncs with acquire in `beginResize`
256266 }
257267
258268 /// Not threadsafe.
259269 fn allocatedSliceUnsafe(node: *Node) []u8 {
260 return @as([*]u8, @ptrCast(node))[0 .. node.size & ~resize_bit];
270 return @as([*]u8, @ptrCast(node))[0..node.size.toInt()];
261271 }
262272};
263273
......@@ -326,19 +336,22 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
326336 retry: while (true) {
327337 const first_node: ?*Node, const prev_size: usize = first_node: {
328338 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);
339 const buf = node.loadBuf();
333340
334 if (aligned_index + n > buf.len) {
335 break :first_node .{ node, buf.len };
336 }
341 // To avoid using a CAS loop in the hot path we atomically increase
342 // `end_index` by a large enough amount to be able to always provide
343 // the required alignment within the reserved memory. To recover the
344 // space this potentially wastes we try to subtract the 'overshoot'
345 // with a single cmpxchg afterwards, which may fail.
337346
338 end_index = node.trySetEndIndex(end_index, aligned_index + n) orelse {
339 return buf[aligned_index..][0..n].ptr;
340 };
341 }
347 const alignable = n + alignment.toByteUnits() - 1;
348 const end_index = @atomicRmw(usize, &node.end_index, .Add, alignable, .monotonic);
349 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
350 assert(end_index + alignable >= aligned_index + n);
351 _ = @cmpxchgStrong(usize, &node.end_index, end_index + alignable, aligned_index + n, .monotonic, .monotonic);
352
353 if (aligned_index + n > buf.len) break :first_node .{ node, buf.len };
354 return buf[aligned_index..][0..n].ptr;
342355 };
343356
344357 resize: {
......@@ -352,7 +365,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
352365 defer node.endResize(size);
353366
354367 const buf = allocated_slice[@sizeOf(Node)..];
355 const end_index = node.loadEndIndex();
368 const end_index = @atomicLoad(usize, &node.end_index, .monotonic);
356369 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
357370 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
358371
......@@ -403,55 +416,59 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
403416 }
404417 }
405418
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) {
419 const candidate: ?*Node, const prev: ?*Node = candidate: {
420 var best_fit_prev: ?*Node = null;
421 var best_fit: ?*Node = null;
422 var best_fit_diff: usize = std.math.maxInt(usize);
423
424 var it_prev: ?*Node = null;
425 var it = free_list;
426 while (it) |node| : ({
427 it_prev = it;
428 it = node.next;
429 }) {
430 last_free = node;
431 assert(!node.size.resizing);
432 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
433 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
434
435 if (aligned_index + n <= buf.len) {
436 break :candidate .{ node, it_prev };
437 }
438
421439 const diff = aligned_index + n - buf.len;
422440 if (diff <= best_fit_diff) {
423441 best_fit_prev = it_prev;
424442 best_fit = node;
425443 best_fit_diff = diff;
426444 }
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 };
445 } else {
446 // Ideally we want to use all nodes in `free_list` eventually,
447 // so even if none fit we'll try to resize the one that was the
448 // closest to being large enough.
449 if (best_fit) |node| {
450 const allocated_slice = node.allocatedSliceUnsafe();
451 const buf = allocated_slice[@sizeOf(Node)..];
452 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
453 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
454
455 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
456 node.size = .fromInt(new_size);
457 break :candidate .{ node, best_fit_prev };
458 }
443459 }
460 break :from_free_list;
444461 }
445 break :from_free_list;
446462 };
447463
448 it = last_free;
449 while (it) |node| : (it = node.next) {
450 last_free = node;
464 {
465 var it = last_free;
466 while (it) |node| : (it = node.next) {
467 last_free = node;
468 }
451469 }
452470
453471 const node = candidate orelse break :from_free_list;
454
455472 const old_next = node.next;
456473
457474 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
......@@ -489,12 +506,11 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
489506 const big_enough_size = prev_size + min_size + 16;
490507 break :size mem.alignForward(usize, big_enough_size + big_enough_size / 2, 2);
491508 };
492 assert(size & Node.resize_bit == 0);
493509 const ptr = arena.child_allocator.rawAlloc(size, .of(Node), @returnAddress()) orelse
494510 return null;
495511 const new_node: *Node = @ptrCast(@alignCast(ptr));
496512 new_node.* = .{
497 .size = size,
513 .size = .fromInt(size),
498514 .end_index = undefined, // set below
499515 .next = undefined, // set below
500516 };
......@@ -504,7 +520,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
504520
505521 const buf = new_node.allocatedSliceUnsafe()[@sizeOf(Node)..];
506522 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
507 assert(new_node.size >= @sizeOf(Node) + aligned_index + n);
523 assert(new_node.size.toInt() >= @sizeOf(Node) + aligned_index + n);
508524
509525 new_node.end_index = aligned_index + n;
510526 new_node.next = first_node;
......@@ -533,7 +549,7 @@ fn resize(ctx: *anyopaque, buf: []u8, alignment: Alignment, new_len: usize, ret_
533549 const node = arena.loadFirstNode().?;
534550 const cur_buf_ptr = @as([*]u8, @ptrCast(node)) + @sizeOf(Node);
535551
536 var cur_end_index = node.loadEndIndex();
552 var cur_end_index = @atomicLoad(usize, &node.end_index, .monotonic);
537553 while (true) {
538554 if (cur_buf_ptr + cur_end_index != buf.ptr + buf.len) {
539555 // It's not the most recent allocation, so it cannot be expanded,
......@@ -554,7 +570,14 @@ fn resize(ctx: *anyopaque, buf: []u8, alignment: Alignment, new_len: usize, ret_
554570 return false;
555571 };
556572
557 cur_end_index = node.trySetEndIndex(cur_end_index, new_end_index) orelse {
573 cur_end_index = @cmpxchgWeak(
574 usize,
575 &node.end_index,
576 cur_end_index,
577 new_end_index,
578 .monotonic,
579 .monotonic,
580 ) orelse {
558581 return true;
559582 };
560583 }
......@@ -580,14 +603,22 @@ fn free(ctx: *anyopaque, buf: []u8, alignment: Alignment, ret_addr: usize) void
580603 const node = arena.loadFirstNode().?;
581604 const cur_buf_ptr: [*]u8 = @as([*]u8, @ptrCast(node)) + @sizeOf(Node);
582605
583 var cur_end_index = node.loadEndIndex();
606 var cur_end_index = @atomicLoad(usize, &node.end_index, .monotonic);
584607 while (true) {
585608 if (cur_buf_ptr + cur_end_index != buf.ptr + buf.len) {
586609 // Not the most recent allocation; we cannot free it.
587610 return;
588611 }
589612 const new_end_index = cur_end_index - buf.len;
590 cur_end_index = node.trySetEndIndex(cur_end_index, new_end_index) orelse {
613
614 cur_end_index = @cmpxchgWeak(
615 usize,
616 &node.end_index,
617 cur_end_index,
618 new_end_index,
619 .monotonic,
620 .monotonic,
621 ) orelse {
591622 return;
592623 };
593624 }
......@@ -637,6 +668,7 @@ test "reset while retaining a buffer" {
637668 try std.testing.expect(arena_allocator.state.used_list.?.next != null);
638669
639670 // This retains the first allocated buffer
640 try std.testing.expect(arena_allocator.reset(.{ .retain_with_limit = 1 }));
671 try std.testing.expect(arena_allocator.reset(.{ .retain_with_limit = 2 }));
641672 try std.testing.expect(arena_allocator.state.used_list.?.next == null);
673 try std.testing.expectEqual(2, arena_allocator.queryCapacity());
642674}