Add memory allocation tracking

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.

build.zig

@@ -63,6 +63,7 @@ fn setupExamples(b: *std.Build, target: std.zig.CrossTarget, optimize: std.built
         "bubble_sort",
         "hooks",
         "json",
+        "memory_tracking",
         "parameterised",
         "progress",
         "sleep",

examples/json.zig

@@ -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();

examples/memory_tracking.zig

@@ -0,0 +1,25 @@
+const std = @import("std");
+const zbench = @import("zbench");
+
+fn myBenchmark(allocator: std.mem.Allocator) void {
+    for (0..2000) |_| {
+        const buf = allocator.alloc(u8, 512) catch @panic("OOM");
+        defer allocator.free(buf);
+    }
+}
+
+test "bench test basic" {
+    const stdout = std.io.getStdOut().writer();
+    var bench = zbench.Benchmark.init(std.testing.allocator, .{
+        .iterations = 64,
+    });
+    defer bench.deinit();
+
+    try bench.add("My Benchmark 1", myBenchmark, .{});
+    try bench.add("My Benchmark 2", myBenchmark, .{
+        .track_allocations = true,
+    });
+
+    try stdout.writeAll("\n");
+    try bench.run(stdout);
+}

util/runner.zig

@@ -4,6 +4,8 @@ const expectEqSlices = std.testing.expectEqualSlices;
 
 pub const Error = @import("./runner/types.zig").Error;
 pub const Step = @import("./runner/types.zig").Step;
+pub const AllocationReading = @import("./runner/types.zig").AllocationReading;
+pub const AllocationReadings = @import("./runner/types.zig").AllocationReadings;
 pub const Reading = @import("./runner/types.zig").Reading;
 pub const Readings = @import("./runner/types.zig").Readings;
 
@@ -45,32 +47,40 @@ const State = union(enum) {
         /// Number of timings still to be performed in the benchmark.
         iterations_remaining: usize,
 
-        /// Array of timings collected.
-        timings_ns: []u64,
+        /// Readings collected during the run.
+        readings: Readings,
     };
 };
 
 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),
+            .readings = try Readings.init(
+                allocator,
+                iterations,
+                track_allocations,
+            ),
         } },
     };
 }
