Fix event loop issue with `Bun.connect` (#4157)

* Fix event loop issue with Bun.connect

* sorry

* add test and fix another edge case

3 files changed, 73 insertions, 3 deletions

diff --git a/src/bun.js/api/bun/socket.zig b/src/bun.js/api/bun/socket.zig
index ce43e3574..05c382359 100644
--- a/src/bun.js/api/bun/socket.zig
+++ b/src/bun.js/api/bun/socket.zig
@@ -985,6 +985,7 @@ pub const Listener = struct {
                 exception.* = ZigString.static("Failed to connect").toErrorInstance(globalObject).asObjectRef();
                 return .zero;
             };
+            tls.poll_ref.ref(handlers.vm);
 
             return promise_value;
         } else {
@@ -1010,6 +1011,7 @@ pub const Listener = struct {
                 exception.* = ZigString.static("Failed to connect").toErrorInstance(globalObject).asObjectRef();
                 return .zero;
             };
+            tcp.poll_ref.ref(handlers.vm);
 
             return promise_value;
         }
@@ -1175,7 +1177,7 @@ fn NewSocket(comptime ssl: bool) type {
             defer this.markInactive();
 
             const handlers = this.handlers;
-            this.poll_ref.unref(handlers.vm);
+            this.poll_ref.unrefOnNextTick(handlers.vm);
 
             const callback = handlers.onConnectError;
             var globalObject = handlers.globalObject;
@@ -1212,7 +1214,6 @@ fn NewSocket(comptime ssl: bool) type {
                 const err_ = err.toErrorInstance(globalObject);
                 promise.rejectOnNextTickAsHandled(globalObject, err_);
                 this.has_pending_activity.store(false, .Release);
-                this.poll_ref.unref(handlers.vm);
             }
         }
 
@@ -1280,7 +1281,6 @@ fn NewSocket(comptime ssl: bool) type {
                 }
             }
 
-            this.poll_ref.ref(this.handlers.vm);
             this.detached = false;
             this.socket = socket;
 
diff --git a/test/js/bun/net/keep-event-loop-alive.js b/test/js/bun/net/keep-event-loop-alive.js
new file mode 100644
index 000000000..28f8941fd
--- /dev/null
+++ b/test/js/bun/net/keep-event-loop-alive.js
@@ -0,0 +1,47 @@
+(async () => {
+  const port = process.argv[2] ? parseInt(process.argv[2]) : null;
+  await Bun.sleep(10);
+  // failed connection
+  console.log("test 1: failed connection");
+  try {
+    const socket = await Bun.connect({
+      hostname: "localhost",
+      port: 9999,
+      socket: { data() {} },
+    });
+    socket.end();
+  } catch (error) {}
+  // failed connection tls
+  console.log("test 2: failed connection [tls]");
+  try {
+    const socket = await Bun.connect({
+      hostname: "localhost",
+      port: 9999,
+      socket: { data() {} },
+      tls: true,
+    });
+    socket.end();
+  } catch (error) {}
+  if (port) {
+    // successful connection
+    console.log("test 3: successful connection");
+    const socket = await Bun.connect({
+      hostname: "localhost",
+      port,
+      socket: { data() {} },
+    });
+    socket.end();
+
+    // successful connection tls
+    console.log("test 4: successful connection [tls]");
+    const socket2 = await Bun.connect({
+      hostname: "localhost",
+      port,
+      socket: { data() {} },
+    });
+    socket2.end();
+  } else {
+    console.log("run with a port as an argument to try the success situation");
+  }
+  console.log("success: event loop was not killed");
+})();
diff --git a/test/js/bun/net/socket.test.ts b/test/js/bun/net/socket.test.ts
index 5126067e6..cf6625dc5 100644
--- a/test/js/bun/net/socket.test.ts
+++ b/test/js/bun/net/socket.test.ts
@@ -27,6 +27,29 @@ it("should keep process alive only when active", async () => {
   ).toEqual(["[Client] OPENED", "[Client] GOT response", "[Client] CLOSED"]);
 });
 
+it("connect without top level await should keep process alive", async () => {
+  const server = Bun.listen({
+    socket: {
+      open(socket) {},
+      data(socket, data) {},
+    },
+    hostname: "localhost",
+    port: 0,
+  });
+  const proc = Bun.spawn({
+    cmd: [bunExe(), "keep-event-loop-alive.js", String(server.port)],
+    cwd: import.meta.dir,
+    env: bunEnv,
+  });
+  await proc.exited;
+  try {
+    expect(proc.exitCode).toBe(0);
+    expect(await new Response(proc.stdout).text()).toContain("event loop was not killed");
+  } finally {
+    server.stop();
+  }
+});
+
 it("connect() should return the socket object", async () => {
   const { exited, stdout, stderr } = spawn({
     cmd: [bunExe(), "connect-returns-socket.js"],