|
| 1 | +//! Budget-enforcing allocator for fail2zig components. |
| 2 | +//! |
| 3 | +//! `BudgetAllocator` wraps a backing allocator and enforces a hard byte |
| 4 | +//! ceiling. Allocation requests that would exceed the ceiling are rejected |
| 5 | +//! with `error.OutOfMemory` — never silently serviced. This is the memory |
| 6 | +//! safety backbone for every fail2zig component (state tracker, parser |
| 7 | +//! buffers, event queue, log buffers): under adversarial load, we run out |
| 8 | +//! of budget, not out of system memory. |
| 9 | +//! |
| 10 | +//! Backing storage is a single page-allocated buffer fed into a |
| 11 | +//! `std.heap.FixedBufferAllocator`. The budget layer sits on top, so every |
| 12 | +//! `alloc`/`resize`/`remap`/`free` goes through accounting before touching |
| 13 | +//! the fixed buffer. This gives two guarantees at once: |
| 14 | +//! 1. Hard byte ceiling (no unbounded growth). |
| 15 | +//! 2. Bounded physical backing (no fragmentation across the heap). |
| 16 | + |
| 17 | +const std = @import("std"); |
| 18 | +const Allocator = std.mem.Allocator; |
| 19 | +const Alignment = std.mem.Alignment; |
| 20 | + |
| 21 | +/// Statistics snapshot for a `BudgetAllocator`. |
| 22 | +pub const Stats = struct { |
| 23 | + /// Bytes currently allocated (live). Strictly ≤ `capacity`. |
| 24 | + bytes_allocated: usize, |
| 25 | + /// High-water mark of `bytes_allocated` since init. |
| 26 | + peak_bytes: usize, |
| 27 | + /// Monotonic counter of successful `alloc` calls (for debugging). |
| 28 | + allocation_count: u64, |
| 29 | + /// Configured byte ceiling (budget). |
| 30 | + capacity: usize, |
| 31 | +}; |
| 32 | + |
| 33 | +/// Budget-enforcing allocator. Not thread-safe by design — each component |
| 34 | +/// owns its own `BudgetAllocator` and may wrap it in a mutex if it shares |
| 35 | +/// the allocator across threads. |
| 36 | +pub const BudgetAllocator = struct { |
| 37 | + /// Hard byte ceiling. Allocations past this return `error.OutOfMemory`. |
| 38 | + capacity: usize, |
| 39 | + /// Live bytes outstanding. Incremented on alloc, decremented on free. |
| 40 | + bytes_allocated: usize, |
| 41 | + /// Max observed `bytes_allocated`. |
| 42 | + peak_bytes: usize, |
| 43 | + /// Count of successful `alloc` calls. |
| 44 | + allocation_count: u64, |
| 45 | + |
| 46 | + /// Backing page-allocated buffer. Owned by the `BudgetAllocator` and |
| 47 | + /// freed on `deinit`. Sized to `capacity`. |
| 48 | + backing_buffer: []u8, |
| 49 | + /// FixedBufferAllocator that serves physical memory from `backing_buffer`. |
| 50 | + fba: std.heap.FixedBufferAllocator, |
| 51 | + |
| 52 | + pub const Error = error{OutOfMemory}; |
| 53 | + |
| 54 | + /// Initialize a budget allocator with the given ceiling. Allocates the |
| 55 | + /// full backing buffer from `std.heap.page_allocator` up front — there |
| 56 | + /// is no lazy growth. Caller must `deinit` to release the buffer. |
| 57 | + pub fn init(capacity: usize) Error!BudgetAllocator { |
| 58 | + const buf = std.heap.page_allocator.alloc(u8, capacity) catch |
| 59 | + return error.OutOfMemory; |
| 60 | + return .{ |
| 61 | + .capacity = capacity, |
| 62 | + .bytes_allocated = 0, |
| 63 | + .peak_bytes = 0, |
| 64 | + .allocation_count = 0, |
| 65 | + .backing_buffer = buf, |
| 66 | + .fba = std.heap.FixedBufferAllocator.init(buf), |
| 67 | + }; |
| 68 | + } |
| 69 | + |
| 70 | + /// Release the backing buffer. After this call the allocator is |
| 71 | + /// invalid and any outstanding allocations from it dangle. |
| 72 | + pub fn deinit(self: *BudgetAllocator) void { |
| 73 | + std.heap.page_allocator.free(self.backing_buffer); |
| 74 | + self.* = undefined; |
| 75 | + } |
| 76 | + |
| 77 | + /// Zig Allocator interface for this budget. |
| 78 | + pub fn allocator(self: *BudgetAllocator) Allocator { |
| 79 | + return .{ |
| 80 | + .ptr = self, |
| 81 | + .vtable = &.{ |
| 82 | + .alloc = alloc, |
| 83 | + .resize = resize, |
| 84 | + .remap = remap, |
| 85 | + .free = free, |
| 86 | + }, |
| 87 | + }; |
| 88 | + } |
| 89 | + |
| 90 | + /// Snapshot current statistics. |
| 91 | + pub fn stats(self: *const BudgetAllocator) Stats { |
| 92 | + return .{ |
| 93 | + .bytes_allocated = self.bytes_allocated, |
| 94 | + .peak_bytes = self.peak_bytes, |
| 95 | + .allocation_count = self.allocation_count, |
| 96 | + .capacity = self.capacity, |
| 97 | + }; |
| 98 | + } |
| 99 | + |
| 100 | + /// Bytes available before the ceiling is hit. |
| 101 | + pub fn available(self: *const BudgetAllocator) usize { |
| 102 | + return self.capacity - self.bytes_allocated; |
| 103 | + } |
| 104 | + |
| 105 | + // ------------------------------------------------------------------ |
| 106 | + // Allocator vtable implementation |
| 107 | + // ------------------------------------------------------------------ |
| 108 | + |
| 109 | + fn alloc(ctx: *anyopaque, n: usize, alignment: Alignment, ret_addr: usize) ?[*]u8 { |
| 110 | + const self: *BudgetAllocator = @ptrCast(@alignCast(ctx)); |
| 111 | + |
| 112 | + // Budget check up front — refuse before we even ask the fixed |
| 113 | + // buffer for memory. Use saturating addition to defend against |
| 114 | + // pathological `n` values that would wrap around. |
| 115 | + const new_total = std.math.add(usize, self.bytes_allocated, n) catch { |
| 116 | + @branchHint(.unlikely); |
| 117 | + return null; |
| 118 | + }; |
| 119 | + if (new_total > self.capacity) { |
| 120 | + @branchHint(.unlikely); |
| 121 | + return null; |
| 122 | + } |
| 123 | + |
| 124 | + const fba_alloc = self.fba.allocator(); |
| 125 | + const ptr = fba_alloc.rawAlloc(n, alignment, ret_addr) orelse { |
| 126 | + @branchHint(.unlikely); |
| 127 | + return null; |
| 128 | + }; |
| 129 | + |
| 130 | + self.bytes_allocated = new_total; |
| 131 | + if (self.bytes_allocated > self.peak_bytes) { |
| 132 | + self.peak_bytes = self.bytes_allocated; |
| 133 | + } |
| 134 | + self.allocation_count += 1; |
| 135 | + return ptr; |
| 136 | + } |
| 137 | + |
| 138 | + fn resize( |
| 139 | + ctx: *anyopaque, |
| 140 | + buf: []u8, |
| 141 | + alignment: Alignment, |
| 142 | + new_len: usize, |
| 143 | + ret_addr: usize, |
| 144 | + ) bool { |
| 145 | + const self: *BudgetAllocator = @ptrCast(@alignCast(ctx)); |
| 146 | + |
| 147 | + // A resize may grow or shrink. If growing, ensure the growth fits |
| 148 | + // the budget before asking the FBA. |
| 149 | + if (new_len > buf.len) { |
| 150 | + const delta = new_len - buf.len; |
| 151 | + const new_total = std.math.add(usize, self.bytes_allocated, delta) catch { |
| 152 | + @branchHint(.unlikely); |
| 153 | + return false; |
| 154 | + }; |
| 155 | + if (new_total > self.capacity) { |
| 156 | + @branchHint(.unlikely); |
| 157 | + return false; |
| 158 | + } |
| 159 | + const fba_alloc = self.fba.allocator(); |
| 160 | + if (!fba_alloc.rawResize(buf, alignment, new_len, ret_addr)) return false; |
| 161 | + self.bytes_allocated = new_total; |
| 162 | + if (self.bytes_allocated > self.peak_bytes) { |
| 163 | + self.peak_bytes = self.bytes_allocated; |
| 164 | + } |
| 165 | + return true; |
| 166 | + } |
| 167 | + |
| 168 | + // Shrink or same-size: always accounting-safe. Ask FBA; if it |
| 169 | + // refuses we leave the accounting untouched. |
| 170 | + const fba_alloc = self.fba.allocator(); |
| 171 | + if (!fba_alloc.rawResize(buf, alignment, new_len, ret_addr)) return false; |
| 172 | + const delta = buf.len - new_len; |
| 173 | + self.bytes_allocated -= delta; |
| 174 | + return true; |
| 175 | + } |
| 176 | + |
| 177 | + fn remap( |
| 178 | + ctx: *anyopaque, |
| 179 | + buf: []u8, |
| 180 | + alignment: Alignment, |
| 181 | + new_len: usize, |
| 182 | + ret_addr: usize, |
| 183 | + ) ?[*]u8 { |
| 184 | + const self: *BudgetAllocator = @ptrCast(@alignCast(ctx)); |
| 185 | + |
| 186 | + // Growing remap must pre-check budget. Shrinking is accounting- |
| 187 | + // safe. If the FBA can serve the remap in place, we adjust |
| 188 | + // accounting; otherwise we return null and the caller will do an |
| 189 | + // alloc+copy+free pair that will each pass through our vtable. |
| 190 | + if (new_len > buf.len) { |
| 191 | + const delta = new_len - buf.len; |
| 192 | + const new_total = std.math.add(usize, self.bytes_allocated, delta) catch { |
| 193 | + @branchHint(.unlikely); |
| 194 | + return null; |
| 195 | + }; |
| 196 | + if (new_total > self.capacity) { |
| 197 | + @branchHint(.unlikely); |
| 198 | + return null; |
| 199 | + } |
| 200 | + const fba_alloc = self.fba.allocator(); |
| 201 | + const ptr = fba_alloc.rawRemap(buf, alignment, new_len, ret_addr) orelse return null; |
| 202 | + self.bytes_allocated = new_total; |
| 203 | + if (self.bytes_allocated > self.peak_bytes) { |
| 204 | + self.peak_bytes = self.bytes_allocated; |
| 205 | + } |
| 206 | + return ptr; |
| 207 | + } |
| 208 | + |
| 209 | + const fba_alloc = self.fba.allocator(); |
| 210 | + const ptr = fba_alloc.rawRemap(buf, alignment, new_len, ret_addr) orelse return null; |
| 211 | + const delta = buf.len - new_len; |
| 212 | + self.bytes_allocated -= delta; |
| 213 | + return ptr; |
| 214 | + } |
| 215 | + |
| 216 | + fn free(ctx: *anyopaque, buf: []u8, alignment: Alignment, ret_addr: usize) void { |
| 217 | + const self: *BudgetAllocator = @ptrCast(@alignCast(ctx)); |
| 218 | + const fba_alloc = self.fba.allocator(); |
| 219 | + fba_alloc.rawFree(buf, alignment, ret_addr); |
| 220 | + // FixedBufferAllocator reclaims only the last allocation; we still |
| 221 | + // decrement the budget counter so caps are measured by logical |
| 222 | + // lifetime rather than physical reclamation. |
| 223 | + self.bytes_allocated -= buf.len; |
| 224 | + } |
| 225 | +}; |
| 226 | + |
| 227 | +// ============================================================================ |
| 228 | +// Tests |
| 229 | +// ============================================================================ |
| 230 | + |
| 231 | +const testing = std.testing; |
| 232 | + |
| 233 | +test "BudgetAllocator: init and deinit do not leak" { |
| 234 | + var budget = try BudgetAllocator.init(4096); |
| 235 | + defer budget.deinit(); |
| 236 | + const s = budget.stats(); |
| 237 | + try testing.expectEqual(@as(usize, 4096), s.capacity); |
| 238 | + try testing.expectEqual(@as(usize, 0), s.bytes_allocated); |
| 239 | + try testing.expectEqual(@as(usize, 0), s.peak_bytes); |
| 240 | + try testing.expectEqual(@as(u64, 0), s.allocation_count); |
| 241 | +} |
| 242 | + |
| 243 | +test "BudgetAllocator: allocate within budget succeeds" { |
| 244 | + var budget = try BudgetAllocator.init(1024); |
| 245 | + defer budget.deinit(); |
| 246 | + const a = budget.allocator(); |
| 247 | + |
| 248 | + const buf = try a.alloc(u8, 256); |
| 249 | + defer a.free(buf); |
| 250 | + |
| 251 | + const s = budget.stats(); |
| 252 | + try testing.expectEqual(@as(usize, 256), s.bytes_allocated); |
| 253 | + try testing.expectEqual(@as(usize, 256), s.peak_bytes); |
| 254 | + try testing.expectEqual(@as(u64, 1), s.allocation_count); |
| 255 | +} |
| 256 | + |
| 257 | +test "BudgetAllocator: allocate beyond budget returns OOM" { |
| 258 | + var budget = try BudgetAllocator.init(512); |
| 259 | + defer budget.deinit(); |
| 260 | + const a = budget.allocator(); |
| 261 | + |
| 262 | + // Request more than the ceiling — must fail. |
| 263 | + try testing.expectError(error.OutOfMemory, a.alloc(u8, 1024)); |
| 264 | + |
| 265 | + // Fill exactly to cap, then the next byte fails. |
| 266 | + const buf = try a.alloc(u8, 512); |
| 267 | + defer a.free(buf); |
| 268 | + try testing.expectError(error.OutOfMemory, a.alloc(u8, 1)); |
| 269 | +} |
| 270 | + |
| 271 | +test "BudgetAllocator: peak bytes tracked across alloc/free cycles" { |
| 272 | + var budget = try BudgetAllocator.init(4096); |
| 273 | + defer budget.deinit(); |
| 274 | + const a = budget.allocator(); |
| 275 | + |
| 276 | + const buf1 = try a.alloc(u8, 100); |
| 277 | + const buf2 = try a.alloc(u8, 200); |
| 278 | + const buf3 = try a.alloc(u8, 300); |
| 279 | + try testing.expectEqual(@as(usize, 600), budget.stats().bytes_allocated); |
| 280 | + try testing.expectEqual(@as(usize, 600), budget.stats().peak_bytes); |
| 281 | + |
| 282 | + // Free in reverse order so the FixedBufferAllocator can actually |
| 283 | + // reclaim the storage. Our accounting decrements regardless of |
| 284 | + // reclamation, which is the point of the budget layer. |
| 285 | + a.free(buf3); |
| 286 | + a.free(buf2); |
| 287 | + a.free(buf1); |
| 288 | + try testing.expectEqual(@as(usize, 0), budget.stats().bytes_allocated); |
| 289 | + try testing.expectEqual(@as(usize, 600), budget.stats().peak_bytes); |
| 290 | +} |
| 291 | + |
| 292 | +test "BudgetAllocator: allocation_count increments" { |
| 293 | + var budget = try BudgetAllocator.init(4096); |
| 294 | + defer budget.deinit(); |
| 295 | + const a = budget.allocator(); |
| 296 | + |
| 297 | + const b1 = try a.alloc(u8, 16); |
| 298 | + const b2 = try a.alloc(u8, 32); |
| 299 | + const b3 = try a.alloc(u8, 64); |
| 300 | + try testing.expectEqual(@as(u64, 3), budget.stats().allocation_count); |
| 301 | + |
| 302 | + // Free in reverse order so the FBA reclaims each. |
| 303 | + a.free(b3); |
| 304 | + a.free(b2); |
| 305 | + a.free(b1); |
| 306 | +} |
| 307 | + |
| 308 | +test "BudgetAllocator: available reports remaining budget" { |
| 309 | + var budget = try BudgetAllocator.init(1000); |
| 310 | + defer budget.deinit(); |
| 311 | + const a = budget.allocator(); |
| 312 | + |
| 313 | + try testing.expectEqual(@as(usize, 1000), budget.available()); |
| 314 | + const buf = try a.alloc(u8, 400); |
| 315 | + defer a.free(buf); |
| 316 | + try testing.expectEqual(@as(usize, 600), budget.available()); |
| 317 | +} |
| 318 | + |
| 319 | +test "BudgetAllocator: rejects pathological huge request" { |
| 320 | + var budget = try BudgetAllocator.init(1024); |
| 321 | + defer budget.deinit(); |
| 322 | + const a = budget.allocator(); |
| 323 | + |
| 324 | + // A request that would overflow if naively added must be rejected |
| 325 | + // without crashing. |
| 326 | + try testing.expectError(error.OutOfMemory, a.alloc(u8, std.math.maxInt(usize))); |
| 327 | +} |
| 328 | + |
| 329 | +test "BudgetAllocator: arena on top of budget stays bounded" { |
| 330 | + // Exercise the common real pattern: an ArenaAllocator backed by a |
| 331 | + // BudgetAllocator. The arena can grab up to budget then OOMs, and |
| 332 | + // `arena.deinit()` frees everything through the budget. |
| 333 | + var budget = try BudgetAllocator.init(2048); |
| 334 | + defer budget.deinit(); |
| 335 | + |
| 336 | + var arena = std.heap.ArenaAllocator.init(budget.allocator()); |
| 337 | + defer arena.deinit(); |
| 338 | + const aa = arena.allocator(); |
| 339 | + |
| 340 | + _ = try aa.alloc(u8, 512); |
| 341 | + _ = try aa.alloc(u8, 512); |
| 342 | + _ = try aa.alloc(u8, 512); |
| 343 | + |
| 344 | + // Arena may round up; the 4th 512-byte alloc should exceed cap. |
| 345 | + try testing.expectError(error.OutOfMemory, aa.alloc(u8, 1024)); |
| 346 | +} |
0 commit comments