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 {...@@ -213,12 +213,12 @@ pub fn reset(arena: *ArenaAllocator, mode: ResetMode) bool {
213const Node = struct {213const Node = struct {
214 /// Only meant to be accessed indirectly via the methods supplied by this type,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.215 /// except if the node is owned by the thread accessing it.
216 /// Must always be an even number to accomodate `resize_bit`.216 /// Must always be an even number to accomodate `resize` bit.
217 size: usize,217 size: Size,
218 /// Concurrent accesses to `end_index` can be monotonic since it is only ever218 /// Concurrent accesses to `end_index` can be monotonic as long as its value
219 /// incremented in `alloc` and `resize` after being compared to `size`.219 /// is compared to a version of `size` before using it to access memory.
220 /// Since `size` can only grow and never shrink, memory access depending on220 /// 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.
222 end_index: usize,222 end_index: usize,
223 /// This field should only be accessed if the node is owned by the thread223 /// This field should only be accessed if the node is owned by the thread
224 /// accessing it.224 /// accessing it.
...@@ -226,16 +226,6 @@ const Node = struct {...@@ -226,16 +226,6 @@ const Node = struct {
226226
227 const resize_bit: usize = 1;227 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
239 fn loadBuf(node: *Node) []u8 {229 fn loadBuf(node: *Node) []u8 {
240 // monotonic is fine since `size` can only ever grow, so the buffer returned230 // monotonic is fine since `size` can only ever grow, so the buffer returned
241 // by this function is always valid memory.231 // by this function is always valid memory.
...@@ -326,19 +316,22 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u...@@ -326,19 +316,22 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
326 retry: while (true) {316 retry: while (true) {
327 const first_node: ?*Node, const prev_size: usize = first_node: {317 const first_node: ?*Node, const prev_size: usize = first_node: {
328 const node = cur_first_node orelse break :first_node .{ null, 0 };318 const node = cur_first_node orelse break :first_node .{ null, 0 };
329 var end_index = node.loadEndIndex();319 const buf = node.loadBuf();
330 while (true) {
331 const buf = node.loadBuf();
332 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
333320
334 if (aligned_index + n > buf.len) {321 // To avoid using a CAS loop in the hot path we atomically increase
335 break :first_node .{ node, buf.len };322 // `end_index` by a large enough amount to be able to always provide
336 }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 {327 const alignable = n + alignment.toByteUnits() - 1;
339 return buf[aligned_index..][0..n].ptr;328 const end_index = @atomicRmw(usize, &node.end_index, .Add, alignable, .monotonic);
340 };329 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
341 }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;
342 };335 };
343336
344 resize: {337 resize: {
...@@ -352,7 +345,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u...@@ -352,7 +345,7 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
352 defer node.endResize(size);345 defer node.endResize(size);
353346
354 const buf = allocated_slice[@sizeOf(Node)..];347 const buf = allocated_slice[@sizeOf(Node)..];
355 const end_index = node.loadEndIndex();348 const end_index = @atomicLoad(usize, &node.end_index, .monotonic);
356 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);349 const aligned_index = alignedIndex(buf.ptr, end_index, alignment);
357 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);350 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...@@ -403,55 +396,59 @@ fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u
403 }396 }
404 }397 }
405398
406 var best_fit_prev: ?*Node = null;399 const candidate: ?*Node, const prev: ?*Node = candidate: {
407 var best_fit: ?*Node = null;400 var best_fit_prev: ?*Node = null;
408 var best_fit_diff: usize = std.math.maxInt(usize);401 var best_fit: ?*Node = null;
409402 var best_fit_diff: usize = std.math.maxInt(usize);
410 var it_prev: ?*Node = null;403
411 var it = free_list;404 var it_prev: ?*Node = null;
412 const candidate: ?*Node, const prev: ?*Node = find: while (it) |node| : ({405 var it = free_list;
413 it_prev = it;406 while (it) |node| : ({
414 it = node.next;407 it_prev = it;
415 }) {408 it = node.next;
416 last_free = node;409 }) {
417 assert(node.size & Node.resize_bit == 0);410 last_free = node;
418 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];411 assert(node.size & Node.resize_bit == 0);
419 const aligned_index = alignedIndex(buf.ptr, 0, alignment);412 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
420 if (buf.len < aligned_index + n) {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
421 const diff = aligned_index + n - buf.len;419 const diff = aligned_index + n - buf.len;
422 if (diff <= best_fit_diff) {420 if (diff <= best_fit_diff) {
423 best_fit_prev = it_prev;421 best_fit_prev = it_prev;
424 best_fit = node;422 best_fit = node;
425 best_fit_diff = diff;423 best_fit_diff = diff;
426 }424 }
427 continue :find;425 } else {
428 }426 // Ideally we want to use all nodes in `free_list` eventually,
429 break :find .{ node, it_prev };427 // so even if none fit we'll try to resize the one that was the
430 } else {428 // closest to being large enough.
431 // Ideally we want to use all nodes in `free_list` eventually,429 if (best_fit) |node| {
432 // so even if none fit we'll try to resize the one that was the430 const allocated_slice = node.allocatedSliceUnsafe();
433 // closest to being large enough.431 const buf = allocated_slice[@sizeOf(Node)..];
434 if (best_fit) |node| {432 const aligned_index = alignedIndex(buf.ptr, 0, alignment);
435 const allocated_slice = node.allocatedSliceUnsafe();433 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);
436 const buf = allocated_slice[@sizeOf(Node)..];434
437 const aligned_index = alignedIndex(buf.ptr, 0, alignment);435 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {
438 const new_size = mem.alignForward(usize, @sizeOf(Node) + aligned_index + n, 2);436 node.size = new_size;
439437 break :candidate .{ node, best_fit_prev };
440 if (arena.child_allocator.rawResize(allocated_slice, .of(Node), new_size, @returnAddress())) {438 }
441 node.size = new_size;
442 break :find .{ node, best_fit_prev };
443 }439 }
440 break :from_free_list;
444 }441 }
445 break :from_free_list;
446 };442 };
447443
448 it = last_free;444 {
449 while (it) |node| : (it = node.next) {445 var it = last_free;
450 last_free = node;446 while (it) |node| : (it = node.next) {
447 last_free = node;
448 }
451 }449 }
452450
453 const node = candidate orelse break :from_free_list;451 const node = candidate orelse break :from_free_list;
454
455 const old_next = node.next;452 const old_next = node.next;
456453
457 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];454 const buf = node.allocatedSliceUnsafe()[@sizeOf(Node)..];
...@@ -533,7 +530,7 @@ fn resize(ctx: *anyopaque, buf: []u8, alignment: Alignment, new_len: usize, ret_...@@ -533,7 +530,7 @@ fn resize(ctx: *anyopaque, buf: []u8, alignment: Alignment, new_len: usize, ret_
533 const node = arena.loadFirstNode().?;530 const node = arena.loadFirstNode().?;
534 const cur_buf_ptr = @as([*]u8, @ptrCast(node)) + @sizeOf(Node);531 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);
537 while (true) {534 while (true) {
538 if (cur_buf_ptr + cur_end_index != buf.ptr + buf.len) {535 if (cur_buf_ptr + cur_end_index != buf.ptr + buf.len) {
539 // It's not the most recent allocation, so it cannot be expanded,536 // 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_...@@ -554,7 +551,14 @@ fn resize(ctx: *anyopaque, buf: []u8, alignment: Alignment, new_len: usize, ret_
554 return false;551 return false;
555 };552 };
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 {
558 return true;562 return true;
559 };563 };
560 }564 }
...@@ -580,14 +584,22 @@ fn free(ctx: *anyopaque, buf: []u8, alignment: Alignment, ret_addr: usize) void...@@ -580,14 +584,22 @@ fn free(ctx: *anyopaque, buf: []u8, alignment: Alignment, ret_addr: usize) void
580 const node = arena.loadFirstNode().?;584 const node = arena.loadFirstNode().?;
581 const cur_buf_ptr: [*]u8 = @as([*]u8, @ptrCast(node)) + @sizeOf(Node);585 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);
584 while (true) {588 while (true) {
585 if (cur_buf_ptr + cur_end_index != buf.ptr + buf.len) {589 if (cur_buf_ptr + cur_end_index != buf.ptr + buf.len) {
586 // Not the most recent allocation; we cannot free it.590 // Not the most recent allocation; we cannot free it.
587 return;591 return;
588 }592 }
589 const new_end_index = cur_end_index - buf.len;593 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 {
591 return;603 return;
592 };604 };
593 }605 }