lib/sys/src/allocator.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const builtin = @import("builtin");
  3 const capabilities = @import("capabilities.zig");
  4 const memory = @import("memory.zig");
  5 
  6 const Allocator = std.mem.Allocator;
  7 
  8 pub const required_capabilities = capabilities.noLibc(&.{.allocator});
  9 
 10 var process_allocator_override = std.atomic.Value(usize).init(0);
 11 var process_allocator_readers = std.atomic.Value(usize).init(0);
 12 var process_allocator_context: u8 = 0;
 13 
 14 const transition_address = 1;
 15 
 16 const process_allocator: Allocator = .{
 17     .ptr = &process_allocator_context,
 18     .vtable = &process_allocator_vtable,
 19 };
 20 
 21 pub const ProcessAllocatorOverride = struct {
 22     handle: *const Allocator,
 23     active: bool = true,
 24 
 25     pub fn deinit(self: *ProcessAllocatorOverride) void {
 26         if (!self.active) return;
 27         const address = @intFromPtr(self.handle);
 28         const previous = process_allocator_override.cmpxchgStrong(
 29             address,
 30             transition_address,
 31             .acq_rel,
 32             .acquire,
 33         );
 34         std.debug.assert(previous == null);
 35         while (process_allocator_readers.load(.acquire) != 0) {
 36             std.atomic.spinLoopHint();
 37         }
 38         process_allocator_override.store(0, .release);
 39         self.active = false;
 40     }
 41 };
 42 
 43 pub fn processAllocator() Allocator {
 44     return process_allocator;
 45 }
 46 
 47 pub fn observerControlAllocator() Allocator {
 48     return baseProcessAllocator();
 49 }
 50 
 51 pub fn installProcessAllocatorOverride(
 52     handle: *const Allocator,
 53 ) error{AlreadyInstalled}!ProcessAllocatorOverride {
 54     const address = @intFromPtr(handle);
 55     std.debug.assert(address > transition_address);
 56     if (process_allocator_override.cmpxchgStrong(
 57         0,
 58         transition_address,
 59         .acq_rel,
 60         .acquire,
 61     ) != null) {
 62         return error.AlreadyInstalled;
 63     }
 64     while (process_allocator_readers.load(.acquire) != 0) {
 65         std.atomic.spinLoopHint();
 66     }
 67     process_allocator_override.store(address, .release);
 68     return .{ .handle = handle };
 69 }
 70 
 71 pub fn hostInteropAllocator() Allocator {
 72     return processAllocator();
 73 }
 74 
 75 pub fn hostInteropFreeRaw(ptr: ?*anyopaque) void {
 76     _ = ptr;
 77 }
 78 
 79 pub fn hostInteropRawFreeAvailable() bool {
 80     return false;
 81 }
 82 
 83 pub fn wasmHostAllocator(fallback: Allocator) Allocator {
 84     if (comptime builtin.cpu.arch.isWasm()) {
 85         return fallback;
 86     }
 87     return processAllocator();
 88 }
 89 
 90 pub fn hostOwnedAllocator() Allocator {
 91     if (comptime builtin.cpu.arch.isWasm() and builtin.os.tag == .freestanding) {
 92         return std.heap.wasm_allocator;
 93     }
 94     return hostInteropAllocator();
 95 }
 96 
 97 pub fn requiredHostInteropAllocator(comptime _: []const u8) Allocator {
 98     return hostInteropAllocator();
 99 }
