Merge pull request #1605 from cloudflare/jsnell/basics-memory-tracker

src/workerd/api/basics.c++

@@ -744,4 +744,45 @@ CustomEvent::CustomEventInit::operator Event::Init() {
   };
 }
 
+size_t EventTarget::EventHandler::JavaScriptHandler::jsgGetMemorySelfSize() const {
+  return sizeof(JavaScriptHandler);
+}
+
+void EventTarget::EventHandler::JavaScriptHandler::jsgGetMemoryInfo(
+    jsg::MemoryTracker& tracker) const {
+  tracker.trackField("identity", identity);
+  tracker.trackField("callback", callback);
+  if (abortHandler != kj::none) {
+    tracker.trackFieldWithSize("abortHandler", sizeof(kj::Own<NativeHandler>) +
+                                                sizeof(NativeHandler));
+  }
+}
+
+size_t EventTarget::EventHandler::jsgGetMemorySelfSize() const {
+  return sizeof(EventHandler);
+}
+
+void EventTarget::EventHandler::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {
+  KJ_SWITCH_ONEOF(handler) {
+    KJ_CASE_ONEOF(js, JavaScriptHandler) {
+      tracker.trackField("js", js);
+    }
+    KJ_CASE_ONEOF(native, NativeHandlerRef) {
+      tracker.trackFieldWithSize("native", sizeof(NativeHandlerRef));
+    }
+  }
+}
+
+size_t EventTarget::EventHandlerSet::jsgGetMemorySelfSize() const {
+  return sizeof(EventHandlerSet);
+}
+
+void EventTarget::EventHandlerSet::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {
+  tracker.trackField("handlers", handlers, kj::none, kj::none, false);
+}
+
+void EventTarget::visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+  tracker.trackField("typeMap", typeMap);
+}
+
 }  // namespace workerd::api

src/workerd/api/basics.h

@@ -148,6 +148,11 @@ class Event: public jsg::Object {
     JSG_STATIC_CONSTANT(BUBBLING_PHASE);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("type", ownType);
+    tracker.trackField("target", target);
+  }
+
 private:
   kj::StringPtr type;
   kj::String ownType;
@@ -215,6 +220,10 @@ class CustomEvent: public Event {
     JSG_READONLY_PROTOTYPE_PROPERTY(detail, getDetail);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("detail", detail);
+  }
+
 private:
   jsg::Optional<jsg::JsRef<jsg::JsValue>> detail;
 };
@@ -334,6 +343,8 @@ class EventTarget: public jsg::Object {
       jsg::Function<void(jsg::Ref<Event>)> func,
       bool once = false);
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const;
+
 private:
   void addNativeListener(jsg::Lock& js, NativeHandler& handler);
   bool removeNativeListener(NativeHandler& handler);
@@ -352,6 +363,10 @@ class EventTarget: public jsg::Object {
           handler->visitForGc(visitor);
         }
       }
