| ... | ... | @@ -39,20 +39,14 @@ const Allocator = std.mem.Allocator; |
| 39 | 39 | const SmpAllocator = @This(); |
| 40 | 40 | const PageAllocator = std.heap.PageAllocator; |
| 41 | 41 | |
| 42 | | /// Protects the state in this struct (global state), except for `threads` |
| 43 | | /// which each have their own mutex. |
| 44 | | mutex: std.Thread.Mutex, |
| 45 | | next_thread_index: u32, |
| 46 | 42 | cpu_count: u32, |
| 47 | 43 | threads: [max_thread_count]Thread, |
| 48 | 44 | |
| 49 | 45 | var global: SmpAllocator = .{ |
| 50 | | .mutex = .{}, |
| 51 | | .next_thread_index = 0, |
| 52 | 46 | .threads = @splat(.{}), |
| 53 | 47 | .cpu_count = 0, |
| 54 | 48 | }; |
| 55 | | threadlocal var thread_id: Thread.Id = .none; |
| 49 | threadlocal var thread_index: u32 = 0; |
| 56 | 50 | |
| 57 | 51 | const max_thread_count = 128; |
| 58 | 52 | const slab_len: usize = @max(std.heap.page_size_max, 256 * 1024); |
| ... | ... | @@ -74,60 +68,22 @@ const Thread = struct { |
| 74 | 68 | /// For each size class, points to the freed pointer. |
| 75 | 69 | frees: [size_class_count]usize = @splat(0), |
| 76 | 70 | |
| 77 | | /// Index into `SmpAllocator.threads`. |
| 78 | | const Id = enum(usize) { |
| 79 | | none = 0, |
| 80 | | first = 1, |
| 81 | | _, |
| 82 | | |
| 83 | | fn fromIndex(index: usize) Id { |
| 84 | | return @enumFromInt(index + 1); |
| 85 | | } |
| 86 | | |
| 87 | | fn toIndex(id: Id) usize { |
| 88 | | return @intFromEnum(id) - 1; |
| 89 | | } |
| 90 | | }; |
| 91 | | |
| 92 | 71 | fn lock() *Thread { |
| 93 | | const id = thread_id; |
| 94 | | if (id != .none) { |
| 95 | | var index = id.toIndex(); |
| 96 | | { |
| 97 | | const t = &global.threads[index]; |
| 98 | | if (t.mutex.tryLock()) return t; |
| 99 | | } |
| 100 | | const cpu_count = global.cpu_count; |
| 101 | | assert(cpu_count != 0); |
| 102 | | while (true) { |
| 103 | | index = (index + 1) % cpu_count; |
| 104 | | const t = &global.threads[index]; |
| 105 | | if (t.mutex.tryLock()) { |
| 106 | | thread_id = .fromIndex(index); |
| 107 | | return t; |
| 108 | | } |
| 72 | var index = thread_index; |
| 73 | { |
| 74 | const t = &global.threads[index]; |
| 75 | if (t.mutex.tryLock()) { |
| 76 | @branchHint(.likely); |
| 77 | return t; |
| 109 | 78 | } |
| 110 | 79 | } |
| 80 | const cpu_count = getCpuCount(); |
| 81 | assert(cpu_count != 0); |
| 111 | 82 | while (true) { |
| 112 | | const thread_index = i: { |
| 113 | | global.mutex.lock(); |
| 114 | | defer global.mutex.unlock(); |
| 115 | | const cpu_count = c: { |
| 116 | | const cpu_count = global.cpu_count; |
| 117 | | if (cpu_count == 0) { |
| 118 | | const n: u32 = @intCast(@max(std.Thread.getCpuCount() catch max_thread_count, max_thread_count)); |
| 119 | | global.cpu_count = n; |
| 120 | | break :c n; |
| 121 | | } |
| 122 | | break :c cpu_count; |
| 123 | | }; |
| 124 | | const thread_index = global.next_thread_index; |
| 125 | | global.next_thread_index = @intCast((thread_index + 1) % cpu_count); |
| 126 | | break :i thread_index; |
| 127 | | }; |
| 128 | | const t = &global.threads[thread_index]; |
| 83 | index = (index + 1) % cpu_count; |
| 84 | const t = &global.threads[index]; |
| 129 | 85 | if (t.mutex.tryLock()) { |
| 130 | | thread_id = .fromIndex(thread_index); |
| 86 | thread_index = index; |
| 131 | 87 | return t; |
| 132 | 88 | } |
| 133 | 89 | } |
| ... | ... | @@ -138,6 +94,13 @@ const Thread = struct { |
| 138 | 94 | } |
| 139 | 95 | }; |
| 140 | 96 | |
| 97 | fn getCpuCount() u32 { |
| 98 | const cpu_count = @atomicLoad(u32, &global.cpu_count, .unordered); |
| 99 | if (cpu_count != 0) return cpu_count; |
| 100 | const n: u32 = @intCast(@max(std.Thread.getCpuCount() catch max_thread_count, max_thread_count)); |
| 101 | return if (@cmpxchgStrong(u32, &global.cpu_count, 0, n, .monotonic, .monotonic)) |other| other else n; |
| 102 | } |
| 103 | |
| 141 | 104 | pub const vtable: Allocator.VTable = .{ |
| 142 | 105 | .alloc = alloc, |
| 143 | 106 | .resize = resize, |
| ... | ... | @@ -159,8 +122,8 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ? |
| 159 | 122 | } |
| 160 | 123 | |
| 161 | 124 | const slot_size = slotSize(class); |
| 162 | | const max_search = 2; |
| 163 | | var search_count: u32 = 0; |
| 125 | const max_search = 1; |
| 126 | var search_count: u8 = 0; |
| 164 | 127 | |
| 165 | 128 | var t = Thread.lock(); |
| 166 | 129 | |
| ... | ... | @@ -191,15 +154,14 @@ fn alloc(context: *anyopaque, len: usize, alignment: mem.Alignment, ra: usize) ? |
| 191 | 154 | } |
| 192 | 155 | |
| 193 | 156 | t.unlock(); |
| 194 | | t = undefined; |
| 195 | | const cpu_count = global.cpu_count; |
| 157 | const cpu_count = getCpuCount(); |
| 196 | 158 | assert(cpu_count != 0); |
| 197 | | var index = thread_id.toIndex(); |
| 159 | var index = thread_index; |
| 198 | 160 | while (true) { |
| 199 | 161 | index = (index + 1) % cpu_count; |
| 200 | 162 | t = &global.threads[index]; |
| 201 | 163 | if (t.mutex.tryLock()) { |
| 202 | | thread_id = .fromIndex(index); |
| 164 | thread_index = index; |
| 203 | 165 | search_count += 1; |
| 204 | 166 | continue :outer; |
| 205 | 167 | } |
Andrew Kelley <andrew@ziglang.org> 2025-02-06 21:47:46-08:00
lib