Drain microtasks in more places

oven-sh · Feb 2, 2024 · 7fc0bce · 7fc0bce
1 parent c0c7991
commit 7fc0bce
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 4 deletions.
diff --git a/src/bun.js/api/server.zig b/src/bun.js/api/server.zig
@@ -3750,6 +3750,7 @@ pub const ServerWebSocket = struct {
         if (onOpenHandler.isEmptyOrUndefinedOrNull()) return;
         const this_value = this.getThisValue();
         var args = [_]JSValue{this_value};
+        defer globalObject.bunVM().drainMicrotasks();
 
         var corker = Corker{
             .args = &args,
@@ -3893,7 +3894,7 @@ pub const ServerWebSocket = struct {
                 .globalObject = globalObject,
                 .callback = handler.onDrain,
             };
-
+            defer globalObject.bunVM().drainMicrotasks();
             this.websocket.cork(&corker, Corker.run);
             const result = corker.result;
 
@@ -4004,15 +4005,17 @@ pub const ServerWebSocket = struct {
 
         if (!handler.onClose.isEmptyOrUndefinedOrNull()) {
             var str = ZigString.init(message);
+            const globalObject = handler.globalObject;
+            defer globalObject.bunVM().drainMicrotasks();
             str.markUTF8();
             const result = handler.onClose.call(
-                handler.globalObject,
-                &[_]JSC.JSValue{ this.this_value, JSValue.jsNumber(code), str.toValueGC(handler.globalObject) },
+                globalObject,
+                &[_]JSC.JSValue{ this.this_value, JSValue.jsNumber(code), str.toValueGC(globalObject) },
             );
 
             if (result.toError()) |err| {
                 log("onClose error", .{});
-                handler.globalObject.bunVM().runErrorHandler(err, null);
+                globalObject.bunVM().runErrorHandler(err, null);
             }
         }
 

diff --git a/src/http/websocket_http_client.zig b/src/http/websocket_http_client.zig
@@ -916,6 +916,7 @@ pub fn NewWebSocketClient(comptime ssl: bool) type {
         header_fragment: ?u8 = null,
 
         initial_data_handler: ?*InitialDataHandler = null,
+        event_loop: *JSC.EventLoop = undefined,
 
         pub const name = if (ssl) "WebSocketClientTLS" else "WebSocketClient";
 
@@ -1057,8 +1058,19 @@ pub fn NewWebSocketClient(comptime ssl: bool) type {
                 this.clearData();
                 return;
             };
+            if (comptime Environment.isDebug) {
+                this.event_loop.debug.enter();
+            }
+            defer {
+                if (comptime Environment.isDebug) {
+                    this.event_loop.debug.enter();
+                }
+            }
+
             switch (kind) {
                 .Text => {
+                    defer this.event_loop.drainMicrotasks();
+
                     // this function encodes to UTF-16 if > 127
                     // so we don't need to worry about latin1 non-ascii code points
                     // we avoid trim since we wanna keep the utf8 validation intact
@@ -1079,6 +1091,7 @@ pub fn NewWebSocketClient(comptime ssl: bool) type {
                     }
                 },
                 .Binary, .Ping, .Pong => {
+                    defer this.event_loop.drainMicrotasks();
                     JSC.markBinding(@src());
                     out.didReceiveBytes(data_.ptr, data_.len, @as(u8, @intFromEnum(kind)));
                 },
@@ -1768,6 +1781,7 @@ pub fn NewWebSocketClient(comptime ssl: bool) type {
                 .globalThis = globalThis,
                 .send_buffer = bun.LinearFifo(u8, .Dynamic).init(bun.default_allocator),
                 .receive_buffer = bun.LinearFifo(u8, .Dynamic).init(bun.default_allocator),
+                .event_loop = globalThis.bunVM().eventLoop(),
             });
             if (!Socket.adoptPtr(
                 tcp,