lib/alloc/observe/src/allocator.zig

daab053ee43316e1809a84551d573ddd1e5bf3d2

  1 const std = @import("std");
  2 const protocol = @import("protocol.zig");
  3 
  4 const Allocator = std.mem.Allocator;
  5 const Alignment = std.mem.Alignment;
  6 
  7 /// A generic adapter wrapping a supplied backing allocator, instrumenting
  8 /// allocation, resizing, remapping, and deallocation virtual table calls while
  9 /// returning backing allocator outcomes unchanged. The adapter emits no
 10 /// lifecycle records and enforces no independent memory reclamation policy.
 11 /// Callers must keep both the adapter and its backing allocator at stable
 12 /// memory addresses throughout use, and must ensure the installed sink can
 13 /// handle concurrent invocations when the backing allocator is accessed across
 14 /// multiple threads.
 15 pub const ObservedAllocator = struct {
 16     backing: Allocator,
 17     producer_id: u64,
 18     producer: protocol.Producer,
 19 
 20     /// Initializes the adapter by storing the backing allocator handle, an
 21     /// explicit producer identifier, and a producer kind. The constructor
 22     /// always asserts that the supplied producer identifier is nonzero, even in
 23     /// disabled builds. Callers must therefore avoid unconditionally passing
 24     /// `producerId`(), which evaluates to zero when observation is disabled.
 25     /// Initialization emits no observation events.
 26     pub fn init(
 27         backing: Allocator,
 28         producer_id: u64,
 29         producer: protocol.Producer,
 30     ) ObservedAllocator {
 31         std.debug.assert(producer_id != 0);
 32         return .{
 33             .backing = backing,
 34             .producer_id = producer_id,
 35             .producer = producer,
 36         };
 37     }
 38 
 39     /// Returns an allocator interface handle borrowing a pointer to this
 40     /// adapter and its virtual table. The adapter and its underlying backing
 41     /// allocator must remain alive and at stable memory addresses for the
 42     /// lifetime of all issued handles. Relocating the owning adapter in memory
 43     /// does not update existing borrowed pointers.
 44     pub fn allocator(self: *ObservedAllocator) Allocator {
 45         return .{ .ptr = self, .vtable = &vtable };
 46     }
 47 
 48     fn rawAlloc(
 49         context: *anyopaque,
 50         len: usize,
 51         alignment: Alignment,
 52         return_address: usize,
 53     ) ?[*]u8 {
 54         const self: *ObservedAllocator = @ptrCast(@alignCast(context));
 55         var span = self.begin(.alloc, 0, 0, len, alignment, return_address);
 56         const result = self.backing.rawAlloc(len, alignment, return_address);
 57         span.finish(.{
 58             .address = if (result) |ptr| @intFromPtr(ptr) else 0,
 59             .succeeded = result != null,
 60         });
 61         return result;
 62     }
 63 
 64     fn rawResize(
 65         context: *anyopaque,
 66         memory: []u8,
 67         alignment: Alignment,
 68         new_len: usize,
 69         return_address: usize,
 70     ) bool {
 71         const self: *ObservedAllocator = @ptrCast(@alignCast(context));
 72         var span = self.begin(
 73             .resize,
 74             @intFromPtr(memory.ptr),
 75             memory.len,
 76             new_len,
 77             alignment,
 78             return_address,
 79         );
 80         const succeeded = self.backing.rawResize(
 81             memory,
 82             alignment,
 83             new_len,
 84             return_address,
 85         );
 86         span.finish(.{
 87             .address = if (succeeded) @intFromPtr(memory.ptr) else 0,
 88             .succeeded = succeeded,
 89         });
 90         return succeeded;
 91     }
 92 
 93     fn rawRemap(
 94         context: *anyopaque,
 95         memory: []u8,
 96         alignment: Alignment,
 97         new_len: usize,
 98         return_address: usize,
 99     ) ?[*]u8 {
100         const self: *ObservedAllocator = @ptrCast(@alignCast(context));
101         var span = self.begin(
102             .remap,
103             @intFromPtr(memory.ptr),
104             memory.len,
105             new_len,
106             alignment,
107             return_address,
108         );
109         const result = self.backing.rawRemap(
110             memory,
111             alignment,
112             new_len,
113             return_address,
114         );
115         span.finish(.{
116             .address = if (result) |ptr| @intFromPtr(ptr) else 0,
117             .succeeded = result != null,
118         });
119         return result;
120     }
121 
122     fn rawFree(
123         context: *anyopaque,
124         memory: []u8,
125         alignment: Alignment,
126         return_address: usize,
127     ) void {
128         const self: *ObservedAllocator = @ptrCast(@alignCast(context));
129         var span = self.begin(
130             .free,
131             @intFromPtr(memory.ptr),
132             memory.len,
133             0,
134             alignment,
135             return_address,
136         );
137         self.backing.rawFree(memory, alignment, return_address);
138         span.finish(.{
139             .address = @intFromPtr(memory.ptr),
140             .succeeded = true,
141         });
142     }
143 
144     fn begin(
145         self: *ObservedAllocator,
146         operation: protocol.Operation,
147         old_address: usize,
148         old_len: usize,
149         len: usize,
150         alignment: Alignment,
151         return_address: usize,
152     ) protocol.Span {
153         return protocol.begin(
154             self.producer_id,
155             self.producer,
156             operation,
157             old_address,
158             old_len,
159             len,
160             alignment.toByteUnits(),
161             return_address,
162         );
163     }
164 
165     const vtable: Allocator.VTable = .{
166         .alloc = rawAlloc,
167         .resize = rawResize,
168         .remap = rawRemap,
169         .free = rawFree,
170     };
171 };