tiny.bench.CountingAllocator

1pub const CountingAllocator = struct {2    backing: Allocator,3    counts: Counts = .{},4    mutex: std.atomic.Mutex = .unlocked,56    pub fn init(backing: Allocator) CountingAllocator {7        return .{ .backing = backing };8    }910    pub fn allocator(self: *CountingAllocator) Allocator {11        return .{ .ptr = self, .vtable = &vtable };12    }1314    /// Zeroes the traffic counters and sets the high water back to the live15    /// total. The live-byte total survives the reset, because memory handed out16    /// earlier is still outstanding. The high water restarts at the live total,17    /// so the next window reports the peak that window itself reached.18    pub fn reset(self: *CountingAllocator) void {19        self.lock();20        defer self.mutex.unlock();21        const live = self.counts.live_bytes;22        self.counts = .{ .live_bytes = live, .peak_live_bytes = live };23    }2425    const vtable: Allocator.VTable = .{26        .alloc = rawAlloc,27        .resize = rawResize,28        .remap = rawRemap,29        .free = rawFree,30    };3132    fn rawAlloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {33        const self: *CountingAllocator = @ptrCast(@alignCast(ctx));34        const ptr = self.backing.rawAlloc(len, alignment, ret_addr) orelse return null;35        self.lock();36        defer self.mutex.unlock();37        self.counts.alloc_count +|= 1;38        self.counts.alloc_bytes +|= timing.saturatingU64(len);39        self.growLive(timing.saturatingU64(len));40        return ptr;41    }4243    fn rawResize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {44        const self: *CountingAllocator = @ptrCast(@alignCast(ctx));45        const resized = self.backing.rawResize(memory, alignment, new_len, ret_addr);46        if (resized) self.recordLength(memory.len, new_len);47        return resized;48    }4950    fn rawRemap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {51        const self: *CountingAllocator = @ptrCast(@alignCast(ctx));52        const ptr = self.backing.rawRemap(memory, alignment, new_len, ret_addr) orelse return null;53        self.recordLength(memory.len, new_len);54        return ptr;55    }5657    fn rawFree(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {58        const self: *CountingAllocator = @ptrCast(@alignCast(ctx));59        self.backing.rawFree(memory, alignment, ret_addr);60        self.lock();61        defer self.mutex.unlock();62        self.counts.free_count +|= 1;63        self.counts.live_bytes -|= timing.saturatingU64(memory.len);64    }6566    fn recordLength(self: *CountingAllocator, old_len: usize, new_len: usize) void {67        if (new_len > old_len) {68            const grew = timing.saturatingU64(new_len - old_len);69            self.lock();70            defer self.mutex.unlock();71            self.counts.alloc_bytes +|= grew;72            self.counts.live_bytes +|= grew;73            self.counts.peak_live_bytes = @max(74                self.counts.peak_live_bytes,75                self.counts.live_bytes,76            );77            return;78        }79        if (new_len == old_len) return;80        self.lock();81        defer self.mutex.unlock();82        self.counts.live_bytes -|= timing.saturatingU64(old_len - new_len);83    }8485    fn growLive(self: *CountingAllocator, bytes: u64) void {86        self.counts.live_bytes +|= bytes;87        self.counts.peak_live_bytes = @max(88            self.counts.peak_live_bytes,89            self.counts.live_bytes,90        );91    }9293    fn lock(self: *CountingAllocator) void {94        while (!self.mutex.tryLock()) std.atomic.spinLoopHint();95    }96};

1pub const CountingAllocator = allocation.CountingAllocator;