1//! An allocator that is designed for ReleaseFast optimization mode, with
2//! multi-threading enabled.
3//!
4//! This allocator is a singleton; it uses global state and only one should be
5//! instantiated for the entire process.
6//!
7//! ## Basic Design
8//!
9//! Each thread gets a separate freelist, however, the data must be recoverable
10//! when the thread exits. We do not directly learn when a thread exits, so
11//! occasionally, one thread must attempt to reclaim another thread's
12//! resources.
13//!
14//! Above a certain size, those allocations are memory mapped directly, with no
15//! storage of allocation metadata. This works because the implementation
16//! refuses resizes that would move an allocation from small category to large
17//! category or vice versa.
18//!
19//! Each allocator operation checks the thread identifier from a threadlocal
20//! variable to find out which metadata in the global state to access, and
21//! attempts to grab its lock. This will usually succeed without contention,
22//! unless another thread has been assigned the same id. In the case of such
23//! contention, the thread moves on to the next thread metadata slot and
24//! repeats the process of attempting to obtain the lock.
25//!
26//! By limiting the thread-local metadata array to the same number as the CPU
27//! count, ensures that as threads are created and destroyed, they cycle
28//! through the full set of freelists.
29const SmpAllocator = @This();
30
31const builtin = @import("builtin");
32
33const std = @import("../std.zig");
34const assert = std.debug.assert;
35const mem = std.mem;
36const math = std.math;
37const Allocator = std.mem.Allocator;
38const Alignment = std.mem.Alignment;
39const PageAllocator = std.heap.PageAllocator;
40
41cpu_count: u32,
42threads: [max_thread_count]Thread,
43
44var global: SmpAllocator = .{
45 .threads = @splat(.{}),
46 .cpu_count = 0,
47};
48threadlocal var thread_index: u32 = 0;
49
50const max_thread_count = 128;
51const slab_len: usize = @max(std.heap.page_size_max, 64 * 1024);
52/// Because of storing free list pointers, the minimum size class is 3.
53const min_class = math.log2(@sizeOf(usize));
54const size_class_count = math.log2(slab_len) - min_class;
55/// Before mapping a fresh page, `alloc` will rotate this many times.
56const max_alloc_search = 1;
57
58const Thread = struct {
59 /// Avoid false sharing.
60 _: void align(std.atomic.cache_line) = {},
61
62 /// Protects the state in this struct (per-thread state).
63 ///
64 /// Threads lock this before accessing their own state in order
65 /// to support freelist reclamation.
66 mutex: std.atomic.Mutex = .unlocked,
67
68 /// For each size class, tracks the next address to be returned from
69 /// `alloc` when the freelist is empty.
70 next_addrs: [size_class_count]usize = @splat(0),
71 /// For each size class, points to the freed pointer.
72 frees: [size_class_count]usize = @splat(0),
73
74 fn lock() *Thread {
75 var index = thread_index;
76 {
77 const t = &global.threads[index];
78 if (t.mutex.tryLock()) {
79 @branchHint(.likely);
80 return t;
81 }
82 }
83 const cpu_count = getCpuCount();
84 assert(cpu_count != 0);
85 while (true) {
86 index = (index + 1) % cpu_count;
87 const t = &global.threads[index];
88 if (t.mutex.tryLock()) {
89 thread_index = index;
90 return t;
91 }
92 }
93 }
94
95 fn unlock(t: *Thread) void {
96 t.mutex.unlock();
97 }
98};
99
100fn getCpuCount() u32 {
101 const cpu_count = @atomicLoad(u32, &global.cpu_count, .unordered);
102 if (cpu_count != 0) return cpu_count;
103 const n: u32 = @min(std.Thread.getCpuCount() catch max_thread_count, max_thread_count);
104 return if (@cmpxchgStrong(u32, &global.cpu_count, 0, n, .monotonic, .monotonic)) |other| other else n;
105}
106
107pub const vtable: Allocator.VTable = .{
108 .alloc = alloc,
109 .resize = resize,
110 .remap = remap,
111 .free = free,
112};
113
114comptime {
115 assert(!builtin.single_threaded); // you're holding it wrong
116}
117
118fn alloc(context: *anyopaque, len: usize, alignment: Alignment, ra: usize) ?[*]u8 {
119 _ = context;
120 _ = ra;
121 const class = sizeClassIndex(len, alignment);
122 if (class >= size_class_count) {
123 @branchHint(.unlikely);
124 return PageAllocator.map(len, alignment);
125 }
126
127 const slot_size = slotSize(class);
128 assert(slab_len % slot_size == 0);
129 var search_count: u8 = 0;
130
131 var t = Thread.lock();
132
133 outer: while (true) {
134 const top_free_ptr = t.frees[class];
135 if (top_free_ptr != 0) {
136 @branchHint(.likely);
137 defer t.unlock();
138 const node: *usize = @ptrFromInt(top_free_ptr);
139 t.frees[class] = node.*;
140 return @ptrFromInt(top_free_ptr);
141 }
142
143 const next_addr = t.next_addrs[class];
144 if ((next_addr % slab_len) != 0) {
145 @branchHint(.likely);
146 defer t.unlock();
147 t.next_addrs[class] = next_addr + slot_size;
148 return @ptrFromInt(next_addr);
149 }
150
151 if (search_count >= max_alloc_search) {
152 @branchHint(.likely);
153 defer t.unlock();
154 // slab alignment here ensures the % slab len earlier catches the end of slots.
155 const slab = PageAllocator.map(slab_len, .fromByteUnits(slab_len)) orelse return null;
156 t.next_addrs[class] = @intFromPtr(slab) + slot_size;
157 return slab;
158 }
159
160 t.unlock();
161 const cpu_count = getCpuCount();
162 assert(cpu_count != 0);
163 var index = thread_index;
164 while (true) {
165 index = (index + 1) % cpu_count;
166 t = &global.threads[index];
167 if (t.mutex.tryLock()) {
168 thread_index = index;
169 search_count += 1;
170 continue :outer;
171 }
172 }
173 }
174}
175
176fn resize(context: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ra: usize) bool {
177 _ = context;
178 _ = ra;
179 const class = sizeClassIndex(memory.len, alignment);
180 const new_class = sizeClassIndex(new_len, alignment);
181 if (class >= size_class_count) {
182 if (new_class < size_class_count) return false;
183 return PageAllocator.realloc(memory, alignment, new_len, false) != null;
184 }
185 return new_class == class;
186}
187
188fn remap(context: *anyopaque, memory: []u8, alignment: Alignment, new_len: usize, ra: usize) ?[*]u8 {
189 _ = context;
190 _ = ra;
191 const class = sizeClassIndex(memory.len, alignment);
192 const new_class = sizeClassIndex(new_len, alignment);
193 if (class >= size_class_count) {
194 if (new_class < size_class_count) return null;
195 return PageAllocator.realloc(memory, alignment, new_len, true);
196 }
197 return if (new_class == class) memory.ptr else null;
198}
199
200fn free(context: *anyopaque, memory: []u8, alignment: Alignment, ra: usize) void {
201 _ = context;
202 _ = ra;
203 const class = sizeClassIndex(memory.len, alignment);
204 if (class >= size_class_count) {
205 @branchHint(.unlikely);
206 return PageAllocator.unmap(@alignCast(memory));
207 }
208
209 const node: *usize = @ptrCast(@alignCast(memory.ptr));
210
211 const t = Thread.lock();
212 defer t.unlock();
213
214 node.* = t.frees[class];
215 t.frees[class] = @intFromPtr(node);
216}
217
218fn sizeClassIndex(len: usize, alignment: Alignment) usize {
219 return @max(@bitSizeOf(usize) - @clz(len - 1), @backingInt(alignment), min_class) - min_class;
220}
221
222fn slotSize(class: usize) usize {
223 return @as(usize, 1) << @intCast(class + min_class);
224}