Implement server.unref() and server.ref() (#8675)

* Implement server.ref() and server.unref() * Update dns_resolver.zig * Update server.zig --------- Co-authored-by: Jarred Sumner <[email protected]>
oven-sh · Feb 4, 2024 · 1009b07 · 1009b07
1 parent 1a695f1
commit 1009b07
Show file tree

Hide file tree

Showing 12 changed files with 167 additions and 23 deletions.
diff --git a/packages/bun-types/bun.d.ts b/packages/bun-types/bun.d.ts
@@ -2473,6 +2473,25 @@ declare module "bun" {
      */
     requestIP(request: Request): SocketAddress | null;
 
+    /**
+     * Undo a call to {@link Server.unref}
+     *
+     * If the Server has already been stopped, this does nothing.
+     *
+     * If {@link Server.ref} is called multiple times, this does nothing. Think of it as a boolean toggle.
+     */
+    ref(): void;
+
+    /**
+     * Don't keep the process alive if this server is the only thing left.
+     * Active connections may continue to keep the process alive.
+     *
+     * By default, the server is ref'd.
+     *
+     * To prevent new connections from being accepted, use {@link Server.stop}
+     */
+    unref(): void;
+
     /**
      * How many requests are in-flight right now?
      */

diff --git a/src/bun.js/api/bun/dns_resolver.zig b/src/bun.js/api/bun/dns_resolver.zig
@@ -956,7 +956,7 @@ pub const CAresNameInfo = struct {
     }
 
     pub fn deinit(this: *@This()) void {
-        this.poll_ref.unrefOnNextTick(this.globalThis.bunVM());
+        this.poll_ref.unref(this.globalThis.bunVM());
         // freed
         bun.default_allocator.free(this.name);
 
@@ -1374,7 +1374,7 @@ pub const CAresReverse = struct {
     }
 
     pub fn deinit(this: *@This()) void {
-        this.poll_ref.unrefOnNextTick(this.globalThis.bunVM());
+        this.poll_ref.unref(this.globalThis.bunVM());
         bun.default_allocator.free(this.name);
 
         if (this.allocated)
@@ -1393,12 +1393,20 @@ pub fn CAresLookup(comptime cares_type: type, comptime type_name: []const u8) ty
         next: ?*@This() = null,
         name: []const u8,
 
-        pub fn init(globalThis: *JSC.JSGlobalObject, allocator: std.mem.Allocator, name: []const u8) !*@This() {
-            const this = try allocator.create(@This());
+        pub usingnamespace bun.New(@This());
+
+        pub fn init(globalThis: *JSC.JSGlobalObject, _: std.mem.Allocator, name: []const u8) !*@This() {
             var poll_ref = Async.KeepAlive.init();
             poll_ref.ref(globalThis.bunVM());
-            this.* = .{ .globalThis = globalThis, .promise = JSC.JSPromise.Strong.init(globalThis), .poll_ref = poll_ref, .allocated = true, .name = name };
-            return this;
+            return @This().new(
+                .{
+                    .globalThis = globalThis,
+                    .promise = JSC.JSPromise.Strong.init(globalThis),
+                    .poll_ref = poll_ref,
+                    .allocated = true,
+                    .name = name,
+                },
+            );
         }
 
         pub fn processResolve(this: *@This(), err_: ?c_ares.Error, _: i32, result: ?*cares_type) void {
@@ -1445,11 +1453,11 @@ pub fn CAresLookup(comptime cares_type: type, comptime type_name: []const u8) ty
         }
 
         pub fn deinit(this: *@This()) void {
-            this.poll_ref.unrefOnNextTick(this.globalThis.bunVM());
+            this.poll_ref.unref(this.globalThis.bunVM());
             bun.default_allocator.free(this.name);
 
             if (this.allocated)
-                this.globalThis.allocator().destroy(this);
+                this.destroy();
         }
     };
 }
@@ -1562,7 +1570,8 @@ pub const DNSLookup = struct {
 
     pub fn deinit(this: *DNSLookup) void {
         log("deinit", .{});
-        this.poll_ref.unrefOnNextTick(this.globalThis.bunVM());
+        this.poll_ref.unref(this.globalThis.bunVM());
+
         if (this.allocated)
             this.globalThis.allocator().destroy(this);
     }

diff --git a/src/bun.js/api/server.classes.ts b/src/bun.js/api/server.classes.ts
@@ -41,6 +41,12 @@ function generate(name) {
       pendingWebSockets: {
         getter: "getPendingWebSockets",
       },
+      ref: {
+        fn: "doRef",
+      },
+      unref: {
+        fn: "doUnref",
+      },
       hostname: {
         getter: "getHostname",
         cache: true,

diff --git a/src/bun.js/api/server.zig b/src/bun.js/api/server.zig
@@ -5755,14 +5755,26 @@ pub fn NewServer(comptime NamespaceType: type, comptime ssl_enabled_: bool, comp
             if (this.poll_ref.isActive()) return;
 
             this.poll_ref.ref(this.vm);
-            this.vm.eventLoop().start_server_on_next_tick = true;
         }
 
         pub fn unref(this: *ThisServer) void {
             if (!this.poll_ref.isActive()) return;
 
-            this.poll_ref.unrefOnNextTick(this.vm);
-            this.vm.eventLoop().start_server_on_next_tick = false;
+            this.poll_ref.unref(this.vm);
+        }
+
+        pub fn doRef(this: *ThisServer, _: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) callconv(.C) JSC.JSValue {
+            const this_value = callframe.this();
+            this.ref();
+
+            return this_value;
+        }
+
+        pub fn doUnref(this: *ThisServer, _: *JSC.JSGlobalObject, callframe: *JSC.CallFrame) callconv(.C) JSC.JSValue {
+            const this_value = callframe.this();
+            this.unref();
+
+            return this_value;
         }
 
         pub fn onBunInfoRequest(this: *ThisServer, req: *uws.Request, resp: *App.Response) void {

diff --git a/src/bun.js/event_loop.zig b/src/bun.js/event_loop.zig
@@ -654,7 +654,6 @@ pub const EventLoop = struct {
     global: *JSGlobalObject = undefined,
     virtual_machine: *JSC.VirtualMachine = undefined,
     waker: ?Waker = null,
-    start_server_on_next_tick: bool = false,
     defer_count: std.atomic.Value(usize) = std.atomic.Value(usize).init(0),
     forever_timer: ?*uws.Timer = null,
     deferred_microtask_map: std.AutoArrayHashMapUnmanaged(?*anyopaque, DeferredRepeatingTask) = .{},

diff --git a/src/http/websocket_http_client.zig b/src/http/websocket_http_client.zig
@@ -258,7 +258,7 @@ pub fn NewHTTPUpgradeClient(comptime ssl: bool) type {
                 &client_protocol_hash,
                 NonUTF8Headers.init(header_names, header_values, header_count),
             ) catch return null;
-            var vm = global.bunVM();
+            const vm = global.bunVM();
 
             var client = HTTPClient.new(.{
                 .tcp = undefined,
@@ -269,8 +269,7 @@ pub fn NewHTTPUpgradeClient(comptime ssl: bool) type {
 
             var host_ = host.toSlice(bun.default_allocator);
             defer host_.deinit();
-            const prev_start_server_on_next_tick = vm.eventLoop().start_server_on_next_tick;
-            vm.eventLoop().start_server_on_next_tick = true;
+
             client.poll_ref.ref(vm);
             const display_host_ = host_.slice();
             const display_host = if (bun.FeatureFlags.hardcode_localhost_to_127_0_0_1 and strings.eqlComptime(display_host_, "localhost"))
@@ -295,8 +294,6 @@ pub fn NewHTTPUpgradeClient(comptime ssl: bool) type {
                 out.tcp.?.timeout(120);
                 return out;
             } else {
-                vm.eventLoop().start_server_on_next_tick = prev_start_server_on_next_tick;
-
                 client.clearData();
                 client.destroy();
             }
@@ -309,7 +306,7 @@ pub fn NewHTTPUpgradeClient(comptime ssl: bool) type {
             this.input_body_buf.len = 0;
         }
         pub fn clearData(this: *HTTPClient) void {
-            this.poll_ref.unrefOnNextTick(JSC.VirtualMachine.get());
+            this.poll_ref.unref(JSC.VirtualMachine.get());
 
             this.clearInput();
             this.body.clearAndFree(bun.default_allocator);
@@ -957,7 +954,7 @@ pub fn NewWebSocketClient(comptime ssl: bool) type {
         }
 
         pub fn clearData(this: *WebSocket) void {
-            this.poll_ref.unrefOnNextTick(this.globalThis.bunVM());
+            this.poll_ref.unref(this.globalThis.bunVM());
             this.clearReceiveBuffers(true);
             this.clearSendBuffers(true);
             this.ping_received = false;
@@ -1706,15 +1703,15 @@ pub fn NewWebSocketClient(comptime ssl: bool) type {
 
         fn dispatchAbruptClose(this: *WebSocket) void {
             var out = this.outgoing_websocket orelse return;
-            this.poll_ref.unrefOnNextTick(this.globalThis.bunVM());
+            this.poll_ref.unref(this.globalThis.bunVM());
             JSC.markBinding(@src());
             this.outgoing_websocket = null;
             out.didAbruptClose(ErrorCode.closed);
         }
 
         fn dispatchClose(this: *WebSocket, code: u16, reason: *const bun.String) void {
             var out = this.outgoing_websocket orelse return;
-            this.poll_ref.unrefOnNextTick(this.globalThis.bunVM());
+            this.poll_ref.unref(this.globalThis.bunVM());
             JSC.markBinding(@src());
             this.outgoing_websocket = null;
             out.didClose(code, reason);

diff --git a/src/js/node/http.ts b/src/js/node/http.ts
@@ -351,9 +351,12 @@ var optionsSymbol = Symbol("options");
 var serverSymbol = Symbol("server");
 function Server(options, callback) {
   if (!(this instanceof Server)) return new Server(options, callback);
+  EventEmitter.$call(this);
 
   this.listening = false;
-  EventEmitter.$call(this);
+  this._unref = false;
+  this[serverSymbol] = undefined;
+
   if (typeof options === "function") {
     callback = options;
     options = {};
@@ -427,6 +430,18 @@ function Server(options, callback) {
 Object.setPrototypeOf((Server.prototype = {}), EventEmitter.prototype);
 Object.setPrototypeOf(Server, EventEmitter);
 
+Server.prototype.ref = function () {
+  this._unref = false;
+  this[serverSymbol]?.ref?.();
+  return this;
+};
+
+Server.prototype.unref = function () {
+  this._unref = true;
+  this[serverSymbol]?.unref?.();
+  return this;
+};
+
 Server.prototype.closeAllConnections = function () {
   const server = this[serverSymbol];
   if (!server) {
@@ -571,6 +586,11 @@ Server.prototype.listen = function (port, host, backlog, onListen) {
       },
     });
     isHTTPS = this[serverSymbol].protocol === "https";
+
+    if (this?._unref) {
+      this[serverSymbol]?.unref?.();
+    }
+
     setTimeout(emitListeningNextTick, 1, this, onListen, null, this[serverSymbol].hostname, this[serverSymbol].port);
   } catch (err) {
     server.emit("error", err);

diff --git a/test/js/bun/http/bun-server.test.ts b/test/js/bun/http/bun-server.test.ts
@@ -423,3 +423,12 @@ describe("Server", () => {
     expect(exitCode).toBe(0);
   });
 });
+
+// By not timing out, this test passes.
+test("Bun.serve().unref() works", async () => {
+  expect([path.join(import.meta.dir, "unref-fixture.ts")]).toRun();
+});
+
+test("unref keeps process alive for ongoing connections", async () => {
+  expect([path.join(import.meta.dir, "unref-fixture-2.ts")]).toRun();
+});
diff --git a/test/js/bun/http/unref-fixture-2.ts b/test/js/bun/http/unref-fixture-2.ts
diff --git a/test/js/bun/http/unref-fixture.ts b/test/js/bun/http/unref-fixture.ts
diff --git a/test/js/node/http/node-http-ref-fixture.js b/test/js/node/http/node-http-ref-fixture.js
diff --git a/test/js/node/http/node-http.test.ts b/test/js/node/http/node-http.test.ts
@@ -1675,3 +1675,11 @@ it("#4415.4 IncomingMessage es5", () => {
   IncomingMessage.call(im, { url: "/foo" });
   expect(im.url).toBe("/foo");
 });
+
+// Windows doesnt support SIGUSR1
+if (process.platform !== "win32") {
+  // By not timing out, this test passes.
+  test(".unref() works", async () => {
+    expect([path.join(import.meta.dir, "node-http-ref-fixture.js")]).toRun();
+  });
+}