Fixes oven-sh#7827 (oven-sh#7858)

* Fixes oven-sh#7827

* Add some assertions and errors

---------

Co-authored-by: Jarred Sumner <709451+Jarred-Sumner@users.noreply.github.com>

src/bun.js/api/html_rewriter.zig

@@ -16,38 +16,43 @@ pub const LOLHTMLContext = struct {
     selectors: SelectorMap = .{},
     element_handlers: std.ArrayListUnmanaged(*ElementHandler) = .{},
     document_handlers: std.ArrayListUnmanaged(*DocumentHandler) = .{},
+    ref_count: u32 = 1,
 
-    pub fn deinit(this: *LOLHTMLContext, allocator: std.mem.Allocator) void {
+    pub usingnamespace bun.NewRefCounted(@This(), deinit);
+
+    fn deinit(this: *LOLHTMLContext) void {
         for (this.selectors.items) |selector| {
             selector.deinit();
         }
-        this.selectors.deinit(allocator);
+        this.selectors.deinit(bun.default_allocator);
         this.selectors = .{};
 
         for (this.element_handlers.items) |handler| {
             handler.deinit();
         }
-        this.element_handlers.deinit(allocator);
+        this.element_handlers.deinit(bun.default_allocator);
         this.element_handlers = .{};
 
         for (this.document_handlers.items) |handler| {
             handler.deinit();
         }
-        this.document_handlers.deinit(allocator);
+        this.document_handlers.deinit(bun.default_allocator);
         this.document_handlers = .{};
+
+        this.destroy();
     }
 };
 pub const HTMLRewriter = struct {
     builder: *LOLHTML.HTMLRewriter.Builder,
-    context: LOLHTMLContext,
+    context: *LOLHTMLContext,
 
     pub usingnamespace JSC.Codegen.JSHTMLRewriter;
 
     pub fn constructor(_: *JSGlobalObject, _: *JSC.CallFrame) callconv(.C) ?*HTMLRewriter {
         const rewriter = bun.default_allocator.create(HTMLRewriter) catch unreachable;
         rewriter.* = HTMLRewriter{
             .builder = LOLHTML.HTMLRewriter.Builder.init(),
-            .context = .{},
+            .context = LOLHTMLContext.new(.{}),
         };
         return rewriter;
     }
@@ -152,13 +157,13 @@ pub const HTMLRewriter = struct {
     }
 
     pub fn finalizeWithoutDestroy(this: *HTMLRewriter) void {
-        this.context.deinit(bun.default_allocator);
+        this.context.deref();
         this.builder.deinit();
     }
 
     pub fn beginTransform(this: *HTMLRewriter, global: *JSGlobalObject, response: *Response) JSValue {
         const new_context = this.context;
-        this.context = .{};
+        new_context.ref();
         return BufferOutputSink.init(new_context, global, response, this.builder);
     }
 
@@ -298,17 +303,10 @@ pub const HTMLRewriter = struct {
         pub fn setup(
             this: *HTMLRewriterLoader,
             builder: *LOLHTML.HTMLRewriter.Builder,
-            context: LOLHTMLContext,
+            context: *LOLHTMLContext,
             size_hint: ?usize,
             output: JSC.WebCore.Sink,
         ) ?[]const u8 {
-            for (context.document_handlers.items) |doc| {
-                doc.ctx = this;
-            }
-            for (context.element_handlers.items) |doc| {
-                doc.ctx = this;
-            }
-
             const chunk_size = @max(size_hint orelse 16384, 1024);
             this.rewriter = builder.build(
                 .UTF8,
@@ -392,13 +390,13 @@ pub const HTMLRewriter = struct {
         global: *JSGlobalObject,
         bytes: bun.MutableString,
         rewriter: ?*LOLHTML.HTMLRewriter = null,
-        context: LOLHTMLContext,
+        context: *LOLHTMLContext,
         response: *Response,
         response_value: JSC.Strong = .{},
         bodyValueBufferer: ?JSC.WebCore.BodyValueBufferer = null,
         tmp_sync_error: ?*JSC.JSValue = null,
         // const log = bun.Output.scoped(.BufferOutputSink, false);
-        pub fn init(context: LOLHTMLContext, global: *JSGlobalObject, original: *Response, builder: *LOLHTML.HTMLRewriter.Builder) JSC.JSValue {
+        pub fn init(context: *LOLHTMLContext, global: *JSGlobalObject, original: *Response, builder: *LOLHTML.HTMLRewriter.Builder) JSC.JSValue {
             var sink = bun.new(BufferOutputSink, BufferOutputSink{
                 .global = global,
                 .bytes = bun.MutableString.initEmpty(bun.default_allocator),
@@ -422,13 +420,6 @@ pub const HTMLRewriter = struct {
 
             sink.response = result;
 
-            for (sink.context.document_handlers.items) |doc| {
-                doc.ctx = sink;
-            }
-            for (sink.context.element_handlers.items) |doc| {
-                doc.ctx = sink;
-            }
-
             const input_size = original.body.len();
             sink.rewriter = builder.build(
                 .UTF8,
@@ -616,7 +607,7 @@ pub const HTMLRewriter = struct {
                 bufferer.deinit();
             }
 
-            this.context.deinit(bun.default_allocator);
+            this.context.deref();
             this.response_value.deinit();
             if (this.rewriter) |rewriter| {
                 rewriter.deinit();
@@ -752,7 +743,6 @@ const DocumentHandler = struct {
     onEndCallback: ?JSValue = null,
     thisObject: JSValue,
     global: *JSGlobalObject,
-    ctx: ?*HTMLRewriter.BufferOutputSink = null,
 
     pub const onDocType = HandlerCallback(
         DocumentHandler,
@@ -928,7 +918,6 @@ const ElementHandler = struct {
     onTextCallback: ?JSValue = null,
     thisObject: JSValue,
     global: *JSGlobalObject,
-    ctx: ?*HTMLRewriter.BufferOutputSink = null,
 
     pub fn init(global: *JSGlobalObject, thisObject: JSValue) !ElementHandler {
         var handler = ElementHandler{

src/bun.zig

@@ -2862,6 +2862,75 @@ pub fn New(comptime T: type) type {
     };
 }
 
+/// Reference-counted heap-allocated instance value.
+///
+/// `ref_count` is expected to be defined on `T` with a default value set to `1`
+pub fn NewRefCounted(comptime T: type, comptime deinit_fn: ?fn (self: *T) void) type {
+    if (!@hasField(T, "ref_count")) {
+        @compileError("Expected a field named \"ref_count\" with a default value of 1 on " ++ @typeName(T));
+    }
+
+    for (std.meta.fields(T)) |field| {
+        if (strings.eqlComptime(field.name, "ref_count")) {
+            if (field.default_value == null) {
+                @compileError("Expected a field named \"ref_count\" with a default value of 1 on " ++ @typeName(T));
+            }
+        }
+    }
+
+    return struct {
+        pub fn destroy(self: *T) void {
+            if (comptime Environment.allow_assert) {
+                std.debug.assert(self.ref_count == 0);
+            }
+
+            if (comptime is_heap_breakdown_enabled) {
+                HeapBreakdown.allocator(T).destroy(self);
+            } else {
+                default_allocator.destroy(self);
+            }
+        }
+
+        pub fn ref(self: *T) void {
+            self.ref_count += 1;
+        }
+
+        pub fn deref(self: *T) void {
+            self.ref_count -= 1;
+
+            if (self.ref_count == 0) {
+                if (comptime deinit_fn) |deinit| {
+                    deinit(self);
+                } else {
+                    self.destroy();
+                }
+            }
+        }
+
+        pub inline fn new(t: T) *T {
+            if (comptime is_heap_breakdown_enabled) {
+                const ptr = HeapBreakdown.allocator(T).create(T) catch outOfMemory();
+                ptr.* = t;
+
+                if (comptime Environment.allow_assert) {
+                    std.debug.assert(ptr.ref_count == 1);
+                }
+
+                return ptr;
+            }
+
+            const ptr = default_allocator.create(T) catch outOfMemory();
+            ptr.* = t;
+
+            if (comptime Environment.allow_assert) {
+                std.debug.assert(ptr.ref_count == 1);
+            }
+
+            return ptr;
+        }
+    };
+}
+
 /// Free a globally-allocated a value.
 ///
 /// Must have used `new` to allocate the value.

test/regression/issue/07827.test.ts

@@ -0,0 +1,22 @@
+import { test, expect, jest } from "bun:test";
+
+test("#7827", () => {
+  for (let i = 0; i < 10; i++)
+    (function () {
+      const element = jest.fn(element => {
+        element.tagName;
+      });
+      const rewriter = new HTMLRewriter().on("p", {
+        element,
+      });
+
+      const content = "<p>Lorem ipsum!</p>";
+
+      rewriter.transform(new Response(content));
+      rewriter.transform(new Response(content));
+
+      expect(element).toHaveBeenCalledTimes(2);
+    })();
+
+  Bun.gc(true);
+});