From 3bfeb83e7e2d114f6cfd4a341e1823555d20aed5 Mon Sep 17 00:00:00 2001
From: Jarred Sumner <jarred@jarredsumner.com>
Date: Thu, 25 Jul 2024 18:49:35 -0700
Subject: [PATCH] Fix [Symbol.dispose] on Bun.listen() & Bun.connect() + add
 types (#12739)

Co-authored-by: Dylan Conway <35280289+dylan-conway@users.noreply.github.com>
---
 packages/bun-types/bun.d.ts       |  4 ++--
 src/bun.js/api/bun/socket.zig     | 18 +++++++++++++-----
 src/bun.js/api/sockets.classes.ts |  4 ++--
 3 files changed, 17 insertions(+), 9 deletions(-)

diff --git a/packages/bun-types/bun.d.ts b/packages/bun-types/bun.d.ts
index 6e7309a89d..4a88ba74d2 100644
--- a/packages/bun-types/bun.d.ts
+++ b/packages/bun-types/bun.d.ts
@@ -3847,7 +3847,7 @@ declare module "bun" {
    */
   const isMainThread: boolean;
 
-  interface Socket<Data = undefined> {
+  interface Socket<Data = undefined> extends Disposable {
     /**
      * Write `data` to the socket
      *
@@ -4129,7 +4129,7 @@ declare module "bun" {
     setMaxSendFragment(size: number): boolean;
   }
 
-  interface SocketListener<Data = undefined> {
+  interface SocketListener<Data = undefined> extends Disposable {
     stop(closeActiveConnections?: boolean): void;
     ref(): void;
     unref(): void;
diff --git a/src/bun.js/api/bun/socket.zig b/src/bun.js/api/bun/socket.zig
index de14ab3c3f..d3ba902e8e 100644
--- a/src/bun.js/api/bun/socket.zig
+++ b/src/bun.js/api/bun/socket.zig
@@ -874,11 +874,22 @@ pub const Listener = struct {
         return JSValue.jsUndefined();
     }
 
+    pub fn dispose(this: *Listener, _: *JSC.JSGlobalObject, _: *JSC.CallFrame) JSValue {
+        this.doStop(true);
+        return .undefined;
+    }
+
     pub fn stop(this: *Listener, _: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) JSValue {
         const arguments = callframe.arguments(1);
         log("close", .{});
 
-        var listener = this.listener orelse return JSValue.jsUndefined();
+        this.doStop(if (arguments.len > 0 and arguments.ptr[0].isBoolean()) arguments.ptr[0].toBoolean() else false);
+
+        return .undefined;
+    }
+
+    fn doStop(this: *Listener, force_close: bool) void {
+        var listener = this.listener orelse return;
         this.listener = null;
 
         this.poll_ref.unref(this.handlers.vm);
@@ -891,8 +902,7 @@ pub const Listener = struct {
             this.strong_self.clear();
             this.strong_data.clear();
         } else {
-            const forceClose = arguments.len > 0 and arguments.ptr[0].isBoolean() and arguments.ptr[0].toBoolean() and this.socket_context != null;
-            if (forceClose) {
+            if (force_close) {
                 // close all connections in this context and wait for them to close
                 this.socket_context.?.close(this.ssl);
             } else {
@@ -900,8 +910,6 @@ pub const Listener = struct {
                 listener.close(this.ssl);
             }
         }
-
-        return JSValue.jsUndefined();
     }
 
     pub fn finalize(this: *Listener) callconv(.C) void {
diff --git a/src/bun.js/api/sockets.classes.ts b/src/bun.js/api/sockets.classes.ts
index 192ba84665..d36ca4629c 100644
--- a/src/bun.js/api/sockets.classes.ts
+++ b/src/bun.js/api/sockets.classes.ts
@@ -116,7 +116,7 @@ function generate(ssl) {
       },
 
       "@@dispose": {
-        fn: "shutdown",
+        fn: "end",
         length: 0,
       },
 
@@ -191,7 +191,7 @@ export default [
         length: 1,
       },
       "@@dispose": {
-        fn: "stop",
+        fn: "dispose",
         length: 0,
       },