@@ -101,15 +111,19 @@ pub fn next(self: *Runner, reading: Reading) Error!?Step {
                 self.state = .{ .running = .{
                     .iterations_count = N,
                     .iterations_remaining = N,
-                    .timings_ns = try self.allocator.alloc(u64, N),
+                    .readings = try Readings.init(
+                        self.allocator,
+                        N,
+                        self.track_allocations,
+                    ),
                 } };
             }
             return .more;
         },
         .running => |*st| {
             if (0 < st.iterations_remaining) {
-                const i = st.timings_ns.len - st.iterations_remaining;
-                st.timings_ns[i] = reading.timing_ns;
+                const i = st.readings.iterations - st.iterations_remaining;
+                st.readings.set(i, reading);
                 st.iterations_remaining -= 1;
             }
             return if (st.iterations_remaining == 0) null else .more;
@@ -121,20 +135,21 @@ pub fn next(self: *Runner, reading: Reading) Error!?Step {
 /// complete, so get the timing results.
 pub fn finish(self: *Runner) Error!Readings {
     return switch (self.state) {
-        .preparing => .{
+        .preparing => Readings{
+            .allocator = self.allocator,
+            .iterations = 0,
             .timings_ns = &.{},
+            .allocations = null,
         },
-        .running => |st| .{
-            .timings_ns = st.timings_ns,
-        },
+        .running => |st| st.readings,
     };
 }
 
 /// Clean up after an error.
 pub fn abort(self: *Runner) void {
     return switch (self.state) {
         .preparing => {},
-        .running => |st| self.allocator.free(st.timings_ns),
+        .running => |st| st.readings.deinit(),
     };
 }
 
@@ -157,33 +172,76 @@ 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(100_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(300_000_000, null)));
+        try expectEq(Step.more, try r.next(Reading.init(100_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(300_000_000, null)));
+        try expectEq(Step.more, try r.next(Reading.init(100_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(300_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(100_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(300_000_000, null)));
+        try expectEq(Step.more, try r.next(Reading.init(400_000_000, null)));
+        try expectEq(Step.more, try r.next(Reading.init(100_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(300_000_000, null)));
+        try expectEq(@as(?Step, null), try r.next(Reading.init(400_000_000, null)));
+    }
+    const result = try r.finish();
+    defer result.deinit();
+    try expectEqSlices(u64, &.{
+        100_000_000, 200_000_000, 300_000_000, 400_000_000,
+        100_000_000, 200_000_000, 300_000_000, 400_000_000,
+    }, result.timings_ns);
+}
+
+test "Runner - memory tracking" {
+    var r = try Runner.init(std.testing.allocator, 0, 16384, 2e9, true);
     {
         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(100_000_000, .{ .max = 256, .count = 1 })));
+        try expectEq(Step.more, try r.next(Reading.init(200_000_000, .{ .max = 256, .count = 1 })));
+        try expectEq(Step.more, try r.next(Reading.init(300_000_000, .{ .max = 512, .count = 2 })));
+        try expectEq(Step.more, try r.next(Reading.init(100_000_000, .{ .max = 512, .count = 2 })));
+        try expectEq(Step.more, try r.next(Reading.init(200_000_000, .{ .max = 1024, .count = 4 })));
+        try expectEq(Step.more, try r.next(Reading.init(300_000_000, .{ .max = 1024, .count = 4 })));
+        try expectEq(Step.more, try r.next(Reading.init(100_000_000, .{ .max = 2048, .count = 8 })));
+        try expectEq(Step.more, try r.next(Reading.init(200_000_000, .{ .max = 2045, .count = 8 })));
+        try expectEq(Step.more, try r.next(Reading.init(300_000_000, .{ .max = 4096, .count = 16 })));
+        try expectEq(Step.more, try r.next(Reading.init(200_000_000, .{ .max = 4096, .count = 16 })));
+        try expectEq(Step.more, try r.next(Reading.init(200_000_000, .{ .max = 8192, .count = 32 })));
+
+        try expectEq(Step.more, try r.next(Reading.init(100, .{ .max = 1, .count = 2 })));
+        try expectEq(Step.more, try r.next(Reading.init(200, .{ .max = 2, .count = 4 })));
+        try expectEq(Step.more, try r.next(Reading.init(300, .{ .max = 4, .count = 8 })));
+        try expectEq(Step.more, try r.next(Reading.init(400, .{ .max = 8, .count = 16 })));
+        try expectEq(Step.more, try r.next(Reading.init(100, .{ .max = 16, .count = 32 })));
+        try expectEq(Step.more, try r.next(Reading.init(200, .{ .max = 32, .count = 64 })));
+        try expectEq(Step.more, try r.next(Reading.init(300, .{ .max = 64, .count = 128 })));
+        try expectEq(@as(?Step, null), try r.next(Reading.init(400, .{ .max = 128, .count = 256 })));
     }
     const result = try r.finish();
-    defer std.testing.allocator.free(result.timings_ns);
+    defer result.deinit();
     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,
+        100, 200, 300, 400, 100, 200, 300, 400,
     }, result.timings_ns);
+    try expectEqSlices(
+        usize,
+        &.{ 1, 2, 4, 8, 16, 32, 64, 128 },
+        result.allocations.?.maxes,
+    );
+    try expectEqSlices(
+        usize,
+        &.{ 2, 4, 8, 16, 32, 64, 128, 256 },
+        result.allocations.?.counts,
+    );
 }

util/runner/types.zig

@@ -6,14 +6,86 @@ pub const Step = enum { more };
 
 pub const Reading = struct {
     timing_ns: u64,
+    allocation: ?AllocationReading,
 
-    pub fn init(timing_ns: u64) Reading {
+    pub fn init(
+        timing_ns: u64,
+        allocation: ?AllocationReading,
+    ) Reading {
         return .{
             .timing_ns = timing_ns,
+            .allocation = allocation,
         };
     }
 };
 
 pub const Readings = struct {
+    allocator: std.mem.Allocator,
+    iterations: usize,
     timings_ns: []u64,
+    allocations: ?AllocationReadings,
+
+    pub fn init(
+        allocator: std.mem.Allocator,
+        n: usize,
+        track_allocations: bool,
+    ) !Readings {
+        return Readings{
+            .allocator = allocator,
+            .iterations = n,
+            .timings_ns = try allocator.alloc(u64, n),
+            .allocations = if (track_allocations)
+                try AllocationReadings.init(allocator, n)
+            else
+                null,
+        };
+    }
+
+    pub fn deinit(self: Readings) void {
+        self.allocator.free(self.timings_ns);
+        if (self.allocations) |allocs| allocs.deinit(self.allocator);
+    }
+
+    pub fn set(self: *Readings, i: usize, reading: Reading) void {
+        self.timings_ns[i] = reading.timing_ns;
+        if (self.allocations) |allocs| {
+            if (reading.allocation) |x| {
+                allocs.maxes[i] = x.max;
+                allocs.counts[i] = x.count;
+            } else {
+                allocs.deinit(self.allocator);
+                self.allocations = null;
+            }
+        }
+    }
+
+    pub fn sort(self: *Readings) void {
+        std.sort.heap(u64, self.timings_ns, {}, std.sort.asc(u64));
+        if (self.allocations) |allocs| {
+            std.sort.heap(usize, allocs.maxes, {}, std.sort.asc(usize));
+            std.sort.heap(usize, allocs.counts, {}, std.sort.asc(usize));
+        }
+    }
+};
+
+pub const AllocationReading = struct {
+    max: usize,
+    count: usize,
+};
+
+pub const AllocationReadings = struct {
+    maxes: []usize,
+    counts: []usize,
+
+    pub fn init(allocator: std.mem.Allocator, n: usize) !AllocationReadings {
+        return AllocationReadings{
+            .maxes = try allocator.alloc(usize, n),
+            .counts = try allocator.alloc(usize, n),
+        };
+    }
+
+    pub fn deinit(self: AllocationReadings, allocator: std.mem.Allocator) void {
+        allocator.free(self.maxes);
+        allocator.free(self.counts);
+    }
 };

util/statistics.zig

@@ -1,5 +1,6 @@
 const std = @import("std");
 
+/// Collect common statistical calculations together.
 pub fn Statistics(comptime T: type) type {
     return struct {
         total: T,