Skip to content

Commit

Permalink
Add memory allocation tracking
Browse files Browse the repository at this point in the history
Adds a custom Allocator wrapper which keeps track of the allocated memory, then
record the maximum allocation readings in the Runner alongside timing readings
and carry them into the Result.
  • Loading branch information
bens committed Mar 21, 2024
1 parent 9aeedd3 commit b3b2b91
Show file tree
Hide file tree
Showing 6 changed files with 227 additions and 35 deletions.
9 changes: 8 additions & 1 deletion examples/json.zig
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,14 @@ test "bench test json" {
var bench = zbench.Benchmark.init(test_allocator, .{});
defer bench.deinit();

try bench.add("My Benchmark 1", myBenchmark, .{ .iterations = 10 });
try bench.add("My Benchmark 1", myBenchmark, .{
.iterations = 10,
.track_allocations = false,
});
try bench.add("My Benchmark 2", myBenchmark, .{
.iterations = 10,
.track_allocations = true,
});

try stdout.writeAll("[");
var iter = try bench.iterator();
Expand Down
69 changes: 48 additions & 21 deletions util/runner.zig
Original file line number Diff line number Diff line change
Expand Up @@ -47,30 +47,41 @@ const State = union(enum) {

/// Array of timings collected.
timings_ns: []u64,

/// Array of memory usages collected.
allocations: ?[]usize,
};
};

allocator: std.mem.Allocator,
track_allocations: bool,
state: State,

