Fix heap snapshots memory usage stats. Introduce estimateDirectMemoryUsageOf function in "bun:jsc" (#15790)

Author: Jarred-Sumner

docs/api/utils.md

@@ -771,3 +771,28 @@ console.log(obj); // => { foo: "bar" }
 ```
 
 Internally, [`structuredClone`](https://developer.mozilla.org/en-US/docs/Web/API/structuredClone) and [`postMessage`](https://developer.mozilla.org/en-US/docs/Web/API/Window/postMessage) serialize and deserialize the same way. This exposes the underlying [HTML Structured Clone Algorithm](https://developer.mozilla.org/en-US/docs/Web/API/Web_Workers_API/Structured_clone_algorithm) to JavaScript as an ArrayBuffer.
+
+## `estimateDirectMemoryUsageOf` in `bun:jsc`
+
+The `estimateDirectMemoryUsageOf` function returns a best-effort estimate of the memory usage of an object in bytes, excluding the memory usage of properties or other objects it references. For accurate per-object memory usage, use `Bun.generateHeapSnapshot`.
+
+```js
+import { estimateDirectMemoryUsageOf } from "bun:jsc";
+
+const obj = { foo: "bar" };
+const usage = estimateDirectMemoryUsageOf(obj);
+console.log(usage); // => 16
+
+const buffer = Buffer.alloc(1024 * 1024);
+estimateDirectMemoryUsageOf(buffer);
+// => 1048624
+
+const req = new Request("https://bun.sh");
+estimateDirectMemoryUsageOf(req);
+// => 167
+
+const array = Array(1024).fill({ a: 1 });
+// Arrays are usually not stored contiguously in memory, so this will not return a useful value (which isn't a bug).
+estimateDirectMemoryUsageOf(array);
+// => 16
+```

packages/bun-types/bun.d.ts

@@ -2151,6 +2151,8 @@ declare module "bun" {
      * });
      */
     data: T;
+
+    getBufferedAmount(): number;
   }
 
   /**

packages/bun-types/jsc.d.ts

@@ -214,4 +214,16 @@ declare module "bun:jsc" {
    * Run JavaScriptCore's sampling profiler
    */
   function startSamplingProfiler(optionalDirectory?: string): void;
+
+  /**
+   * Non-recursively estimate the memory usage of an object, excluding the memory usage of
+   * properties or other objects it references. For more accurate per-object
+   * memory usage, use {@link Bun.generateHeapSnapshot}.
+   *
+   * This is a best-effort estimate. It may not be 100% accurate. When it's
+   * wrong, it may mean the memory is non-contiguous (such as a large array).
+   *
+   * Passing a primitive type that isn't heap allocated returns 0.
+   */
+  function estimateDirectMemoryUsageOf(value: object | CallableFunction | bigint | symbol | string): number;
 }

packages/bun-uws/src/WebSocket.h

@@ -73,6 +73,10 @@ struct WebSocket : AsyncSocket<SSL> {
         DROPPED
     };
 
+    size_t memoryCost() {
+        return getBufferedAmount() + sizeof(WebSocket);
+    }
+
     /* Sending fragmented messages puts a bit of effort on the user; you must not interleave regular sends
      * with fragmented sends and you must sendFirstFragment, [sendFragment], then finally sendLastFragment. */
     SendStatus sendFirstFragment(std::string_view message, OpCode opCode = OpCode::BINARY, bool compress = false) {

src/bun.js/api/BunObject.classes.ts

@@ -48,6 +48,7 @@ export default [
     finalize: true,
     hasPendingActivity: true,
     configurable: false,
+    memoryCost: true,
     klass: {},
     JSType: "0b11101110",
     proto: {

src/bun.js/api/bun/process.zig

@@ -153,6 +153,10 @@ pub const Process = struct {
     sync: bool = false,
     event_loop: JSC.EventLoopHandle,
 
+    pub fn memoryCost(_: *const Process) usize {
+        return @sizeOf(@This());
+    }
+
     pub usingnamespace bun.NewRefCounted(Process, deinit);
 
     pub fn setExitHandler(this: *Process, handler: anytype) void {

src/bun.js/api/bun/socket.zig

@@ -1394,6 +1394,10 @@ fn NewSocket(comptime ssl: bool) type {
             return this.has_pending_activity.load(.acquire);
         }
 
+        pub fn memoryCost(this: *This) usize {
+            return @sizeOf(This) + this.buffered_data_for_node_net.cap;
+        }
+
         pub fn attachNativeCallback(this: *This, callback: NativeCallbacks) bool {
             if (this.native_callback != .none) return false;
             this.native_callback = callback;

src/bun.js/api/bun/subprocess.zig

@@ -395,6 +395,14 @@ pub const Subprocess = struct {
         closed: void,
         buffer: []u8,
 
+        pub fn memoryCost(this: *const Readable) usize {
+            return switch (this.*) {
+                .pipe => @sizeOf(PipeReader) + this.pipe.memoryCost(),
+                .buffer => this.buffer.len,
+                else => 0,
+            };
+        }
+
         pub fn hasPendingActivity(this: *const Readable) bool {
             return switch (this.*) {
                 .pipe => this.pipe.hasPendingActivity(),
@@ -794,6 +802,16 @@ pub const Subprocess = struct {
         array_buffer: JSC.ArrayBuffer.Strong,
         detached: void,
 
+        pub fn memoryCost(this: *const Source) usize {
+            // Memory cost of Source and each of the particular fields is covered by @sizeOf(Subprocess).
+            return switch (this.*) {
+                .blob => this.blob.memoryCost(),
+                // ArrayBuffer is owned by GC.
+                .array_buffer => 0,
+                .detached => 0,
+            };
+        }
+
         pub fn slice(this: *const Source) []const u8 {
             return switch (this.*) {
                 .blob => this.blob.slice(),
@@ -921,6 +939,10 @@ pub const Subprocess = struct {
                 this.destroy();
             }
 
+            pub fn memoryCost(this: *const This) usize {
+                return @sizeOf(@This()) + this.source.memoryCost() + this.writer.memoryCost();
+            }
+
             pub fn loop(this: *This) *uws.Loop {
                 return this.event_loop.loop();
             }
@@ -954,6 +976,10 @@ pub const Subprocess = struct {
 
         pub usingnamespace bun.NewRefCounted(PipeReader, PipeReader.deinit);
 
+        pub fn memoryCost(this: *const PipeReader) usize {
+            return this.reader.memoryCost();
+        }
+
         pub fn hasPendingActivity(this: *const PipeReader) bool {
             if (this.state == .pending)
                 return true;
@@ -1141,6 +1167,15 @@ pub const Subprocess = struct {
         inherit: void,
         ignore: void,
 
+        pub fn memoryCost(this: *const Writable) usize {
+            return switch (this.*) {
+                .pipe => |pipe| pipe.memoryCost(),
+                .buffer => |buffer| buffer.memoryCost(),
+                // TODO: memfd
+                else => 0,
+            };
+        }
+
         pub fn hasPendingActivity(this: *const Writable) bool {
             return switch (this.*) {
                 .pipe => false,
@@ -1415,6 +1450,14 @@ pub const Subprocess = struct {
         }
     };
 
+    pub fn memoryCost(this: *const Subprocess) usize {
+        return @sizeOf(@This()) +
+            this.process.memoryCost() +
+            this.stdin.memoryCost() +
+            this.stdout.memoryCost() +
+            this.stderr.memoryCost();
+    }
+
     pub fn onProcessExit(this: *Subprocess, process: *Process, status: bun.spawn.Status, rusage: *const Rusage) void {
         log("onProcessExit()", .{});
         const this_jsvalue = this.this_jsvalue;

src/bun.js/api/server.classes.ts

@@ -93,6 +93,7 @@ export default [
   define({
     name: "ServerWebSocket",
     JSType: "0b11101110",
+    memoryCost: true,
     proto: {
       send: {
         fn: "send",

src/bun.js/api/server.zig

@@ -1686,6 +1686,29 @@ pub const AnyRequestContext = struct {
     pub fn init(request_ctx: anytype) AnyRequestContext {
         return .{ .tagged_pointer = Pointer.init(request_ctx) };
     }
+
+    pub fn memoryCost(self: AnyRequestContext) usize {
+        if (self.tagged_pointer.isNull()) {
+            return 0;
+        }
+
+        switch (self.tagged_pointer.tag()) {
+            @field(Pointer.Tag, bun.meta.typeBaseName(@typeName(HTTPServer.RequestContext))) => {
+                return self.tagged_pointer.as(HTTPServer.RequestContext).memoryCost();
+            },
+            @field(Pointer.Tag, bun.meta.typeBaseName(@typeName(HTTPSServer.RequestContext))) => {
+                return self.tagged_pointer.as(HTTPSServer.RequestContext).memoryCost();
+            },
+            @field(Pointer.Tag, bun.meta.typeBaseName(@typeName(DebugHTTPServer.RequestContext))) => {
+                return self.tagged_pointer.as(DebugHTTPServer.RequestContext).memoryCost();
+            },
+            @field(Pointer.Tag, bun.meta.typeBaseName(@typeName(DebugHTTPSServer.RequestContext))) => {
+                return self.tagged_pointer.as(DebugHTTPSServer.RequestContext).memoryCost();
+            },
+            else => @panic("Unexpected AnyRequestContext tag"),
+        }
+    }
+
     pub fn get(self: AnyRequestContext, comptime T: type) ?*T {
         return self.tagged_pointer.get(T);
     }
@@ -1909,6 +1932,11 @@ fn NewRequestContext(comptime ssl_enabled: bool, comptime debug_mode: bool, comp
         // TODO: support builtin compression
         const can_sendfile = !ssl_enabled and !Environment.isWindows;
 
+        pub fn memoryCost(this: *const RequestContext) usize {
+            // The Sink and ByteStream aren't owned by this.
+            return @sizeOf(RequestContext) + this.request_body_buf.capacity + this.response_buf_owned.capacity + this.blob.memoryCost();
+        }
+
         pub inline fn isAsync(this: *const RequestContext) bool {
             return this.defer_deinit_until_callback_completes == null;
         }
@@ -4464,6 +4492,13 @@ pub const ServerWebSocket = struct {
     pub usingnamespace JSC.Codegen.JSServerWebSocket;
     pub usingnamespace bun.New(ServerWebSocket);
 
+    pub fn memoryCost(this: *const ServerWebSocket) usize {
+        if (this.flags.closed) {
+            return @sizeOf(ServerWebSocket);
+        }
+        return this.websocket().memoryCost() + @sizeOf(ServerWebSocket);
+    }
+
     const log = Output.scoped(.WebSocketServer, false);
 
     pub fn onOpen(this: *ServerWebSocket, ws: uws.AnyWebSocket) void {

src/bun.js/api/sockets.classes.ts

@@ -7,6 +7,7 @@ function generate(ssl) {
     hasPendingActivity: true,
     noConstructor: true,
     configurable: false,
+    memoryCost: true,
     proto: {
       getAuthorizationError: {
         fn: "getAuthorizationError",

src/bun.js/api/streams.classes.ts

@@ -7,6 +7,7 @@ function source(name) {
     noConstructor: true,
     finalize: true,
     configurable: false,
+    memoryCost: true,
     proto: {
       drain: {
         fn: "drainFromJS",

src/bun.js/bindings/CommonJSModuleRecord.cpp

@@ -727,6 +727,19 @@ JSCommonJSModule* JSCommonJSModule::create(
     return JSCommonJSModule::create(globalObject, requireMapKey, exportsObject, hasEvaluated, parent);
 }
 
+size_t JSCommonJSModule::estimatedSize(JSC::JSCell* cell, JSC::VM& vm)
+{
+    auto* thisObject = jsCast<JSCommonJSModule*>(cell);
+    size_t additionalSize = 0;
+    if (!thisObject->sourceCode.isNull() && !thisObject->sourceCode.view().isEmpty()) {
+        additionalSize += thisObject->sourceCode.view().length();
+        if (!thisObject->sourceCode.view().is8Bit()) {
+            additionalSize *= 2;
+        }
+    }
+    return Base::estimatedSize(cell, vm) + additionalSize;
+}
+
 void JSCommonJSModule::destroy(JSC::JSCell* cell)
 {
     static_cast<JSCommonJSModule*>(cell)->JSCommonJSModule::~JSCommonJSModule();
@@ -999,9 +1012,14 @@ void JSCommonJSModule::analyzeHeap(JSCell* cell, HeapAnalyzer& analyzer)
     if (auto* id = thisObject->m_id.get()) {
         if (!id->isRope()) {
             auto label = id->tryGetValue(false);
-            analyzer.setLabelForCell(cell, label);
+            analyzer.setLabelForCell(cell, makeString("CommonJS Module: "_s, StringView(label)));
+        } else {
+            analyzer.setLabelForCell(cell, "CommonJS Module"_s);
         }
+    } else {
+        analyzer.setLabelForCell(cell, "CommonJS Module"_s);
     }
+
     Base::analyzeHeap(cell, analyzer);
 }
 

src/bun.js/bindings/CommonJSModuleRecord.h

@@ -56,6 +56,8 @@ class JSCommonJSModule final : public JSC::JSDestructibleObject {
     bool ignoreESModuleAnnotation { false };
     JSC::SourceCode sourceCode = JSC::SourceCode();
 
+    static size_t estimatedSize(JSC::JSCell* cell, JSC::VM& vm);
+
     void setSourceCode(JSC::SourceCode&& sourceCode);
 
     static void destroy(JSC::JSCell*);

src/bun.js/bindings/DOMFormData.cpp

@@ -200,4 +200,19 @@ std::optional<KeyValuePair<String, DOMFormData::FormDataEntryValue>> DOMFormData
     return makeKeyValuePair(item.name, item.data);
 }
 
+size_t DOMFormData::memoryCost() const
+{
+    size_t cost = m_items.sizeInBytes();
+    for (auto& item : m_items) {
+        cost += item.name.sizeInBytes();
+        if (auto value = std::get_if<RefPtr<Blob>>(&item.data)) {
+            cost += value->get()->memoryCost();
+        } else if (auto value = std::get_if<String>(&item.data)) {
+            cost += value->sizeInBytes();
+        }
+    }
+
+    return cost;
+}
+
 } // namespace WebCore

src/bun.js/bindings/DOMFormData.h

@@ -74,6 +74,7 @@ class DOMFormData : public RefCounted<DOMFormData>, public ContextDestructionObs
     Ref<DOMFormData> clone() const;
 
     size_t count() const { return m_items.size(); }
+    size_t memoryCost() const;
 
     String toURLEncodedString();
 

src/bun.js/bindings/DOMURL.h

@@ -61,6 +61,11 @@ class DOMURL final : public RefCounted<DOMURL>, public CanMakeWeakPtr<DOMURL>, p
 
     static String createPublicURL(ScriptExecutionContext&, URLRegistrable&);
 
+    size_t memoryCost() const
+    {
+        return sizeof(DOMURL) + m_url.string().sizeInBytes();
+    }
+
 private:
     static ExceptionOr<Ref<DOMURL>> create(const String& url, const URL& base);
     DOMURL(URL&& completeURL);

src/bun.js/bindings/URLSearchParams.cpp

@@ -192,4 +192,13 @@ URLSearchParams::Iterator::Iterator(URLSearchParams& params)
 {
 }
 
+size_t URLSearchParams::memoryCost() const
+{
+    size_t cost = sizeof(URLSearchParams);
+    for (const auto& pair : m_pairs) {
+        cost += pair.key.sizeInBytes();
+        cost += pair.value.sizeInBytes();
+    }
+    return cost;
+}
 }

src/bun.js/bindings/URLSearchParams.h

@@ -56,6 +56,7 @@ class URLSearchParams : public RefCounted<URLSearchParams> {
     void updateFromAssociatedURL();
     void sort();
     size_t size() const { return m_pairs.size(); }
+    size_t memoryCost() const;
 
     class Iterator {
     public:

src/bun.js/bindings/blob.cpp

@@ -26,4 +26,9 @@ JSC::JSValue toJSNewlyCreated(JSC::JSGlobalObject* lexicalGlobalObject, JSDOMGlo
     return JSC::JSValue::decode(encoded);
 }
 
+size_t Blob::memoryCost() const
+{
+    return sizeof(Blob) + JSBlob::memoryCost(m_impl);
+}
+
 }

src/bun.js/bindings/blob.h

@@ -51,6 +51,8 @@ class Blob : public RefCounted<Blob> {
     }
     void* m_impl;
 
+    size_t memoryCost() const;
+
 private:
     Blob(void* impl, String fileName = String())
     {