Implement MemoryTracker for Heap Snapshots

src/workerd/jsg/BUILD.bazel

@@ -18,6 +18,7 @@ wd_cc_library(
         exclude = [
             "exception.h",
             "observer.h",
+            "memory.h",
             "rtti.h",
             "url-test-corpus-failures.h",
             "url-test-corpus-success.h",
@@ -26,6 +27,7 @@ wd_cc_library(
     visibility = ["//visibility:public"],
     deps = [
         ":exception",
+        ":memory-tracker",
         ":modules_capnp",
         ":observer",
         ":url",
@@ -38,6 +40,16 @@ wd_cc_library(
     ],
 )
 
+wd_cc_library(
+    name = "memory-tracker",
+    hdrs = ["memory.h"],
+    visibility = ["//visibility:public"],
+    deps = [
+        "@capnp-cpp//src/kj",
+        "@workerd-v8//:v8",
+    ],
+)
+
 wd_cc_library(
     name = "url",
     srcs = ["url.c++"],
@@ -48,6 +60,7 @@ wd_cc_library(
     deps = [
         "@capnp-cpp//src/kj",
         "@ada-url",
+        ":memory-tracker",
     ],
 )
 

src/workerd/jsg/async-context.c++

@@ -4,6 +4,7 @@
 #include "async-context.h"
 #include "jsg.h"
 #include "setup.h"
+#include <workerd/jsg/memory.h>
 #include <v8.h>
 
 namespace workerd::jsg {
@@ -42,6 +43,13 @@ AsyncContextFrame::AsyncContextFrame(Lock& js, StorageEntry storageEntry) {
   });
 }
 
+AsyncContextFrame::StorageEntry::StorageEntry(kj::Own<StorageKey> key, Value value)
+    : key(kj::mv(key)), value(kj::mv(value)) {}
+
+AsyncContextFrame::StorageEntry AsyncContextFrame::StorageEntry::clone(Lock& js) {
+  return StorageEntry(kj::addRef(*key), value.addRef(js));
+}
+
 kj::Maybe<AsyncContextFrame&> AsyncContextFrame::current(Lock& js) {
   return current(js.v8Isolate);
 }
@@ -172,10 +180,7 @@ AsyncContextFrame::StorageScope::StorageScope(
     Lock& js,
     StorageKey& key,
     Value store)
-    : frame(AsyncContextFrame::create(js, StorageEntry {
-        .key = kj::addRef(key),
-        .value = kj::mv(store)
-      })),
+    : frame(AsyncContextFrame::create(js, StorageEntry(kj::addRef(key), kj::mv(store)))),
       scope(js, *frame) {}
 
 v8::Local<v8::Object> AsyncContextFrame::getJSWrapper(v8::Isolate* isolate) {

src/workerd/jsg/async-context.h

@@ -90,6 +90,8 @@ class AsyncContextFrame final: public Wrappable {
       return this == &other;
     }
 
+    JSG_MEMORY_INFO(StorageKey) {}
+
   private:
     uint hash;
     bool dead = false;
@@ -98,12 +100,12 @@ class AsyncContextFrame final: public Wrappable {
   struct StorageEntry {
     kj::Own<StorageKey> key;
     Value value;
+    StorageEntry(kj::Own<StorageKey> key, Value value);
+    StorageEntry clone(Lock& js);
 
-    inline StorageEntry clone(Lock& js) {
-      return {
-        .key = kj::addRef(*key),
-        .value = value.addRef(js)
-      };
+    JSG_MEMORY_INFO(StorageEntry) {
+      tracker.trackField("key", key);
+      tracker.trackField("value", value);
     }
   };
 
@@ -189,6 +191,13 @@ class AsyncContextFrame final: public Wrappable {
     KJ_DISALLOW_COPY(StorageScope);
   };
 
+  kj::StringPtr jsgGetMemoryName() const override { return "AsyncContextFrame"_kjc; }
+  size_t jsgGetMemorySelfSize() const override { return sizeof(AsyncContextFrame); }
+  void jsgGetMemoryInfo(MemoryTracker& tracker) const override {
+    Wrappable::jsgGetMemoryInfo(tracker);
+    tracker.trackField("storage", storage);
+  }
+
 private:
   struct StorageEntryCallbacks {
     StorageKey& keyForRow(StorageEntry& entry) const {

src/workerd/jsg/buffersource.h

@@ -207,6 +207,10 @@ class BackingStore {
     return BackingStore(backingStore, byteLength, byteOffset, elementSize, ctor, integerType);
   }
 
+  JSG_MEMORY_INFO(BackingStore) {
+    tracker.trackFieldWithSize("buffer", size());
+  }
+
 private:
   std::shared_ptr<v8::BackingStore> backingStore;
   size_t byteLength;
@@ -363,6 +367,13 @@ class BufferSource {
   // to successfully detach the backing store.
   void setDetachKey(Lock& js, v8::Local<v8::Value> key);
 
+  JSG_MEMORY_INFO(BufferSource) {
+    tracker.trackField("handle", handle);
+    KJ_IF_SOME(backing, maybeBackingStore) {
+      tracker.trackField("backing", backing);
+    }
+  }
+
 private:
   Value handle;
   kj::Maybe<BackingStore> maybeBackingStore;

src/workerd/jsg/dom-exception.c++

@@ -3,6 +3,7 @@
 //     https://opensource.org/licenses/Apache-2.0
 
 #include "dom-exception.h"
+#include <workerd/jsg/memory.h>
 #include <kj/string.h>
 #include <map>
 

src/workerd/jsg/dom-exception.h

@@ -106,6 +106,12 @@ class DOMException: public Object {
 #undef JSG_DOM_EXCEPTION_CONSTANT_CXX
 #undef JSG_DOM_EXCEPTION_CONSTANT_JS
 
+  void visitForMemoryInfo(MemoryTracker& tracker) const {
+    tracker.trackField("message", message);
+    tracker.trackField("name", name);
+    tracker.trackField("errorForStack", errorForStack);
+  }
+
 private:
   kj::String message;
   kj::String name;

src/workerd/jsg/function-test.c++

@@ -8,7 +8,7 @@ namespace workerd::jsg::test {
 namespace {
 
 V8System v8System;
-class ContextGlobalObject: public Object, public ContextGlobal { };
+class ContextGlobalObject: public Object, public ContextGlobal {};
 
 struct CallbackContext: public ContextGlobalObject {
   kj::String callCallback(Lock& js, jsg::Function<kj::String(kj::StringPtr, double)> function) {

src/workerd/jsg/function.h

@@ -23,6 +23,20 @@ class WrappableFunction<Ret(Args...)>: public Wrappable {
   virtual Ret operator()(Lock& js, Args&&... args) = 0;
 
   const bool needsGcTracing;
+
+  kj::StringPtr jsgGetMemoryName() const override {
+    return "WrappableFunction"_kjc;
+  }
+  size_t jsgGetMemorySelfSize() const override {
+    return sizeof(WrappableFunction<Ret(Args...)>);
+  }
+  void jsgGetMemoryInfo(MemoryTracker& tracker) const override {
+    Wrappable::jsgGetMemoryInfo(tracker);
+    visitForMemoryInfo(tracker);
+  }
+  virtual void visitForMemoryInfo(MemoryTracker& tracker) const {
+    // TODO(soon): Implement tracking for WrappableFunction.
+  }
 };
 
 template <typename Signature, typename Impl, bool = hasPublicVisitForGc<Impl>()>
@@ -240,16 +254,33 @@ class Function<Ret(Args...)> {
     }
   }
 
+  JSG_MEMORY_INFO(Function) {
+    KJ_SWITCH_ONEOF(impl) {
+      KJ_CASE_ONEOF(ref, Ref<NativeFunction>) {
+        tracker.trackField("native", ref);
+      }
+      KJ_CASE_ONEOF(impl, JsImpl) {
+        tracker.trackField("impl", impl);
+      }
+    }
+  }
+
 private:
   Function(Ref<NativeFunction>&& func) : impl(kj::mv(func)) {}
 
   struct JsImpl {
     Wrapper* wrapper;
     Value receiver;
     V8Ref<v8::Function> handle;
+
+    JSG_MEMORY_INFO(JsImpl) {
+      tracker.trackField("receiver", receiver);
+      tracker.trackField("handle", handle);
+    }
   };
 
   kj::OneOf<Ref<NativeFunction>, JsImpl> impl;
+  friend class MemoryTracker;
 };
 
 template <typename T>

src/workerd/jsg/iterator-test.c++

@@ -10,6 +10,7 @@ namespace {
 V8System v8System;
 
 struct GeneratorContext: public Object, public ContextGlobal {
+
   uint generatorTest(Lock& js, Generator<kj::String> generator) {
 
     KJ_DEFER(generator.forEach(js, [](auto& js, auto, auto&) {

src/workerd/jsg/iterator.h

@@ -6,6 +6,7 @@
 
 #include "jsg.h"
 #include "struct.h"
+#include <workerd/jsg/memory.h>
 #include <concepts>
 #include <deque>
 
@@ -670,6 +671,14 @@ class IteratorBase: public Object {
     }
   }
 
+  JSG_MEMORY_INFO(IteratorBase) {
+    if constexpr (MemoryRetainer<State>) {
+      tracker.trackField("state", state);
+    } else {
+      tracker.trackFieldWithSize("state", sizeof(State));
+    }
+  }
+
 private:
   State state;
 
@@ -704,6 +713,10 @@ class AsyncIteratorImpl {
     JSG_STRUCT(done, value);
   };
 
+  JSG_MEMORY_INFO(AsyncIteratorImpl) {
+    // TODO(soon): Implement memory tracking
+  }
+
 private:
   std::deque<Promise<void>> pendingStack;
 };
@@ -780,10 +793,30 @@ class AsyncIteratorBase: public Object {
     }
   }
 
+  JSG_MEMORY_INFO(AsyncIteratorBase) {
+    KJ_SWITCH_ONEOF(state) {
+      KJ_CASE_ONEOF(fin, Finished) {
+        tracker.trackFieldWithSize("state", sizeof(Finished));
+      }
+      KJ_CASE_ONEOF(state, InnerState) {
+        tracker.trackField("state", state);
+      }
+    }
+  }
+
 private:
   struct InnerState {
     State state;
     AsyncIteratorImpl impl;
+
+    JSG_MEMORY_INFO(InnerState) {
+      if constexpr (MemoryRetainer<State>) {
+        tracker.trackField("state", state);
+      } else {
+        tracker.trackField("state", sizeof(State));
+      }
+      tracker.trackField("impl", impl);
+    }
   };
 
   kj::OneOf<Finished, InnerState> state;

src/workerd/jsg/jsg-test.c++

@@ -35,7 +35,7 @@ static_assert(kj::_::isDisallowedInCoroutine<Lock*>());
 // ========================================================================================
 
 V8System v8System;
-class ContextGlobalObject: public Object, public ContextGlobal { };
+class ContextGlobalObject: public Object, public ContextGlobal {};
 
 struct TestContext: public ContextGlobalObject {
   JSG_RESOURCE_TYPE(TestContext) {}