authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-02-07 14:05:28-08:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-02-07 14:41:49-08:00
log1754e014f5da09bf83a7ee1e53132325fd78d1c1
tree840242c06840a052c963349702df4a9566feb0f7
parenta9d30056167cbe54c0c144199f01a56f6cdcdafc

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+48-7
...@@ -47,10 +47,16 @@ var global: SmpAllocator = .{...@@ -47,10 +47,16 @@ var global: SmpAllocator = .{
47threadlocal var thread_index: u32 = 0;47threadlocal var thread_index: u32 = 0;
4848
49const max_thread_count = 128;49const max_thread_count = 128;
50const slab_len: usize = @max(std.heap.page_size_max, 256 * 1024);50const slab_len: usize = @max(std.heap.page_size_max, 64 * 1024);
51/// Because of storing free list pointers, the minimum size class is 3.51/// Because of storing free list pointers, the minimum size class is 3.
52const min_class = math.log2(@sizeOf(usize));52const min_class = math.log2(@sizeOf(usize));
53const size_class_count = math.log2(slab_len) - min_class;53const 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.
56const max_freelist_len: u8 = 16;
57const max_free_search = 1;
58/// Before mapping a fresh page, `alloc` will rotate this many times.
59const max_alloc_search = 1;
5460
55const Thread = struct {61const Thread = struct {
56 /// Avoid false sharing.62 /// Avoid false sharing.
...@@ -62,9 +68,13 @@ const Thread = struct {...@@ -62,9 +68,13 @@ const Thread = struct {
62 /// to support freelist reclamation.68 /// to support freelist reclamation.
63 mutex: std.Thread.Mutex = .{},69 mutex: std.Thread.Mutex = .{},
6470
71 /// For each size class, tracks the next address to be returned from
72 /// `alloc` when the freelist is empty.
65 next_addrs: [size_class_count]usize = @splat(0),73 next_addrs: [size_class_count]usize = @splat(0),
66 /// For each size class, points to the freed pointer.74 /// For each size class, points to the freed pointer.
67 frees: [size_class_count]usize = @splat(0),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),
6878
69 fn lock() *Thread {79 fn lock() *Thread {
70 var index = thread_index;80 var index = thread_index;
...@@ -121,7 +131,6 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ?...@@ -121,7 +131,6 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ?
121131
122 const slot_size = slotSize(class);132 const slot_size = slotSize(class);
123 assert(slab_len % slot_size == 0);133 assert(slab_len % slot_size == 0);
124 const max_search = 1;
125 var search_count: u8 = 0;134 var search_count: u8 = 0;
126135
127 var t = Thread.lock();136 var t = Thread.lock();
...@@ -133,6 +142,7 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ?...@@ -133,6 +142,7 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ?
133 defer t.unlock();142 defer t.unlock();
134 const node: *usize = @ptrFromInt(top_free_ptr);143 const node: *usize = @ptrFromInt(top_free_ptr);
135 t.frees[class] = node.*;144 t.frees[class] = node.*;
145 t.freelist_lens[class] -|= 1;
136 return @ptrFromInt(top_free_ptr);146 return @ptrFromInt(top_free_ptr);
137 }147 }
138148
...@@ -144,12 +154,13 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ?...@@ -144,12 +154,13 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ?
144 return @ptrFromInt(next_addr);154 return @ptrFromInt(next_addr);
145 }155 }
146156
147 if (search_count >= max_search) {157 if (search_count >= max_alloc_search) {
148 @branchHint(.likely);158 @branchHint(.likely);
149 defer t.unlock();159 defer t.unlock();
150 // slab alignment here ensures the % slab len earlier catches the end of slots.160 // slab alignment here ensures the % slab len earlier catches the end of slots.
151 const slab = PageAllocator.map(slab_len, .fromByteUnits(slab_len)) orelse return null;161 const slab = PageAllocator.map(slab_len, .fromByteUnits(slab_len)) orelse return null;
152 t.next_addrs[class] = @intFromPtr(slab) + slot_size;162 t.next_addrs[class] = @intFromPtr(slab) + slot_size;
163 t.freelist_lens[class] = 0;
153 return slab;164 return slab;
154 }165 }
155166
...@@ -203,12 +214,42 @@ fn free(context: *anyopaque, memory: []u8, alignment: mem.Alignment, ra: usize)...@@ -203,12 +214,42 @@ fn free(context: *anyopaque, memory: []u8, alignment: mem.Alignment, ra: usize)
203 }214 }
204215
205 const node: *usize = @alignCast(@ptrCast(memory.ptr));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 }
206230
207 const t = Thread.lock();231 if (search_count >= max_free_search) {
208 defer t.unlock();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 }
209238
210 node.* = t.frees[class];239 t.unlock();
211 t.frees[class] = @intFromPtr(node);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}
213254
214fn sizeClassIndex(len: usize, alignment: mem.Alignment) usize {255fn sizeClassIndex(len: usize, alignment: mem.Alignment) usize {