Implement heap snapshot memory tracking for TransformStream

cloudflare · Feb 2, 2024 · a9982b1 · a9982b1
1 parent ad520e4
commit a9982b1
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 6 deletions.
diff --git a/src/workerd/api/streams/queue.h b/src/workerd/api/streams/queue.h
@@ -1041,7 +1041,7 @@ kj::StringPtr ValueQueue::Consumer::jsgGetMemoryName() const {
 }
 
 size_t ValueQueue::Consumer::jsgGetMemorySelfSize() const {
-  return sizeof(ValueQueue::Consumer) - sizeof(ConsumerImpl);
+  return sizeof(ValueQueue::Consumer);
 }
 
 void ValueQueue::Consumer::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {
@@ -1053,7 +1053,7 @@ kj::StringPtr ValueQueue::jsgGetMemoryName() const {
 }
 
 size_t ValueQueue::jsgGetMemorySelfSize() const {
-  return sizeof(ValueQueue) - sizeof(QueueImpl);
+  return sizeof(ValueQueue);
 }
 
 void ValueQueue::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {
@@ -1065,7 +1065,7 @@ kj::StringPtr ByteQueue::Consumer::jsgGetMemoryName() const {
 }
 
 size_t ByteQueue::Consumer::jsgGetMemorySelfSize() const {
-  return sizeof(ByteQueue::Consumer) - sizeof(ConsumerImpl);
+  return sizeof(ByteQueue::Consumer);
 }
 
 void ByteQueue::Consumer::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {
@@ -1077,7 +1077,7 @@ kj::StringPtr ByteQueue::jsgGetMemoryName() const {
 }
 
 size_t ByteQueue::jsgGetMemorySelfSize() const {
-  return sizeof(ByteQueue) - sizeof(QueueImpl);
+  return sizeof(ByteQueue);
 }
 
 void ByteQueue::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {

diff --git a/src/workerd/api/streams/standard.c++ b/src/workerd/api/streams/standard.c++
@@ -4185,7 +4185,7 @@ kj::StringPtr WritableStreamJsController::jsgGetMemoryName() const {
 }
 
 size_t WritableStreamJsController::jsgGetMemorySelfSize() const {
-  return sizeof(WritableStreamJsController) - sizeof(WritableLockImpl);
+  return sizeof(WritableStreamJsController);
 }
 
 void WritableStreamJsController::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {
@@ -4211,7 +4211,7 @@ kj::StringPtr ReadableStreamJsController::jsgGetMemoryName() const {
 }
 
 size_t ReadableStreamJsController::jsgGetMemorySelfSize() const {
-  return sizeof(ReadableStreamJsController) - sizeof(ReadableLockImpl);
+  return sizeof(ReadableStreamJsController);
 }
 
 void ReadableStreamJsController::jsgGetMemoryInfo(jsg::MemoryTracker& tracker) const {
@@ -4277,4 +4277,11 @@ void ReadableStreamBYOBRequest::visitForMemoryInfo(jsg::MemoryTracker& tracker)
   }
 }
 
+void TransformStreamDefaultController::visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+  tracker.trackField("startPromise", startPromise);
+  tracker.trackField("maybeBackpressureChange", maybeBackpressureChange);
+  tracker.trackField("transformAlgorithm", algorithms.transform);
+  tracker.trackField("flushAlgorithm", algorithms.flush);
+}
+
 }  // namespace workerd::api
diff --git a/src/workerd/api/streams/standard.h b/src/workerd/api/streams/standard.h
@@ -679,6 +679,8 @@ class TransformStreamDefaultController: public jsg::Object {
   jsg::Promise<void> pull(jsg::Lock& js);
   jsg::Promise<void> cancel(jsg::Lock& js, v8::Local<v8::Value> reason);
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const;
+
 private:
   struct Algorithms {
     kj::Maybe<jsg::Function<Transformer::TransformAlgorithm>> transform;

diff --git a/src/workerd/api/streams/transform.h b/src/workerd/api/streams/transform.h
@@ -63,6 +63,11 @@ class TransformStream: public jsg::Object {
     }
   }
 
+  void visitForMemoryInfo(jsg::MemoryTracker& tracker) const {
+    tracker.trackField("readable", readable);
+    tracker.trackField("writable", writable);
+  }
+
 private:
   jsg::Ref<ReadableStream> readable;
   jsg::Ref<WritableStream> writable;