authorgravatar for justus@klausecker.deJustus Klausecker <justus@klausecker.de> 2026-02-26 11:51:23+01:00
committergravatar for justus@klausecker.deJustus Klausecker <justus@klausecker.de> 2026-02-26 15:30:55+01:00
log2fa2300ba46f35a9073475bd24a8aa9180e0b228
tree7729433688b91aa33e7a800aec27c8285c7f97b4
parent56253d9e31c0576f024d95929a8fe26428b35176

std.heap.ArenaAllocator: Get rid of cmpxchg loop in hot path

This is achieved by bumping `end_index` by a large enough amount so that a suitably aligned region of memory can always be provided. The potential wasted space this creates is then recovered by a single cmpxchg. This is always successful for single-threaded arenas which means that this version still behaves exactly the same as the old single-threaded implementation when only being accessed by one thread at a time. It can however fail when another thread bumps `end_index` in the meantime. The observerd failure rates under extreme load are: 2 Threads: 4-5% 3 Threads: 13-15% 4 Threads: 15-17% 5 Threads: 17-18% 6 Threads: 19-20% 7 Threads: 18-21% This version offers ~25% faster performance under extreme load from 7 threads, with diminishing speedups for less threads. The performance for 1 and 2 threads is nearly identical.

1 files changed, 79 insertions(+), 67 deletions(-)