100 
101 pub fn benchmarkAllocator() Allocator {
102     return processAllocator();
103 }
104 
105 pub fn moduleLocalAbiAllocator() Allocator {
106     return processAllocator();
107 }
108 
109 pub fn preloadBackingAllocator() Allocator {
110     return processAllocator();
111 }
112 
113 pub fn debugInfoScratchAllocator() Allocator {
114     return memory.page_allocator;
115 }
116 
117 const ProcessAllocatorSelection = struct {
118     allocator: Allocator,
119 
120     fn release(_: @This()) void {
121         _ = process_allocator_readers.fetchSub(1, .release);
122     }
123 };
124 
125 fn selectProcessAllocator() ProcessAllocatorSelection {
126     while (true) {
127         const address = process_allocator_override.load(.acquire);
128         if (address == transition_address) {
129             std.atomic.spinLoopHint();
130             continue;
131         }
132         _ = process_allocator_readers.fetchAdd(1, .acquire);
133         if (process_allocator_override.load(.acquire) == address) {
134             return .{
135                 .allocator = if (address == 0)
136                     baseProcessAllocator()
137                 else
138                     @as(*const Allocator, @ptrFromInt(address)).*,
139             };
140         }
141         _ = process_allocator_readers.fetchSub(1, .release);
142     }
143 }
144 
145 fn processAllocatorAlloc(
146     _: *anyopaque,
147     len: usize,
148     alignment: std.mem.Alignment,
149     return_address: usize,
150 ) ?[*]u8 {
151     const selected = selectProcessAllocator();
152     defer selected.release();
153     return selected.allocator.rawAlloc(len, alignment, return_address);
154 }
155 
156 fn processAllocatorResize(
157     _: *anyopaque,
158     memory_slice: []u8,
159     alignment: std.mem.Alignment,
160     new_len: usize,
161     return_address: usize,
162 ) bool {
163     const selected = selectProcessAllocator();
164     defer selected.release();
165     return selected.allocator.rawResize(
166         memory_slice,
167         alignment,
168         new_len,
169         return_address,
170     );
171 }
172 
173 fn processAllocatorRemap(
174     _: *anyopaque,
175     memory_slice: []u8,
176     alignment: std.mem.Alignment,
177     new_len: usize,
178     return_address: usize,
179 ) ?[*]u8 {
180     const selected = selectProcessAllocator();
181     defer selected.release();
182     return selected.allocator.rawRemap(
183         memory_slice,
184         alignment,
185         new_len,
186         return_address,
187     );
188 }
189 
190 fn processAllocatorFree(
191     _: *anyopaque,
192     memory_slice: []u8,
193     alignment: std.mem.Alignment,
194     return_address: usize,
195 ) void {
196     const selected = selectProcessAllocator();
197     defer selected.release();
198     selected.allocator.rawFree(
199         memory_slice,
200         alignment,
201         return_address,
202     );
203 }
204 
205 const process_allocator_vtable: Allocator.VTable = .{
206     .alloc = processAllocatorAlloc,
207     .resize = processAllocatorResize,
208     .remap = processAllocatorRemap,
209     .free = processAllocatorFree,
210 };
211 
212 fn baseProcessAllocator() Allocator {
213     if (comptime builtin.cpu.arch.isWasm()) return std.heap.wasm_allocator;
214     return std.heap.smp_allocator;
215 }
216 
217 fn expectSameAllocator(expected: Allocator, actual: Allocator) !void {
218     try std.testing.expectEqual(expected.ptr, actual.ptr);
219     try std.testing.expectEqual(expected.vtable, actual.vtable);
220 }
221 
222 test "process and host interop allocators use libc-free sys policy" {
223     const expected = process_allocator;
224     try expectSameAllocator(expected, processAllocator());
225     try expectSameAllocator(expected, hostInteropAllocator());
226     try expectSameAllocator(
227         baseProcessAllocator(),
228         observerControlAllocator(),
229     );
230 }
231 
232 test "process allocator override is scoped and exclusive" {
233     var storage: [64]u8 = undefined;
234     var fixed = std.heap.FixedBufferAllocator.init(&storage);
235     var handle = fixed.allocator();
236     const process = processAllocator();
237     var session = try installProcessAllocatorOverride(&handle);
238     defer session.deinit();
239 
240     const bytes = try process.alloc(u8, 16);
241     try std.testing.expect(@intFromPtr(bytes.ptr) >= @intFromPtr(&storage));
242     try std.testing.expect(
243         @intFromPtr(bytes.ptr) + bytes.len <=
244             @intFromPtr(&storage) + storage.len,
245     );
246     process.free(bytes);
247     try std.testing.expectError(
248         error.AlreadyInstalled,
249         installProcessAllocatorOverride(&handle),
250     );
251     session.deinit();
252     try expectSameAllocator(process_allocator, processAllocator());
253 }
254 
255 test "host interop raw free is unavailable without libc ownership" {
256     hostInteropFreeRaw(null);
257     try std.testing.expect(!hostInteropRawFreeAvailable());
258 }
259 
260 test "wasm host allocator uses fallback for wasm hosts" {
261     const fallback = std.testing.allocator;
262     const expected = if (comptime builtin.cpu.arch.isWasm())
263         fallback
264     else
265         processAllocator();
266     try expectSameAllocator(expected, wasmHostAllocator(fallback));
267 }
268 
269 test "host-owned allocator uses wasm allocation for browser-style targets" {
270     const expected = if (comptime builtin.cpu.arch.isWasm() and builtin.os.tag == .freestanding)
271         std.heap.wasm_allocator
272     else
273         hostInteropAllocator();
274     try expectSameAllocator(expected, hostOwnedAllocator());
275 }
276 
277 test "module-local ABI and preload policies use the process seam" {
278     try expectSameAllocator(processAllocator(), moduleLocalAbiAllocator());
279     try expectSameAllocator(processAllocator(), preloadBackingAllocator());
280 }
281 
282 test "debug-info scratch policy avoids userspace allocator locks" {
283     try expectSameAllocator(memory.page_allocator, debugInfoScratchAllocator());
284 }
285 
286 test "required host interop allocator uses sys policy" {
287     try expectSameAllocator(hostInteropAllocator(), requiredHostInteropAllocator("test owner"));
288 }