Commit 1754e014f5da09bf83a7ee1e53132325fd78d1c1 | zig

author	Andrew Kelley <andrew@ziglang.org> 2025-02-07 14:05:28-08:00
committer	Andrew Kelley <andrew@ziglang.org> 2025-02-07 14:41:49-08:00
log	1754e014f5da09bf83a7ee1e53132325fd78d1c1
tree	840242c06840a052c963349702df4a9566feb0f7
parent	a9d30056167cbe54c0c144199f01a56f6cdcdafc

std.heap.SmpAllocator: rotate on free sometimes

* slab length reduced to 64K
* track freelist length with u8s
* on free(), rotate if freelist length exceeds max_freelist_len

Prevents memory leakage in the scenario where one thread only allocates
and another thread only frees.

1 files changed, 48 insertions(+), 7 deletions(-)

lib
- std
  - heap
    SmpAllocator.zig

lib/std/heap/SmpAllocator.zig+48-7

...	...	@@ -47,10 +47,16 @@ var global: SmpAllocator = .{
47	47	threadlocal var thread_index: u32 = 0;
48	48
49	49	const max_thread_count = 128;
50		const slab_len: usize = @max(std.heap.page_size_max, 256 * 1024);
	50	const slab_len: usize = @max(std.heap.page_size_max, 64 * 1024);
51	51	/// Because of storing free list pointers, the minimum size class is 3.
52	52	const min_class = math.log2(@sizeOf(usize));
53	53	const size_class_count = math.log2(slab_len) - min_class;
	54	/// When a freelist length exceeds this number, a `free` will rotate up to
	55	/// `max_free_search` times before pushing.
	56	const max_freelist_len: u8 = 16;
	57	const max_free_search = 1;
	58	/// Before mapping a fresh page, `alloc` will rotate this many times.
	59	const max_alloc_search = 1;
54	60
55	61	const Thread = struct {
56	62	/// Avoid false sharing.
...	...	@@ -62,9 +68,13 @@ const Thread = struct {
62	68	/// to support freelist reclamation.
63	69	mutex: std.Thread.Mutex = .{},
64	70
	71	/// For each size class, tracks the next address to be returned from
	72	/// `alloc` when the freelist is empty.
65	73	next_addrs: [size_class_count]usize = @splat(0),
66	74	/// For each size class, points to the freed pointer.
67	75	frees: [size_class_count]usize = @splat(0),
	76	/// For each size class, tracks the number of items in the freelist.
	77	freelist_lens: [size_class_count]u8 = @splat(0),
68	78
69	79	fn lock() *Thread {
70	80	var index = thread_index;
...	...	@@ -121,7 +131,6 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ?
121	131
122	132	const slot_size = slotSize(class);
123	133	assert(slab_len % slot_size == 0);
124		const max_search = 1;
125	134	var search_count: u8 = 0;
126	135
127	136	var t = Thread.lock();
...	...	@@ -133,6 +142,7 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ?
133	142	defer t.unlock();
134	143	const node: *usize = @ptrFromInt(top_free_ptr);
135	144	t.frees[class] = node.*;
	145	t.freelist_lens[class] -\|= 1;
136	146	return @ptrFromInt(top_free_ptr);
137	147	}
138	148
...	...	@@ -144,12 +154,13 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ?
144	154	return @ptrFromInt(next_addr);
145	155	}
146	156
147		if (search_count >= max_search) {
	157	if (search_count >= max_alloc_search) {
148	158	@branchHint(.likely);
149	159	defer t.unlock();
150	160	// slab alignment here ensures the % slab len earlier catches the end of slots.
151	161	const slab = PageAllocator.map(slab_len, .fromByteUnits(slab_len)) orelse return null;
152	162	t.next_addrs[class] = @intFromPtr(slab) + slot_size;
	163	t.freelist_lens[class] = 0;
153	164	return slab;
154	165	}
155	166
...	...	@@ -203,12 +214,42 @@ fn free(context: *anyopaque, memory: []u8, alignment: mem.Alignment, ra: usize)
203	214	}
204	215
205	216	const node: *usize = @alignCast(@ptrCast(memory.ptr));
	217	var search_count: u8 = 0;
	218
	219	var t = Thread.lock();
	220
	221	outer: while (true) {
	222	const freelist_len = t.freelist_lens[class];
	223	if (freelist_len < max_freelist_len) {
	224	@branchHint(.likely);
	225	defer t.unlock();
	226	node.* = t.frees[class];
	227	t.frees[class] = @intFromPtr(node);
	228	return;
	229	}
206	230
207		const t = Thread.lock();
208		defer t.unlock();
	231	if (search_count >= max_free_search) {
	232	defer t.unlock();
	233	t.freelist_lens[class] = freelist_len +\| 1;
	234	node.* = t.frees[class];
	235	t.frees[class] = @intFromPtr(node);
	236	return;
	237	}
209	238
210		node.* = t.frees[class];
211		t.frees[class] = @intFromPtr(node);
	239	t.unlock();
	240	const cpu_count = getCpuCount();
	241	assert(cpu_count != 0);
	242	var index = thread_index;
	243	while (true) {
	244	index = (index + 1) % cpu_count;
	245	t = &global.threads[index];
	246	if (t.mutex.tryLock()) {
	247	thread_index = index;
	248	search_count += 1;
	249	continue :outer;
	250	}
	251	}
	252	}
212	253	}
213	254
214	255	fn sizeClassIndex(len: usize, alignment: mem.Alignment) usize {