From 546fcd8b4c0cf4f8b7c87ec6873d40ae5c5da192 Mon Sep 17 00:00:00 2001 From: Sayan- <1415138+Sayan-@users.noreply.github.com> Date: Wed, 24 Jun 2026 20:13:22 +0000 Subject: [PATCH 1/3] Close hijacked WebSockets with Going Away on shutdown MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit http.Server.Shutdown does not touch hijacked connections, so the CDP proxy, WebDriver/BiDi, ChromeDriver, and PTY attach WebSockets were cut abruptly when the server stops — clients saw a 1006 abnormal closure. Add a small wsdrain.Registry that tracks live client connections; the shutdown path closes them all with a 1001 Going Away before draining the HTTP servers. The first close wins, so the proxies' own normal-closure teardown no longer overrides the code clients observe. Co-Authored-By: Claude Opus 4.7 --- server/cmd/api/api/process.go | 6 ++- server/cmd/api/main.go | 20 ++++++- server/lib/chromedriverproxy/proxy.go | 6 +++ server/lib/devtoolsproxy/proxy.go | 6 ++- server/lib/devtoolsproxy/proxy_test.go | 69 ++++++++++++++++++++++-- server/lib/wsdrain/wsdrain.go | 74 ++++++++++++++++++++++++++ server/lib/wsdrain/wsdrain_test.go | 56 +++++++++++++++++++ server/lib/wsproxy/wsproxy.go | 7 +++ 8 files changed, 235 insertions(+), 9 deletions(-) create mode 100644 server/lib/wsdrain/wsdrain.go create mode 100644 server/lib/wsdrain/wsdrain_test.go diff --git a/server/cmd/api/api/process.go b/server/cmd/api/api/process.go index 12eb448a..38a458f4 100644 --- a/server/cmd/api/api/process.go +++ b/server/cmd/api/api/process.go @@ -26,6 +26,7 @@ import ( "github.com/kernel/kernel-images/server/lib/logger" oapi "github.com/kernel/kernel-images/server/lib/oapi" "github.com/kernel/kernel-images/server/lib/ptyio" + "github.com/kernel/kernel-images/server/lib/wsdrain" openapi_types "github.com/oapi-codegen/runtime/types" ) @@ -659,7 +660,7 @@ func writeJSON(w http.ResponseWriter, status int, body string) { // - Server sends TextMessage with JSON for events (e.g., exit code) // // This endpoint is intentionally not defined in OpenAPI. -func (s *ApiService) HandleProcessAttachWS(w http.ResponseWriter, r *http.Request, id string) { +func (s *ApiService) HandleProcessAttachWS(w http.ResponseWriter, r *http.Request, id string, reg *wsdrain.Registry) { ctx := r.Context() log := logger.FromContext(ctx) @@ -702,6 +703,9 @@ func (s *ApiService) HandleProcessAttachWS(w http.ResponseWriter, r *http.Reques } defer wsConn.CloseNow() + untrack := reg.Track(wsConn) + defer untrack() + // Set a generous read limit for PTY data wsConn.SetReadLimit(1024 * 1024) // 1MB diff --git a/server/cmd/api/main.go b/server/cmd/api/main.go index b3500a52..0b40fff0 100644 --- a/server/cmd/api/main.go +++ b/server/cmd/api/main.go @@ -14,6 +14,7 @@ import ( "syscall" "time" + "github.com/coder/websocket" "github.com/ghodss/yaml" "github.com/go-chi/chi/v5" chiMiddleware "github.com/go-chi/chi/v5/middleware" @@ -32,6 +33,7 @@ import ( "github.com/kernel/kernel-images/server/lib/scaletozero" "github.com/kernel/kernel-images/server/lib/sysmon" "github.com/kernel/kernel-images/server/lib/telemetry" + "github.com/kernel/kernel-images/server/lib/wsdrain" ) func main() { @@ -45,6 +47,11 @@ func main() { } slogger.Info("server configuration", "config", config) + // Tracks live hijacked WebSocket connections (CDP, WebDriver/BiDi, + // ChromeDriver, PTY attach) so they can be closed with a clean Going Away + // frame on shutdown — http.Server.Shutdown does not touch hijacked conns. + wsRegistry := wsdrain.New() + // context cancellation on SIGINT/SIGTERM ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() @@ -166,7 +173,7 @@ func main() { // Uses WebSocket for bidirectional streaming, which works well through proxies. r.Get("/process/{process_id}/attach", func(w http.ResponseWriter, r *http.Request) { id := chi.URLParam(r, "process_id") - apiService.HandleProcessAttachWS(w, r, id) + apiService.HandleProcessAttachWS(w, r, id, wsRegistry) }) // Serve extension files for Chrome policy-installed extensions @@ -214,7 +221,7 @@ func main() { rDevtools.Get("/json/list", jsonTargetHandler) rDevtools.Get("/json/list/", jsonTargetHandler) rDevtools.Get("/*", func(w http.ResponseWriter, r *http.Request) { - devtoolsproxy.WebSocketProxyHandler(upstreamMgr, slogger, config.LogCDPMessages, stz, telemetrySession.Publish).ServeHTTP(w, r) + devtoolsproxy.WebSocketProxyHandler(upstreamMgr, slogger, config.LogCDPMessages, stz, telemetrySession.Publish, wsRegistry).ServeHTTP(w, r) }) srvDevtools := &http.Server{ @@ -240,6 +247,7 @@ func main() { rChromeDriver.Handle("/*", chromedriverproxy.Handler(slogger, &chromedriverproxy.Options{ ChromeDriverUpstream: config.ChromeDriverUpstreamAddr, DevToolsProxyAddr: config.DevToolsProxyAddr, + Registry: wsRegistry, })) srvChromeDriver := &http.Server{ @@ -275,6 +283,14 @@ func main() { <-ctx.Done() slogger.Info("shutdown signal received") + // Close hijacked WebSockets with Going Away before draining the HTTP + // servers. http.Server.Shutdown leaves hijacked conns untouched, so without + // this CDP/WebDriver clients would see a 1006 abnormal closure when the VM + // is destroyed instead of a clean 1001. + if n := wsRegistry.CloseAll(websocket.StatusGoingAway, "browser shutting down"); n > 0 { + slogger.Info("closed active websocket connections for shutdown", "count", n) + } + shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second) defer shutdownCancel() g, _ := errgroup.WithContext(shutdownCtx) diff --git a/server/lib/chromedriverproxy/proxy.go b/server/lib/chromedriverproxy/proxy.go index 0b83a5c7..612f34b6 100644 --- a/server/lib/chromedriverproxy/proxy.go +++ b/server/lib/chromedriverproxy/proxy.go @@ -12,6 +12,7 @@ import ( "strings" "github.com/coder/websocket" + "github.com/kernel/kernel-images/server/lib/wsdrain" "github.com/kernel/kernel-images/server/lib/wsproxy" ) @@ -25,6 +26,9 @@ const ( type Options struct { ChromeDriverUpstream string DevToolsProxyAddr string + // Registry, when set, tracks accepted WebDriver/BiDi connections so they + // are closed with a Going Away frame on server shutdown. + Registry *wsdrain.Registry } func resolveOptions(opts *Options) Options { @@ -41,6 +45,7 @@ func resolveOptions(opts *Options) Options { if opts.DevToolsProxyAddr != "" { resolved.DevToolsProxyAddr = opts.DevToolsProxyAddr } + resolved.Registry = opts.Registry return resolved } @@ -313,6 +318,7 @@ func proxyWebSocket(w http.ResponseWriter, r *http.Request, logger *slog.Logger, DialOptions: dialOpts, Logger: logger, Transform: transform, + Registry: cfg.Registry, }) } diff --git a/server/lib/devtoolsproxy/proxy.go b/server/lib/devtoolsproxy/proxy.go index aded45a8..df8b47eb 100644 --- a/server/lib/devtoolsproxy/proxy.go +++ b/server/lib/devtoolsproxy/proxy.go @@ -22,6 +22,7 @@ import ( "github.com/kernel/kernel-images/server/lib/events" oapi "github.com/kernel/kernel-images/server/lib/oapi" "github.com/kernel/kernel-images/server/lib/scaletozero" + "github.com/kernel/kernel-images/server/lib/wsdrain" "github.com/kernel/kernel-images/server/lib/wsproxy" ) @@ -309,7 +310,7 @@ type EventPublisher func(ev events.Event) (events.Envelope, bool) // If logCDPMessages is true, all CDP messages will be logged with their direction. // publish is invoked on accept (cdp_connect) and on teardown (cdp_disconnect); pass // nil to disable emission. -func WebSocketProxyHandler(mgr *UpstreamManager, logger *slog.Logger, logCDPMessages bool, ctrl scaletozero.Controller, publish EventPublisher) http.Handler { +func WebSocketProxyHandler(mgr *UpstreamManager, logger *slog.Logger, logCDPMessages bool, ctrl scaletozero.Controller, publish EventPublisher, reg *wsdrain.Registry) http.Handler { return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { // Counts every relayed message so cdp_disconnect can report message_count. var msgCount atomic.Int64 @@ -349,6 +350,9 @@ func WebSocketProxyHandler(mgr *UpstreamManager, logger *slog.Logger, logCDPMess } clientConn.SetReadLimit(100 * 1024 * 1024) + untrack := reg.Track(clientConn) + defer untrack() + publishCdpConnect(publish) connectedAt := time.Now() diff --git a/server/lib/devtoolsproxy/proxy_test.go b/server/lib/devtoolsproxy/proxy_test.go index c3de042e..5956a555 100644 --- a/server/lib/devtoolsproxy/proxy_test.go +++ b/server/lib/devtoolsproxy/proxy_test.go @@ -23,6 +23,7 @@ import ( "github.com/kernel/kernel-images/server/lib/events" oapi "github.com/kernel/kernel-images/server/lib/oapi" "github.com/kernel/kernel-images/server/lib/scaletozero" + "github.com/kernel/kernel-images/server/lib/wsdrain" "github.com/kernel/kernel-images/server/lib/wsproxy" ) @@ -132,7 +133,7 @@ func TestWebSocketProxyHandler_ProxiesEcho(t *testing.T) { // seed current upstream to echo server including path/query (bypass tailing) mgr.setCurrent((&url.URL{Scheme: u.Scheme, Host: u.Host, Path: u.Path, RawQuery: u.RawQuery}).String()) - proxy := WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), nil) + proxy := WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), nil, nil) proxySrv := httptest.NewServer(proxy) defer proxySrv.Close() @@ -164,6 +165,64 @@ func TestWebSocketProxyHandler_ProxiesEcho(t *testing.T) { } } +func TestWebSocketProxyHandler_RegistryClosesClientWithGoingAway(t *testing.T) { + echoSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + c, err := websocket.Accept(w, r, &websocket.AcceptOptions{OriginPatterns: []string{"*"}}) + if err != nil { + return + } + defer c.Close(websocket.StatusNormalClosure, "") + ctx := r.Context() + for { + mt, msg, err := c.Read(ctx) + if err != nil { + return + } + if err := c.Write(ctx, mt, msg); err != nil { + return + } + } + })) + defer echoSrv.Close() + + u, _ := url.Parse(echoSrv.URL) + logger := silentLogger() + mgr := NewUpstreamManager("/dev/null", logger) + mgr.setCurrent((&url.URL{Scheme: "ws", Host: u.Host, Path: "/echo"}).String()) + + reg := wsdrain.New() + proxySrv := httptest.NewServer(WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), nil, reg)) + defer proxySrv.Close() + + pu, _ := url.Parse(proxySrv.URL) + pu.Scheme = "ws" + ctx := context.Background() + conn, _, err := websocket.Dial(ctx, pu.String(), nil) + if err != nil { + t.Fatalf("dial proxy failed: %v", err) + } + defer conn.Close(websocket.StatusInternalError, "") + + // Round-trip so the proxy session is fully established and registered. + if err := conn.Write(ctx, websocket.MessageText, []byte("ping")); err != nil { + t.Fatalf("write failed: %v", err) + } + if _, _, err := conn.Read(ctx); err != nil { + t.Fatalf("read failed: %v", err) + } + + if n := reg.CloseAll(websocket.StatusGoingAway, "shutting down"); n != 1 { + t.Fatalf("CloseAll closed %d conns, want 1", n) + } + + // The client should observe a 1001 Going Away, not the 1000 the proxy's + // own cleanup would otherwise send. + _, _, err = conn.Read(ctx) + if got := websocket.CloseStatus(err); got != websocket.StatusGoingAway { + t.Fatalf("client close status = %v (err %v), want StatusGoingAway", got, err) + } +} + func TestDialUpstreamWithRetry_RechecksCurrentAfterMissedUpdate(t *testing.T) { // Start a working websocket upstream. upstreamSrv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { @@ -461,7 +520,7 @@ func TestWebSocketProxyHandler_EmitsConnectAndDisconnect(t *testing.T) { mgr.setCurrent(u.String()) rp := &recordingPublisher{} - proxySrv := httptest.NewServer(WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), rp.publish)) + proxySrv := httptest.NewServer(WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), rp.publish, nil)) defer proxySrv.Close() pu, _ := url.Parse(proxySrv.URL) @@ -636,7 +695,7 @@ func TestWebSocketProxyHandler_EmitsUpstreamChangedOnMidStreamRestart(t *testing mgr.setCurrent(urlA.String()) rp := &recordingPublisher{} - proxySrv := httptest.NewServer(WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), rp.publish)) + proxySrv := httptest.NewServer(WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), rp.publish, nil)) defer proxySrv.Close() pu, _ := url.Parse(proxySrv.URL) @@ -720,7 +779,7 @@ func TestWebSocketProxyHandler_KicksClientOffStaleUpstreamOnURLChange(t *testing mgr.setCurrent(urlA.String()) rp := &recordingPublisher{} - proxySrv := httptest.NewServer(WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), rp.publish)) + proxySrv := httptest.NewServer(WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), rp.publish, nil)) defer proxySrv.Close() pu, _ := url.Parse(proxySrv.URL) @@ -772,7 +831,7 @@ func TestWebSocketProxyHandler_EmitsUpstreamErrorOnDialFailure(t *testing.T) { mgr.setCurrent(deadURL) rp := &recordingPublisher{} - proxySrv := httptest.NewServer(WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), rp.publish)) + proxySrv := httptest.NewServer(WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), rp.publish, nil)) defer proxySrv.Close() pu, _ := url.Parse(proxySrv.URL) diff --git a/server/lib/wsdrain/wsdrain.go b/server/lib/wsdrain/wsdrain.go new file mode 100644 index 00000000..826c70f2 --- /dev/null +++ b/server/lib/wsdrain/wsdrain.go @@ -0,0 +1,74 @@ +// Package wsdrain tracks live WebSocket connections so they can be closed with +// a single status code when the server shuts down. +// +// Browser VMs are torn down on session destroy. The hijacked WebSocket flows +// (CDP proxy, WebDriver/BiDi, ChromeDriver, PTY attach) are not touched by +// http.Server.Shutdown, so without a proactive close the client sees a 1006 +// abnormal closure when the VM dies. Registering each connection here lets the +// shutdown path send every client a clean 1001 Going Away first. +package wsdrain + +import ( + "sync" + + "github.com/coder/websocket" +) + +// Conn is the subset of *websocket.Conn the registry needs. +type Conn interface { + Close(code websocket.StatusCode, reason string) error +} + +// Registry tracks active connections. Construct one with New. All methods are +// safe for concurrent use and tolerate a nil receiver, so callers may pass a +// nil *Registry to disable tracking. +type Registry struct { + mu sync.Mutex + conns map[Conn]struct{} +} + +func New() *Registry { + return &Registry{conns: make(map[Conn]struct{})} +} + +// Track registers conn and returns a function that removes it. The returned +// function is idempotent; call it (e.g. via defer) when the connection ends. +func (r *Registry) Track(conn Conn) func() { + if r == nil || conn == nil { + return func() {} + } + r.mu.Lock() + r.conns[conn] = struct{}{} + r.mu.Unlock() + + var once sync.Once + return func() { + once.Do(func() { + r.mu.Lock() + delete(r.conns, conn) + r.mu.Unlock() + }) + } +} + +// CloseAll closes every tracked connection with the given code and reason and +// returns how many it closed. Connections are snapshotted under the lock and +// closed outside it. Close errors are ignored: the connection is being +// discarded regardless, and the first Close wins, so a later normal-closure +// from the connection's own teardown does not override the code sent here. +func (r *Registry) CloseAll(code websocket.StatusCode, reason string) int { + if r == nil { + return 0 + } + r.mu.Lock() + conns := make([]Conn, 0, len(r.conns)) + for c := range r.conns { + conns = append(conns, c) + } + r.mu.Unlock() + + for _, c := range conns { + _ = c.Close(code, reason) + } + return len(conns) +} diff --git a/server/lib/wsdrain/wsdrain_test.go b/server/lib/wsdrain/wsdrain_test.go new file mode 100644 index 00000000..d48b4bdb --- /dev/null +++ b/server/lib/wsdrain/wsdrain_test.go @@ -0,0 +1,56 @@ +package wsdrain + +import ( + "testing" + + "github.com/coder/websocket" +) + +type fakeConn struct { + closes []websocket.StatusCode +} + +func (c *fakeConn) Close(code websocket.StatusCode, _ string) error { + c.closes = append(c.closes, code) + return nil +} + +func TestCloseAllClosesTrackedConns(t *testing.T) { + r := New() + a, b := &fakeConn{}, &fakeConn{} + r.Track(a) + r.Track(b) + + if n := r.CloseAll(websocket.StatusGoingAway, "bye"); n != 2 { + t.Fatalf("CloseAll returned %d, want 2", n) + } + for _, c := range []*fakeConn{a, b} { + if len(c.closes) != 1 || c.closes[0] != websocket.StatusGoingAway { + t.Fatalf("conn closed with %v, want one StatusGoingAway", c.closes) + } + } +} + +func TestUntrackRemovesConn(t *testing.T) { + r := New() + c := &fakeConn{} + untrack := r.Track(c) + untrack() + untrack() // idempotent + + if n := r.CloseAll(websocket.StatusGoingAway, "bye"); n != 0 { + t.Fatalf("CloseAll returned %d after untrack, want 0", n) + } + if len(c.closes) != 0 { + t.Fatalf("untracked conn was closed: %v", c.closes) + } +} + +func TestNilRegistryIsNoop(t *testing.T) { + var r *Registry + untrack := r.Track(&fakeConn{}) + untrack() + if n := r.CloseAll(websocket.StatusGoingAway, "bye"); n != 0 { + t.Fatalf("nil registry CloseAll returned %d, want 0", n) + } +} diff --git a/server/lib/wsproxy/wsproxy.go b/server/lib/wsproxy/wsproxy.go index 8a14e490..8a7e5940 100644 --- a/server/lib/wsproxy/wsproxy.go +++ b/server/lib/wsproxy/wsproxy.go @@ -7,6 +7,7 @@ import ( "sync" "github.com/coder/websocket" + "github.com/kernel/kernel-images/server/lib/wsdrain" ) // Conn abstracts a WebSocket connection for testing and flexibility. @@ -28,6 +29,9 @@ type ProxyOptions struct { DialOptions *websocket.DialOptions Logger *slog.Logger Transform MessageTransform + // Registry, when set, tracks the accepted client connection so it is + // closed with a Going Away frame on server shutdown. + Registry *wsdrain.Registry } // PumpExitCause names which side caused Pump to return. Callers use this to @@ -121,6 +125,9 @@ func Proxy(w http.ResponseWriter, r *http.Request, upstreamURL string, opts Prox } clientConn.SetReadLimit(100 * 1024 * 1024) + untrack := opts.Registry.Track(clientConn) + defer untrack() + upstreamConn, _, err := websocket.Dial(r.Context(), upstreamURL, opts.DialOptions) if err != nil { logger.Error("dial upstream failed", slog.String("err", err.Error()), slog.String("url", upstreamURL)) From df433781b71921911b21fb29722b0a5408d527d1 Mon Sep 17 00:00:00 2001 From: Sayan- <1415138+Sayan-@users.noreply.github.com> Date: Wed, 24 Jun 2026 20:43:16 +0000 Subject: [PATCH 2/3] Close shutdown WebSockets concurrently with server drain CloseAll touches the hijacked connections that http.Server.Shutdown ignores, so the two drain disjoint connection sets. Run CloseAll as a member of the shutdown errgroup instead of a serialized step before it, overlapping its close handshakes with the servers' drain. Co-Authored-By: Claude Opus 4.7 --- server/cmd/api/main.go | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/server/cmd/api/main.go b/server/cmd/api/main.go index 0b40fff0..aae7d301 100644 --- a/server/cmd/api/main.go +++ b/server/cmd/api/main.go @@ -283,18 +283,21 @@ func main() { <-ctx.Done() slogger.Info("shutdown signal received") - // Close hijacked WebSockets with Going Away before draining the HTTP - // servers. http.Server.Shutdown leaves hijacked conns untouched, so without - // this CDP/WebDriver clients would see a 1006 abnormal closure when the VM - // is destroyed instead of a clean 1001. - if n := wsRegistry.CloseAll(websocket.StatusGoingAway, "browser shutting down"); n > 0 { - slogger.Info("closed active websocket connections for shutdown", "count", n) - } - shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 10*time.Second) defer shutdownCancel() g, _ := errgroup.WithContext(shutdownCtx) + // Close hijacked WebSockets with Going Away. http.Server.Shutdown leaves + // hijacked conns untouched, so without this CDP/WebDriver clients would see + // a 1006 abnormal closure when the VM is destroyed instead of a clean 1001. + // This drains a disjoint set of connections from the servers below, so it + // runs alongside them. + g.Go(func() error { + if n := wsRegistry.CloseAll(websocket.StatusGoingAway, "browser shutting down"); n > 0 { + slogger.Info("closed active websocket connections for shutdown", "count", n) + } + return nil + }) g.Go(func() error { return srv.Shutdown(shutdownCtx) }) From 0eb856d826200b083ca9f8292714eb2bc25f62ff Mon Sep 17 00:00:00 2001 From: Sayan Samanta Date: Wed, 24 Jun 2026 13:52:50 -0700 Subject: [PATCH 3/3] vibe code --- server/cmd/api/main.go | 25 +++++++++---------------- server/lib/wsdrain/wsdrain.go | 6 ------ 2 files changed, 9 insertions(+), 22 deletions(-) diff --git a/server/cmd/api/main.go b/server/cmd/api/main.go index aae7d301..d839f4fb 100644 --- a/server/cmd/api/main.go +++ b/server/cmd/api/main.go @@ -47,11 +47,6 @@ func main() { } slogger.Info("server configuration", "config", config) - // Tracks live hijacked WebSocket connections (CDP, WebDriver/BiDi, - // ChromeDriver, PTY attach) so they can be closed with a clean Going Away - // frame on shutdown — http.Server.Shutdown does not touch hijacked conns. - wsRegistry := wsdrain.New() - // context cancellation on SIGINT/SIGTERM ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() @@ -87,6 +82,9 @@ func main() { os.Exit(1) } + // ws conn tracker + wsRegistry := wsdrain.New() + // DevTools WebSocket upstream manager: tail Chromium supervisord log const chromiumLogPath = "/var/log/supervisord/chromium" upstreamMgr := devtoolsproxy.NewUpstreamManager(chromiumLogPath, slogger) @@ -287,23 +285,18 @@ func main() { defer shutdownCancel() g, _ := errgroup.WithContext(shutdownCtx) - // Close hijacked WebSockets with Going Away. http.Server.Shutdown leaves - // hijacked conns untouched, so without this CDP/WebDriver clients would see - // a 1006 abnormal closure when the VM is destroyed instead of a clean 1001. - // This drains a disjoint set of connections from the servers below, so it - // runs alongside them. - g.Go(func() error { - if n := wsRegistry.CloseAll(websocket.StatusGoingAway, "browser shutting down"); n > 0 { - slogger.Info("closed active websocket connections for shutdown", "count", n) - } - return nil - }) g.Go(func() error { return srv.Shutdown(shutdownCtx) }) g.Go(func() error { return apiService.Shutdown(shutdownCtx) }) + g.Go(func() error { + if n := wsRegistry.CloseAll(websocket.StatusGoingAway, "browser shutting down"); n > 0 { + slogger.Info("closed active websocket connections for shutdown", "count", n) + } + return nil + }) g.Go(func() error { upstreamMgr.Stop() return srvDevtools.Shutdown(shutdownCtx) diff --git a/server/lib/wsdrain/wsdrain.go b/server/lib/wsdrain/wsdrain.go index 826c70f2..0df8a093 100644 --- a/server/lib/wsdrain/wsdrain.go +++ b/server/lib/wsdrain/wsdrain.go @@ -1,11 +1,5 @@ // Package wsdrain tracks live WebSocket connections so they can be closed with // a single status code when the server shuts down. -// -// Browser VMs are torn down on session destroy. The hijacked WebSocket flows -// (CDP proxy, WebDriver/BiDi, ChromeDriver, PTY attach) are not touched by -// http.Server.Shutdown, so without a proactive close the client sees a 1006 -// abnormal closure when the VM dies. Registering each connection here lets the -// shutdown path send every client a clean 1001 Going Away first. package wsdrain import (