+
+      kj::StringPtr jsgGetMemoryName() const { return "JavaScriptHandler"_kjc; }
+      size_t jsgGetMemorySelfSize() const;
+      void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const;
     };
 
     struct NativeHandlerRef {
@@ -367,6 +382,10 @@ class EventTarget: public jsg::Object {
 
     // When once is true, the handler will be removed after it is invoked one time.
     bool once = false;
+
+    kj::StringPtr jsgGetMemoryName() const { return "EventHandler"_kjc; }
+    size_t jsgGetMemorySelfSize() const;
+    void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const;
   };
 
   struct EventHandlerHashCallbacks {
@@ -389,6 +408,10 @@ class EventTarget: public jsg::Object {
 
     EventHandlerSet()
         : handlers(EventHandlerHashCallbacks(), {}) {}
+
+    kj::StringPtr jsgGetMemoryName() const { return "EventHandlerSet"_kjc; }
+    size_t jsgGetMemorySelfSize() const;
+    void jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const;
   };
 
   EventHandlerSet& getOrCreate(kj::StringPtr str) KJ_LIFETIMEBOUND;
@@ -492,6 +515,13 @@ class AbortSignal final: public EventTarget {
 
   RefcountedCanceler& getCanceler();
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    EventTarget::visitForMemoryInfo(tracker);
+    tracker.trackInlineFieldWithSize("IoOwn<RefcountedCanceler>",
+        sizeof(IoOwn<RefcountedCanceler>));
+    tracker.trackField("reason", reason);
+  }
+
 private:
   IoOwn<RefcountedCanceler> canceler;
   Flag flag;
@@ -525,6 +555,10 @@ class AbortController final: public jsg::Object {
     JSG_METHOD(abort);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("signal", signal);
+  }
+
 private:
   jsg::Ref<AbortSignal> signal;
 

src/workerd/api/global-scope.h

@@ -122,6 +122,11 @@ class PromiseRejectionEvent: public Event {
     JSG_READONLY_INSTANCE_PROPERTY(reason, getReason);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("promise", promise);
+    tracker.trackField("reason", reason);
+  }
+
 private:
   jsg::V8Ref<v8::Promise> promise;
   jsg::Value reason;
@@ -656,6 +661,10 @@ class ServiceWorkerGlobalScope: public WorkerGlobalScope {
   TimeoutId::Generator timeoutIdGenerator;
   // The generator for all timeout IDs associated with this scope.
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("unhandledRejections", unhandledRejections);
+  }
+
 private:
   jsg::UnhandledRejectionHandler unhandledRejections;
 

src/workerd/api/gpu/gpu-device.h

@@ -120,6 +120,10 @@ class GPUUncapturedErrorEvent : public Event {
     JSG_READONLY_INSTANCE_PROPERTY(error, getError);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("error", error_);
+  }
+
 private:
   jsg::Ref<GPUError> error_;
 

src/workerd/api/web-socket.h

@@ -67,6 +67,10 @@ class MessageEvent: public Event {
     });
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("data", data);
+  }
+
 private:
   jsg::JsRef<jsg::JsValue> data;
 
@@ -112,6 +116,10 @@ class CloseEvent: public Event {
     // CloseEvent will be referenced from the `WebSocketEventMap` define
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("reason", reason);
+  }
+
 private:
   int code;
   kj::String reason;
@@ -147,6 +155,11 @@ class ErrorEvent: public Event {
     // ErrorEvent will be referenced from the `WebSocketEventMap` define
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("message", message);
+    tracker.trackField("error", error);
+  }
+
 private:
   kj::String message;
   jsg::JsRef<jsg::JsValue> error;

src/workerd/jsg/BUILD.bazel

@@ -10,6 +10,7 @@ wd_cc_library(
         ["*.c++"],
         exclude = [
             "exception.c++",
+            "memory.c++",
             "*-test.c++",
         ],
     ),
@@ -43,6 +44,7 @@ wd_cc_library(
 wd_cc_library(
     name = "memory-tracker",
     hdrs = ["memory.h"],
+    srcs = ["memory.c++"],
     visibility = ["//visibility:public"],
     deps = [
         "@capnp-cpp//src/kj",

src/workerd/jsg/iterator.h

@@ -813,7 +813,7 @@ class AsyncIteratorBase: public Object {
       if constexpr (MemoryRetainer<State>) {
         tracker.trackField("state", state);
       } else {
-        tracker.trackField("state", sizeof(State));
+        tracker.trackFieldWithSize("state", sizeof(State));
       }
       tracker.trackField("impl", impl);
     }

src/workerd/jsg/jsg.h

@@ -833,12 +833,12 @@ class HashableV8Ref: public V8Ref<T> {
 };
 
 template <V8Value T>
-MemoryTracker& MemoryTracker::trackField(
+void MemoryTracker::trackField(
     kj::StringPtr edgeName,
     const V8Ref<T>& value,
     kj::Maybe<kj::StringPtr> nodeName) {
   // Even though we're passing in a template T, casting to a v8::Value is sufficient here.
-  return trackField(edgeName, value.handle.Get(isolate_)
+  trackField(edgeName, value.handle.Get(isolate_)
       .template As<v8::Value>(), nodeName);
 }
 
@@ -1221,11 +1221,11 @@ class Ref {
 };
 
 template <MemoryRetainer T>
-MemoryTracker& MemoryTracker::trackField(
+void MemoryTracker::trackField(
     kj::StringPtr edgeName,
     const Ref<T>& value,
     kj::Maybe<kj::StringPtr> nodeName) {
-  return trackField(edgeName, value.get(), nodeName);
+  trackField(edgeName, value.get(), nodeName);
 }
 
 template <typename T, typename... Params>