pub fn init(
allocator: std.mem.Allocator,
iterations: u16,
max_iterations: u16,
time_budget_ns: u64,
track_allocations: bool,
) Error!Runner {
return if (iterations == 0) .{
.allocator = allocator,
.track_allocations = track_allocations,
.state = .{ .preparing = .{
.max_iterations = max_iterations,
.time_budget_ns = time_budget_ns,
} },
} else .{
.allocator = allocator,
.track_allocations = track_allocations,
.state = .{ .running = .{
.iterations_count = iterations,
.iterations_remaining = iterations,
.timings_ns = try allocator.alloc(u64, iterations),
.allocations = if (track_allocations)
try allocator.alloc(usize, iterations)
else
null,
} },
};
}
Expand Down Expand Up @@ -102,6 +113,10 @@ pub fn next(self: *Runner, reading: Reading) Error!?Step {
.iterations_count = N,
.iterations_remaining = N,
.timings_ns = try self.allocator.alloc(u64, N),
.allocations = if (self.track_allocations)
try self.allocator.alloc(usize, N)
else
null,
} };
}
return .more;
Expand All @@ -110,6 +125,12 @@ pub fn next(self: *Runner, reading: Reading) Error!?Step {
if (0 < st.iterations_remaining) {
const i = st.timings_ns.len - st.iterations_remaining;
st.timings_ns[i] = reading.timing_ns;
if (st.allocations) |allocs| {
if (reading.max_allocated) |m| allocs[i] = m else {
self.allocator.free(allocs);
st.allocations = null;
}
}
st.iterations_remaining -= 1;
}
return if (st.iterations_remaining == 0) null else .more;
Expand All @@ -123,9 +144,11 @@ pub fn finish(self: *Runner) Error!Readings {
return switch (self.state) {
.preparing => .{
.timings_ns = &.{},
.max_allocations = null,
},
.running => |st| .{
.timings_ns = st.timings_ns,
.max_allocations = st.allocations,
},
};
}
Expand All @@ -134,7 +157,10 @@ pub fn finish(self: *Runner) Error!Readings {
pub fn abort(self: *Runner) void {
return switch (self.state) {
.preparing => {},
.running => |st| self.allocator.free(st.timings_ns),
.running => |st| {
self.allocator.free(st.timings_ns);
if (st.allocations) |allocs| self.allocator.free(allocs);
},
};
}

Expand All @@ -157,31 +183,32 @@ pub fn status(self: Runner) Status {
}

test "Runner" {
var r = try Runner.init(std.testing.allocator, 0, 16384, 2e9);
var r = try Runner.init(std.testing.allocator, 0, 16384, 2e9, false);
{
errdefer r.abort();
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000)));
try expectEq(@as(?Step, null), try r.next(Reading.init(200_000_000)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(Step.more, try r.next(Reading.init(200_000_000, null)));
try expectEq(@as(?Step, null), try r.next(Reading.init(200_000_000, null)));
}
const result = try r.finish();
defer std.testing.allocator.free(result.timings_ns);
defer if (result.max_allocations) |m| std.testing.allocator.free(m);
try expectEqSlices(u64, &.{
200_000_000, 200_000_000, 200_000_000, 200_000_000,
200_000_000, 200_000_000, 200_000_000, 200_000_000,
Expand Down
5 changes: 4 additions & 1 deletion util/runner/types.zig
Original file line number Diff line number Diff line change
Expand Up @@ -6,14 +6,17 @@ pub const Step = enum { more };

pub const Reading = struct {
timing_ns: u64,
max_allocated: ?usize,

pub fn init(timing_ns: u64) Reading {
pub fn init(timing_ns: u64, max_allocated: ?usize) Reading {
return .{
.timing_ns = timing_ns,
.max_allocated = max_allocated,
};
}
};

pub const Readings = struct {
timings_ns: []u64,
max_allocations: ?[]usize,
};
1 change: 1 addition & 0 deletions util/statistics.zig
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
const std = @import("std");

/// Collect common statistical calculations together.
pub fn Statistics(comptime T: type) type {
return struct {
total: T,
Expand Down
81 changes: 81 additions & 0 deletions util/tracking_allocator.zig
Original file line number Diff line number Diff line change
@@ -0,0 +1,81 @@
const std = @import("std");
const Allocator = std.mem.Allocator;

const TrackingAllocator = @This();

parent_allocator: Allocator,
current_allocated: usize = 0,
max_allocated: usize = 0,

pub fn init(parent_allocator: Allocator) TrackingAllocator {
return .{
.parent_allocator = parent_allocator,
};
}

pub fn allocator(self: *TrackingAllocator) Allocator {
return .{
.ptr = self,
.vtable = &.{
.alloc = alloc,
.resize = resize,
.free = free,
},
};
}

pub fn maxAllocated(self: TrackingAllocator) usize {
return self.max_allocated;
}

fn alloc(
ctx: *anyopaque,
len: usize,
log2_ptr_align: u8,
ra: usize,
) ?[*]u8 {
const self: *TrackingAllocator = @ptrCast(@alignCast(ctx));
const result = self.parent_allocator.rawAlloc(len, log2_ptr_align, ra);
if (result) |_| {
self.current_allocated += len;
if (self.max_allocated < self.current_allocated)
self.max_allocated = self.current_allocated;
}
return result;
}

fn resize(
ctx: *anyopaque,
buf: []u8,
log2_buf_align: u8,
new_len: usize,
ra: usize,
) bool {
const self: *TrackingAllocator = @ptrCast(@alignCast(ctx));
const result = self.parent_allocator.rawResize(buf, log2_buf_align, new_len, ra);
if (result) {
if (buf.len < new_len) {
self.current_allocated += new_len - buf.len;
if (self.max_allocated < self.current_allocated)
self.max_allocated = self.current_allocated;
} else self.current_allocated -= buf.len - new_len;
}
return result;
}

fn free(
ctx: *anyopaque,
buf: []u8,
log2_buf_align: u8,
ra: usize,
) void {
const self: *TrackingAllocator = @ptrCast(@alignCast(ctx));
self.parent_allocator.rawFree(buf, log2_buf_align, ra);
self.current_allocated -= buf.len;
}

/// This allocator is used in front of another allocator and tracks the maximum
/// memory usage on every call to the allocator.
pub fn trackingAllocator(parent_allocator: Allocator) TrackingAllocator {
return TrackingAllocator.init(parent_allocator);
}
Loading

0 comments on commit b3b2b91

Please sign in to comment.