From 983b4ecb5394af5f7e53120dfdef1b9059174713 Mon Sep 17 00:00:00 2001
From: James M Snell <jasnell@gmail.com>
Date: Wed, 7 Aug 2024 12:12:43 -0700
Subject: [PATCH] Adds APIs to jsg::Lock to perform manual external memory
 accounting.

---
 src/workerd/jsg/jsg.c++ | 23 +++++++++++++++++++++++
 src/workerd/jsg/jsg.h   | 14 ++++++++++++++
 2 files changed, 37 insertions(+)

diff --git a/src/workerd/jsg/jsg.c++ b/src/workerd/jsg/jsg.c++
index cc4601b34b6..12ac288d2e9 100644
--- a/src/workerd/jsg/jsg.c++
+++ b/src/workerd/jsg/jsg.c++
@@ -278,6 +278,29 @@ JsSymbol Lock::symbolAsyncDispose() {
   return IsolateBase::from(v8Isolate).getSymbolAsyncDispose();
 }
 
+void Lock::adjustExternalMemory(ssize_t amount) {
+  v8Isolate->AdjustAmountOfExternalAllocatedMemory(static_cast<int64_t>(amount));
+}
+
+namespace {
+struct ExternalMemoryAdjuster {
+  int64_t size;
+  v8::Isolate* isolate;
+  ExternalMemoryAdjuster(v8::Isolate* isolate, size_t size)
+      : size(static_cast<int64_t>(size)), isolate(isolate) {
+    isolate->AdjustAmountOfExternalAllocatedMemory(size);
+  }
+
+  ~ExternalMemoryAdjuster() noexcept(false) {
+    isolate->AdjustAmountOfExternalAllocatedMemory(-size);
+  }
+};
+}  // namespace
+
+kj::Own<void> Lock::getExternalMemoryAdjuster(size_t amount) {
+  return kj::heap<ExternalMemoryAdjuster>(v8Isolate, amount);
+}
+
 Name::Name(kj::String string)
     : hash(kj::hashCode(string)),
       inner(kj::mv(string)) {}
diff --git a/src/workerd/jsg/jsg.h b/src/workerd/jsg/jsg.h
index eda9ff57793..580e2fd5e16 100644
--- a/src/workerd/jsg/jsg.h
+++ b/src/workerd/jsg/jsg.h
@@ -2085,6 +2085,20 @@ class Lock {
     return *reinterpret_cast<Lock*>(v8Isolate->GetData(2));
   }
 
+  // Manually make adjustments to the amount of external memory reported to V8.
+  // This is useful when we have a large amount of external memory allocated that
+  // typically would not be visible to v8's memory tracking. In the future we hope
+  // to make most memory accounting automatic, making most direct uses of this
+  // API unnecessary but there will always be times when manual adjustments are
+  // necessary.
+  void adjustExternalMemory(ssize_t amount);
+
+  // Reports amount of external memory to be manually attributed to the isolate.
+  // When the returned kj::Own<void> is dropped, the amount will be subtracted
+  // from the isolate's external memory accounting. It is important that these
+  // be held onto only during the lifetime of the isolate.
+  kj::Own<void> getExternalMemoryAdjuster(size_t amount);
+
   Value parseJson(kj::ArrayPtr<const char> data);
   Value parseJson(v8::Local<v8::String> text);
   template <typename T>