use v8::ExternalMemoryAccounter more (#3485)

src/workerd/jsg/jsg.c++

@@ -289,10 +289,11 @@ JsSymbol Lock::symbolAsyncDispose() {
   return IsolateBase::from(v8Isolate).getSymbolAsyncDispose();
 }
 
-void ExternalMemoryAdjustment::maybeDeferAdjustment(v8::Isolate* isolate, size_t amount) {
+void ExternalMemoryAdjustment::maybeDeferAdjustment(
+    v8::ExternalMemoryAccounter& externalMemoryAccounter, v8::Isolate* isolate, size_t amount) {
   if (isolate == nullptr) return;
   if (v8::Locker::IsLocked(isolate)) {
-    isolate->AdjustAmountOfExternalAllocatedMemory(static_cast<int64_t>(-amount));
+    externalMemoryAccounter.Decrease(isolate, amount);
   } else {
     // Otherwise, if we don't have the isolate locked, defer the adjustment to the next
     // time that we do.
@@ -301,16 +302,19 @@ void ExternalMemoryAdjustment::maybeDeferAdjustment(v8::Isolate* isolate, size_t
   }
 }
 
-ExternalMemoryAdjustment::ExternalMemoryAdjustment(v8::Isolate* isolate, size_t amount)
-    : amount(amount),
-      isolate(isolate) {
+ExternalMemoryAdjustment::ExternalMemoryAdjustment(
+    v8::ExternalMemoryAccounter& externalMemoryAccounter, v8::Isolate* isolate, size_t amount)
+    : externalMemoryAccounter(externalMemoryAccounter),
+      isolate(isolate),
+      amount(amount) {
   KJ_DASSERT(isolate != nullptr);
-  isolate->AdjustAmountOfExternalAllocatedMemory(amount);
+  externalMemoryAccounter.Increase(isolate, amount);
 }
 
 ExternalMemoryAdjustment::ExternalMemoryAdjustment(ExternalMemoryAdjustment&& other)
-    : amount(other.amount),
-      isolate(other.isolate) {
+    : externalMemoryAccounter(other.externalMemoryAccounter),
+      isolate(other.isolate),
+      amount(other.amount) {
   other.amount = 0;
   other.isolate = nullptr;
 }
@@ -319,7 +323,8 @@ ExternalMemoryAdjustment& ExternalMemoryAdjustment::operator=(ExternalMemoryAdju
   // If we currently have an amount, adjust it back to zero.
   // In the case we don't have the isolate lock here, the adjustment
   // will be deferred until the next time we do.
-  if (amount > 0) maybeDeferAdjustment(isolate, amount);
+  if (amount > 0) maybeDeferAdjustment(externalMemoryAccounter, isolate, amount);
+  externalMemoryAccounter = kj::mv(other.externalMemoryAccounter);
   amount = other.amount;
   isolate = other.isolate;
   other.amount = 0;
@@ -329,22 +334,23 @@ ExternalMemoryAdjustment& ExternalMemoryAdjustment::operator=(ExternalMemoryAdju
 
 ExternalMemoryAdjustment::~ExternalMemoryAdjustment() noexcept(false) {
   if (amount != 0) {
-    maybeDeferAdjustment(isolate, amount);
+    maybeDeferAdjustment(externalMemoryAccounter, isolate, amount);
   }
 }
 
 void ExternalMemoryAdjustment::adjust(Lock& js, ssize_t amount) {
   amount = kj::max(amount, -static_cast<ssize_t>(this->amount));
   this->amount += amount;
-  isolate->AdjustAmountOfExternalAllocatedMemory(amount);
+  externalMemoryAccounter.Update(isolate, amount);
 }
 
 void ExternalMemoryAdjustment::set(Lock& js, size_t amount) {
   adjust(js, amount - this->amount);
 }
 
 ExternalMemoryAdjustment Lock::getExternalMemoryAdjustment(int64_t amount) {
-  return ExternalMemoryAdjustment(v8Isolate, amount);
+  return ExternalMemoryAdjustment(
+      IsolateBase::from(v8Isolate).getExternalMemoryAccounter(), v8Isolate, amount);
 }
 
 Name::Name(kj::String string): hash(kj::hashCode(string)), inner(kj::mv(string)) {}

src/workerd/jsg/jsg.h

@@ -14,6 +14,7 @@
 #include <workerd/jsg/exception.h>
 #include <workerd/jsg/memory.h>
 
+#include <v8-external-memory-accounter.h>
 #include <v8-profiler.h>
 #include <v8.h>
 
@@ -2220,7 +2221,8 @@ class DOMException;
 // The allocation amount can be adjusted up or down during the lifetime of an object.
 class ExternalMemoryAdjustment final {
  public:
-  ExternalMemoryAdjustment() = default;
+  ExternalMemoryAdjustment(
+      v8::ExternalMemoryAccounter& externalMemoryAccounter, v8::Isolate* isolate, size_t amount);
   ExternalMemoryAdjustment(v8::Isolate* isolate, size_t amount);
   ExternalMemoryAdjustment(ExternalMemoryAdjustment&& other);
   ExternalMemoryAdjustment& operator=(ExternalMemoryAdjustment&& other);
@@ -2243,10 +2245,12 @@ class ExternalMemoryAdjustment final {
   }
 
  private:
-  size_t amount = 0;
+  v8::ExternalMemoryAccounter& externalMemoryAccounter;
   v8::Isolate* isolate = nullptr;
+  size_t amount = 0;
 
-  static void maybeDeferAdjustment(v8::Isolate* isolate, size_t amount);
+  static void maybeDeferAdjustment(
+      v8::ExternalMemoryAccounter& externalMemoryAccounter, v8::Isolate* isolate, size_t amount);
 };
 
 // Represents an isolate lock, which allows the current thread to execute JavaScript code within

src/workerd/jsg/setup.c++

@@ -193,7 +193,9 @@ void IsolateBase::deferExternalMemoryDecrement(int64_t size) {
 void IsolateBase::clearPendingExternalMemoryDecrement() {
   KJ_ASSERT(v8::Locker::IsLocked(ptr));
   int64_t amount = pendingExternalMemoryDecrement.exchange(0, std::memory_order_relaxed);
-  if (amount > 0) ptr->AdjustAmountOfExternalAllocatedMemory(-amount);
+  if (amount > 0) {
+    externalMemoryAccounter.Decrease(ptr, amount);
+  }
 }
 
 void IsolateBase::terminateExecution() const {

src/workerd/jsg/setup.h

@@ -202,6 +202,10 @@ class IsolateBase {
     return JsSymbol(symbolAsyncDispose.Get(ptr));
   }
 
+  v8::ExternalMemoryAccounter& getExternalMemoryAccounter() {
+    return externalMemoryAccounter;
+  }
+
  private:
   template <typename TypeWrapper>
   friend class Isolate;
@@ -266,6 +270,9 @@ class IsolateBase {
   // Polyfilled Symbol.asyncDispose.
   v8::Global<v8::Symbol> symbolAsyncDispose;
 
+  // Used to account for external memory
+  v8::ExternalMemoryAccounter externalMemoryAccounter;
+
   // We expect queues to remain relatively small -- 8 is the largest size I have observed from local
   // testing.
   static constexpr auto DESTRUCTION_QUEUE_INITIAL_SIZE = 8;