Merge pull request #1652 from cloudflare/jsnell/even-moar-memory-tracker

src/workerd/api/actor-state.h

@@ -373,6 +373,21 @@ class ActorState: public jsg::Object {
     JSG_TS_OVERRIDE(type ActorState = never);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    KJ_SWITCH_ONEOF(id) {
+      KJ_CASE_ONEOF(str, kj::String) {
+        tracker.trackField("id", str);
+      }
+      KJ_CASE_ONEOF(id, kj::Own<ActorIdFactory::ActorId>) {
+        // TODO(later): This only yields the shallow size of the ActorId and not the
+        // size of the actual value. Should probably make ActorID a MemoryRetainer.
+        tracker.trackFieldWithSize("id", sizeof(ActorIdFactory::ActorId));
+      }
+    }
+    tracker.trackField("transient", transient);
+    tracker.trackField("persistent", persistent);
+  }
+
 private:
   Worker::Actor::Id id;
   kj::Maybe<jsg::JsRef<jsg::JsValue>> transient;
@@ -396,6 +411,11 @@ class WebSocketRequestResponsePair: public jsg::Object {
     JSG_READONLY_PROTOTYPE_PROPERTY(response, getResponse);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("request", request);
+    tracker.trackField("response", response);
+  }
+
 private:
   kj::String request;
   kj::String response;
@@ -504,6 +524,20 @@ class DurableObjectState: public jsg::Object {
     // Make `storage` non-optional
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    KJ_SWITCH_ONEOF(id) {
+      KJ_CASE_ONEOF(str, kj::String) {
+        tracker.trackField("id", str);
+      }
+      KJ_CASE_ONEOF(id, kj::Own<ActorIdFactory::ActorId>) {
+        // TODO(later): This only yields the shallow size of the ActorId and not the
+        // size of the actual value. Should probably make ActorID a MemoryRetainer.
+        tracker.trackFieldWithSize("id", sizeof(ActorIdFactory::ActorId));
+      }
+    }
+    tracker.trackField("storage", storage);
+  }
+
 private:
   Worker::Actor::Id id;
   kj::Maybe<jsg::Ref<DurableObjectStorage>> storage;

src/workerd/api/actor.h

@@ -64,6 +64,10 @@ class DurableObjectId: public jsg::Object {
     JSG_READONLY_INSTANCE_PROPERTY(name, getName);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackFieldWithSize("id", sizeof(ActorIdFactory::ActorId));
+  }
+
 private:
   kj::Own<ActorIdFactory::ActorId> id;
 
@@ -100,6 +104,10 @@ class DurableObject final: public Fetcher {
     // the interface implemented by users' Durable Object classes.
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("id", id);
+  }
+
 private:
   jsg::Ref<DurableObjectId> id;
 

src/workerd/api/html-rewriter.c++

@@ -177,6 +177,12 @@ struct UnregisteredElementHandlers {
   void visitForGc(jsg::GcVisitor& visitor) {
     visitor.visit(element, comments, text);
   }
+
+  JSG_MEMORY_INFO(UnregisteredElementHandlers) {
+    tracker.trackField("element", element);
+    tracker.trackField("comments", comments);
+    tracker.trackField("text", text);
+  }
 };
 
 struct UnregisteredDocumentHandlers {
@@ -193,6 +199,13 @@ struct UnregisteredDocumentHandlers {
   void visitForGc(jsg::GcVisitor& visitor) {
     visitor.visit(doctype, comments, text, end);
   }
+
+  JSG_MEMORY_INFO(UnregisteredDocumentHandlers) {
+    tracker.trackField("doctype", doctype);
+    tracker.trackField("comments", comments);
+    tracker.trackField("text", text);
+    tracker.trackField("end", end);
+  }
 };
 
 using UnregisteredElementOrDocumentHandlers =
@@ -988,11 +1001,28 @@ struct HTMLRewriter::Impl {
   //
   //   In the meantime, we keep this list of handlers around and "replay" their registration, in
   //   order, on the builder object that we create inside of .transform().
+
+  JSG_MEMORY_INFO(HTMLRewriter::Impl) {
+    for (const auto& handlers : unregisteredHandlers) {
+      KJ_SWITCH_ONEOF(handlers) {
+        KJ_CASE_ONEOF(h, UnregisteredElementHandlers) {
+          tracker.trackField(nullptr, h);
+        }
+        KJ_CASE_ONEOF(h, UnregisteredDocumentHandlers) {
+          tracker.trackField(nullptr, h);
+        }
+      }
+    }
+  }
 };
 
 HTMLRewriter::HTMLRewriter(): impl(kj::heap<Impl>()) {}
 HTMLRewriter::~HTMLRewriter() noexcept(false) {}
 
+void HTMLRewriter::visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+  tracker.trackField("impl", impl);
+}
+
 jsg::Ref<HTMLRewriter> HTMLRewriter::constructor() {
   return jsg::alloc<HTMLRewriter>();
 }

src/workerd/api/html-rewriter.h

@@ -108,6 +108,8 @@ class HTMLRewriter: public jsg::Object {
     JSG_METHOD(transform);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const;
+
 private:
   void visitForGc(jsg::GcVisitor& visitor);
 

src/workerd/api/url.h

@@ -106,6 +106,30 @@ class URL: public jsg::Object {
   // Treat as private -- needs to be public for jsg::alloc<T>()...
   explicit URL(kj::Url&& u);
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    size_t size = 0;
+    size += url->scheme.size();
+    KJ_IF_SOME(ui, url->userInfo) {
+      size += ui.username.size();
+      KJ_IF_SOME(pwd, ui.password) {
+        size += pwd.size();
+      }
+    }
+    size += url->host.size();
+    for (const auto& p: url->path) {
+      size += p.size();
+    }
+    for (const auto& param : url->query) {
+      size += param.name.size();
+      size += param.value.size();
+    }
+    KJ_IF_SOME(frag, url->fragment) {
+      size += frag.size();
+    }
+    tracker.trackFieldWithSize("inner", size);
+    tracker.trackField("searchParams", searchParams);
+  }
+
 private:
   friend class URLSearchParams;
 
@@ -213,6 +237,10 @@ class URLSearchParams: public jsg::Object {
     }
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackFieldWithSize("url", url->toString().size());
+  }
+
 private:
   kj::Own<URL::RefcountedUrl> url;