lib/alloc/arena/src/arena.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const observe = @import("alloc_observe");
  3 
  4 const Allocator = std.mem.Allocator;
  5 const Alignment = std.mem.Alignment;
  6 
  7 /// `Arena` groups allocations under one caller-controlled lifetime and releases
  8 /// them together. When allocation observation is disabled, `Arena` aliases
  9 /// `std.heap.ArenaAllocator`. When observation is enabled through
 10 /// `-Dobserve-allocations=true`, it aliases `ObservedArena`, which wraps a
 11 /// standard arena with an observation identity. Both variants provide `init`,
 12 /// `allocator`, `deinit`, `queryCapacity`, `reset`, and the `ResetMode` union.
 13 /// Standard library details such as `std.heap.ArenaAllocator.State`, `promote`,
 14 /// and concrete struct fields are not guaranteed when observation is active.
 15 ///
 16 /// The owner enforces no fixed byte or node capacity cap of its own, drawing
 17 /// fresh blocks on demand from a backing allocator. An `Allocator` handle
 18 /// returned by `allocator` stores a pointer back to the arena, requiring the
 19 /// caller to keep the arena owner at a stable memory address while using that
 20 /// handle. The caller must keep the backing allocator state alive until the
 21 /// arena and all allocated memory are torn down, but the small handle value
 22 /// itself can be copied freely.
 23 ///
 24 /// Calling `reset` or `deinit` ends the lifetime of all existing allocations at
 25 /// once and invalidates every slice previously handed out. The arena does not
 26 /// invoke destructors on individual objects before reclaiming storage, and
 27 /// physical release of memory to the operating system depends on the backing
 28 /// allocator.
 29 ///
 30 /// The allocation vtable is thread-safe when the backing allocator is
 31 /// thread-safe. In contrast, `reset` and `deinit` mutate internal state, and
 32 /// `queryCapacity` reads shared state without synchronization, requiring
 33 /// callers to exclude concurrent operations during those calls. When
 34 /// observation is active, the allocation vtable remains thread-safe if both the
 35 /// backing allocator and observation sinks support concurrent calls, but
 36 /// observation callbacks require their own synchronization.
 37 pub const Arena = if (observe.enabled)
 38     ObservedArena
 39 else
 40     std.heap.ArenaAllocator;
 41 
 42 /// This wrapper embeds a standard library arena allocator and emits logical
 43 /// memory operations and lifecycle events to the alloc_observe subsystem. The
 44 /// wrapper assigns an immutable producer identity upon initialization that
 45 /// survives moving the wrapper value in memory, though existing borrowed
 46 /// Allocator handles continue to point to the owner's original memory address.
 47 /// The wrapper acts as the sole owner of its inner storage and cannot be cloned
 48 /// into an independent copy. Its observation generation counter begins at zero
 49 /// and increments after every call to reset.
 50 const ObservedArena = struct {
 51     /// The underlying standard library arena allocator that owns backing memory
 52     /// blocks and internal chunk bookkeeping, maintained directly by
 53     /// ObservedArena. Mutating this instance directly bypasses event emission
 54     /// and breaks lifecycle observation.
 55     inner: std.heap.ArenaAllocator,
 56     /// The `producerId` is a numeric identifier assigned at initialization from
 57     /// a process counter to identify the arena's operation stream. It is
 58     /// allocated whether or not an observation sink session is installed. The
 59     /// value derives an `owner_cookie` that is computed from this identifier
 60     /// rather than a memory address. This derivation keeps the cookie
 61     /// consistent when the owner moves in memory, but does not guarantee global
 62     /// collision freedom.
 63     producer_id: u64,
 64     /// The lifecycle generation counter tracks arena observation epochs,
 65     /// starting at 0 on initialization. The counter increments after each call
 66     /// to `reset`, including calls that return `false`, and panics on integer
 67     /// overflow.
 68     generation: u64 = 0,
 69 
 70     const Self = @This();
 71 
 72     /// Selects how the arena manages previously allocated storage during
 73     /// `reset`, aliasing `std.heap.ArenaAllocator.ResetMode`. The `free_all`
 74     /// mode returns all tracked memory nodes to the backing allocator and
 75     /// clears internal state to empty. The `retain_capacity` mode requests
 76     /// retention and preheating of existing buffer capacity for reuse, but does
 77     /// not guarantee that the immediate next allocation will avoid a backing
 78     /// request. The `retain_with_limit` mode requests an upper bound on
 79     /// retained payload capacity following the reset rather than establishing a
 80     /// lifetime memory cap. When `reset` fails and returns `false`, the arena
 81     /// may leave an old rewound node whose capacity exceeds the requested
 82     /// limit. The mode does not enforce a hard limit on future allocation
 83     /// requests or guarantee that unneeded storage is physically released to
 84     /// the operating system.
 85     pub const ResetMode = std.heap.ArenaAllocator.ResetMode;
 86 
 87     /// Creates an empty arena owner backed by the provided allocator,
 88     /// performing no backing block allocation during initialization. In the
 89     /// observed variant, this constructor assigns a new producer identity and
 90     /// sets the observation generation counter to 0.
 91     pub fn init(backing: Allocator) Self {
 92         return .{
 93             .inner = std.heap.ArenaAllocator.init(backing),
 94             .producer_id = observe.producerId(),
 95         };
 96     }
 97 
 98     /// Returns a borrowed std.mem.Allocator handle containing a pointer to this
 99     /// owner. The caller must keep the owner alive at a stable memory address
100     /// for the duration of the handle's use, and retains ownership of the
101     /// underlying backing allocator.
102     pub fn allocator(self: *Self) Allocator {
103         return .{
104             .ptr = self,
105             .vtable = &vtable,
106         };
107     }
108 
109     /// Releases all tracked backing memory blocks by calling deinit on the
110     /// inner arena, and emits a lifecycle deinit event for the current
111     /// observation generation. Every allocation handed out by the arena is
112     /// immediately invalidated, and physical memory reclamation depends on the
113     /// backing allocator. Any required destructor calls or object cleanup must
114     /// occur before tearing down the arena storage. The operation cannot run
115     /// concurrently with allocation, capacity queries, or resets, and does not
116     /// reset internal state to permit reusing the instance afterward.
117     pub fn deinit(self: *Self) void {
118         var span = observe.beginLifecycle(
119             self.identity(),
120             .end,
121             .deinit,
122             @returnAddress(),
123         );
124         self.inner.deinit();
125         span.finish(.{ .succeeded = true });
126     }
127 
128     /// Returns the total payload capacity in bytes currently held across the
129     /// arena's used and free chunk lists, excluding node header structures.
130     /// This figure includes alignment padding and unallocated space within
131     /// retained blocks, differing from the count of active payload bytes. The
132     /// calculation provides no saturation guarantee against numeric overflow,
133     /// and callers must serialize capacity queries with any concurrent
134     /// mutations.
135     pub fn queryCapacity(self: Self) usize {
136         return self.inner.queryCapacity();
137     }
138 
139     /// Ends the lifetime of all existing allocations, invalidating previously
140     /// issued slices even if subsequent allocations reuse the same memory
141     /// addresses. The return value indicates whether the requested retention or
142     /// reallocation policy succeeded. When returning false, the arena remains
143     /// functional but may leave an existing rewound chunk in place whose
144     /// capacity exceeds the requested limit, while free_all always succeeds and
145     /// returns true. Under observation, the method emits an invalidation event
146     /// for the old generation reporting whether retention succeeded, then
147     /// increments the generation counter, panicking on counter overflow.
148     /// Callers must exclude all concurrent access while resetting the arena.
149     pub fn reset(self: *Self, mode: ResetMode) bool {
150         var span = observe.beginLifecycle(
151             self.identity(),
152             .invalidate,
153             .reset,
154             @returnAddress(),
155         );
156         const succeeded = self.inner.reset(mode);
157         span.finish(.{ .succeeded = succeeded });
158         self.generation = nextGeneration(self.generation);
159         return succeeded;
160     }
161 
162     fn rawAlloc(
163         context: *anyopaque,
164         len: usize,
165         alignment: Alignment,
166         return_address: usize,
167     ) ?[*]u8 {
168         const self: *Self = @ptrCast(@alignCast(context));
169         var span = observe.beginOwned(
170             self.identity(),
171             .alloc,
172             0,
173             0,
174             len,
175             alignment.toByteUnits(),
176             return_address,
177         );
178         const result = self.inner.allocator().rawAlloc(
179             len,
180             alignment,
181             return_address,
182         );
183         span.finish(.{
184             .address = if (result) |ptr| @intFromPtr(ptr) else 0,
185             .succeeded = result != null,
186         });
187         return result;
188     }
189 
190     fn rawResize(
191         context: *anyopaque,
192         memory: []u8,
193         alignment: Alignment,
194         new_len: usize,
195         return_address: usize,
196     ) bool {
197         const self: *Self = @ptrCast(@alignCast(context));
198         var span = observe.beginOwned(
199             self.identity(),
200             .resize,
201             @intFromPtr(memory.ptr),
202             memory.len,
203             new_len,
204             alignment.toByteUnits(),
205             return_address,
206         );
207         const succeeded = self.inner.allocator().rawResize(
208             memory,
209             alignment,
210             new_len,
211             return_address,
212         );
213         span.finish(.{
214             .address = if (succeeded) @intFromPtr(memory.ptr) else 0,
215             .succeeded = succeeded,
216         });
217         return succeeded;
218     }
219 
220     fn rawRemap(
221         context: *anyopaque,
222         memory: []u8,
223         alignment: Alignment,
224         new_len: usize,
225         return_address: usize,
226     ) ?[*]u8 {
227         const self: *Self = @ptrCast(@alignCast(context));
228         var span = observe.beginOwned(
229             self.identity(),
230             .remap,
231             @intFromPtr(memory.ptr),
232             memory.len,
233             new_len,
234             alignment.toByteUnits(),
235             return_address,
236         );
237         const result = self.inner.allocator().rawRemap(
238             memory,
239             alignment,
240             new_len,
241             return_address,
242         );
243         span.finish(.{
244             .address = if (result) |ptr| @intFromPtr(ptr) else 0,
245             .succeeded = result != null,
246         });
247         return result;
248     }
249 
250     fn rawFree(
251         context: *anyopaque,
252         memory: []u8,
253         alignment: Alignment,
254         return_address: usize,
255     ) void {
256         const self: *Self = @ptrCast(@alignCast(context));
257         var span = observe.beginOwned(
258             self.identity(),
259             .free,
260             @intFromPtr(memory.ptr),
261             memory.len,
262             0,
263             alignment.toByteUnits(),
264             return_address,
265         );
266         self.inner.allocator().rawFree(
267             memory,
268             alignment,
269             return_address,
270         );
271         span.finish(.{
272             .address = @intFromPtr(memory.ptr),
273             .succeeded = true,
274         });
275     }
276 
277     const vtable: Allocator.VTable = .{
278         .alloc = rawAlloc,
279         .resize = rawResize,
280         .remap = rawRemap,
281         .free = rawFree,
282     };
283 
284     fn identity(self: *const Self) observe.Identity {
285         return observe.Identity.movable(
286             self.producer_id,
287             .arena,
288             self.generation,
289         );
290     }
291 };
292 
293 fn nextGeneration(generation: u64) u64 {
294     return std.math.add(u64, generation, 1) catch
295         @panic("allocator observation generation exhausted");
296 }