fix websocket hang (#1910)

* ref and deref, and some tests

* subprocess file

* remove deref

* use flag in test

3 files changed, 69 insertions, 1 deletions

diff --git a/src/bun.js/bindings/webcore/WebSocket.cpp b/src/bun.js/bindings/webcore/WebSocket.cpp
index 015b706d9..c2acb0b98 100644
--- a/src/bun.js/bindings/webcore/WebSocket.cpp
+++ b/src/bun.js/bindings/webcore/WebSocket.cpp
@@ -159,6 +159,7 @@ WebSocket::WebSocket(ScriptExecutionContext& context)
 {
     m_state = CONNECTING;
     m_hasPendingActivity.store(true);
+    ref();
 }
 
 WebSocket::~WebSocket()
@@ -803,7 +804,9 @@ void WebSocket::didReceiveMessage(String&& message)
 
     if (this->hasEventListeners("message"_s)) {
         // the main reason for dispatching on a separate tick is to handle when you haven't yet attached an event listener
+        this->incPendingActivityCount();
         dispatchEvent(MessageEvent::create(WTFMove(message), m_url.string()));
+        this->decPendingActivityCount();
         return;
     }
 
@@ -839,7 +842,9 @@ void WebSocket::didReceiveBinaryData(Vector<uint8_t>&& binaryData)
     case BinaryType::ArrayBuffer: {
         if (this->hasEventListeners("message"_s)) {
             // the main reason for dispatching on a separate tick is to handle when you haven't yet attached an event listener
+            this->incPendingActivityCount();
             dispatchEvent(MessageEvent::create(ArrayBuffer::create(binaryData.data(), binaryData.size()), m_url.string()));
+            this->decPendingActivityCount();
             return;
         }
 
diff --git a/test/bun.js/websocket-subprocess.ts b/test/bun.js/websocket-subprocess.ts
new file mode 100644
index 000000000..4700f75b7
--- /dev/null
+++ b/test/bun.js/websocket-subprocess.ts
@@ -0,0 +1,12 @@
+const hostname = process.argv[2];
+const port = process.argv[3];
+
+const host = port ? `http://${hostname}:${port}` : hostname;
+
+const ws = new WebSocket(host);
+
+ws.onmessage = (message) => {
+  if (message.data == "hello websocket") {
+    ws.send("hello");
+  }
+};
diff --git a/test/bun.js/websocket.test.js b/test/bun.js/websocket.test.js
index ff5362050..0d4a0ce06 100644
--- a/test/bun.js/websocket.test.js
+++ b/test/bun.js/websocket.test.js
@@ -1,5 +1,7 @@
 import { describe, it, expect } from "bun:test";
-import { unsafe } from "bun";
+import { unsafe, spawn, readableStreamToText } from "bun";
+import { bunExe } from "bunExe";
+
 import { gc } from "./gc";
 
 const TEST_WEBSOCKET_HOST =
@@ -131,3 +133,52 @@ describe("WebSocket", () => {
     gc(true);
   });
 });
+
+describe("websocket in subprocess", () => {
+  it("should exit", async () => {
+    let messageReceived = false;
+    const server = Bun.serve({
+      port: 8765,
+      fetch(req, server) {
+        if (server.upgrade(req)) {
+          return;
+        }
+
+        return new Response("http response");
+      },
+      websocket: {
+        open(ws) {
+          ws.send("hello websocket");
+        },
+        message(ws) {
+          messageReceived = true;
+          ws.close();
+        },
+        close(ws) {},
+      },
+    });
+    const subprocess = Bun.spawn({
+      cmd: [bunExe(), "websocket-subprocess.ts", server.hostname, server.port],
+      stderr: "pipe",
+      stdin: "pipe",
+      stdout: "pipe",
+    });
+
+    expect(await subprocess.exited).toBe(0);
+    expect(messageReceived).toBe(true);
+    server.stop();
+  });
+
+  it("should exit after killed", async () => {
+    const subprocess = Bun.spawn({
+      cmd: [bunExe(), "websocket-subprocess.ts", TEST_WEBSOCKET_HOST],
+      stderr: "pipe",
+      stdin: "pipe",
+      stdout: "pipe",
+    });
+
+    subprocess.kill();
+
+    expect(await subprocess.exited).toBe("SIGHUP");
+  });
+});