diff --git a/cmd/serve.go b/cmd/serve.go index 067e893..3e57002 100644 --- a/cmd/serve.go +++ b/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,7 +120,7 @@ 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 @@ -125,7 +128,7 @@ func serve(cmd *cobra.Command, args []string) { 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,7 +390,7 @@ 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() @@ -309,7 +398,7 @@ func buildBrowserEngines(baseOpts core.BrowserOpts, proxyCfg core.ProxyConfig) ( 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 { diff --git a/core/browser.go b/core/browser.go index cbec094..c817b7e 100644 --- a/core/browser.go +++ b/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. diff --git a/core/server.go b/core/server.go index 364e0de..386e32a 100644 --- a/core/server.go +++ b/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) } diff --git a/core/server_test.go b/core/server_test.go index dd446f6..2189da7 100644 --- a/core/server_test.go +++ b/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", diff --git a/docs/openapi.yaml b/docs/openapi.yaml index 1287d90..ae9daa8 100644 --- a/docs/openapi.yaml +++ b/docs/openapi.yaml @@ -303,6 +303,34 @@ paths: go_version: go1.24.6 "404": $ref: "#/components/responses/NotFoundError" + /ready: + get: + tags: [Health] + operationId: readinessCheck + summary: Service readiness status + responses: + "200": + description: Instance is ready to receive traffic + content: + application/json: + schema: + $ref: "#/components/schemas/ReadinessStatus" + examples: + ready: + value: + status: ready + "503": + description: Instance is draining and should not receive new traffic + content: + application/json: + schema: + $ref: "#/components/schemas/ReadinessStatus" + examples: + draining: + value: + status: draining + "404": + $ref: "#/components/responses/NotFoundError" /stats: get: tags: [Stats] @@ -730,6 +758,13 @@ components: type: integer go_version: type: string + ReadinessStatus: + type: object + required: [status] + properties: + status: + type: string + enum: [ready, draining] CacheStatsEnabled: type: object required: [status, entries, hits, misses, bypasses, evictions, ttl_seconds, max_size]