feat: graceful shutdown with drain signal and browser cleanup

Author: karust

cmd/serve.go

@@ -2,10 +2,13 @@ package cmd
 
 import (
 	"context"
+	"errors"
 	"fmt"
 	"os"
+	"os/signal"
 	"strings"
 	"sync"
+	"syscall"
 	"time"
 
 	"github.com/karust/openserp/baidu"
@@ -96,7 +99,7 @@ func serve(cmd *cobra.Command, args []string) {
 			&rawEngine{name: "yandex"},
 			&rawEngine{name: "baidu"},
 		)
-		if err := serv.Listen(); err != nil {
+		if err := listenWithGracefulShutdown(serv, nil); err != nil {
 			logrus.Error(err)
 		}
 		return
@@ -117,15 +120,15 @@ func serve(cmd *cobra.Command, args []string) {
 		baseOpts.IsHeadless = false
 	}
 
-	engines, err := buildBrowserEngines(baseOpts, proxyCfg)
+	engines, closeBrowsers, err := buildBrowserEngines(baseOpts, proxyCfg)
 	if err != nil {
 		logrus.Error(err)
 		return
 	}
 
 	serverOpts := buildServerOptions(corsCfg, proxyCfg)
 	serv := core.NewServerWithOptions(config.Server.Host, config.Server.Port, serverOpts, engines...)
-	if err := serv.Listen(); err != nil {
+	if err := listenWithGracefulShutdown(serv, closeBrowsers); err != nil {
 		logrus.Error(err)
 	}
 }
@@ -154,6 +157,71 @@ func buildServerOptions(corsCfg core.CORSConfig, proxyCfg core.ProxyConfig) core
 	}
 }
 
