Merge pull request #2726 from cloudflare/yagiz/enforce-isolate-lock

enforce isolate lock and track allocations on v8

src/workerd/api/node/zlib-util.h

@@ -293,7 +293,6 @@ class ZlibUtil final: public jsg::Object {
   class CompressionStream: public jsg::Object {
   public:
     explicit CompressionStream(ZlibMode _mode): context_(_mode) {}
-    CompressionStream() = default;
     // TODO(soon): Find a way to add noexcept(false) to this destructor.
     ~CompressionStream();
     KJ_DISALLOW_COPY_AND_MOVE(CompressionStream);

src/workerd/api/streams/compression.c++

@@ -26,20 +26,42 @@ void* CompressionAllocator::AllocForZlib(void* data, uInt items, uInt size) {
   return AllocForBrotli(data, real_size);
 }
 
-// TODO(soon): Enabling V8 reporting breaks CompressionStreamImpl on edgeworker. Investigate.
-// This is mostly likely due to running CompressionStream without an isolate lock.
 void* CompressionAllocator::AllocForBrotli(void* opaque, size_t size) {
-  auto* thisAllocator = static_cast<CompressionAllocator*>(opaque);
-  auto data = kj::heapArray<uint8_t>(size);
+  auto* allocator = static_cast<CompressionAllocator*>(opaque);
+  auto data = kj::heapArray<kj::byte>(size);
   auto begin = data.begin();
-  thisAllocator->allocations.insert(begin, kj::mv(data));
+  auto isolate = v8::Isolate::TryGetCurrent();
+  kj::Maybe<jsg::ExternalMemoryAdjustment> maybeMemoryAdjustment;
+  // TODO(soon): Improve this. We want to track external memory allocations
+  // with the v8 isolate so we can account for these as part of the isolate
+  // heap memory limits. However, we don't always have an isolate lock or
+  // current isolate when this is called so we can't just blindly try
+  // grabbing the isolate. For now we'll only be able to account for the
+  // allocations when we actually have an isolate. It's a bit tricky but
+  // we could possibly try implementing a deferred accounting adjustment?
+  // Basically, defer incrementing the memory allocation reported to the
+  // isolate until we have the isolate lock again? But that's a bit tricky
+  // if the adjustment is dropped before that happens. Will have to think
+  // through how best to approach that.
+  if (isolate != nullptr) {
+    auto& js = jsg::Lock::from(isolate);
+    maybeMemoryAdjustment = js.getExternalMemoryAdjustment(size);
+  }
+  allocator->allocations.insert(begin,
+      {
+        .data = kj::mv(data),
+        .memoryAdjustment = kj::mv(maybeMemoryAdjustment),
+      });
   return begin;
 }
 
 void CompressionAllocator::FreeForZlib(void* opaque, void* pointer) {
   if (KJ_UNLIKELY(pointer == nullptr)) return;
-  auto* thisAllocator = static_cast<CompressionAllocator*>(opaque);
-  JSG_REQUIRE(thisAllocator->allocations.erase(pointer), Error, "Zlib allocation should exist"_kj);
+  auto* allocator = static_cast<CompressionAllocator*>(opaque);
+  // No need to destroy memoryAdjustment here.
+  // Dropping the allocation from the hashmap will defer the adjustment
+  // until the isolate lock is held.
+  JSG_REQUIRE(allocator->allocations.erase(pointer), Error, "Zlib allocation should exist"_kj);
 }
 
 namespace {
@@ -138,6 +160,9 @@ public:
     };
   }
 
+protected:
+  CompressionAllocator allocator;
+
 private:
   static int getWindowBits(kj::StringPtr format) {
     // We use a windowBits value of 15 combined with the magic value
@@ -161,7 +186,6 @@ private:
   Mode mode;
   z_stream ctx = {};
   kj::byte buffer[16384];
-  CompressionAllocator allocator;
 
   // For the eponymous compatibility flag
   ContextFlags strictCompression;
@@ -469,7 +493,7 @@ private:
 };
 }  // namespace
 
-jsg::Ref<CompressionStream> CompressionStream::constructor(jsg::Lock& js, kj::String format) {
+jsg::Ref<CompressionStream> CompressionStream::constructor(kj::String format) {
   JSG_REQUIRE(format == "deflate" || format == "gzip" || format == "deflate-raw", TypeError,
       "The compression format must be either 'deflate', 'deflate-raw' or 'gzip'.");
 

src/workerd/api/streams/compression.h

@@ -18,28 +18,25 @@ namespace workerd::api {
 // isolate pointer and use that to get the external memory adjustment.
 class CompressionAllocator final {
 public:
-  CompressionAllocator() = default;
   void configure(z_stream* stream);
 
   static void* AllocForZlib(void* data, uInt items, uInt size);
   static void* AllocForBrotli(void* data, size_t size);
   static void FreeForZlib(void* data, void* pointer);
 
 private:
-  // TODO(soon): Use this instead of array<byte> after fixing edgeworker.
-  // struct Allocation {
-  //   kj::Array<kj::byte> data;
-  //   kj::Maybe<jsg::ExternalMemoryAdjustment> memoryAdjustment;
-  // };
-
-  kj::HashMap<void*, kj::Array<kj::byte>> allocations;
+  struct Allocation {
+    kj::Array<kj::byte> data;
+    kj::Maybe<jsg::ExternalMemoryAdjustment> memoryAdjustment = kj::none;
+  };
+  kj::HashMap<void*, Allocation> allocations;
 };
 
 class CompressionStream: public TransformStream {
 public:
   using TransformStream::TransformStream;
 
-  static jsg::Ref<CompressionStream> constructor(jsg::Lock& js, kj::String format);
+  static jsg::Ref<CompressionStream> constructor(kj::String format);
 
   JSG_RESOURCE_TYPE(CompressionStream) {
     JSG_INHERIT(TransformStream);