From d181608175a703537e75db07026cd3eeaf808c34 Mon Sep 17 00:00:00 2001
From: Yagiz Nizipli <yagiz@nizipli.com>
Date: Fri, 13 Sep 2024 12:52:20 -0400
Subject: [PATCH] use `ExternalMemoryAdjuster` on zlib allocations

---
 src/workerd/api/node/zlib-util.c++ | 16 +++++++++++++---
 src/workerd/api/node/zlib-util.h   | 14 ++++++++++----
 2 files changed, 23 insertions(+), 7 deletions(-)

diff --git a/src/workerd/api/node/zlib-util.c++ b/src/workerd/api/node/zlib-util.c++
index 45ed4591af3..1976df1277d 100644
--- a/src/workerd/api/node/zlib-util.c++
+++ b/src/workerd/api/node/zlib-util.c++
@@ -802,9 +802,19 @@ void* ZlibUtil::Allocator::AllocForZlib(void* data, uInt items, uInt size) {
 
 void* ZlibUtil::Allocator::AllocForBrotli(void* opaque, size_t size) {
   auto* thisAllocator = static_cast<Allocator*>(opaque);
-  auto memory = kj::heapArray<uint8_t>(size);
-  auto begin = memory.begin();
-  thisAllocator->allocations.insert(begin, kj::mv(memory));
+  kj::Maybe<jsg::ExternalMemoryAdjustment> memory{};
+  auto isolate = v8::Isolate::GetCurrent();
+  if (isolate != nullptr) {
+    auto& js = jsg::Lock::from(isolate);
+    memory = js.getExternalMemoryAdjustment(size);
+  }
+  auto data = kj::heapArray<uint8_t>(size);
+  auto begin = data.begin();
+  thisAllocator->allocations.insert(begin,
+      {
+        .data = kj::mv(data),
+        .memory = kj::mv(memory),
+      });
   return begin;
 }
 
diff --git a/src/workerd/api/node/zlib-util.h b/src/workerd/api/node/zlib-util.h
index defb9db137f..07c7e77cc61 100644
--- a/src/workerd/api/node/zlib-util.h
+++ b/src/workerd/api/node/zlib-util.h
@@ -291,17 +291,23 @@ class ZlibUtil final: public jsg::Object {
   ZlibUtil() = default;
   ZlibUtil(jsg::Lock&, const jsg::Url&) {}
 
-  // A custom allocator to be used by the zlib and brotli libraries
-  // The current implementation stores allocations in a hash map.
-  // TODO: Use an arena allocator implementation instead of hashing pointers in order to improve performance
+  // A custom allocator to be used by the zlib and brotli libraries.
+  // The allocator should not and can not safely hold a reference to the jsg::Lock
+  // instance. Therefore, we lookup the current jsg::Lock instance from the
+  // isolate pointer and use that to get the external memory adjustment.
   class Allocator final {
   public:
     static void* AllocForZlib(void* data, uInt items, uInt size);
     static void* AllocForBrotli(void* data, size_t size);
     static void FreeForZlib(void* data, void* pointer);
 
+    struct Allocation {
+      kj::Array<kj::byte> data;
+      kj::Maybe<jsg::ExternalMemoryAdjustment> memory;
+    };
+
   private:
-    kj::HashMap<void*, kj::Array<kj::byte>> allocations;
+    kj::HashMap<void*, Allocation> allocations;
   };
 
   template <class CompressionContext>