From a9982b1e6953b11bf113126fb2f90bd5e1a830f8 Mon Sep 17 00:00:00 2001
From: James M Snell <jsnell@cloudflare.com>
Date: Thu, 1 Feb 2024 18:18:14 -0800
Subject: [PATCH] Implement heap snapshot memory tracking for TransformStream

---
 src/workerd/api/streams/queue.h      |  8 ++++----
 src/workerd/api/streams/standard.c++ | 11 +++++++++--
 src/workerd/api/streams/standard.h   |  2 ++
 src/workerd/api/streams/transform.h  |  5 +++++
 4 files changed, 20 insertions(+), 6 deletions(-)

diff --git a/src/workerd/api/streams/queue.h b/src/workerd/api/streams/queue.h
index 76e7a07b471..203672dd1ca 100644
--- 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++
index ebbd113ce51..54b85c7bdc8 100644
--- 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
index 2db25b1cd7e..218e2f42d8a 100644
--- 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
index 334d27fa5f5..7549271bd45 100644
--- 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;