lib/std/heap/ArenaAllocator.zig+79-67
......@@ -213,12 +213,12 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
213213const Node = struct {
214214 /// Only meant to be accessed indirectly via the methods supplied by this type,
215215 /// 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`.
216 /// Must always be an even number to accomodate `resize` bit.
217 size: Size,
218 /// Concurrent accesses to `end_index` can be monotonic as long as its value
219 /// is compared to a version of `size` before using it to access memory.
220220 /// Since `size` can only grow and never shrink, memory access depending on
221 /// `end_index` can never be OOB.
221 /// any `end_index` <= any `size` can never be OOB.
222222 end_index: usize,
223223 /// This field should only be accessed if the node is owned by the thread
224224 /// accessing it.
......@@ -226,16 +226,6 @@ const Node = struct {
226226
227227 const resize_bit: usize = 1;
228228
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
239229 fn loadBuf(node: *Node) []u8 {
240230 // monotonic is fine since `size` can only ever grow, so the buffer returned
241231 // by this function is always valid memory.
......@@ -326,19 +316,22 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
326316 retry: while (true) {
327317 const first_node: ?*Node, const prev_size: usize = first_node: {
328318 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);
319 const buf = node.loadBuf();
333320
334 if (aligned_index + n > buf.len) {
335 break :first_node .{ node, buf.len };
336 }
321 // To avoid using a CAS loop in the hot path we atomically increase
322 // `end_index` by a large enough amount to be able to always provide
323 // the required alignment within the reserved memory. To recover the
324 // space this potentially wastes we try to subtract the 'overshoot'
325 // with a single cmpxchg afterwards, which may fail.
337326
338 end_index = node.trySetEndIndex(end_index, aligned_index + n) orelse {
339 return buf[aligned_index..][0..n].ptr;
340 };
341 }
327 const alignable = n + alignment.toByteUnits() - 1;
328 const end_index = @atomicRmw(usize, &node.end_index, .Add, alignable, .monotonic);
329 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
330 assert(end_index + alignable >= aligned_index + n);
331 _ = @cmpxchgStrong(usize, &node.end_index, end_index + alignable, aligned_index + n, .monotonic, .monotonic);
332
333 if (aligned_index + n > buf.len) break :first_node .{ node, buf.len };
334 return buf[aligned_index..][0..n].ptr;
342335 };
343336
344337 resize: {
......@@ -352,7 +345,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
352345 defer node.endResize(size);
353346
354347 const buf = allocated_slice[@sizeOf(Node)..];
355 const end_index = node.loadEndIndex();
348 const end_index = @atomicLoad(usize, &node.end_index, .monotonic);
356349 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
357350 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
358351
......@@ -403,55 +396,59 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
403396 }
404397 }
405398
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) {
399 const candidate: ?*Node, const prev: ?*Node = candidate: {
400 var best_fit_prev: ?*Node = null;
401 var best_fit: ?*Node = null;
402 var best_fit_diff: usize = std.math.maxInt(usize);
403
404 var it_prev: ?*Node = null;
405 var it = free_list;
406 while (it) |node| : ({
407 it_prev = it;
408 it = node.next;
409 }) {
410 last_free = node;
411 assert(node.size & Node.resize_bit == 0);
412 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
413 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
414
415 if (aligned_index + n <= buf.len) {
416 break :candidate .{ node, it_prev };
417 }
418
421419 const diff = aligned_index + n - buf.len;
422420 if (diff <= best_fit_diff) {
423421 best_fit_prev = it_prev;
424422 best_fit = node;
425423 best_fit_diff = diff;
426424 }
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 };
425 } else {
426 // Ideally we want to use all nodes in `free_list` eventually,
427 // so even if none fit we'll try to resize the one that was the
428 // closest to being large enough.
429 if (best_fit) |node| {
430 const allocated_slice = node.allocatedSliceUnsafe();
431 const buf = allocated_slice[@sizeOf(Node)..];
432 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
433 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
434
435 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
436 node.size = new_size;
437 break :candidate .{ node, best_fit_prev };
438 }
443439 }
440 break :from_free_list;
444441 }
445 break :from_free_list;
446442 };
447443
448 it = last_free;
449 while (it) |node| : (it = node.next) {
450 last_free = node;
444 {
445 var it = last_free;
446 while (it) |node| : (it = node.next) {
447 last_free = node;
448 }
451449 }
452450
453451 const node = candidate orelse break :from_free_list;
454
455452 const old_next = node.next;
456453
457454 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
......@@ -533,7 +530,7 @@ fn resize(ctx: *anyopaque, buf: []u8, alignment: Alignment, new_len: usize, ret_
533530 const node = arena.loadFirstNode().?;
534531 const cur_buf_ptr = @as([*]u8, @ptrCast(node)) + @sizeOf(Node);
535532
536 var cur_end_index = node.loadEndIndex();
533 var cur_end_index = @atomicLoad(usize, &node.end_index, .monotonic);
537534 while (true) {
538535 if (cur_buf_ptr + cur_end_index != buf.ptr + buf.len) {
539536 // It's not the most recent allocation, so it cannot be expanded,
......@@ -554,7 +551,14 @@ fn resize(ctx: *anyopaque, buf: []u8, alignment: Alignment, new_len: usize, ret_
554551 return false;
555552 };
556553
557 cur_end_index = node.trySetEndIndex(cur_end_index, new_end_index) orelse {
554 cur_end_index = @cmpxchgWeak(
555 usize,
556 &node.end_index,
557 cur_end_index,
558 new_end_index,
559 .monotonic,
560 .monotonic,
561 ) orelse {
558562 return true;
559563 };
560564 }
......@@ -580,14 +584,22 @@ fn free(ctx: *anyopaque, buf: []u8, alignment: Alignment, ret_addr: usize) void
580584 const node = arena.loadFirstNode().?;
581585 const cur_buf_ptr: [*]u8 = @as([*]u8, @ptrCast(node)) + @sizeOf(Node);
582586
583 var cur_end_index = node.loadEndIndex();
587 var cur_end_index = @atomicLoad(usize, &node.end_index, .monotonic);
584588 while (true) {
585589 if (cur_buf_ptr + cur_end_index != buf.ptr + buf.len) {
586590 // Not the most recent allocation; we cannot free it.
587591 return;
588592 }
589593 const new_end_index = cur_end_index - buf.len;
590 cur_end_index = node.trySetEndIndex(cur_end_index, new_end_index) orelse {
594
595 cur_end_index = @cmpxchgWeak(
596 usize,
597 &node.end_index,
598 cur_end_index,
599 new_end_index,
600 .monotonic,
601 .monotonic,
602 ) orelse {
591603 return;
592604 };
593605 }