Fix promise memory leak

5 files changed, 87 insertions, 23 deletions

diff --git a/src/bun.js/event_loop.zig b/src/bun.js/event_loop.zig
index d6b7f75da..95bf5456c 100644
--- a/src/bun.js/event_loop.zig
+++ b/src/bun.js/event_loop.zig
@@ -42,7 +42,7 @@ pub fn ConcurrentPromiseTask(comptime Context: type) type {
                 .promise = JSValue.createInternalPromise(globalThis),
                 .globalThis = globalThis,
             };
-            js.JSValueProtect(globalThis.ref(), this.promise.asObjectRef());
+            this.promise.protect();
             VirtualMachine.vm.active_tasks +|= 1;
             return this;
         }
@@ -55,6 +55,8 @@ pub fn ConcurrentPromiseTask(comptime Context: type) type {
 
         pub fn runFromJS(this: This) void {
             var promise_value = this.promise;
+            promise_value.ensureStillAlive();
+            promise_value.unprotect();
             var promise = promise_value.asInternalPromise() orelse {
                 if (comptime @hasDecl(Context, "deinit")) {
                     @call(.{}, Context.deinit, .{this.ctx});
@@ -64,7 +66,6 @@ pub fn ConcurrentPromiseTask(comptime Context: type) type {
 
             var ctx = this.ctx;
 
-            js.JSValueUnprotect(this.globalThis.ref(), promise_value.asObjectRef());
             ctx.then(promise);
         }
 
diff --git a/src/bun.js/javascript_core_c_api.zig b/src/bun.js/javascript_core_c_api.zig
index c2bcc45cf..aa9af2b7e 100644
--- a/src/bun.js/javascript_core_c_api.zig
+++ b/src/bun.js/javascript_core_c_api.zig
@@ -101,7 +101,15 @@ pub extern fn JSValueToBoolean(ctx: JSContextRef, value: JSValueRef) bool;
 pub extern fn JSValueToNumber(ctx: JSContextRef, value: JSValueRef, exception: ExceptionRef) f64;
 pub extern fn JSValueToStringCopy(ctx: JSContextRef, value: JSValueRef, exception: ExceptionRef) JSStringRef;
 pub extern fn JSValueToObject(ctx: JSContextRef, value: JSValueRef, exception: ExceptionRef) JSObjectRef;
-pub extern fn JSValueProtect(ctx: JSContextRef, value: JSValueRef) void;
+
+pub inline fn JSValueProtect(ctx: JSContextRef, value: JSValueRef) void {
+    const Wrapped = struct {
+        pub extern fn JSValueProtect(ctx: JSContextRef, value: JSValueRef) void;
+    };
+    // wrapper exists to make it easier to set a breakpoint
+    Wrapped.JSValueProtect(ctx, value);
+}
+
 pub extern fn JSValueUnprotect(ctx: JSContextRef, value: JSValueRef) void;
 pub const JSPropertyAttributes = enum(c_uint) {
     kJSPropertyAttributeNone = 0,
diff --git a/src/bun.js/node/types.zig b/src/bun.js/node/types.zig
index 52a35699b..5fea3af8d 100644
--- a/src/bun.js/node/types.zig
+++ b/src/bun.js/node/types.zig
@@ -415,7 +415,8 @@ pub const PathLike = union(Tag) {
 
                 if (!Valid.pathString(zig_str, ctx, exception)) return null;
 
-                arguments.protectEat();
+                arguments.eat();
+                arg.ensureStillAlive();
 
                 if (zig_str.is16Bit()) {
                     var printed = std.mem.span(std.fmt.allocPrintZ(arguments.arena.allocator(), "{}", .{zig_str}) catch unreachable);
diff --git a/src/bun.js/webcore/response.zig b/src/bun.js/webcore/response.zig
index 4d2f92e59..d8bb01008 100644
--- a/src/bun.js/webcore/response.zig
+++ b/src/bun.js/webcore/response.zig
@@ -1283,6 +1283,7 @@ pub const Blob = struct {
         exception: js.ExceptionRef,
     ) js.JSObjectRef {
         var args = JSC.Node.ArgumentsSlice.from(ctx.bunVM(), arguments);
+        defer args.deinit();
         // accept a path or a blob
         var path_or_blob = PathOrBlob.fromJS(ctx, &args, exception) orelse {
             exception.* = JSC.toInvalidArguments("Bun.write expects a path, file descriptor or a blob", .{}, ctx).asObjectRef();
@@ -3177,11 +3178,15 @@ pub const Blob = struct {
             var promise = handler.promise;
             var globalThis = handler.globalThis;
             bun.default_allocator.destroy(handler);
+            const value = promise.asValue(globalThis);
+            value.ensureStillAlive();
             switch (count) {
                 .err => |err| {
+                    value.unprotect();
                     promise.reject(globalThis, err.toErrorInstance(globalThis));
                 },
                 .result => |wrote| {
+                    value.unprotect();
                     promise.resolve(globalThis, JSC.JSValue.jsNumberFromUint64(wrote));
                 },
             }
diff --git a/src/bun.js/webcore/streams.zig b/src/bun.js/webcore/streams.zig
index 603694f60..1bbd6cdba 100644
--- a/src/bun.js/webcore/streams.zig
+++ b/src/bun.js/webcore/streams.zig
@@ -57,16 +57,19 @@ pub const ReadableStream = struct {
 
     pub fn cancel(this: *const ReadableStream, globalThis: *JSGlobalObject) void {
         JSC.markBinding();
+        this.value.unprotect();
         ReadableStream__cancel(this.value, globalThis);
     }
 
     pub fn abort(this: *const ReadableStream, globalThis: *JSGlobalObject) void {
         JSC.markBinding();
+        this.value.unprotect();
         ReadableStream__abort(this.value, globalThis);
     }
 
     pub fn detach(this: *const ReadableStream, globalThis: *JSGlobalObject) void {
         JSC.markBinding();
+        this.value.unprotect();
         ReadableStream__detach(this.value, globalThis);
     }
 
@@ -1115,7 +1118,9 @@ pub fn NewJSSink(comptime SinkType: type, comptime name_: []const u8) type {
             if (this.sink.signal.isDead())
                 return;
             this.sink.signal.clear();
-            detachPtr(@intToEnum(JSValue, @bitCast(JSC.JSValueReprInt, @ptrToInt(ptr))));
+            const value = @intToEnum(JSValue, @bitCast(JSC.JSValueReprInt, @ptrToInt(ptr)));
+            value.unprotect();
+            detachPtr(value);
         }
 
         pub fn detachPtr(ptr: JSValue) callconv(.C) void {
@@ -1251,6 +1256,12 @@ pub fn NewJSSink(comptime SinkType: type, comptime name_: []const u8) type {
                 }
             }
 
+            defer {
+                if (comptime @hasField(SinkType, "done") and this.sink.done) {
+                    callframe.this().unprotect();
+                }
+            }
+
             if (comptime @hasDecl(SinkType, "drainFromJS")) {
                 return this.sink.drainFromJS(globalThis).result;
             }
@@ -1303,6 +1314,12 @@ pub fn NewJSSink(comptime SinkType: type, comptime name_: []const u8) type {
                 }
             }
 
+            defer {
+                if (comptime @hasField(SinkType, "done") and this.sink.done) {
+                    callframe.this().unprotect();
+                }
+            }
+
             return this.sink.endFromJS(globalThis).toJS(globalThis);
         }
 
@@ -1368,17 +1385,25 @@ pub fn NewJSSink(comptime SinkType: type, comptime name_: []const u8) type {
 //                 globalThis.vm().throwError(globalThis, err);
 //                 return JSC.JSValue.jsUndefined();
 //             }));
-
-//             this.socket.connect()
 //         }
 //     };
 // }
 
+// TODO: make this JSGlobalObject local
+// for better security
+const ByteListPool = ObjectPool(
+    bun.ByteList,
+    null,
+    true,
+    8,
+);
+
 pub fn HTTPServerWritable(comptime ssl: bool) type {
     return struct {
         const UWSResponse = uws.NewApp(ssl).Response;
         res: *UWSResponse,
         buffer: bun.ByteList,
+        pooled_buffer: ?*ByteListPool.Node = null,
         offset: Blob.SizeType = 0,
 
         is_listening_for_abort: bool = false,
@@ -1439,6 +1464,7 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
 
             if (this.done) {
                 this.res.endStream(false);
+                this.finalize();
                 return false;
             }
 
@@ -1463,13 +1489,7 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
             }
 
             // flush the javascript promise from calling .drain()
-            if (this.pending_drain) |prom| {
-                this.pending_drain = null;
-                log("flush promise ({d})", .{readable.len});
-
-                prom.asValue(this.globalThis).unprotect();
-                prom.resolve(this.globalThis, JSC.jsNumber(readable.len));
-            }
+            this.flushPromise();
 
             if (this.has_callback) {
                 this.has_callback = false;
@@ -1502,6 +1522,14 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
                 return .{ .result = {} };
             }
 
+            if (this.buffer.cap == 0) {
+                std.debug.assert(this.pooled_buffer == null);
+                if (ByteListPool.has()) {
+                    this.pooled_buffer = ByteListPool.get(this.allocator);
+                    this.buffer = this.pooled_buffer.?.data;
+                }
+            }
+
             this.buffer.len = 0;
 
             switch (stream_start) {
@@ -1731,6 +1759,7 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
             if (this.done or this.res.hasResponded()) {
                 this.signal.close(err);
                 this.done = true;
+                this.finalize();
                 return .{ .result = {} };
             }
 
@@ -1739,8 +1768,10 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
             this.end_len = readable.len;
 
             if (readable.len == 0) {
+                this.signal.close(err);
                 this.done = true;
                 this.res.endStream(false);
+                this.finalize();
                 return .{ .result = {} };
             }
 
@@ -1766,6 +1797,7 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
             if (this.done or this.res.hasResponded()) {
                 this.signal.close(null);
                 this.done = true;
+                this.finalize();
                 return .{ .result = JSC.JSValue.jsNumber(0) };
             }
 
@@ -1777,8 +1809,9 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
                 this.done = true;
                 this.res.endStream(false);
                 this.signal.close(null);
-                this.done = true;
-                return .{ .result = JSC.JSValue.jsNumber(this.wrote) };
+                const wrote = this.wrote;
+                this.finalize();
+                return .{ .result = JSC.JSValue.jsNumber(wrote) };
             }
 
             if (!this.hasBackpressure()) {
@@ -1786,7 +1819,9 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
                     this.handleWrote(readable.len);
                     this.signal.close(null);
                     this.done = true;
-                    return .{ .result = JSC.JSValue.jsNumber(this.wrote) };
+                    const wrote = this.wrote;
+                    this.finalize();
+                    return .{ .result = JSC.JSValue.jsNumber(wrote) };
                 }
 
                 this.res.onWritable(*@This(), onWritable, this);
@@ -1814,6 +1849,7 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
             this.done = true;
             this.aborted = true;
             this.flushPromise();
+            this.finalize();
         }
 
         pub fn finalize(this: *@This()) void {
@@ -1824,10 +1860,22 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
                 this.res.endStream(false);
             }
 
-            var bytes = this.buffer.listManaged(this.allocator);
-            bytes.deinit();
-            this.buffer.len = 0;
-            this.flushPromise();
+            if (this.pooled_buffer) |pooled| {
+                this.buffer.len = 0;
+                pooled.data = this.buffer;
+                this.buffer = bun.ByteList.init("");
+                this.pooled_buffer = null;
+                pooled.release();
+            } else if (!ByteListPool.full()) {
+                var entry = ByteListPool.get(this.allocator);
+                entry.data = this.buffer;
+                this.buffer = bun.ByteList.init("");
+                entry.release();
+            } else {
+                var bytes = this.buffer.listManaged(this.allocator);
+                bytes.deinit();
+                this.buffer.len = 0;
+            }
         }
 
         pub fn flushPromise(this: *@This()) void {
@@ -1835,8 +1883,9 @@ pub fn HTTPServerWritable(comptime ssl: bool) type {
                 log("flushPromise()", .{});
 
                 this.pending_drain = null;
-                prom.asValue(this.globalThis).unprotect();
-                prom.resolve(this.globalThis, JSC.JSValue.jsNumber(0));
+                const globalThis = this.globalThis;
+                prom.asValue(globalThis).unprotect();
+                prom.resolve(globalThis, JSC.JSValue.jsNumber(0));
             }
         }