Reduce Chromium restart wait time

server/cmd/api/api/process.go

@@ -26,7 +26,6 @@ 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"
 )
 
@@ -660,7 +659,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, reg *wsdrain.Registry) {
+func (s *ApiService) HandleProcessAttachWS(w http.ResponseWriter, r *http.Request, id string) {
 	ctx := r.Context()
 	log := logger.FromContext(ctx)
 
@@ -703,9 +702,6 @@ 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
 

server/cmd/api/main.go

@@ -14,7 +14,6 @@ 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"
@@ -34,7 +33,6 @@ 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() {
@@ -83,10 +81,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	// ws conn tracker
-	wsRegistry := wsdrain.New()
-
-	// DevTools WebSocket upstream manager: tail Chromium supervisord log
+		// DevTools WebSocket upstream manager: tail Chromium supervisord log
 	const chromiumLogPath = "/var/log/supervisord/chromium"
 	upstreamMgr := devtoolsproxy.NewUpstreamManager(chromiumLogPath, slogger)
 	upstreamMgr.Start(ctx)
@@ -176,7 +171,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, wsRegistry)
+		apiService.HandleProcessAttachWS(w, r, id)
 	})
 
 	// Serve extension files for Chrome policy-installed extensions
@@ -224,7 +219,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, wsRegistry).ServeHTTP(w, r)
+		devtoolsproxy.WebSocketProxyHandler(upstreamMgr, slogger, config.LogCDPMessages, stz, telemetrySession.Publish).ServeHTTP(w, r)
 	})
 
 	srvDevtools := &http.Server{
@@ -250,7 +245,6 @@ func main() {
 	rChromeDriver.Handle("/*", chromedriverproxy.Handler(slogger, &chromedriverproxy.Options{
 		ChromeDriverUpstream: config.ChromeDriverUpstreamAddr,
 		DevToolsProxyAddr:    config.DevToolsProxyAddr,
-		Registry:             wsRegistry,
 	}))
 
 	srvChromeDriver := &http.Server{
@@ -320,12 +314,6 @@ func main() {
 	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)

server/cmd/chromium-launcher/main.go

@@ -1,7 +1,6 @@
 package main
 
 import (
-	"context"
 	"flag"
 	"fmt"
 	"net"
@@ -167,41 +166,19 @@ func execLookPath(file string) (string, error) {
 	return exec.LookPath(file)
 }
 
-// waitForPort waits until the given port is available for binding on IPv4.
-// This handles the delay after SIGKILL before the kernel releases the socket.
-// We disable SO_REUSEADDR to get an accurate check matching chromium's bind behavior.
+// waitForPort waits until the devtools port can be bound on IPv4. After SIGKILL
+// the old listener may linger briefly; this loop waits for that to clear.
+// Go's net.Listen sets SO_REUSEADDR, matching chromium's DevTools bind — so
+// TIME_WAIT sockets from prior CDP client connections do not block the probe.
 // Only IPv4 is checked because IPv6 is disabled at the kernel level in the VM.
 func waitForPort(port string, timeout time.Duration) {
 	deadline := time.Now().Add(timeout)
-	addrs := []string{"127.0.0.1:" + port}
-
-	// ListenConfig with Control to disable SO_REUSEADDR for accurate port availability check
-	lc := &net.ListenConfig{
-		Control: func(network, address string, c syscall.RawConn) error {
-			var sockErr error
-			err := c.Control(func(fd uintptr) {
-				// Disable SO_REUSEADDR to match chromium's behavior
-				sockErr = syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_REUSEADDR, 0)
-			})
-			if err != nil {
-				return err
-			}
-			return sockErr
-		},
-	}
+	addr := "127.0.0.1:" + port
 
-	ctx := context.Background()
 	for time.Now().Before(deadline) {
-		allFree := true
-		for _, addr := range addrs {
-			ln, err := lc.Listen(ctx, "tcp", addr)
-			if err != nil {
-				allFree = false
-				break
-			}
+		ln, err := net.Listen("tcp", addr)
+		if err == nil {
 			ln.Close()
-		}
-		if allFree {
 			return
 		}
 		time.Sleep(50 * time.Millisecond)

server/lib/chromedriverproxy/proxy.go

@@ -12,7 +12,6 @@ import (
 	"strings"
 
 	"github.com/coder/websocket"
-	"github.com/kernel/kernel-images/server/lib/wsdrain"
 	"github.com/kernel/kernel-images/server/lib/wsproxy"
 )
 
@@ -26,9 +25,6 @@ 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 {
@@ -44,8 +40,6 @@ func resolveOptions(opts *Options) Options {
 	}
 	if opts.DevToolsProxyAddr != "" {
 		resolved.DevToolsProxyAddr = opts.DevToolsProxyAddr
-	}
-	resolved.Registry = opts.Registry
 	return resolved
 }
 
@@ -318,7 +312,6 @@ func proxyWebSocket(w http.ResponseWriter, r *http.Request, logger *slog.Logger,
 		DialOptions:   dialOpts,
 		Logger:        logger,
 		Transform:     transform,
-		Registry:      cfg.Registry,
 	})
 }
 

server/lib/devtoolsproxy/proxy.go

@@ -22,7 +22,6 @@ 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"
 )
 
@@ -310,7 +309,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, reg *wsdrain.Registry) http.Handler {
+func WebSocketProxyHandler(mgr *UpstreamManager, logger *slog.Logger, logCDPMessages bool, ctrl scaletozero.Controller, publish EventPublisher) 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
@@ -350,9 +349,6 @@ func WebSocketProxyHandler(mgr *UpstreamManager, logger *slog.Logger, logCDPMess
 		}
 		clientConn.SetReadLimit(100 * 1024 * 1024)
 
-		untrack := reg.Track(clientConn)
-		defer untrack()
-
 		publishCdpConnect(publish)
 		connectedAt := time.Now()
 

server/lib/devtoolsproxy/proxy_test.go

@@ -23,7 +23,6 @@ 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"
 )
 
@@ -133,7 +132,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, nil)
+	proxy := WebSocketProxyHandler(mgr, logger, false, scaletozero.NewNoopController(), nil)
 	proxySrv := httptest.NewServer(proxy)
 	defer proxySrv.Close()
 
@@ -165,64 +164,6 @@ 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) {