Merge pull request #1625 from cloudflare/jsnell/even-moar-memorytracker

cloudflare · Feb 6, 2024 · 15f4640 · 15f4640
2 parents ec3a570 + 3911304
commit 15f4640
Show file tree

Hide file tree

Showing 7 changed files with 78 additions and 2 deletions.
diff --git a/src/workerd/api/cache.h b/src/workerd/api/cache.h
@@ -74,6 +74,10 @@ class Cache: public jsg::Object {
     // Use RequestInfo type alias to allow `URL`s as cache keys
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("cacheName", cacheName);
+  }
+
 private:
   kj::Maybe<kj::String> cacheName;
 
@@ -109,6 +113,10 @@ class CacheStorage: public jsg::Object {
     JSG_READONLY_INSTANCE_PROPERTY(default, getDefault);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("default", default_);
+  }
+
 private:
   jsg::Ref<Cache> default_;
 };

diff --git a/src/workerd/api/cf-property.h b/src/workerd/api/cf-property.h
@@ -42,6 +42,19 @@ class CfProperty {
 
   void visitForGc(jsg::GcVisitor& visitor);
 
+  JSG_MEMORY_INFO(CfProperty) {
+    KJ_IF_SOME(v, value) {
+      KJ_SWITCH_ONEOF(v) {
+        KJ_CASE_ONEOF(str, kj::String) {
+          tracker.trackField("value", str);
+        }
+        KJ_CASE_ONEOF(obj, jsg::JsRef<jsg::JsObject>) {
+          tracker.trackField("value", obj);
+        }
+      }
+    }
+  }
+
 private:
   kj::Maybe<kj::OneOf<kj::String, jsg::JsRef<jsg::JsObject>>> value;
 };

diff --git a/src/workerd/api/form-data.h b/src/workerd/api/form-data.h
@@ -32,6 +32,10 @@ class FormData: public jsg::Object {
     void visitForGc(jsg::GcVisitor& visitor) {
       visitor.visit(parent);
     }
+
+    JSG_MEMORY_INFO(IteratorState) {
+      tracker.trackField("parent", parent);
+    }
   };
 
 public:
@@ -51,6 +55,18 @@ class FormData: public jsg::Object {
   struct Entry {
     kj::String name;
     kj::OneOf<jsg::Ref<File>, kj::String> value;
+
+    JSG_MEMORY_INFO(Entry) {
+      tracker.trackField("name", name);
+      KJ_SWITCH_ONEOF(value) {
+        KJ_CASE_ONEOF(file, jsg::Ref<File>) {
+          tracker.trackField("value", file);
+        }
+        KJ_CASE_ONEOF(str, kj::String) {
+          tracker.trackField("value", str);
+        }
+      }
+    }
   };
 
   kj::ArrayPtr<const Entry> getData() { return data; }
@@ -142,6 +158,10 @@ class FormData: public jsg::Object {
     }
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("data", data.asPtr());
+  }
+
 private:
   kj::Vector<Entry> data;
 

diff --git a/src/workerd/api/hyperdrive.h b/src/workerd/api/hyperdrive.h
@@ -43,6 +43,14 @@ class Hyperdrive : public jsg::Object {
     JSG_METHOD(connect);
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("randomHost", randomHost);
+    tracker.trackField("database", database);
+    tracker.trackField("user", user);
+    tracker.trackField("password", password);
+    tracker.trackField("scheme", scheme);
+  }
+
 private:
   uint clientIndex;
   kj::String randomHost;

diff --git a/src/workerd/api/kv.h b/src/workerd/api/kv.h
@@ -17,6 +17,11 @@ class KvNamespace: public jsg::Object {
   struct AdditionalHeader {
     kj::String name;
     kj::String value;
+
+    JSG_MEMORY_INFO(AdditionalHeader) {
+      tracker.trackField("name", name);
+      tracker.trackField("value", value);
+    }
   };
 
   // `subrequestChannel` is what to pass to IoContext::getHttpClient() to get an HttpClient
@@ -155,6 +160,10 @@ class KvNamespace: public jsg::Object {
     });
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const  {
+    tracker.trackField("additionalHeaders", additionalHeaders.asPtr());
+  }
+
 protected:
   // Do the boilerplate work of constructing an HTTP client to KV. Setting a KvOptType causes
   // the limiter for that op type to be checked. If a string is used, that's used as the operation

diff --git a/src/workerd/api/queue.h b/src/workerd/api/queue.h
@@ -136,6 +136,12 @@ class QueueMessage final: public jsg::Object {
     });
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("id", id);
+    tracker.trackField("body", body);
+    tracker.trackFieldWithSize("IoPtr<QueueEventResult>", sizeof(IoPtr<QueueEventResult>));
+  }
+
 private:
   kj::String id;
   kj::Date timestamp;
@@ -182,6 +188,14 @@ class QueueEvent final: public ExtendableEvent {
     });
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    for (auto& message: messages) {
+      tracker.trackField("message", message);
+    }
+    tracker.trackField("queueName", queueName);
+    tracker.trackFieldWithSize("IoPtr<QueueEventResult>", sizeof(IoPtr<QueueEventResult>));
+  }
+
 private:
   // TODO(perf): Should we store these in a v8 array directly rather than this intermediate kj
   // array to avoid one intermediate copy?
@@ -218,6 +232,10 @@ class QueueController final: public jsg::Object {
     });
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("event", event);
+  }
+
 private:
   jsg::Ref<QueueEvent> event;
 

diff --git a/src/workerd/jsg/memory.h b/src/workerd/jsg/memory.h
@@ -536,7 +536,7 @@ void MemoryTracker::trackField(
     decCurrentNodeSize(sizeof(T));
   }
   pushNode(nodeName.orDefault(edgeName), sizeof(T), edgeName);
-  for (const auto item : value) {
+  for (const auto& item : value) {
     trackField(nullptr, item, elementName);
   }
   nodeStack_.pop();
@@ -554,7 +554,7 @@ void MemoryTracker::trackField(
     decCurrentNodeSize(sizeof(T));
   }
   pushNode(nodeName.orDefault(edgeName), sizeof(T), edgeName);
-  for (const auto item : value) {
+  for (const auto& item : value) {
     trackField(nullptr, item, elementName);
   }
   nodeStack_.pop();