| ... | ... | @@ -0,0 +1,288 @@ |
| 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 | //! Avoid locking the global mutex as much as possible. |
| 10 | //! |
| 11 | //! Each thread gets a separate freelist, however, the data must be recoverable |
| 12 | //! when the thread exits. We do not directly learn when a thread exits, so |
| 13 | //! occasionally, one thread must attempt to reclaim another thread's |
| 14 | //! resources. |
| 15 | //! |
| 16 | //! Above a certain size, those allocations are memory mapped directly, with no |
| 17 | //! storage of allocation metadata. This works because the implementation |
| 18 | //! refuses resizes that would move an allocation from small category to large |
| 19 | //! category or vice versa. |
| 20 | //! |
| 21 | //! Each allocator operation checks the thread identifier from a threadlocal |
| 22 | //! variable to find out which metadata in the global state to access, and |
| 23 | //! attempts to grab its lock. This will usually succeed without contention, |
| 24 | //! unless another thread has been assigned the same id. In the case of such |
| 25 | //! contention, the thread moves on to the next thread metadata slot and |
| 26 | //! repeats the process of attempting to obtain the lock. |
| 27 | //! |
| 28 | //! By limiting the thread-local metadata array to the same number as the CPU |
| 29 | //! count, ensures that as threads are created and destroyed, they cycle |
| 30 | //! through the full set of freelists. |
| 31 | |
| 32 | const builtin = @import("builtin"); |
| 33 | const native_os = builtin.os.tag; |
| 34 | |
| 35 | const std = @import("../std.zig"); |
| 36 | const assert = std.debug.assert; |
| 37 | const mem = std.mem; |
| 38 | const math = std.math; |
| 39 | const Allocator = std.mem.Allocator; |
| 40 | const SmpAllocator = @This(); |
| 41 | const PageAllocator = std.heap.PageAllocator; |
| 42 | |
| 43 | /// Protects the state in this struct (global state), except for `threads` |
| 44 | /// which each have their own mutex. |
| 45 | mutex: std.Thread.Mutex, |
| 46 | next_thread_index: u32, |
| 47 | cpu_count: u32, |
| 48 | threads: [max_thread_count]Thread, |
| 49 | |
| 50 | var global: SmpAllocator = .{ |
| 51 | .mutex = .{}, |
| 52 | .next_thread_index = 0, |
| 53 | .threads = @splat(.{}), |
| 54 | .cpu_count = 0, |
| 55 | }; |
| 56 | threadlocal var thread_id: Thread.Id = .none; |
| 57 | |
| 58 | const max_thread_count = 128; |
| 59 | const slab_len: usize = @max(std.heap.page_size_max, switch (builtin.os.tag) { |
| 60 | .windows => 64 * 1024, // Makes `std.heap.PageAllocator` take the happy path. |
| 61 | .wasi => 64 * 1024, // Max alignment supported by `std.heap.WasmAllocator`. |
| 62 | else => 256 * 1024, // Avoids too many active mappings when `page_size_max` is low. |
| 63 | }); |
| 64 | /// Because of storing free list pointers, the minimum size class is 3. |
| 65 | const min_class = math.log2(math.ceilPowerOfTwoAssert(usize, 1 + @sizeOf(usize))); |
| 66 | const size_class_count = math.log2(slab_len) - min_class; |
| 67 | |
| 68 | const Thread = struct { |
| 69 | /// Avoid false sharing. |
| 70 | _: void align(std.atomic.cache_line) = {}, |
| 71 | |
| 72 | /// Protects the state in this struct (per-thread state). |
| 73 | /// |
| 74 | /// Threads lock this before accessing their own state in order |
| 75 | /// to support freelist reclamation. |
| 76 | mutex: std.Thread.Mutex = .{}, |
| 77 | |
| 78 | next_addrs: [size_class_count]usize = @splat(0), |
| 79 | /// For each size class, points to the freed pointer. |
| 80 | frees: [size_class_count]usize = @splat(0), |
| 81 | |
| 82 | /// Index into `SmpAllocator.threads`. |
| 83 | const Id = enum(usize) { |
| 84 | none = 0, |
| 85 | first = 1, |
| 86 | _, |
| 87 | |
| 88 | fn fromIndex(index: usize) Id { |
| 89 | return @enumFromInt(index + 1); |
| 90 | } |
| 91 | |
| 92 | fn toIndex(id: Id) usize { |
| 93 | return @intFromEnum(id) - 1; |
| 94 | } |
| 95 | }; |
| 96 | |
| 97 | fn lock() *Thread { |
| 98 | const id = thread_id; |
| 99 | if (id != .none) { |
| 100 | var index = id.toIndex(); |
| 101 | { |
| 102 | const t = &global.threads[index]; |
| 103 | if (t.mutex.tryLock()) return t; |
| 104 | } |
| 105 | const cpu_count = global.cpu_count; |
| 106 | assert(cpu_count != 0); |
| 107 | while (true) { |
| 108 | index = (index + 1) % cpu_count; |
| 109 | const t = &global.threads[index]; |
| 110 | if (t.mutex.tryLock()) { |
| 111 | thread_id = .fromIndex(index); |
| 112 | return t; |
| 113 | } |
| 114 | } |
| 115 | } |
| 116 | while (true) { |
| 117 | const thread_index = i: { |
| 118 | global.mutex.lock(); |
| 119 | defer global.mutex.unlock(); |
| 120 | const cpu_count = c: { |
| 121 | const cpu_count = global.cpu_count; |
| 122 | if (cpu_count == 0) { |
| 123 | const n: u32 = @intCast(@max(std.Thread.getCpuCount() catch max_thread_count, max_thread_count)); |
| 124 | global.cpu_count = n; |
| 125 | break :c n; |
| 126 | } |
| 127 | break :c cpu_count; |
| 128 | }; |
| 129 | const thread_index = global.next_thread_index; |
| 130 | global.next_thread_index = @intCast((thread_index + 1) % cpu_count); |
| 131 | break :i thread_index; |
| 132 | }; |
| 133 | const t = &global.threads[thread_index]; |
| 134 | if (t.mutex.tryLock()) { |
| 135 | thread_id = .fromIndex(thread_index); |
| 136 | return t; |
| 137 | } |
| 138 | } |
| 139 | } |
| 140 | |
| 141 | fn unlock(t: *Thread) void { |
| 142 | t.mutex.unlock(); |
| 143 | } |
| 144 | }; |
| 145 | |
| 146 | pub const vtable: Allocator.VTable = .{ |
| 147 | .alloc = alloc, |
| 148 | .resize = resize, |
| 149 | .remap = remap, |
| 150 | .free = free, |
| 151 | }; |
| 152 | |
| 153 | comptime { |
| 154 | assert(!builtin.single_threaded); // you're holding it wrong |
| 155 | } |
| 156 | |
| 157 | fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ?[*]u8 { |
| 158 | _ = context; |
| 159 | _ = ra; |
| 160 | const class = sizeClassIndex(len, alignment); |
| 161 | if (class >= size_class_count) { |
| 162 | @branchHint(.unlikely); |
| 163 | return PageAllocator.map(len, alignment); |
| 164 | } |
| 165 | |
| 166 | const t = Thread.lock(); |
| 167 | defer t.unlock(); |
| 168 | |
| 169 | const slot_size = slotSize(class); |
| 170 | |
| 171 | const top_free_ptr = t.frees[class]; |
| 172 | if (top_free_ptr != 0) { |
| 173 | const node: *usize = @ptrFromInt(top_free_ptr + (slot_size - @sizeOf(usize))); |
| 174 | t.frees[class] = node.*; |
| 175 | return @ptrFromInt(top_free_ptr); |
| 176 | } |
| 177 | |
| 178 | const next_addr = t.next_addrs[class]; |
| 179 | if (next_addr % slab_len == 0) { |
| 180 | const slab = PageAllocator.map(slab_len, .fromByteUnits(std.heap.pageSize())) orelse return null; |
| 181 | t.next_addrs[class] = @intFromPtr(slab) + slot_size; |
| 182 | return slab; |
| 183 | } |
| 184 | |
| 185 | t.next_addrs[class] = next_addr + slot_size; |
| 186 | return @ptrFromInt(next_addr); |
| 187 | } |
| 188 | |
| 189 | fn resize(context: *anyopaque, memory: []u8, alignment: mem.Alignment, new_len: usize, ra: usize) bool { |
| 190 | _ = context; |
| 191 | _ = ra; |
| 192 | const class = sizeClassIndex(memory.len, alignment); |
| 193 | const new_class = sizeClassIndex(new_len, alignment); |
| 194 | if (class >= size_class_count) { |
| 195 | if (new_class < size_class_count) return false; |
| 196 | return PageAllocator.realloc(memory, new_len, false) != null; |
| 197 | } |
| 198 | return new_class == class; |
| 199 | } |
| 200 | |
| 201 | fn remap(context: *anyopaque, memory: []u8, alignment: mem.Alignment, new_len: usize, ra: usize) ?[*]u8 { |
| 202 | _ = context; |
| 203 | _ = ra; |
| 204 | const class = sizeClassIndex(memory.len, alignment); |
| 205 | const new_class = sizeClassIndex(new_len, alignment); |
| 206 | if (class >= size_class_count) { |
| 207 | if (new_class < size_class_count) return null; |
| 208 | return PageAllocator.realloc(memory, new_len, true); |
| 209 | } |
| 210 | return if (new_class == class) memory.ptr else null; |
| 211 | } |
| 212 | |
| 213 | fn free(context: *anyopaque, memory: []u8, alignment: mem.Alignment, ra: usize) void { |
| 214 | _ = context; |
| 215 | _ = ra; |
| 216 | const class = sizeClassIndex(memory.len, alignment); |
| 217 | if (class >= size_class_count) { |
| 218 | @branchHint(.unlikely); |
| 219 | return PageAllocator.unmap(@alignCast(memory)); |
| 220 | } |
| 221 | |
| 222 | const t = Thread.lock(); |
| 223 | defer t.unlock(); |
| 224 | |
| 225 | const slot_size = slotSize(class); |
| 226 | const addr = @intFromPtr(memory.ptr); |
| 227 | const node: *usize = @ptrFromInt(addr + (slot_size - @sizeOf(usize))); |
| 228 | node.* = t.frees[class]; |
| 229 | t.frees[class] = addr; |
| 230 | } |
| 231 | |
| 232 | fn sizeClassIndex(len: usize, alignment: mem.Alignment) usize { |
| 233 | return @max( |
| 234 | @bitSizeOf(usize) - @clz(len - 1), |
| 235 | @intFromEnum(alignment), |
| 236 | min_class, |
| 237 | ); |
| 238 | } |
| 239 | |
| 240 | fn slotSize(class: usize) usize { |
| 241 | const Log2USize = std.math.Log2Int(usize); |
| 242 | return @as(usize, 1) << @as(Log2USize, @intCast(class)); |
| 243 | } |
| 244 | |
| 245 | test "large alloc, resize, remap, free" { |
| 246 | const gpa = std.heap.smp_allocator; |
| 247 | |
| 248 | const ptr1 = try gpa.alloc(u64, 42768); |
| 249 | const ptr2 = try gpa.alloc(u64, 52768); |
| 250 | gpa.free(ptr1); |
| 251 | const ptr3 = try gpa.alloc(u64, 62768); |
| 252 | gpa.free(ptr3); |
| 253 | gpa.free(ptr2); |
| 254 | } |
| 255 | |
| 256 | test "small allocations - free in same order" { |
| 257 | const gpa = std.heap.smp_allocator; |
| 258 | |
| 259 | var list = std.ArrayList(*u64).init(std.testing.allocator); |
| 260 | defer list.deinit(); |
| 261 | |
| 262 | var i: usize = 0; |
| 263 | while (i < 513) : (i += 1) { |
| 264 | const ptr = try gpa.create(u64); |
| 265 | try list.append(ptr); |
| 266 | } |
| 267 | |
| 268 | for (list.items) |ptr| { |
| 269 | gpa.destroy(ptr); |
| 270 | } |
| 271 | } |
| 272 | |
| 273 | test "small allocations - free in reverse order" { |
| 274 | const gpa = std.heap.smp_allocator; |
| 275 | |
| 276 | var list = std.ArrayList(*u64).init(std.testing.allocator); |
| 277 | defer list.deinit(); |
| 278 | |
| 279 | var i: usize = 0; |
| 280 | while (i < 513) : (i += 1) { |
| 281 | const ptr = try gpa.create(u64); |
| 282 | try list.append(ptr); |
| 283 | } |
| 284 | |
| 285 | while (list.popOrNull()) |ptr| { |
| 286 | gpa.destroy(ptr); |
| 287 | } |
| 288 | } |
Andrew Kelley <andrew@ziglang.org> 2025-02-06 03:31:32-08:00