+const gracefulShutdownTimeout = 30 * time.Second
+
+func listenWithGracefulShutdown(serv *core.Server, onShutdown func() error) error {
+	listenErrCh := make(chan error, 1)
+	go func() {
+		listenErrCh <- serv.Listen()
+	}()
+
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
+	defer signal.Stop(sigCh)
+
+	select {
+	case err := <-listenErrCh:
+		return err
+	case sig := <-sigCh:
+		logrus.WithField("signal", sig.String()).Info("Shutdown signal received, draining traffic")
+	}
+
+	serv.SetDraining(true)
+
+	shutdownErr := serv.ShutdownWithTimeout(gracefulShutdownTimeout)
+	if isServerNotRunningError(shutdownErr) {
+		shutdownErr = nil
+	}
+
+	if onShutdown != nil {
+		resourceErr := onShutdown()
+		if resourceErr != nil {
+			shutdownErr = errors.Join(shutdownErr, resourceErr)
+		}
+	}
+
+	if listenErr := waitForListenExit(listenErrCh); listenErr != nil && !isExpectedListenShutdownError(listenErr) {
+		shutdownErr = errors.Join(shutdownErr, listenErr)
+	}
+
+	return shutdownErr
+}
+
+func waitForListenExit(listenErrCh <-chan error) error {
+	select {
+	case err := <-listenErrCh:
+		return err
+	case <-time.After(time.Second):
+		return nil
+	}
+}
+
+func isExpectedListenShutdownError(err error) bool {
+	if err == nil {
+		return true
+	}
+	msg := strings.ToLower(err.Error())
+	return strings.Contains(msg, "server closed") ||
+		strings.Contains(msg, "closed network connection")
+}
+
+func isServerNotRunningError(err error) bool {
+	if err == nil {
+		return false
+	}
+	return strings.Contains(strings.ToLower(err.Error()), "server is not running")
+}
+
 type browserPool struct {
 	mu      sync.Mutex
 	base    core.BrowserOpts
@@ -193,6 +261,27 @@ func (p *browserPool) get(proxyURL string) (*core.Browser, error) {
 	return b, nil
 }
 
+func (p *browserPool) close() error {
+	p.mu.Lock()
+	browsers := make([]*core.Browser, 0, len(p.browser))
+	for key, b := range p.browser {
+		browsers = append(browsers, b)
+		delete(p.browser, key)
+	}
+	p.mu.Unlock()
+
+	var closeErr error
+	for _, browser := range browsers {
+		if browser == nil {
+			continue
+		}
+		if err := browser.Close(); err != nil {
+			closeErr = errors.Join(closeErr, err)
+		}
+	}
+	return closeErr
+}
+
 type pooledBrowserEngine struct {
 	name    string
 	limiter *rate.Limiter
@@ -301,15 +390,15 @@ func browserEngineSpecs() []browserEngineSpec {
 	}
 }
 
-func buildBrowserEngines(baseOpts core.BrowserOpts, proxyCfg core.ProxyConfig) ([]core.SearchEngine, error) {
+func buildBrowserEngines(baseOpts core.BrowserOpts, proxyCfg core.ProxyConfig) ([]core.SearchEngine, func() error, error) {
 	pool := newBrowserPool(baseOpts)
 	specs := browserEngineSpecs()
 
 	engines := make([]core.SearchEngine, 0, len(specs))
 	for _, spec := range specs {
 		policy := resolveEngineProxyPolicy(proxyCfg, spec.name)
 		if err := validateBrowserProxyPolicy(proxyCfg, policy); err != nil {
-			return nil, fmt.Errorf("browser proxy validation failed for engine %s: %w", spec.name, err)
+			return nil, nil, fmt.Errorf("browser proxy validation failed for engine %s: %w", spec.name, err)
 		}
 
 		opts := spec.opts
@@ -324,7 +413,7 @@ func buildBrowserEngines(baseOpts core.BrowserOpts, proxyCfg core.ProxyConfig) (
 		})
 	}
 
-	return engines, nil
+	return engines, pool.close, nil
 }
 
 func validateBrowserProxyPolicy(proxyCfg core.ProxyConfig, policy core.ProxyPolicy) error {

core/browser.go

@@ -174,11 +174,7 @@ func resolveBrowserBinaryPath(browserPath string, lookPath func() (string, bool)
 
 // IsInitialized reports whether the browser launcher has been created.
 func (b *Browser) IsInitialized() bool {
-	if b.browserAddr != "" {
-		return true
-	} else {
-		return false
-	}
+	return b.browserAddr != ""
 }
 
 // Navigate connects to Chromium, creates a page, applies stealth/emulation and
@@ -317,7 +313,40 @@ func (b *Browser) Navigate(ctx context.Context, URL string) (*rod.Page, error) {
 
 // Close closes the active browser connection.
 func (b *Browser) Close() error {
-	return b.browser.Close()
+	if b == nil || b.browserAddr == "" {
+		return nil
+	}
+
+	browser := b.browser
+	if browser == nil {
+		browser = rod.New().ControlURL(b.browserAddr)
+		if b.Timeout > 0 {
+			browser = browser.Timeout(b.Timeout)
+		}
+		if err := browser.Connect(); err != nil {
+			if isBrowserClosedError(err) {
+				return nil
+			}
+			return err
+		}
+	}
+
+	b.browser = nil
+	if err := browser.Close(); err != nil && !isBrowserClosedError(err) {
+		return err
+	}
+	return nil
+}
+
+func isBrowserClosedError(err error) bool {
+	if err == nil {
+		return false
+	}
+	msg := strings.ToLower(err.Error())
+	return strings.Contains(msg, "connection refused") ||
+		strings.Contains(msg, "closed network connection") ||
+		strings.Contains(msg, "target closed") ||
+		strings.Contains(msg, "eof")
 }
 
 // ClosePageWithTimeout bounds page close calls so shutdown paths don't hang.

core/server.go

@@ -9,6 +9,7 @@ import (
 	"runtime"
 	"sort"
 	"strings"
+	"sync/atomic"
 	"time"
 
 	"github.com/gofiber/fiber/v2"
@@ -43,6 +44,7 @@ type Server struct {
 	resilient     *ResilientSearcher
 	startTime     time.Time
 	opts          ServerOptions
+	draining      atomic.Bool
 }
 
 // ServerOptions configures HTTP server middleware and resilience behavior.
@@ -98,6 +100,7 @@ func NewServerWithOptions(host string, port int, opts ServerOptions, searchEngin
 		startTime:     time.Now(),
 		opts:          opts,
 	}
+	serv.draining.Store(false)
 	logrus.Info("Resilient search enabled: retry + circuit breaker")
 	if opts.AllowEndpointFallback {
 		logrus.Warn("Dedicated endpoint fallback is enabled")
@@ -120,6 +123,7 @@ func NewServerWithOptions(host string, port int, opts ServerOptions, searchEngin
 	app.Get("/docs", serv.handleSwaggerUI)
 	app.Get("/docs/", serv.handleSwaggerUI)
 	app.Get("/health", serv.handleHealthCheck)
+	app.Get("/ready", serv.handleReadinessCheck)
 	app.Get("/stats", serv.handleStats)
 	app.Get("/stats/cache", serv.handleCacheStats)
 	app.Get("/stats/proxy", serv.handleProxyStats)
@@ -264,6 +268,12 @@ type HealthStatus struct {
 	System  map[string]interface{} `json:"system"`
 }
 
+// ReadinessStatus is returned by /ready to indicate if this instance can
+// receive new traffic.
+type ReadinessStatus struct {
+	Status string `json:"status"`
+}
+
 // EngineHealth describes availability of one configured engine.
 type EngineHealth struct {
 	Name        string `json:"name"`
@@ -338,6 +348,15 @@ func (s *Server) handleHealthCheck(c *fiber.Ctx) error {
 	return c.JSON(health)
 }
 
+func (s *Server) handleReadinessCheck(c *fiber.Ctx) error {
+	status := ReadinessStatus{Status: "ready"}
+	if s.draining.Load() {
+		status.Status = "draining"
+		return c.Status(fiber.StatusServiceUnavailable).JSON(status)
+	}
+	return c.JSON(status)
+}
+
 func (s *Server) handleStats(c *fiber.Ctx) error {
 	return c.JSON(map[string]interface{}{
 		"cache":            s.cacheStatsPayload(),
@@ -672,12 +691,30 @@ func (s *Server) handleSwaggerUI(c *fiber.Ctx) error {
 	return c.SendString(page)
 }
 
+// SetDraining controls readiness state exposed by /ready.
+func (s *Server) SetDraining(draining bool) {
+	s.draining.Store(draining)
+}
+
+const defaultShutdownTimeout = 30 * time.Second
+
 // Listen starts the Fiber HTTP server on the configured address.
 func (s *Server) Listen() error {
+	s.SetDraining(false)
 	return s.app.Listen(s.addr)
 }
 
 // Shutdown gracefully stops the Fiber HTTP server.
 func (s *Server) Shutdown() error {
-	return s.app.Shutdown()
+	return s.ShutdownWithTimeout(defaultShutdownTimeout)
+}
+
+// ShutdownWithTimeout drains the server before force-closing active
+// connections when timeout is exceeded.
+func (s *Server) ShutdownWithTimeout(timeout time.Duration) error {
+	if timeout <= 0 {
+		timeout = defaultShutdownTimeout
+	}
+	s.SetDraining(true)
+	return s.app.ShutdownWithTimeout(timeout)
 }

core/server_test.go

@@ -408,6 +408,37 @@ func TestHealthEndpointStatusSemantics(t *testing.T) {
 	}
 }
 
+func TestReadinessEndpointStatusSemantics(t *testing.T) {
+	engine := &engineMock{name: "google", initialized: true}
+	srv := NewServerWithOptions("127.0.0.1", 7077, DefaultServerOptions(), engine)
+
+	resp := request(t, srv, "/ready")
+	if resp.StatusCode != http.StatusOK {
+		t.Fatalf("expected ready endpoint to return 200 while serving, got %d", resp.StatusCode)
+	}
+
+	var ready ReadinessStatus
+	if err := json.NewDecoder(resp.Body).Decode(&ready); err != nil {
+		t.Fatalf("decode ready response: %v", err)
+	}
+	if ready.Status != "ready" {
+		t.Fatalf("expected readiness status=ready, got %q", ready.Status)
+	}
+
+	srv.SetDraining(true)
+	resp = request(t, srv, "/ready")
+	if resp.StatusCode != http.StatusServiceUnavailable {
+		t.Fatalf("expected ready endpoint to return 503 while draining, got %d", resp.StatusCode)
+	}
+
+	if err := json.NewDecoder(resp.Body).Decode(&ready); err != nil {
+		t.Fatalf("decode draining response: %v", err)
+	}
+	if ready.Status != "draining" {
+		t.Fatalf("expected readiness status=draining, got %q", ready.Status)
+	}
+}
+
 func TestDedicatedEndpointNoFallbackByDefault(t *testing.T) {
 	primary := &engineMock{
 		name:        "google",