Report extra memory more (#4289)

* Report memory allocated in fetch

* Memory size reporting to `Headers`

* Fixup memory reporting allocator

* Make these tests do more

* cleanup some of this

---------

Co-authored-by: Jarred Sumner <709451+Jarred-Sumner@users.noreply.github.com>

1 files changed, 83 insertions, 0 deletions

diff --git a/src/bun.js/base.zig b/src/bun.js/base.zig
index e72e196a3..0ed780a5d 100644
--- a/src/bun.js/base.zig
+++ b/src/bun.js/base.zig
@@ -2484,3 +2484,86 @@ pub const AsyncTaskTracker = struct {
         bun.JSC.Debugger.didDispatchAsyncCall(globalObject, bun.JSC.Debugger.AsyncCallType.EventListener, this.id);
     }
 };
+
+pub const MemoryReportingAllocator = struct {
+    child_allocator: std.mem.Allocator,
+    memory_cost: std.atomic.Atomic(usize) = std.atomic.Atomic(usize).init(0),
+    const log = Output.scoped(.MEM, false);
+
+    fn alloc(this: *MemoryReportingAllocator, n: usize, log2_ptr_align: u8, return_address: usize) ?[*]u8 {
+        var result = this.child_allocator.rawAlloc(n, log2_ptr_align, return_address) orelse return null;
+        _ = this.memory_cost.fetchAdd(n, .Monotonic);
+        if (comptime Environment.allow_assert)
+            log("malloc({d}) = {d}", .{ n, this.memory_cost.loadUnchecked() });
+        return result;
+    }
+
+    pub fn discard(this: *MemoryReportingAllocator, buf: []const u8) void {
+        _ = this.memory_cost.fetchSub(buf.len, .Monotonic);
+        if (comptime Environment.allow_assert)
+            log("discard({d}) = {d}", .{ buf.len, this.memory_cost.loadUnchecked() });
+    }
+
+    fn resize(this: *MemoryReportingAllocator, buf: []u8, buf_align: u8, new_len: usize, ret_addr: usize) bool {
+        if (this.child_allocator.rawResize(buf, buf_align, new_len, ret_addr)) {
+            _ = this.memory_cost.fetchAdd(new_len -| buf.len, .Monotonic);
+            if (comptime Environment.allow_assert)
+                log("resize() = {d}", .{this.memory_cost.loadUnchecked()});
+            return true;
+        } else {
+            return false;
+        }
+    }
+
+    fn free(this: *MemoryReportingAllocator, buf: []u8, buf_align: u8, ret_addr: usize) void {
+        this.child_allocator.rawFree(buf, buf_align, ret_addr);
+
+        const prev = this.memory_cost.fetchSub(buf.len, .Monotonic);
+        if (comptime Environment.allow_assert) {
+            // check for overflow, racily
+            std.debug.assert(prev > this.memory_cost.load(.Monotonic));
+            log("free({d}) = {d}", .{ buf.len, this.memory_cost.loadUnchecked() });
+        }
+    }
+
+    pub fn wrap(this: *MemoryReportingAllocator, allocator_: std.mem.Allocator) std.mem.Allocator {
+        this.* = .{
+            .child_allocator = allocator_,
+        };
+
+        return this.allocator();
+    }
+
+    pub fn allocator(this: *MemoryReportingAllocator) std.mem.Allocator {
+        return std.mem.Allocator{
+            .ptr = this,
+            .vtable = &MemoryReportingAllocator.VTable,
+        };
+    }
+
+    pub fn report(this: *MemoryReportingAllocator, vm: *JSC.VM) void {
+        const mem = this.memory_cost.load(.Monotonic);
+        if (mem > 0) {
+            vm.reportExtraMemory(mem);
+            if (comptime Environment.allow_assert)
+                log("report({d})", .{mem});
+        }
+    }
+
+    pub inline fn assert(this: *const MemoryReportingAllocator) void {
+        if (comptime !Environment.allow_assert) {
+            return;
+        }
+
+        const memory_cost = this.memory_cost.load(.Monotonic);
+        if (memory_cost > 0) {
+            Output.panic("MemoryReportingAllocator still has {d} bytes allocated", .{memory_cost});
+        }
+    }
+
+    pub const VTable = std.mem.Allocator.VTable{
+        .alloc = @ptrCast(&MemoryReportingAllocator.alloc),
+        .resize = @ptrCast(&MemoryReportingAllocator.resize),
+        .free = @ptrCast(&MemoryReportingAllocator.free),
+    };
+};