Implement heap snapshot memory tracking for globals

This is a bit more complicated than it otherwise ideally should be
because context globals are handled so specially by jsg. Specifically,
they are not included in traceable handles and we don't generally
keep track of them. However, in order to trace them for the heap
snapshot we need to.

src/workerd/api/actor-state-test.c++

@@ -27,6 +27,7 @@ jsg::V8System v8System;
 struct ActorStateContext: public jsg::Object, public jsg::ContextGlobal {
   JSG_RESOURCE_TYPE(ActorStateContext) {
   }
+  const jsg::Object& getSelfObject() const override { return *this; }
 };
 JSG_DECLARE_ISOLATE_TYPE(ActorStateIsolate, ActorStateContext);
 

src/workerd/api/basics-test.c++

@@ -53,6 +53,8 @@ struct BasicsContext: public jsg::Object, public jsg::ContextGlobal {
   JSG_RESOURCE_TYPE(BasicsContext) {
     JSG_METHOD(test);
   }
+
+  const jsg::Object& getSelfObject() const override { return *this; }
 };
 JSG_DECLARE_ISOLATE_TYPE(
   BasicsIsolate,

src/workerd/api/crypto-impl-aes-test.c++

@@ -20,6 +20,7 @@ jsg::V8System v8System;
 struct CryptoContext: public jsg::Object, public jsg::ContextGlobal {
   JSG_RESOURCE_TYPE(CryptoContext) {
   }
+  const jsg::Object& getSelfObject() const override { return *this; }
 };
 JSG_DECLARE_ISOLATE_TYPE(CryptoIsolate, CryptoContext);
 

src/workerd/api/global-scope.h

@@ -136,6 +136,8 @@ class PromiseRejectionEvent: public Event {
 
 class WorkerGlobalScope: public EventTarget, public jsg::ContextGlobal {
 public:
+  const jsg::Object& getSelfObject() const override { return *this; }
+
   jsg::Unimplemented importScripts(kj::String s) { return {}; };
 
   JSG_RESOURCE_TYPE(WorkerGlobalScope, CompatibilityFlags::Reader flags) {

src/workerd/api/streams/queue-test.c++

@@ -13,6 +13,7 @@ jsg::V8System v8System;
 
 struct QueueContext: public jsg::Object, public jsg::ContextGlobal {
   JSG_RESOURCE_TYPE(QueueContext) {}
+  const jsg::Object& getSelfObject() const override { return *this; }
 };
 JSG_DECLARE_ISOLATE_TYPE(QueueIsolate, QueueContext);
 

src/workerd/io/promise-wrapper-test.c++

@@ -15,6 +15,7 @@ namespace {
 jsg::V8System v8System;
 
 struct CaptureThrowContext: public jsg::Object, public ContextGlobal {
+  const jsg::Object& getSelfObject() const override { return *this; }
   kj::Promise<int> test1() {
     JSG_FAIL_REQUIRE(TypeError, "boom");
   }

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

@@ -12,6 +12,7 @@ namespace {
 V8System v8System;
 
 struct BufferSourceContext: public jsg::Object, public jsg::ContextGlobal {
+  const jsg::Object& getSelfObject() const override { return *this; }
   BufferSource takeBufferSource(BufferSource buf) {
     auto ptr = buf.asArrayPtr();
     KJ_ASSERT(!buf.isDetached());

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

@@ -12,6 +12,7 @@ namespace {
 V8System v8System;
 
 struct DOMExceptionContext: public Object, public ContextGlobal {
+  const jsg::Object& getSelfObject() const override { return *this; }
   JSG_RESOURCE_TYPE(DOMExceptionContext) {
     JSG_NESTED_TYPE(DOMException);
   }

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

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

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

@@ -10,6 +10,7 @@ namespace {
 V8System v8System;
 
 struct GeneratorContext: public Object, public ContextGlobal {
+  const jsg::Object& getSelfObject() const override { return *this; }
 
   uint generatorTest(Lock& js, Generator<kj::String> generator) {
 

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

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

src/workerd/jsg/jsg.h

@@ -16,6 +16,7 @@
 #include <v8.h>
 #include <v8-profiler.h>
 #include <workerd/jsg/memory.h>
+#include <workerd/util/weak-refs.h>
 #include "macro-meta.h"
 #include "wrappable.h"
 #include "util.h"
@@ -1529,13 +1530,26 @@ class ModuleRegistry;
 // The lifetime of the global object matches the lifetime of the JavaScript context.
 class ContextGlobal {
 public:
-  ContextGlobal() {}
+  inline ContextGlobal()
+      : self(kj::refcounted<WeakRef<ContextGlobal>>(kj::Badge<ContextGlobal> {}, *this)) {}
+  inline virtual ~ContextGlobal() noexcept(false) {
+    KJ_DASSERT(self.get() != nullptr);
+    self->invalidate();
+  }
 
   KJ_DISALLOW_COPY_AND_MOVE(ContextGlobal);
 
   ModuleRegistry& getModuleRegistry() { return *moduleRegistry; }
 
+  inline kj::Own<WeakRef<ContextGlobal>> addWeakRef() {
+    KJ_ASSERT(self.get() != nullptr);
+    return self->addRef();
+  }
+
+  virtual const Object& getSelfObject() const = 0;
+
 private:
+  kj::Own<WeakRef<ContextGlobal>> self;
   kj::Own<ModuleRegistry> moduleRegistry;
 
   void setModuleRegistry(kj::Own<ModuleRegistry> registry) {
@@ -2483,6 +2497,10 @@ inline v8::Local<v8::Context> JsContext<T>::getHandle(Lock& js) {
   return handle.Get(js.v8Isolate);
 }
 
+MemoryTracker& MemoryTracker::trackField(const ContextGlobal& value) {
+ return trackField(nullptr, value.getSelfObject());
+}
+
 }  // namespace workerd::jsg
 
 // These two includes are needed for the JSG type glue macros to work.

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

@@ -8,7 +8,10 @@ namespace workerd::jsg::test {
 namespace {
 
 V8System v8System;
-class ContextGlobalObject: public Object, public ContextGlobal {};
+class ContextGlobalObject: public Object, public ContextGlobal {
+public:
+  const jsg::Object& getSelfObject() const override { return *this; }
+};
 
 struct JsValueContext: public ContextGlobalObject {
   JsRef<JsValue> persisted;

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

@@ -11,7 +11,10 @@ namespace workerd::jsg::test {
 namespace {
 
 V8System v8System;
-class ContextGlobalObject: public Object, public ContextGlobal {};
+class ContextGlobalObject: public Object, public ContextGlobal {
+public:
+  const jsg::Object& getSelfObject() const override { return *this; }
+};
 
 struct Foo: public Object {
   kj::String bar = kj::str("test");

src/workerd/jsg/memory.h

@@ -27,6 +27,7 @@ class MemoryRetainerNode;
 
 template <typename T> class V8Ref;
 template <typename T> class Ref;
+class ContextGlobal;
 
 enum class MemoryInfoDetachedState {
   UNKNOWN,
@@ -90,6 +91,8 @@ class MemoryTracker final {
       kj::StringPtr edgeName, size_t size,
       kj::Maybe<kj::StringPtr> nodeName = kj::none) KJ_LIFETIMEBOUND;
 
+  inline MemoryTracker& trackField(const ContextGlobal& value) KJ_LIFETIMEBOUND;
+
   template <MemoryRetainer T>
   inline MemoryTracker& trackField(
       kj::StringPtr edgeName,

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

@@ -14,6 +14,7 @@ int promiseTestResult = 0;
 kj::String catchTestResult;
 
 struct PromiseContext: public jsg::Object, public jsg::ContextGlobal {
+  const jsg::Object& getSelfObject() const override { return *this; }
   Promise<kj::String> makePromise(jsg::Lock& js) {
     auto [ p, r ] = js.newPromiseAndResolver<int>();
     resolver = kj::mv(r);

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

@@ -9,7 +9,10 @@ namespace workerd::jsg::test {
 namespace {
 
 V8System v8System;
-class ContextGlobalObject: public Object, public ContextGlobal {};
+class ContextGlobalObject: public Object, public ContextGlobal {
+public:
+  const jsg::Object& getSelfObject() const override { return *this; }
+};
 
 struct BoxContext: public ContextGlobalObject {
   JSG_RESOURCE_TYPE(BoxContext) {

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

@@ -10,6 +10,7 @@ namespace {
 V8System v8System;
 
 struct EvalContext: public Object, public ContextGlobal {
+  const jsg::Object& getSelfObject() const override { return *this; }
   JSG_RESOURCE_TYPE(EvalContext) {}
 };
 JSG_DECLARE_ISOLATE_TYPE(EvalIsolate, EvalContext);
@@ -46,6 +47,7 @@ KJ_TEST("eval() is blocked") {
 // ========================================================================================
 
 struct ConfigContext: public Object, public ContextGlobal {
+  const jsg::Object& getSelfObject() const override { return *this; }
   struct Nested: public Object {
     JSG_RESOURCE_TYPE(Nested, int configuration) {
       KJ_EXPECT(configuration == 123, configuration);

src/workerd/jsg/setup.c++

@@ -170,6 +170,12 @@ void IsolateBase::buildEmbedderGraph(v8::Isolate* isolate,
 void IsolateBase::jsgGetMemoryInfo(MemoryTracker& tracker) const {
   tracker.trackField("uuid", uuid);
   tracker.trackField("heapTracer", heapTracer);
+  for (auto& ctx : contextGlobals) {
+    KJ_ASSERT(ctx.get() != nullptr);
+    ctx->runIfAlive([&](ContextGlobal& context) {
+      tracker.trackField(context);
+    });
+  }
 }
 
 void IsolateBase::deferDestruction(Item item) {
@@ -680,4 +686,16 @@ kj::StringPtr IsolateBase::getUuid() {
   return uuid.emplace(randomUUID(kj::none));
 }
 
+void IsolateBase::purgeInvalidContextGlobals() {
+  auto released = contextGlobals.releaseAsArray();
+  contextGlobals.clear();
+  KJ_ASSERT(contextGlobals.empty());
+  for (auto& ctx : released) {
+    KJ_ASSERT(ctx.get() != nullptr);
+    if (ctx->isValid()) {
+      contextGlobals.add(kj::mv(ctx));
+    }
+  }
+}
+
 }  // namespace workerd::jsg

src/workerd/jsg/setup.h

@@ -239,13 +239,15 @@ class IsolateBase {
 
   static void jitCodeEvent(const v8::JitCodeEvent* event) noexcept;
 
-  friend class IsolateBase;
   friend kj::Maybe<kj::StringPtr> getJsStackTrace(void* ucontext, kj::ArrayPtr<char> scratch);
 
   HeapTracer heapTracer;
   std::unique_ptr<v8::CppHeap> cppgcHeap;
   kj::Own<IsolateObserver> observer;
 
+  kj::Vector<kj::Own<WeakRef<ContextGlobal>>> contextGlobals;
+  void purgeInvalidContextGlobals();
+
   friend class Data;
   friend class Wrappable;
   friend class HeapTracer;
@@ -355,6 +357,7 @@ class Isolate: public IsolateBase {
     // order for V8's stack-scanning GC to find them.
     Lock(const Isolate& isolate, V8StackScope&)
         : jsg::Lock(isolate.ptr), jsgIsolate(const_cast<Isolate&>(isolate)) {
+      jsgIsolate.purgeInvalidContextGlobals();
       jsgIsolate.clearDestructionQueue();
     }
     KJ_DISALLOW_COPY_AND_MOVE(Lock);
@@ -436,7 +439,10 @@ class Isolate: public IsolateBase {
       //   allocate the object (forwarding arguments to the constructor) and return something like
       //   a Ref.
 
-      return jsgIsolate.wrapper->newContext(*this, jsgIsolate.getObserver(), (T*)nullptr, kj::fwd<Args>(args)...);
+      auto ctx = jsgIsolate.wrapper->newContext(*this, jsgIsolate.getObserver(),
+          (T*)nullptr, kj::fwd<Args>(args)...);
+      jsgIsolate.contextGlobals.add(ctx->addWeakRef());
+      return kj::mv(ctx);
     }
 
   private:

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

@@ -351,6 +351,7 @@ KJ_TEST("UsvString from kj::String") {
 V8System v8System;
 
 struct UsvStringContext: public jsg::Object, public jsg::ContextGlobal {
+  const jsg::Object& getSelfObject() const override { return *this; }
   UsvString testUsv(jsg::Optional<UsvString> str) {
     return kj::mv(str).orDefault(usv("undefined"));
   }

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

@@ -17,6 +17,7 @@ struct SelfStruct {
 };
 
 struct StructContext: public Object, public ContextGlobal {
+  const jsg::Object& getSelfObject() const override { return *this; }
   kj::String readTestStruct(TestStruct s) {
     return kj::str(s.str, ", ", s.num, ", ", s.box->value);
   }

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

@@ -102,6 +102,7 @@ private:
 };
 
 struct TraceTestContext: public Object, public ContextGlobal {
+  const jsg::Object& getSelfObject() const override { return *this; }
   kj::Maybe<jsg::Ref<NumberBox>> strongRef;
   // A strong reference to a NumberBox which may be get and set.
 

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

@@ -8,7 +8,10 @@ namespace workerd::jsg::test {
 namespace {
 
 V8System v8System;
-class ContextGlobalObject: public Object, public ContextGlobal {};
+class ContextGlobalObject: public Object, public ContextGlobal {
+public:
+  const jsg::Object& getSelfObject() const override { return *this; }
+};
 
 struct InfoContext: public ContextGlobalObject {
   struct WantInfo: public Object {

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

@@ -11,7 +11,10 @@ namespace workerd::jsg::test {
 namespace {
 
 V8System v8System;
-class ContextGlobalObject: public Object, public ContextGlobal {};
+class ContextGlobalObject: public Object, public ContextGlobal {
+public:
+  const jsg::Object& getSelfObject() const override { return *this; }
+};
 
 struct FreezeContext: public ContextGlobalObject {
   void recursivelyFreeze(v8::Local<v8::Value> value, v8::Isolate* isolate) {

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

@@ -8,7 +8,10 @@ namespace workerd::jsg::test {
 namespace {
 
 V8System v8System;
-class ContextGlobalObject: public Object, public ContextGlobal {};
+class ContextGlobalObject: public Object, public ContextGlobal {
+public:
+  const jsg::Object& getSelfObject() const override { return *this; }
+};
 
 struct BoolContext: public ContextGlobalObject {
   kj::String takeBool(bool b) {