serve: keep the JS wrapper alive until no dispatch can fire

src/runtime/server/ServerWebSocket.rs

@@ -388,8 +388,12 @@ impl ServerWebSocket {
 
         let handler = self.handler();
         let vm = handler.vm();
-        // The handler is shared (&), so mutate via the interior-mutability helper.
-        handler.active_connections_saturating_add(1);
+        // Live-socket accounting lives on the server (`Cell`), reached
+        // through the type-erased backref so the shared `&Handler` suffices.
+        let server = handler.server;
+        if let Some(server) = server {
+            server.on_websocket_opened();
+        }
         let global_object = handler.global_object();
         let on_open_handler = handler.on_open;
         if vm.is_shutting_down() {
@@ -426,17 +430,26 @@ impl ServerWebSocket {
         if let Some(err_value) = result.to_error() {
             bun_output::scoped_log!(WebSocketServer, "onOpen exception");
 
+            let mut closed_here = false;
             if !self.flags.get().closed() {
                 self.update_flags(|f| f.set_closed(true));
                 // we un-gracefully close the connection if there was an exception
                 // we don't want any event handlers to fire after this for anything other than error()
                 // https://github.com/oven-sh/bun/issues/1480
+                // (`close()` re-enters `on_close`, which skips its own
+                // accounting because the closed flag is already set.)
                 self.websocket().close();
-                handler.active_connections_saturating_sub(1);
+                closed_here = true;
                 this_value.unprotect();
             }
 
             handler.run_error_callback(vm, global_object, err_value);
+            if closed_here {
+                if let Some(server) = server {
+                    // May run the idle pass; no `&Handler` borrow is live here.
+                    server.on_websocket_closed();
+                }
+            }
         }
     }
 
@@ -633,11 +646,18 @@ impl ServerWebSocket {
         bun_output::scoped_log!(WebSocketServer, "onClose");
         // TODO: Can this called inside finalize?
         let handler = self.handler();
+        // Copy the erased server handle out now: the guard below runs after
+        // every `handler` borrow has expired, and `on_websocket_closed` may
+        // form `&mut NewServer` (which owns the handler storage) to run the
+        // idle pass when this was the last live socket.
+        let server = handler.server;
         let was_closed = self.is_closed();
         self.update_flags(|f| f.set_closed(true));
         scopeguard::defer! {
             if !was_closed {
-                handler.active_connections_saturating_sub(1);
+                if let Some(server) = server {
+                    server.on_websocket_closed();
+                }
             }
         }
         let signal = self.signal.take();
@@ -673,7 +693,10 @@ impl ServerWebSocket {
             return;
         }
 
-        if !handler.on_close.is_empty_or_undefined_or_null() {
+        // Copy to a stack local before `sig.signal()` re-enters JS: a GC
+        // between the test and the `.call(...)` could otherwise collect it.
+        let on_close_handler = handler.on_close;
+        if !on_close_handler.is_empty_or_undefined_or_null() {
             let global_object = handler.global_object();
 
             let _loop_guard = vm.enter_event_loop_scope();
@@ -702,10 +725,7 @@ impl ServerWebSocket {
             };
 
             let call_args = [cached_this, JSValue::js_number(code as f64), message_js];
-            if let Err(e) = handler
-                .on_close
-                .call(global_object, JSValue::UNDEFINED, &call_args)
-            {
+            if let Err(e) = on_close_handler.call(global_object, JSValue::UNDEFINED, &call_args) {
                 let err = global_object.take_exception(e);
                 bun_output::scoped_log!(WebSocketServer, "onClose error {}", was_not_empty);
                 handler.run_error_callback(vm, global_object, err);

src/runtime/server/WebSocketServerContext.rs

@@ -34,10 +34,11 @@ pub struct Handler {
     // LIFETIMES.tsv = STATIC (vm) / JSC_BORROW (global_object) — both outlive the handler.
     pub vm: bun_ptr::BackRef<VirtualMachine>,
     pub global_object: bun_ptr::BackRef<JSGlobalObject>,
-    /// Mutated through `&Handler` (the field is owned by
-    /// `ServerConfig.websocket` and only ever touched on the JS thread), so
-    /// it's a `Cell`.
-    pub active_connections: core::cell::Cell<usize>,
+    /// Type-erased backref to the owning `NewServer`, set alongside `app`
+    /// in `set_routes` (so it is in place before any socket can upgrade and
+    /// refreshed whenever a reload installs a new context). `ServerWebSocket`
+    /// open/close events route the live-socket accounting through it.
+    pub server: Option<super::AnyServer>,
 
     /// used by publish()
     pub flags: HandlerFlags,
@@ -68,19 +69,6 @@ impl Handler {
         self.vm.get()
     }
 
-    #[inline]
-    pub fn active_connections_saturating_add(&self, n: usize) {
-        self.active_connections
-            .set(self.active_connections.get().saturating_add(n));
-    }
-
-    /// See `active_connections_saturating_add`.
-    #[inline]
-    pub fn active_connections_saturating_sub(&self, n: usize) {
-        self.active_connections
-            .set(self.active_connections.get().saturating_sub(n));
-    }
-
     pub fn run_error_callback(
         &self,
         vm: &VirtualMachine,
@@ -119,7 +107,7 @@ impl Handler {
             app: None,
             vm: bun_ptr::BackRef::new(VirtualMachine::get()),
             global_object: bun_ptr::BackRef::new(global_object),
-            active_connections: core::cell::Cell::new(0),
+            server: None,
             flags: HandlerFlags::empty(),
         };
 

src/runtime/server/mod.rs

@@ -265,6 +265,11 @@ pub struct NewServer<const SSL: bool, const DEBUG: bool> {
     pub base_url_string_for_joining: Box<[u8]>,
     pub config: ServerConfig,
     pub pending_requests: usize,
+    /// Live `ServerWebSocket` count. Lives on the server (not the websocket
+    /// context) so a reload's context swap cannot reset it, and sits in a
+    /// `Cell` because the open/close accounting arrives through shared
+    /// `AnyServer` handles on the JS thread.
+    pub active_websocket_count: core::cell::Cell<u32>,
     pub request_pool: *mut request_context::RequestContextStackAllocator<Self, SSL, DEBUG, false>,
     /// Null until the H3 listen path runs (`HAS_H3 && config.http3`); never
     /// allocated when `!SSL`. Kept as a raw nullable pointer rather than a
@@ -565,6 +570,17 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
         self.js_value.try_get().expect("js_value alive")
     }
 
+    /// Returns the wrapper while it is strongly rooted, or None once the
+    /// server has gone idle and downgraded. Dispatch trampolines should
+    /// answer 503+close on None. Checking strong (not just Finalized) closes
+    /// the dead-but-unswept window where `JsRef::Weak` holds a stale address.
+    pub fn js_value_for_dispatch(&self) -> Option<JSValue> {
+        match &self.js_value {
+            jsc::JsRef::Strong(_) => self.js_value.try_get(),
+            _ => None,
+        }
+    }
+
     /// Per-monomorphization static.
     /// Rust statics cannot be const-generic; routed through a
     /// `&'static AtomicBool` so the four (SSL,DEBUG) instantiations share one
@@ -1026,6 +1042,16 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
         req: &mut uws_sys::Request,
         resp: *mut uws_sys::NewAppResponse<SSL>,
     ) {
+        // Idle keep-alive sockets aren't counted in pending_requests, so the
+        // wrapper can have been collected before this fires. Refuse and close
+        // rather than panicking in js_value_assert_alive() below.
+        // SAFETY: `this` is the live server backref for this request.
+        if unsafe { &*this }.js_value_for_dispatch().is_none() {
+            let resp_ref = bun_opaque::opaque_deref_mut(resp);
+            resp_ref.write_status(b"503 Service Unavailable");
+            resp_ref.end_without_body(true);
+            return;
+        }
         let should_deinit_context = core::cell::Cell::new(false);
         let Some(prepared) = Self::prepare_js_request_context(
             this,
@@ -1077,6 +1103,14 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
         let server = user_route.server.cast_mut();
         let index = user_route.id;
 
+        // SAFETY: `server` is the live backref stored in `user_route`.
+        if unsafe { &*server }.js_value_for_dispatch().is_none() {
+            let resp_ref = bun_opaque::opaque_deref_mut(resp);
+            resp_ref.write_status(b"503 Service Unavailable");
+            resp_ref.end_without_body(true);
+            return;
+        }
+
         let should_deinit_context = core::cell::Cell::new(false);
         let Some(mut prepared) = Self::prepare_js_request_context(
             server,
@@ -1462,10 +1496,19 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
     }
 
     pub fn active_sockets_count(&self) -> u32 {
-        self.config
-            .websocket
-            .as_ref()
-            .map_or(0, |ws| ws.handler.active_connections.get() as u32)
+        self.active_websocket_count.get()
+    }
+
+    pub(crate) fn note_websocket_opened(&self) {
+        self.active_websocket_count
+            .set(self.active_websocket_count.get().saturating_add(1));
+    }
+
+    /// Returns true when this close drained the last live websocket.
+    pub(crate) fn note_websocket_closed(&self) -> bool {
+        let remaining = self.active_websocket_count.get().saturating_sub(1);
+        self.active_websocket_count.set(remaining);
+        remaining == 0
     }
 
     pub fn has_active_web_sockets(&self) -> bool {
@@ -1572,6 +1615,15 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
         if !abrupt {
             // S012: `app::ListenSocket<SSL>` is a ZST opaque — safe deref.
             bun_opaque::opaque_deref_mut(listener).close();
+            // Shut idle keep-alive sockets so a late request can't dispatch
+            // after the JS wrapper drops its strong root. `JsRef::Weak` holds
+            // a raw JSValue (not a JSC::Weak), so `try_get` can hand back a
+            // dead-but-unswept cell and the Finalized→503 guard alone leaves a
+            // window. In-flight requests aren't idle and drain normally.
+            if let Some(app) = self.app {
+                // S012: `NewApp<SSL>` is a ZST opaque — safe `*mut → &mut` deref.
+                bun_opaque::opaque_deref_mut(app).close_idle_connections();
+            }
         } else if !self.flags.contains(ServerFlags::TERMINATED) {
             if let Some(ws) = self.config.websocket.as_mut() {
                 ws.handler.app = None;
@@ -1583,9 +1635,6 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
     }
 
     pub fn stop(&mut self, abrupt: bool) {
-        if self.js_value.is_not_empty() {
-            self.js_value.downgrade();
-        }
         if self.config.allow_hot && !self.config.id.is_empty() {
             // `hot_map()` is reached via the thread-local VM singleton (raw ptr
             // deref) and does not borrow `self`, so it cannot overlap with the
@@ -1664,6 +1713,9 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
             );
         }
         if self.pending_requests == 0 && !self.has_listener() && !self.has_active_web_sockets() {
+            // Wrapper-rooted handlers need the wrapper to outlive every
+            // dispatch; downgrade only once nothing can call them.
+            self.js_value.downgrade();
             if let Some(ws) = self.config.websocket.as_mut() {
                 ws.handler.app = None;
             }
@@ -1905,6 +1957,7 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
             h3_alt_svc: Box::<[u8]>::default(),
             js_value: jsc::JsRef::empty(),
             pending_requests: 0,
+            active_websocket_count: core::cell::Cell::new(0),
             request_pool: <Self as ServerPools<SSL, DEBUG>>::request_pool(),
             // Plain HTTP servers never allocate the ~816 KB H3 pool; defer to
             // the H3-listen path (`listen()` below) so HTTPS servers that
@@ -2047,6 +2100,7 @@ impl<const SSL: bool, const DEBUG: bool> NewServer<SSL, DEBUG> {
             websocket.global_object =
                 bun_ptr::BackRef::new(bun_opaque::opaque_deref(self.global_this));
             websocket.handler.app = Some(std::ptr::from_mut(app).cast::<c_void>());
+            websocket.handler.server = Some(any_server);
             websocket
                 .handler
                 .flags
@@ -3453,6 +3507,29 @@ impl AnyServer {
         any_server_dispatch!(self, |s| s.inspector_server_id)
     }
 
+    pub(crate) fn on_websocket_opened(&self) {
+        any_server_dispatch!(self, |s| s.note_websocket_opened());
+    }
+
+    /// Decrement the live-socket count and, when the last socket drained on
+    /// an already-stopped server, run the idle pass so the `JsRef` downgrade
+    /// (and deferred deinit) that was held back by the open sockets fires.
+    ///
+    /// Skipped while still listening (the idle pass would no-op) and while
+    /// TERMINATED: the abrupt-stop path drains every socket synchronously
+    /// from inside `stop_listening` (which holds `&mut self`), and `stop()`
+    /// runs `deinit_if_we_can` itself right after.
+    pub(crate) fn on_websocket_closed(&self) {
+        let drained = any_server_dispatch!(self, |s| {
+            s.note_websocket_closed()
+                && !s.has_listener()
+                && !s.flags.contains(ServerFlags::TERMINATED)
+        });
+        if drained {
+            any_server_dispatch_mut!(self, |s| s.deinit_if_we_can());
+        }
+    }
+
     pub fn set_inspector_server_id(&mut self, id: jsc::DebuggerId) {
         any_server_dispatch_mut!(self, |s| {
             s.inspector_server_id = id;

src/runtime/server/server_body.rs

@@ -2806,6 +2806,12 @@ where
         let server = unsafe { &mut *server_ptr };
         let index = user_route.id;
 
+        if server.js_value_for_dispatch().is_none() {
+            RespLike::write_status(resp, b"503 Service Unavailable");
+            RespLike::end_without_body(resp, true);
+            return;
+        }
+
         let should_deinit_context = core::cell::Cell::new(false);
         let Some(mut prepared) = server.prepare_js_request_context_for::<Ctx>(
             req,
@@ -2890,6 +2896,11 @@ where
         req: &mut Ctx::Req,
         resp: &mut Ctx::Resp,
     ) {
+        if self.js_value_for_dispatch().is_none() {
+            RespLike::write_status(resp, b"503 Service Unavailable");
+            RespLike::end_without_body(resp, true);
+            return;
+        }
         let self_ptr: *mut Self = self;
         let should_deinit_context = core::cell::Cell::new(false);
         let Some(prepared) = self.prepare_js_request_context_for::<Ctx>(
@@ -3189,6 +3200,12 @@ where
         let server_ptr = server_ref.as_ptr();
         let index = this.id;
 
+        if server_ref.js_value_for_dispatch().is_none() {
+            resp.write_status(b"503 Service Unavailable");
+            resp.end_without_body(true);
+            return;
+        }
+
         let should_deinit_context = core::cell::Cell::new(false);
         // SAFETY: `server_ptr` is the live heap server registered for this route;
         // `req`/`resp` are the live uWS handles passed to the route handler.
@@ -3265,6 +3282,11 @@ where
         // (mod.rs `app.ws("/*", self_ptr, 0, ..)`); live for the request's
         // duration.
         let this = unsafe { &mut *self_ptr };
+        if this.js_value_for_dispatch().is_none() {
+            resp.write_status(b"503 Service Unavailable");
+            resp.end_without_body(true);
+            return;
+        }
         if this.config.on_node_http_request.is_some() {
             // NOTE: receiver is `*mut Self` (mod.rs) — the callee re-enters
             // JS, so a long-lived `&mut self` here would alias on callback.