Expose subscriberCount in WebSocket server (#13498)

packages/bun-types/bun.d.ts

@@ -2767,6 +2767,16 @@ declare module "bun" {
       compress?: boolean,
     ): ServerWebSocketSendStatus;
 
+    /**
+     * A count of connections subscribed to a given topic
+     *
+     * This operation will loop through each topic internally to get the count.
+     *
+     * @param topic the websocket topic to check how many subscribers are connected to
+     * @returns the number of subscribers
+     */
+    subscriberCount(topic: string): number;
+
     /**
      * Returns the client IP address and port of the given Request. If the request was closed or is a unix socket, returns null.
      *

src/bun.js/api/server.classes.ts

@@ -16,6 +16,10 @@ function generate(name) {
         fn: "doPublish",
         length: 3,
       },
+      subscriberCount: {
+        fn: "doSubscriberCount",
+        length: 1,
+      },
       reload: {
         fn: "doReload",
         length: 2,

src/bun.js/api/server.zig

@@ -4339,7 +4339,6 @@ pub const ServerWebSocket = struct {
         callframe: *JSC.CallFrame,
     ) JSValue {
         const args = callframe.arguments(4);
-
         if (args.len < 1) {
             log("publish()", .{});
             globalThis.throw("publish requires at least 1 argument", .{});
@@ -5354,6 +5353,31 @@ pub fn NewServer(comptime NamespaceType: type, comptime ssl_enabled_: bool, comp
         pub const doFetch = onFetch;
         pub const doRequestIP = JSC.wrapInstanceMethod(ThisServer, "requestIP", false);
 
+        pub fn doSubscriberCount(this: *ThisServer, globalThis: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) JSC.JSValue {
+            const arguments = callframe.arguments(1);
+            if (arguments.len < 1) {
+                globalThis.throwNotEnoughArguments("subscriberCount", 1, 0);
+                return .zero;
+            }
+
+            if (arguments.ptr[0].isEmptyOrUndefinedOrNull()) {
+                globalThis.throwInvalidArguments("subscriberCount requires a topic name as a string", .{});
+                return .zero;
+            }
+
+            var topic = arguments.ptr[0].toSlice(globalThis, bun.default_allocator);
+            defer topic.deinit();
+            if (globalThis.hasException()) {
+                return .zero;
+            }
+
+            if (topic.len == 0) {
+                return JSValue.jsNumber(0);
+            }
+
+            return JSValue.jsNumber((this.app.num_subscribers(topic.slice())));
+        }
+
         pub usingnamespace NamespaceType;
         pub usingnamespace bun.New(@This());
 

test/js/bun/websocket/websocket-server.test.ts

@@ -492,9 +492,11 @@ describe("ServerWebSocket", () => {
           }
         }
       };
-      test(label, (done, connect) => ({
+      test(label, (done, connect, options) => ({
         async open(ws) {
+          const initial = options.server.subscriberCount(topic);
           ws.subscribe(topic);
+          expect(options.server.subscriberCount(topic)).toBe(initial + 1);
           if (ws.data.id === 0) {
             await connect();
           } else if (ws.data.id === 1) {
@@ -525,10 +527,12 @@ describe("ServerWebSocket", () => {
           }
         }
       };
-      test(label, done => ({
+      test(label, (done, _, options) => ({
         publishToSelf: true,
         async open(ws) {
+          const initial = options.server.subscriberCount(topic);
           ws.subscribe(topic);
+          expect(options.server.subscriberCount(topic)).toBe(initial + 1);
           send(ws);
         },
         drain(ws) {
@@ -690,7 +694,11 @@ describe("ServerWebSocket", () => {
 
 function test(
   label: string,
-  fn: (done: (err?: unknown) => void, connect: () => Promise<void>) => Partial<WebSocketHandler<{ id: number }>>,
+  fn: (
+    done: (err?: unknown) => void,
+    connect: () => Promise<void>,
+    options: { server: Server },
+  ) => Partial<WebSocketHandler<{ id: number }>>,
   timeout?: number,
 ) {
   it(
@@ -705,6 +713,9 @@ function test(
         }
       };
       let id = 0;
+      var options = {
+        server: undefined,
+      };
       const server: Server = serve({
         port: 0,
         fetch(request, server) {
@@ -717,9 +728,11 @@ function test(
         websocket: {
           sendPings: false,
           message() {},
-          ...fn(done, () => connect(server)),
+          ...fn(done, () => connect(server), options as any),
         },
       });
+      options.server = server;
+      expect(server.subscriberCount("empty topic")).toBe(0);
       await connect(server);
     },
     { timeout: timeout ?? 1000 },