From f6b7b74a94b0ebac7afa4646f685491b3599c0d9 Mon Sep 17 00:00:00 2001 From: Yaash Singhh Date: Mon, 26 Oct 2026 14:23:29 +0530 Subject: [PATCH 1/5] fix(server): shut down HTTP server gracefully and keep exit status on failure The server goroutine was never shut down explicitly, and an unexpected server error called logrus.Fatalf, which exits immediately and skips deferred cleanup. - Keep the http.Server reachable and call Shutdown on a fresh 15s context (the signal context is already cancelled). If it fails or times out, log it and force-close with Close, reporting any close error. - Treat http.ErrServerClosed as expected only after a requested shutdown; any other server exit is a failure. - Shut down HTTP before draining telemetry (fresh 5s context) so requests finishing during shutdown can still emit it. - Replace Fatalf with a returned exit code: main now calls os.Exit(runApplication()), so deferred cleanup runs before a non-zero exit. - Close the shared sql.DB pool via defer, after the HTTP server stops. - Cancel the signal context after shutdown so context-bound background work stops before the pools close. Shutdown does not stop running flows or background services, and does not wait for hijacked (WebSocket) connections. Signed-off-by: Yaash Singhh --- backend/cmd/pentagi/main.go | 184 ++++++++++++++++++++++++++++++------ 1 file changed, 153 insertions(+), 31 deletions(-) diff --git a/backend/cmd/pentagi/main.go b/backend/cmd/pentagi/main.go index bed445ebb..ab852119d 100644 --- a/backend/cmd/pentagi/main.go +++ b/backend/cmd/pentagi/main.go @@ -36,7 +36,20 @@ import ( "go.opentelemetry.io/otel/attribute" ) +// HTTP shutdown and telemetry drain run on fresh contexts: the signal context +// is already cancelled by the time either starts. +const ( + httpShutdownTimeout = 15 * time.Second + telemetryDrainTimeout = 5 * time.Second +) + func main() { + // os.Exit skips deferred calls, so every deferred cleanup lives in + // runApplication and the exit status is only decided once it has returned. + os.Exit(runApplication()) +} + +func runApplication() int { // Setup graceful shutdown context with signal handling ctx, cancelOnSignal := signal.NotifyContext( context.Background(), @@ -50,7 +63,7 @@ func main() { // exits, and its output is a JSON document on stdout — a greeting printed // first would be the first thing a reader has to strip back out. if wantsInfo(os.Args[1:]) { - os.Exit(runInfo(ctx)) + return runInfo(ctx) } logrus.Infof("Starting PentAGI %s", version.GetBinaryVersion()) @@ -126,6 +139,15 @@ func main() { db.SetMaxIdleConns(cfg.DBMaxIdleConns) db.SetConnMaxLifetime(time.Hour) + // One pool shared by sqlc and GORM, closed here exactly once. Deferred, so + // it runs only after the HTTP server has stopped (and after the pgxpool + // defer below, which was registered later). + defer func() { + if err := db.Close(); err != nil { + logrus.WithError(err).Warn("failed to close the database connection pool") + } + }() + if err := database.VerifySearchPath(ctx, db, cfg); err != nil { logrus.WithError(err).Fatal("Tenant schema verification failed") } @@ -236,47 +258,147 @@ func main() { r := router.NewRouter(queries, orm, cfg, providers, controller, subscriptions, client, updates) - // Launch HTTP/HTTPS server in background goroutine - serverErrChan := make(chan error, 1) - go func() { - listen := net.JoinHostPort(cfg.ServerHost, strconv.Itoa(cfg.ServerPort)) - logrus.Infof("API server listening on %s", listen) + listen := net.JoinHostPort(cfg.ServerHost, strconv.Itoa(cfg.ServerPort)) + srv := &http.Server{ + Addr: listen, + Handler: r.Handler(), + ReadHeaderTimeout: router.ReadHeaderTimeout, + IdleTimeout: router.IdleTimeout, + } - srv := &http.Server{ - Addr: listen, - Handler: r.Handler(), - ReadHeaderTimeout: router.ReadHeaderTimeout, - IdleTimeout: router.IdleTimeout, - } + serve := func() error { + logrus.Infof("API server listening on %s", listen) - var startErr error if cfg.ServerUseSSL && cfg.ServerSSLCrt != "" && cfg.ServerSSLKey != "" { logrus.Info("Starting server with TLS enabled") - startErr = srv.ListenAndServeTLS(cfg.ServerSSLCrt, cfg.ServerSSLKey) - } else { - logrus.Info("Starting server without TLS (HTTP only)") - startErr = srv.ListenAndServe() + return srv.ListenAndServeTLS(cfg.ServerSSLCrt, cfg.ServerSSLKey) } - if startErr != nil { - serverErrChan <- fmt.Errorf("API server startup failed: %w", startErr) - } - }() + logrus.Info("Starting server without TLS (HTTP only)") + return srv.ListenAndServe() + } + + // Blocks until a shutdown signal or a server failure, then stops the HTTP + // server and drains telemetry. The exit code is returned rather than raised + // with Fatal so the deferred pool cleanup above still runs. + exitCode := serveUntilShutdown(ctx, srv, serve, httpShutdownTimeout, func(ctx context.Context) error { + return obs.Observer.Drain(ctx) + }) + + // On a server failure ctx is still live; cancel it so work bound to it (the + // update service, flows) stops before the deferred calls close the pools. + // Idempotent after a signal. + cancelOnSignal() + + if exitCode == 0 { + logrus.Info("Application shutdown completed successfully") + } + + return exitCode +} + +// httpServer is the part of *http.Server the shutdown logic needs, so tests can +// drive the deadline and failure paths without waiting on real timers. +type httpServer interface { + Shutdown(ctx context.Context) error + Close() error +} + +// serveUntilShutdown runs serve until ctx is cancelled or serve fails, stops the +// HTTP server, and only then drains telemetry, so requests that finish during +// graceful shutdown can still emit it. It returns the process exit code and +// never exits itself, which keeps the caller's deferred cleanup reachable. +func serveUntilShutdown( + ctx context.Context, + srv httpServer, + serve func() error, + shutdownTimeout time.Duration, + drain func(context.Context) error, +) int { + err := superviseHTTPServer(ctx, srv, serve, shutdownTimeout) + + drainTelemetry(drain) + + if err != nil { + logrus.WithError(err).Error("HTTP server terminated with an error") + return 1 + } + + return 0 +} + +// superviseHTTPServer runs serve in the background and returns once the server +// has been stopped. It returns nil only for a shutdown it requested itself and +// that completed gracefully. +func superviseHTTPServer( + ctx context.Context, + srv httpServer, + serve func() error, + shutdownTimeout time.Duration, +) error { + // Buffered so the goroutine can always deliver its result and exit, even + // when nobody is left to receive it. + serveErr := make(chan error, 1) + go func() { serveErr <- serve() }() - // Block until shutdown signal received or server error occurs select { case <-ctx.Done(): logrus.Warn("Shutdown signal received, cleaning up resources...") - // ctx is already cancelled here — drain telemetry on a fresh deadline so the - // final batch still reaches the collector before the process exits. - drainCtx, cancelDrain := context.WithTimeout(context.Background(), 5*time.Second) - if err := obs.Observer.Drain(drainCtx); err != nil { - logrus.WithError(err).Warn("Telemetry drain incomplete") + shutdownErr := shutdownHTTPServer(srv, shutdownTimeout) + + // Shutdown makes serve return http.ErrServerClosed. Wait for it so a real + // error that raced the signal is reported instead of lost. + var failure error + if err := <-serveErr; err != nil && !errors.Is(err, http.ErrServerClosed) { + failure = fmt.Errorf("API server failed: %w", err) + } + + return errors.Join(failure, shutdownErr) + + case err := <-serveErr: + // Nobody asked the server to stop, so even http.ErrServerClosed or a nil + // return means it died on its own. + if err == nil { + err = errors.New("serve returned without an error") } - cancelDrain() - case err := <-serverErrChan: - logrus.Fatalf("Server terminated unexpectedly: %v", err) + failure := fmt.Errorf("API server failed: %w", err) + + // Connections accepted before the failure are still being served; give + // them the same bounded shutdown before their resources are released. + return errors.Join(failure, shutdownHTTPServer(srv, shutdownTimeout)) } +} + +// shutdownHTTPServer gracefully stops srv within timeout and force-closes what +// remains if that fails. Shutdown does not wait for hijacked connections +// (WebSockets), nor does it stop background services or running flows. +func shutdownHTTPServer(srv httpServer, timeout time.Duration) error { + ctx, cancel := context.WithTimeout(context.Background(), timeout) + defer cancel() + + err := srv.Shutdown(ctx) + if err == nil { + return nil + } + + logrus.WithError(err).Warn("Graceful HTTP shutdown failed, closing remaining connections") + shutdownErr := fmt.Errorf("graceful HTTP shutdown: %w", err) - logrus.Info("Application shutdown completed successfully") + if closeErr := srv.Close(); closeErr != nil { + logrus.WithError(closeErr).Error("Forced HTTP server close failed") + return errors.Join(shutdownErr, fmt.Errorf("force close HTTP server: %w", closeErr)) + } + + return shutdownErr +} + +// drainTelemetry flushes buffered telemetry on its own deadline. A failure is +// only logged: losing a final batch must not change the exit status. +func drainTelemetry(drain func(context.Context) error) { + ctx, cancel := context.WithTimeout(context.Background(), telemetryDrainTimeout) + defer cancel() + + if err := drain(ctx); err != nil { + logrus.WithError(err).Warn("Telemetry drain incomplete") + } } From ac513938cdf9bf064e04c466a107449cf3c53eac Mon Sep 17 00:00:00 2001 From: Yaash Singhh Date: Mon, 26 Oct 2026 14:27:46 +0530 Subject: [PATCH 2/5] test(server): deadline expiry and failure exit status - Active request completes before graceful shutdown returns, and Close is not called. - Deadline expiry is reported (context.DeadlineExceeded), Close is called once, and the stuck connection is cut. - A forced-close error is reported alongside the shutdown error. - http.ErrServerClosed after a requested shutdown is not an error; an unrequested one is. A server error racing the signal is not lost. - Unexpected server failure shuts the server down, drains telemetry, runs deferred cleanup in that order, and returns exit code 1. - A real listen failure (address in use) returns 1 and still drains. - A re-executed child process confirms the real exit status is 1 and deferred cleanup ran before exit. No sleeps, external network, PostgreSQL, Docker or LLM providers. Signed-off-by: Yaash Singhh --- backend/cmd/pentagi/main_test.go | 539 +++++++++++++++++++++++++++++++ 1 file changed, 539 insertions(+) create mode 100644 backend/cmd/pentagi/main_test.go diff --git a/backend/cmd/pentagi/main_test.go b/backend/cmd/pentagi/main_test.go new file mode 100644 index 000000000..fdad9a617 --- /dev/null +++ b/backend/cmd/pentagi/main_test.go @@ -0,0 +1,539 @@ +package main + +import ( + "context" + "errors" + "fmt" + "io" + "net" + "net/http" + "net/http/httptest" + "os" + "os/exec" + "reflect" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +// testWaitTimeout only bounds how long a broken test can hang; a passing test +// never waits on it. +const testWaitTimeout = 10 * time.Second + +var errSimulatedServerFailure = errors.New("accept tcp: too many open files") + +// --- helpers --------------------------------------------------------------- + +type shutdownRecorder struct { + mu sync.Mutex + events []string +} + +func (r *shutdownRecorder) add(event string) { + r.mu.Lock() + defer r.mu.Unlock() + + r.events = append(r.events, event) +} + +func (r *shutdownRecorder) snapshot() []string { + r.mu.Lock() + defer r.mu.Unlock() + + return append([]string(nil), r.events...) +} + +func requireEvents(t *testing.T, got, want []string) { + t.Helper() + + if !reflect.DeepEqual(got, want) { + t.Fatalf("events = %v, want %v", got, want) + } +} + +// ctxSnapshot captures a context as the callee saw it. The production code +// cancels its contexts on return, so they must be inspected during the call. +type ctxSnapshot struct { + err error + deadline time.Time + hasDeadline bool + taken time.Time +} + +func snapshotCtx(ctx context.Context) ctxSnapshot { + deadline, ok := ctx.Deadline() + + return ctxSnapshot{err: ctx.Err(), deadline: deadline, hasDeadline: ok, taken: time.Now()} +} + +// requireFresh asserts the context was live (not derived from the cancelled +// signal context) and bounded by a deadline no longer than limit. +func (c ctxSnapshot) requireFresh(t *testing.T, name string, limit time.Duration) { + t.Helper() + + if c.err != nil { + t.Errorf("%s context was already done: %v", name, c.err) + } + if !c.hasDeadline { + t.Errorf("%s context has no deadline", name) + return + } + if remaining := c.deadline.Sub(c.taken); remaining <= 0 || remaining > limit { + t.Errorf("%s deadline is %v away, want within (0, %v]", name, remaining, limit) + } +} + +// stubHTTPServer lets a test decide how Shutdown and Close behave. +type stubHTTPServer struct { + shutdownFn func(ctx context.Context) error + closeFn func() error +} + +func (s *stubHTTPServer) Shutdown(ctx context.Context) error { return s.shutdownFn(ctx) } +func (s *stubHTTPServer) Close() error { return s.closeFn() } + +// newGracefulStub returns a server whose serve func blocks until it is shut +// down or closed and then reports http.ErrServerClosed, like net/http does. +// Shutdown runs on the caller's goroutine, so *shutdownCtx is safe to read once +// the call under test has returned. +func newGracefulStub(emit func(string)) (srv *stubHTTPServer, serve func() error, shutdownCtx *ctxSnapshot) { + stopped := make(chan struct{}) + var once sync.Once + stop := func() { once.Do(func() { close(stopped) }) } + + shutdownCtx = &ctxSnapshot{} + srv = &stubHTTPServer{ + shutdownFn: func(ctx context.Context) error { + *shutdownCtx = snapshotCtx(ctx) + emit("shutdown") + stop() + return nil + }, + closeFn: func() error { + emit("close") + stop() + return nil + }, + } + serve = func() error { + <-stopped + return http.ErrServerClosed + } + + return srv, serve, shutdownCtx +} + +// spyHTTPServer wraps a real *http.Server and observes how it is stopped. +type spyHTTPServer struct { + *http.Server + + shutdownStarted chan struct{} + startOnce sync.Once + closeCalls atomic.Int32 +} + +func (s *spyHTTPServer) Shutdown(ctx context.Context) error { + s.startOnce.Do(func() { close(s.shutdownStarted) }) + return s.Server.Shutdown(ctx) +} + +func (s *spyHTTPServer) Close() error { + s.closeCalls.Add(1) + return s.Server.Close() +} + +// blockedHandler holds every request until it is released. +type blockedHandler struct { + entered chan struct{} + release chan struct{} + enterOnce sync.Once + releaseOnce sync.Once +} + +func (h *blockedHandler) ServeHTTP(w http.ResponseWriter, _ *http.Request) { + h.enterOnce.Do(func() { close(h.entered) }) + <-h.release + _, _ = io.WriteString(w, "done") +} + +func (h *blockedHandler) unblock() { h.releaseOnce.Do(func() { close(h.release) }) } + +// startBlockedServer prepares a real loopback server whose handler blocks. +// The server is not serving yet: the returned serve func starts it, so it can +// be handed to superviseHTTPServer exactly as production does. +func startBlockedServer(t *testing.T) (spy *spyHTTPServer, serve func() error, h *blockedHandler, url string) { + t.Helper() + + h = &blockedHandler{entered: make(chan struct{}), release: make(chan struct{})} + ts := httptest.NewUnstartedServer(h) + + // Cleanups run last-in first-out: release the handler, then close. + t.Cleanup(ts.Close) + t.Cleanup(h.unblock) + + spy = &spyHTTPServer{Server: ts.Config, shutdownStarted: make(chan struct{})} + serve = func() error { return ts.Config.Serve(ts.Listener) } + + return spy, serve, h, "http://" + ts.Listener.Addr().String() +} + +type requestResult struct { + status int + body string + err error +} + +// startRequest issues a GET on its own connection and reports the outcome. +func startRequest(url string) <-chan requestResult { + result := make(chan requestResult, 1) + + go func() { + client := &http.Client{ + Timeout: testWaitTimeout, + Transport: &http.Transport{DisableKeepAlives: true}, + } + + resp, err := client.Get(url) + if err != nil { + result <- requestResult{err: err} + return + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + result <- requestResult{status: resp.StatusCode, body: string(body), err: err} + }() + + return result +} + +func waitClosed(t *testing.T, ch <-chan struct{}, what string) { + t.Helper() + + select { + case <-ch: + case <-time.After(testWaitTimeout): + t.Fatalf("timed out waiting for %s", what) + } +} + +func waitError(t *testing.T, ch <-chan error, what string) error { + t.Helper() + + select { + case err := <-ch: + return err + case <-time.After(testWaitTimeout): + t.Fatalf("timed out waiting for %s", what) + return nil + } +} + +func waitResult(t *testing.T, ch <-chan requestResult, what string) requestResult { + t.Helper() + + select { + case res := <-ch: + return res + case <-time.After(testWaitTimeout): + t.Fatalf("timed out waiting for %s", what) + return requestResult{} + } +} + +// --- graceful shutdown against a real http.Server --------------------------- + +func TestSuperviseHTTPServer_ActiveRequestCompletes(t *testing.T) { + spy, serve, handler, url := startBlockedServer(t) + + ctx, sendSignal := context.WithCancel(context.Background()) + defer sendSignal() + + done := make(chan error, 1) + go func() { done <- superviseHTTPServer(ctx, spy, serve, httpShutdownTimeout) }() + + response := startRequest(url) + waitClosed(t, handler.entered, "the request to reach the handler") + + sendSignal() // stands in for SIGTERM + waitClosed(t, spy.shutdownStarted, "Shutdown to be called") + + // The handler is still running, so a graceful Shutdown cannot be done yet. + select { + case err := <-done: + t.Fatalf("shutdown returned while a request was in flight: %v", err) + default: + } + + handler.unblock() + + if err := waitError(t, done, "graceful shutdown to finish"); err != nil { + t.Fatalf("graceful shutdown failed: %v", err) + } + + res := waitResult(t, response, "the in-flight response") + if res.err != nil || res.status != http.StatusOK || res.body != "done" { + t.Fatalf("in-flight request = (status %d, body %q, err %v), want it to complete normally", + res.status, res.body, res.err) + } + + if n := spy.closeCalls.Load(); n != 0 { + t.Errorf("Close called %d times after a successful graceful shutdown, want 0", n) + } +} + +func TestSuperviseHTTPServer_DeadlineExpiryForcesClose(t *testing.T) { + spy, serve, handler, url := startBlockedServer(t) + + ctx, sendSignal := context.WithCancel(context.Background()) + defer sendSignal() + + // The handler is never released before the deadline, so expiry is certain + // whatever the machine's speed. + done := make(chan error, 1) + go func() { done <- superviseHTTPServer(ctx, spy, serve, 100*time.Millisecond) }() + + response := startRequest(url) + waitClosed(t, handler.entered, "the request to reach the handler") + + sendSignal() + + err := waitError(t, done, "shutdown to give up") + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("err = %v, want it to report context.DeadlineExceeded", err) + } + + if n := spy.closeCalls.Load(); n != 1 { + t.Errorf("Close called %d times, want exactly 1 forced close", n) + } + + // The forced close must really cut the stuck connection. + if res := waitResult(t, response, "the stuck request to be cut"); res.err == nil { + t.Errorf("stuck request finished with status %d, want its connection closed", res.status) + } +} + +func TestShutdownHTTPServer_ForceCloseErrorIsReported(t *testing.T) { + closeErr := errors.New("close failed") + srv := &stubHTTPServer{ + shutdownFn: func(context.Context) error { return context.DeadlineExceeded }, + closeFn: func() error { return closeErr }, + } + + err := shutdownHTTPServer(srv, httpShutdownTimeout) + + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("err = %v, want it to keep the graceful-shutdown failure", err) + } + if !errors.Is(err, closeErr) { + t.Errorf("err = %v, want it to include the forced-close failure", err) + } +} + +// --- expected closure vs. real failure -------------------------------------- + +func TestSuperviseHTTPServer_ErrServerClosedAfterShutdownIsExpected(t *testing.T) { + rec := &shutdownRecorder{} + srv, serve, shutdownCtx := newGracefulStub(rec.add) + + ctx, sendSignal := context.WithCancel(context.Background()) + sendSignal() // the signal context is already cancelled when shutdown starts + + if err := superviseHTTPServer(ctx, srv, serve, httpShutdownTimeout); err != nil { + t.Fatalf("http.ErrServerClosed after a requested shutdown must not be an error, got: %v", err) + } + + requireEvents(t, rec.snapshot(), []string{"shutdown"}) + shutdownCtx.requireFresh(t, "Shutdown", httpShutdownTimeout) +} + +func TestSuperviseHTTPServer_ErrServerClosedWithoutShutdownIsUnexpected(t *testing.T) { + rec := &shutdownRecorder{} + srv, _, _ := newGracefulStub(rec.add) + serve := func() error { return http.ErrServerClosed } + + // Nothing asked the server to stop, so its exit is a failure, not success. + err := superviseHTTPServer(context.Background(), srv, serve, httpShutdownTimeout) + + if !errors.Is(err, http.ErrServerClosed) { + t.Fatalf("err = %v, want the unrequested closure reported", err) + } +} + +func TestSuperviseHTTPServer_ServerErrorRacingSignalIsNotLost(t *testing.T) { + stopped := make(chan struct{}) + srv := &stubHTTPServer{ + shutdownFn: func(context.Context) error { close(stopped); return nil }, + closeFn: func() error { return nil }, + } + serve := func() error { + <-stopped + return errSimulatedServerFailure + } + + ctx, sendSignal := context.WithCancel(context.Background()) + sendSignal() + + err := superviseHTTPServer(ctx, srv, serve, httpShutdownTimeout) + + if !errors.Is(err, errSimulatedServerFailure) { + t.Fatalf("err = %v, want the genuine server error reported", err) + } +} + +func TestSuperviseHTTPServer_UnexpectedFailureIsReportedAndServerStopped(t *testing.T) { + rec := &shutdownRecorder{} + srv, _, shutdownCtx := newGracefulStub(rec.add) + serve := func() error { return errSimulatedServerFailure } + + // context.Background() is never cancelled: this is the failure path, not a + // signal. + err := superviseHTTPServer(context.Background(), srv, serve, httpShutdownTimeout) + + if !errors.Is(err, errSimulatedServerFailure) { + t.Fatalf("err = %v, want the server failure propagated", err) + } + + requireEvents(t, rec.snapshot(), []string{"shutdown"}) + shutdownCtx.requireFresh(t, "Shutdown", httpShutdownTimeout) +} + +// --- exit status and cleanup ordering --------------------------------------- + +func TestServeUntilShutdown_SignalShutsDownThenDrains(t *testing.T) { + rec := &shutdownRecorder{} + srv, serve, _ := newGracefulStub(rec.add) + + var drainCtx ctxSnapshot + drain := func(ctx context.Context) error { + drainCtx = snapshotCtx(ctx) + rec.add("drain") + return nil + } + + ctx, sendSignal := context.WithCancel(context.Background()) + sendSignal() + + if code := serveUntilShutdown(ctx, srv, serve, httpShutdownTimeout, drain); code != 0 { + t.Fatalf("exit code = %d, want 0 after a clean signal shutdown", code) + } + + // HTTP first, so requests finishing during shutdown can still emit telemetry. + requireEvents(t, rec.snapshot(), []string{"shutdown", "drain"}) + drainCtx.requireFresh(t, "telemetry drain", telemetryDrainTimeout) +} + +func TestServeUntilShutdown_DrainFailureDoesNotChangeExitCode(t *testing.T) { + rec := &shutdownRecorder{} + srv, serve, _ := newGracefulStub(rec.add) + drain := func(context.Context) error { return errors.New("collector unreachable") } + + ctx, sendSignal := context.WithCancel(context.Background()) + sendSignal() + + if code := serveUntilShutdown(ctx, srv, serve, httpShutdownTimeout, drain); code != 0 { + t.Fatalf("exit code = %d, want 0: an incomplete drain is only logged", code) + } +} + +// failureScenario has the shape of runApplication: cleanup is deferred, and the +// exit code comes back from serveUntilShutdown instead of from a Fatal call. +// The server fails on its own, with no signal. +func failureScenario(emit func(string)) int { + defer emit("cleanup") + + srv := &stubHTTPServer{ + shutdownFn: func(context.Context) error { emit("shutdown"); return nil }, + closeFn: func() error { emit("close"); return nil }, + } + serve := func() error { return errSimulatedServerFailure } + drain := func(context.Context) error { emit("drain"); return nil } + + return serveUntilShutdown(context.Background(), srv, serve, httpShutdownTimeout, drain) +} + +// This test finishing at all is part of the proof: the old logrus.Fatalf would +// have terminated the test binary here. +func TestServeUntilShutdown_UnexpectedFailureCleansUpAndReturnsFailure(t *testing.T) { + rec := &shutdownRecorder{} + + code := failureScenario(rec.add) + + if code != 1 { + t.Fatalf("exit code = %d, want 1", code) + } + + // HTTP stopped, telemetry drained, and only then the deferred cleanup. + requireEvents(t, rec.snapshot(), []string{"shutdown", "drain", "cleanup"}) +} + +func TestServeUntilShutdown_ListenFailureOnRealServer(t *testing.T) { + occupied, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + defer occupied.Close() + + srv := &http.Server{Addr: occupied.Addr().String(), Handler: http.NotFoundHandler()} + + // The timeout turns a (wrongly) successful bind into a failed assertion + // instead of a hang. + ctx, cancel := context.WithTimeout(context.Background(), testWaitTimeout) + defer cancel() + + drained := false + drain := func(context.Context) error { drained = true; return nil } + + code := serveUntilShutdown(ctx, srv, srv.ListenAndServe, httpShutdownTimeout, drain) + + if code != 1 { + t.Errorf("exit code = %d, want 1 when the address is already in use", code) + } + if !drained { + t.Error("telemetry was not drained after the listener failed") + } +} + +const failureHelperEnv = "PENTAGI_TEST_SERVER_FAILURE_HELPER" + +// TestServerFailureHelperProcess is not a test of its own: it is the child +// process of TestUnexpectedServerFailureExitStatus, and skips otherwise. +func TestServerFailureHelperProcess(t *testing.T) { + if os.Getenv(failureHelperEnv) != "1" { + t.Skip("helper process for TestUnexpectedServerFailureExitStatus") + } + + // The same pattern as main: exit only with the code a returned function chose. + os.Exit(failureScenario(func(event string) { fmt.Println("EVENT:" + event) })) +} + +// Proves the real process exit status, not just a returned integer: the child +// must exit non-zero, and its deferred cleanup must have run before it did. +// It covers serveUntilShutdown and the os.Exit(run()) pattern; runApplication +// itself needs PostgreSQL and Docker and is not executed here. +func TestUnexpectedServerFailureExitStatus(t *testing.T) { + cmd := exec.Command(os.Args[0], "-test.run=^TestServerFailureHelperProcess$") + cmd.Env = append(os.Environ(), failureHelperEnv+"=1") + + out, err := cmd.Output() + + var exitErr *exec.ExitError + if !errors.As(err, &exitErr) { + t.Fatalf("child err = %v, want a non-zero exit\nstdout:\n%s", err, out) + } + if got := exitErr.ExitCode(); got != 1 { + t.Fatalf("child exit status = %d, want 1\nstdout:\n%s\nstderr:\n%s", got, out, exitErr.Stderr) + } + + var events []string + for _, line := range strings.Split(string(out), "\n") { + line = strings.TrimSuffix(line, "\r") + if strings.HasPrefix(line, "EVENT:") { + events = append(events, strings.TrimPrefix(line, "EVENT:")) + } + } + + requireEvents(t, events, []string{"shutdown", "drain", "cleanup"}) +} From cda783dddc4bf30db606f6c090717bf4a382f243 Mon Sep 17 00:00:00 2001 From: Yaash Singhh Date: Mon, 26 Oct 2026 14:36:46 +0530 Subject: [PATCH 3/5] fix(server): shut down HTTP server gracefully and keep exit status on failure The server goroutine was never shut down explicitly, and an unexpected server error called logrus.Fatalf, which exits immediately and skips deferred cleanup. - Keep the http.Server reachable and call Shutdown on a fresh 15s context (the signal context is already cancelled). If it fails or times out, log it and force-close with Close, reporting any close error. - Treat http.ErrServerClosed as expected only after a requested shutdown; any other server exit is a failure. - Shut down HTTP before draining telemetry (fresh 5s context) so requests finishing during shutdown can still emit it. - Replace Fatalf with a returned exit code: main now calls os.Exit(runApplication()), so deferred cleanup runs before a non-zero exit. - Close the shared sql.DB pool via defer, after the HTTP server stops. - Cancel the signal context after shutdown so context-bound background work stops before the pools close. Shutdown does not stop running flows or background services, and does not wait for hijacked (WebSocket) connections. Signed-off-by: Yaash Singhh --- backend/cmd/pentagi/main.go | 713 +++++++++++++++++++++--------------- 1 file changed, 408 insertions(+), 305 deletions(-) diff --git a/backend/cmd/pentagi/main.go b/backend/cmd/pentagi/main.go index ab852119d..01a47dbb1 100644 --- a/backend/cmd/pentagi/main.go +++ b/backend/cmd/pentagi/main.go @@ -2,403 +2,506 @@ package main import ( "context" - "database/sql" "errors" "fmt" + "io" "net" "net/http" + "net/http/httptest" "os" - "os/signal" - "strconv" - "syscall" + "os/exec" + "reflect" + "strings" + "sync" + "sync/atomic" + "testing" "time" - // The runtime image is alpine without the tzdata package, so IANA names the - // dashboard sends would not resolve unless the binary carries the database. - _ "time/tzdata" - - "pentagi/migrations" - "pentagi/pkg/config" - "pentagi/pkg/controller" - "pentagi/pkg/database" - "pentagi/pkg/docker" - "pentagi/pkg/graph/subscriptions" - obs "pentagi/pkg/observability" - "pentagi/pkg/observability/profiling" - "pentagi/pkg/providers" - router "pentagi/pkg/server" - "pentagi/pkg/server/update" - "pentagi/pkg/version" - - "github.com/jackc/pgx/v5/pgxpool" - _ "github.com/lib/pq" - "github.com/pressly/goose/v3" - "github.com/sirupsen/logrus" - "go.opentelemetry.io/otel/attribute" ) -// HTTP shutdown and telemetry drain run on fresh contexts: the signal context -// is already cancelled by the time either starts. const ( - httpShutdownTimeout = 15 * time.Second - telemetryDrainTimeout = 5 * time.Second + testWaitTimeout = 10 * time.Second + testShutdownTimeout = 15 * time.Second + testDrainTimeout = 5 * time.Second ) -func main() { - // os.Exit skips deferred calls, so every deferred cleanup lives in - // runApplication and the exit status is only decided once it has returned. - os.Exit(runApplication()) +var errSimulatedServerFailure = errors.New("accept tcp: too many open files") + +type shutdownRecorder struct { + mu sync.Mutex + events []string +} + +func (r *shutdownRecorder) add(event string) { + r.mu.Lock() + defer r.mu.Unlock() + + r.events = append(r.events, event) } -func runApplication() int { - // Setup graceful shutdown context with signal handling - ctx, cancelOnSignal := signal.NotifyContext( - context.Background(), - os.Interrupt, - syscall.SIGTERM, - syscall.SIGQUIT, - ) - defer cancelOnSignal() +func (r *shutdownRecorder) snapshot() []string { + r.mu.Lock() + defer r.mu.Unlock() - // Before anything is logged or started. `-info` describes this server and - // exits, and its output is a JSON document on stdout — a greeting printed - // first would be the first thing a reader has to strip back out. - if wantsInfo(os.Args[1:]) { - return runInfo(ctx) + return append([]string(nil), r.events...) +} + +func requireEvents(t *testing.T, got, want []string) { + t.Helper() + + if !reflect.DeepEqual(got, want) { + t.Fatalf("events = %v, want %v", got, want) } +} - logrus.Infof("Starting PentAGI %s", version.GetBinaryVersion()) +type ctxSnapshot struct { + err error + deadline time.Time + hasDeadline bool + taken time.Time +} - cfg, err := config.NewConfig() - if err != nil { - // Fatal rather than Error: cfg is nil on failure, so continuing would - // panic anyway, and an invalid TENANT_ID must never boot into a namespace - // it did not intend to own. - logrus.WithError(err).Fatal("Unable to load config") - } - - // Configure logrus log level based on DEBUG env variable - if cfg.Debug { - logrus.SetLevel(logrus.DebugLevel) - logrus.Debug("Debug logging enabled") - } else { - logrus.SetLevel(logrus.InfoLevel) - } - - // Record the identity of this instance for the record: which tenant namespace - // it owns and which data directory it writes to. Both are operator-supplied - // and must differ between instances that share a host; having them in the - // startup log makes an accidental overlap obvious after the fact. - logrus.WithFields(logrus.Fields{ - "tenant_id": cfg.TenantID, - "data_dir": cfg.DataDir, - "schema": cfg.SchemaName(), - "installation_id": cfg.InstallationID, - }).Info("Instance identity") - - // Telemetry is optional — degrade to a no-op observer on init failure so an - // unreachable collector can't take the app down. - lfclient, err := obs.NewLangfuseClient(ctx, cfg) - if err != nil && !errors.Is(err, obs.ErrNotConfigured) { - logrus.WithError(err).Warn("langfuse telemetry disabled: client init failed") - lfclient = nil - } - - otelclient, err := obs.NewTelemetryClient(ctx, cfg) - if err != nil && !errors.Is(err, obs.ErrNotConfigured) { - logrus.WithError(err).Warn("opentelemetry disabled: client init failed") - otelclient = nil - } - - obs.InitObserver(ctx, lfclient, otelclient, []logrus.Level{ - logrus.DebugLevel, - logrus.InfoLevel, - logrus.WarnLevel, - logrus.ErrorLevel, - }) - - _ = obs.Observer.StartProcessMetricCollect(attribute.String("component", "server")) - _ = obs.Observer.StartGoRuntimeMetricCollect(attribute.String("component", "server")) - - // Create this tenant's schema and repoint DATABASE_URL at it before any - // consumer reads the DSN. No-op when TENANT_ID is empty. - if err := database.EnsureTenantSchema(ctx, cfg); err != nil { - logrus.WithError(err).Fatal("Tenant schema initialization failed") - } - - boundedURL, err := database.WithStatementTimeout(cfg.DatabaseURL, database.StatementTimeout) - if err != nil { - logrus.WithError(err).Fatal("Unable to bound database statements") +// Snapshot at call time: production code cancels these contexts on return. +func snapshotCtx(ctx context.Context) ctxSnapshot { + deadline, ok := ctx.Deadline() + + return ctxSnapshot{err: ctx.Err(), deadline: deadline, hasDeadline: ok, taken: time.Now()} +} + +func (c ctxSnapshot) requireFresh(t *testing.T, name string, limit time.Duration) { + t.Helper() + + if c.err != nil { + t.Errorf("%s context was already done: %v", name, c.err) + } + if !c.hasDeadline { + t.Errorf("%s context has no deadline", name) + return } + if remaining := c.deadline.Sub(c.taken); remaining <= 0 || remaining > limit { + t.Errorf("%s deadline is %v away, want within (0, %v]", name, remaining, limit) + } +} - db, err := sql.Open("postgres", boundedURL) - if err != nil { - logrus.WithError(err).Fatal("Unable to open database") +type stubHTTPServer struct { + shutdownFn func(ctx context.Context) error + closeFn func() error +} + +func (s *stubHTTPServer) Shutdown(ctx context.Context) error { return s.shutdownFn(ctx) } +func (s *stubHTTPServer) Close() error { return s.closeFn() } + +func newGracefulStub(emit func(string)) (srv *stubHTTPServer, serve func() error, shutdownCtx *ctxSnapshot) { + stopped := make(chan struct{}) + var once sync.Once + stop := func() { once.Do(func() { close(stopped) }) } + + shutdownCtx = &ctxSnapshot{} + srv = &stubHTTPServer{ + shutdownFn: func(ctx context.Context) error { + *shutdownCtx = snapshotCtx(ctx) + emit("shutdown") + stop() + return nil + }, + closeFn: func() error { + emit("close") + stop() + return nil + }, + } + serve = func() error { + <-stopped + return http.ErrServerClosed } - db.SetMaxOpenConns(cfg.DBMaxOpenConns) - db.SetMaxIdleConns(cfg.DBMaxIdleConns) - db.SetConnMaxLifetime(time.Hour) + return srv, serve, shutdownCtx +} + +type spyHTTPServer struct { + *http.Server + + shutdownStarted chan struct{} + startOnce sync.Once + closeCalls atomic.Int32 +} + +func (s *spyHTTPServer) Shutdown(ctx context.Context) error { + s.startOnce.Do(func() { close(s.shutdownStarted) }) + return s.Server.Shutdown(ctx) +} + +func (s *spyHTTPServer) Close() error { + s.closeCalls.Add(1) + return s.Server.Close() +} + +type blockedHandler struct { + entered chan struct{} + release chan struct{} + enterOnce sync.Once + releaseOnce sync.Once +} + +func (h *blockedHandler) ServeHTTP(w http.ResponseWriter, _ *http.Request) { + h.enterOnce.Do(func() { close(h.entered) }) + <-h.release + _, _ = io.WriteString(w, "done") +} + +func (h *blockedHandler) unblock() { h.releaseOnce.Do(func() { close(h.release) }) } + +func startBlockedServer(t *testing.T) (spy *spyHTTPServer, serve func() error, h *blockedHandler, url string) { + t.Helper() + + h = &blockedHandler{entered: make(chan struct{}), release: make(chan struct{})} + ts := httptest.NewUnstartedServer(h) + + t.Cleanup(ts.Close) + t.Cleanup(h.unblock) + + spy = &spyHTTPServer{Server: ts.Config, shutdownStarted: make(chan struct{})} + serve = func() error { return ts.Config.Serve(ts.Listener) } + + return spy, serve, h, "http://" + ts.Listener.Addr().String() +} + +type requestResult struct { + status int + body string + err error +} + +func startRequest(url string) <-chan requestResult { + result := make(chan requestResult, 1) - // One pool shared by sqlc and GORM, closed here exactly once. Deferred, so - // it runs only after the HTTP server has stopped (and after the pgxpool - // defer below, which was registered later). - defer func() { - if err := db.Close(); err != nil { - logrus.WithError(err).Warn("failed to close the database connection pool") + go func() { + client := &http.Client{ + Timeout: testWaitTimeout, + Transport: &http.Transport{DisableKeepAlives: true}, } + + resp, err := client.Get(url) + if err != nil { + result <- requestResult{err: err} + return + } + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + result <- requestResult{status: resp.StatusCode, body: string(body), err: err} }() - if err := database.VerifySearchPath(ctx, db, cfg); err != nil { - logrus.WithError(err).Fatal("Tenant schema verification failed") + return result +} + +func waitClosed(t *testing.T, ch <-chan struct{}, what string) { + t.Helper() + + select { + case <-ch: + case <-time.After(testWaitTimeout): + t.Fatalf("timed out waiting for %s", what) } +} - queries := database.New(db) +func waitError(t *testing.T, ch <-chan error, what string) error { + t.Helper() - // Pass the same *sql.DB to GORM so both sqlc and GORM share one connection - // pool. Previously each opened its own *sql.DB, consuming up to 40 - // Postgres connections. Now together they consume at most DBMaxOpenConns. - orm, err := database.NewGorm(db, cfg.Debug) - if err != nil { - logrus.WithError(err).Fatal("Unable to open database with gorm") + select { + case err := <-ch: + return err + case <-time.After(testWaitTimeout): + t.Fatalf("timed out waiting for %s", what) + return nil + } +} + +func waitResult(t *testing.T, ch <-chan requestResult, what string) requestResult { + t.Helper() + + select { + case res := <-ch: + return res + case <-time.After(testWaitTimeout): + t.Fatalf("timed out waiting for %s", what) + return requestResult{} } +} - // Create a shared pgxpool for all pgvector stores so that each executor - // reuses pooled connections instead of opening a dedicated pgx.Connect. - pgPoolConfig, err := pgxpool.ParseConfig(boundedURL) - if err != nil { - logrus.WithError(err).Fatal("Failed to parse pgxpool config") +func TestMain_Timeouts_MatchContract(t *testing.T) { + shutdown, drain := httpShutdownTimeout, telemetryDrainTimeout + + if shutdown != 15*time.Second { + t.Errorf("httpShutdownTimeout = %v, want 15s", shutdown) } - pgPoolConfig.MaxConns = int32(cfg.DBVectorMaxConns) - pgPool, err := pgxpool.NewWithConfig(ctx, pgPoolConfig) - if err != nil { - logrus.WithError(err).Fatal("Failed to create pgxpool") + if drain != 5*time.Second { + t.Errorf("telemetryDrainTimeout = %v, want 5s", drain) } - defer pgPool.Close() +} - // Attach the live pool to the config so router and tool executors can use - // pgvector.WithConn(cfg.PgxPool) without any interface changes. - cfg.PgxPool = pgPool +func TestMain_SuperviseHTTPServer_ActiveRequestCompletes(t *testing.T) { + spy, serve, handler, url := startBlockedServer(t) - goose.SetBaseFS(migrations.EmbedMigrations) + ctx, sendSignal := context.WithCancel(context.Background()) + defer sendSignal() - if err := goose.SetDialect("postgres"); err != nil { - logrus.WithError(err).Fatal("Database dialect configuration failed") - } + done := make(chan error, 1) + go func() { done <- superviseHTTPServer(ctx, spy, serve, testShutdownTimeout) }() - // goose's own queries are unqualified, so without this a fresh tenant - // schema silently inherits public's version table via search_path and - // skips its migrations. See pkg/database/tenant.go for the - // schema/search_path setup. - goose.SetTableName(cfg.SchemaName() + ".goose_db_version") + response := startRequest(url) + waitClosed(t, handler.entered, "the request to reach the handler") - // Hold an advisory lock so simultaneous boots cannot execute the same - // migration set concurrently; the initial migration uses bare CREATE TABLE, - // so the loser would otherwise abort on "relation already exists". - // Migrations get a connection of their own, without the statement ceiling: - // one of them may legitimately run longer than any request would. - migrator, err := sql.Open("postgres", cfg.DatabaseURL) - if err != nil { - logrus.WithError(err).Fatal("Unable to open database for migrations") + sendSignal() + waitClosed(t, spy.shutdownStarted, "Shutdown to be called") + + select { + case err := <-done: + t.Fatalf("shutdown returned while a request was in flight: %v", err) + default: } - migrationErr := database.RunMigrations(ctx, migrator, cfg, func(db *sql.DB) error { - return goose.Up(db, "sql") - }) + handler.unblock() + + if err := waitError(t, done, "graceful shutdown to finish"); err != nil { + t.Fatalf("graceful shutdown failed: %v", err) + } - if err := migrator.Close(); err != nil { - logrus.WithError(err).Warn("failed to close the migration connection") + res := waitResult(t, response, "the in-flight response") + if res.err != nil || res.status != http.StatusOK || res.body != "done" { + t.Fatalf("in-flight request = (status %d, body %q, err %v), want it to complete normally", + res.status, res.body, res.err) } - if err := migrationErr; err != nil { - // Fatal: continuing on a half-migrated schema and then serving traffic is - // strictly worse than refusing to start. - logrus.WithError(err).Fatal("Schema migration execution failed") + if n := spy.closeCalls.Load(); n != 0 { + t.Errorf("Close called %d times after a successful graceful shutdown, want 0", n) } +} - logrus.Info("Database schema updated successfully") +func TestMain_SuperviseHTTPServer_DeadlineExpiryForcesClose(t *testing.T) { + spy, serve, handler, url := startBlockedServer(t) - // Keep the update service informed about this installation, periodically. - // After the migrations, deliberately: on a first boot the tables it reads do - // not exist until this point, and a summary of a schema that is not there yet - // describes an installation with nothing in it — which is also what an - // installation nobody uses looks like. - // - // A goroutine that owns nothing: it reads, it sends, and every failure along - // the way is a debug line. A server must not fail to start because something - // it only talks to is unreachable. What it learns is handed to the router, so - // the UI can show whether this build is current; nil means it never runs, and - // the UI says so. - var updates *update.Service - if built, err := update.New(cfg, queries); err != nil { - logrus.WithError(err).Debug("update service integration is not configured") - } else { - updates = built - go updates.Run(ctx) + ctx, sendSignal := context.WithCancel(context.Background()) + defer sendSignal() + + done := make(chan error, 1) + go func() { done <- superviseHTTPServer(ctx, spy, serve, 100*time.Millisecond) }() + + response := startRequest(url) + waitClosed(t, handler.entered, "the request to reach the handler") + + sendSignal() + + err := waitError(t, done, "shutdown to give up") + if !errors.Is(err, context.DeadlineExceeded) { + t.Fatalf("err = %v, want it to report context.DeadlineExceeded", err) } - if cfg.PprofAddr != "" { - go profiling.Start(cfg.PprofAddr) + if n := spy.closeCalls.Load(); n != 1 { + t.Errorf("Close called %d times, want exactly 1 forced close", n) } - client, err := docker.NewDockerClient(ctx, queries, cfg) - if err != nil { - logrus.WithError(err).Fatal("Docker runtime client initialization failed") + if res := waitResult(t, response, "the stuck request to be cut"); res.err == nil { + t.Errorf("stuck request finished with status %d, want its connection closed", res.status) } +} - providers, err := providers.NewProviderController(cfg, queries, client) - if err != nil { - logrus.WithError(err).Fatal("LLM provider controller initialization failed") +func TestMain_ShutdownHTTPServer_ForceCloseErrorIsReported(t *testing.T) { + closeErr := errors.New("close failed") + srv := &stubHTTPServer{ + shutdownFn: func(context.Context) error { return context.DeadlineExceeded }, + closeFn: func() error { return closeErr }, + } + + err := shutdownHTTPServer(srv, testShutdownTimeout) + + if !errors.Is(err, context.DeadlineExceeded) { + t.Errorf("err = %v, want it to keep the graceful-shutdown failure", err) } - subscriptions := subscriptions.NewSubscriptionsController() - controller := controller.NewFlowController(queries, cfg, client, providers, subscriptions) + if !errors.Is(err, closeErr) { + t.Errorf("err = %v, want it to include the forced-close failure", err) + } +} - if err := controller.LoadFlows(ctx); err != nil { - logrus.WithError(err).Fatal("Active flows restoration failed") +func TestMain_SuperviseHTTPServer_ErrServerClosedAfterShutdownIsExpected(t *testing.T) { + rec := &shutdownRecorder{} + srv, serve, shutdownCtx := newGracefulStub(rec.add) + + ctx, sendSignal := context.WithCancel(context.Background()) + + if err := superviseHTTPServer(ctx, srv, serve, testShutdownTimeout); err != nil { + t.Fatalf("http.ErrServerClosed after a requested shutdown must not be an error, got: %v", err) } - r := router.NewRouter(queries, orm, cfg, providers, controller, subscriptions, client, updates) + requireEvents(t, rec.snapshot(), []string{"shutdown"}) + shutdownCtx.requireFresh(t, "Shutdown", testShutdownTimeout) +} + +func TestMain_SuperviseHTTPServer_ErrServerClosedWithoutShutdownIsUnexpected(t *testing.T) { + rec := &shutdownRecorder{} + srv, _, _ := newGracefulStub(rec.add) + serve := func() error { return http.ErrServerClosed } - listen := net.JoinHostPort(cfg.ServerHost, strconv.Itoa(cfg.ServerPort)) - srv := &http.Server{ - Addr: listen, - Handler: r.Handler(), - ReadHeaderTimeout: router.ReadHeaderTimeout, - IdleTimeout: router.IdleTimeout, + err := superviseHTTPServer(context.Background(), srv, serve, testShutdownTimeout) + + if !errors.Is(err, http.ErrServerClosed) { + t.Fatalf("err = %v, want the unrequested closure reported", err) } +} +func TestMain_SuperviseHTTPServer_ServerErrorRacingSignalIsNotLost(t *testing.T) { + stopped := make(chan struct{}) + srv := &stubHTTPServer{ + shutdownFn: func(context.Context) error { close(stopped); return nil }, + closeFn: func() error { return nil }, + } serve := func() error { - logrus.Infof("API server listening on %s", listen) + <-stopped + return errSimulatedServerFailure + } - if cfg.ServerUseSSL && cfg.ServerSSLCrt != "" && cfg.ServerSSLKey != "" { - logrus.Info("Starting server with TLS enabled") - return srv.ListenAndServeTLS(cfg.ServerSSLCrt, cfg.ServerSSLKey) - } + ctx, sendSignal := context.WithCancel(context.Background()) + sendSignal() + + err := superviseHTTPServer(ctx, srv, serve, testShutdownTimeout) - logrus.Info("Starting server without TLS (HTTP only)") - return srv.ListenAndServe() + if !errors.Is(err, errSimulatedServerFailure) { + t.Fatalf("err = %v, want the genuine server error reported", err) } +} - // Blocks until a shutdown signal or a server failure, then stops the HTTP - // server and drains telemetry. The exit code is returned rather than raised - // with Fatal so the deferred pool cleanup above still runs. - exitCode := serveUntilShutdown(ctx, srv, serve, httpShutdownTimeout, func(ctx context.Context) error { - return obs.Observer.Drain(ctx) - }) +func TestMain_SuperviseHTTPServer_UnexpectedFailureIsReportedAndServerStopped(t *testing.T) { + rec := &shutdownRecorder{} + srv, _, shutdownCtx := newGracefulStub(rec.add) + serve := func() error { return errSimulatedServerFailure } - // On a server failure ctx is still live; cancel it so work bound to it (the - // update service, flows) stops before the deferred calls close the pools. - // Idempotent after a signal. - cancelOnSignal() + err := superviseHTTPServer(context.Background(), srv, serve, testShutdownTimeout) - if exitCode == 0 { - logrus.Info("Application shutdown completed successfully") + if !errors.Is(err, errSimulatedServerFailure) { + t.Fatalf("err = %v, want the server failure propagated", err) } - return exitCode + requireEvents(t, rec.snapshot(), []string{"shutdown"}) + shutdownCtx.requireFresh(t, "Shutdown", testShutdownTimeout) } -// httpServer is the part of *http.Server the shutdown logic needs, so tests can -// drive the deadline and failure paths without waiting on real timers. -type httpServer interface { - Shutdown(ctx context.Context) error - Close() error -} +func TestMain_ServeUntilShutdown_SignalShutsDownThenDrains(t *testing.T) { + rec := &shutdownRecorder{} + srv, serve, _ := newGracefulStub(rec.add) -// serveUntilShutdown runs serve until ctx is cancelled or serve fails, stops the -// HTTP server, and only then drains telemetry, so requests that finish during -// graceful shutdown can still emit it. It returns the process exit code and -// never exits itself, which keeps the caller's deferred cleanup reachable. -func serveUntilShutdown( - ctx context.Context, - srv httpServer, - serve func() error, - shutdownTimeout time.Duration, - drain func(context.Context) error, -) int { - err := superviseHTTPServer(ctx, srv, serve, shutdownTimeout) + var drainCtx ctxSnapshot + drain := func(ctx context.Context) error { + drainCtx = snapshotCtx(ctx) + rec.add("drain") + return nil + } - drainTelemetry(drain) + ctx, sendSignal := context.WithCancel(context.Background()) + sendSignal() - if err != nil { - logrus.WithError(err).Error("HTTP server terminated with an error") - return 1 + if code := serveUntilShutdown(ctx, srv, serve, testShutdownTimeout, drain); code != 0 { + t.Fatalf("exit code = %d, want 0 after a clean signal shutdown", code) } - return 0 + requireEvents(t, rec.snapshot(), []string{"shutdown", "drain"}) + drainCtx.requireFresh(t, "telemetry drain", testDrainTimeout) } -// superviseHTTPServer runs serve in the background and returns once the server -// has been stopped. It returns nil only for a shutdown it requested itself and -// that completed gracefully. -func superviseHTTPServer( - ctx context.Context, - srv httpServer, - serve func() error, - shutdownTimeout time.Duration, -) error { - // Buffered so the goroutine can always deliver its result and exit, even - // when nobody is left to receive it. - serveErr := make(chan error, 1) - go func() { serveErr <- serve() }() +func TestMain_ServeUntilShutdown_DrainFailureDoesNotChangeExitCode(t *testing.T) { + rec := &shutdownRecorder{} + srv, serve, _ := newGracefulStub(rec.add) + drain := func(context.Context) error { return errors.New("collector unreachable") } - select { - case <-ctx.Done(): - logrus.Warn("Shutdown signal received, cleaning up resources...") - shutdownErr := shutdownHTTPServer(srv, shutdownTimeout) - - // Shutdown makes serve return http.ErrServerClosed. Wait for it so a real - // error that raced the signal is reported instead of lost. - var failure error - if err := <-serveErr; err != nil && !errors.Is(err, http.ErrServerClosed) { - failure = fmt.Errorf("API server failed: %w", err) - } + ctx, sendSignal := context.WithCancel(context.Background()) + sendSignal() - return errors.Join(failure, shutdownErr) + if code := serveUntilShutdown(ctx, srv, serve, testShutdownTimeout, drain); code != 0 { + t.Fatalf("exit code = %d, want 0: an incomplete drain is only logged", code) + } +} - case err := <-serveErr: - // Nobody asked the server to stop, so even http.ErrServerClosed or a nil - // return means it died on its own. - if err == nil { - err = errors.New("serve returned without an error") - } - failure := fmt.Errorf("API server failed: %w", err) +func failureScenario(emit func(string)) int { + defer emit("cleanup") - // Connections accepted before the failure are still being served; give - // them the same bounded shutdown before their resources are released. - return errors.Join(failure, shutdownHTTPServer(srv, shutdownTimeout)) + srv := &stubHTTPServer{ + shutdownFn: func(context.Context) error { emit("shutdown"); return nil }, + closeFn: func() error { emit("close"); return nil }, } + serve := func() error { return errSimulatedServerFailure } + drain := func(context.Context) error { emit("drain"); return nil } + + return serveUntilShutdown(context.Background(), srv, serve, testShutdownTimeout, drain) +} + +func TestMain_ServeUntilShutdown_UnexpectedFailureCleansUpAndReturnsFailure(t *testing.T) { + rec := &shutdownRecorder{} + + code := failureScenario(rec.add) + + if code != 1 { + t.Fatalf("exit code = %d, want 1", code) + } + + requireEvents(t, rec.snapshot(), []string{"shutdown", "drain", "cleanup"}) } -// shutdownHTTPServer gracefully stops srv within timeout and force-closes what -// remains if that fails. Shutdown does not wait for hijacked connections -// (WebSockets), nor does it stop background services or running flows. -func shutdownHTTPServer(srv httpServer, timeout time.Duration) error { - ctx, cancel := context.WithTimeout(context.Background(), timeout) +func TestMain_ServeUntilShutdown_ListenFailureOnRealServer(t *testing.T) { + occupied, err := net.Listen("tcp", "127.0.0.1:0") + if err != nil { + t.Fatalf("listen: %v", err) + } + defer occupied.Close() + + srv := &http.Server{Addr: occupied.Addr().String(), Handler: http.NotFoundHandler()} + + ctx, cancel := context.WithTimeout(context.Background(), testWaitTimeout) defer cancel() - err := srv.Shutdown(ctx) - if err == nil { - return nil + drained := false + drain := func(context.Context) error { drained = true; return nil } + + code := serveUntilShutdown(ctx, srv, srv.ListenAndServe, testShutdownTimeout, drain) + + if code != 1 { + t.Errorf("exit code = %d, want 1 when the address is already in use", code) } + if !drained { + t.Error("telemetry was not drained after the listener failed") + } +} - logrus.WithError(err).Warn("Graceful HTTP shutdown failed, closing remaining connections") - shutdownErr := fmt.Errorf("graceful HTTP shutdown: %w", err) +const failureHelperEnv = "PENTAGI_TEST_SERVER_FAILURE_HELPER" - if closeErr := srv.Close(); closeErr != nil { - logrus.WithError(closeErr).Error("Forced HTTP server close failed") - return errors.Join(shutdownErr, fmt.Errorf("force close HTTP server: %w", closeErr)) +// Child process of TestMain_ServeUntilShutdown_UnexpectedFailureExitStatus. +func TestMain_ServeUntilShutdown_FailureHelperProcess(t *testing.T) { + if os.Getenv(failureHelperEnv) != "1" { + t.Skip("helper process for TestMain_ServeUntilShutdown_UnexpectedFailureExitStatus") } - return shutdownErr + os.Exit(failureScenario(func(event string) { fmt.Println("EVENT:" + event) })) } -// drainTelemetry flushes buffered telemetry on its own deadline. A failure is -// only logged: losing a final batch must not change the exit status. -func drainTelemetry(drain func(context.Context) error) { - ctx, cancel := context.WithTimeout(context.Background(), telemetryDrainTimeout) - defer cancel() +// runApplication itself needs PostgreSQL and Docker, so a stand-in with the same shape runs instead. +func TestMain_ServeUntilShutdown_UnexpectedFailureExitStatus(t *testing.T) { + cmd := exec.Command(os.Args[0], "-test.run=^TestMain_ServeUntilShutdown_FailureHelperProcess$") + cmd.Env = append(os.Environ(), failureHelperEnv+"=1") + + out, err := cmd.Output() - if err := drain(ctx); err != nil { - logrus.WithError(err).Warn("Telemetry drain incomplete") + var exitErr *exec.ExitError + if !errors.As(err, &exitErr) { + t.Fatalf("child err = %v, want a non-zero exit\nstdout:\n%s", err, out) } + if got := exitErr.ExitCode(); got != 1 { + t.Fatalf("child exit status = %d, want 1\nstdout:\n%s\nstderr:\n%s", got, out, exitErr.Stderr) + } + + var events []string + for _, line := range strings.Split(string(out), "\n") { + line = strings.TrimSuffix(line, "\r") + if strings.HasPrefix(line, "EVENT:") { + events = append(events, strings.TrimPrefix(line, "EVENT:")) + } + } + + requireEvents(t, events, []string{"shutdown", "drain", "cleanup"}) } From 3fcf33f2631efc582c690b975663a0dde9c964fe Mon Sep 17 00:00:00 2001 From: Yaash Singhh Date: Tue, 27 Oct 2026 23:22:35 +0530 Subject: [PATCH 4/5] fix(server): restore HTTP shutdown implementation Signed-off-by: Yaash Singhh --- backend/cmd/pentagi/main.go | 713 +++++++++++++++--------------------- 1 file changed, 305 insertions(+), 408 deletions(-) diff --git a/backend/cmd/pentagi/main.go b/backend/cmd/pentagi/main.go index 01a47dbb1..ab852119d 100644 --- a/backend/cmd/pentagi/main.go +++ b/backend/cmd/pentagi/main.go @@ -2,506 +2,403 @@ package main import ( "context" + "database/sql" "errors" "fmt" - "io" "net" "net/http" - "net/http/httptest" "os" - "os/exec" - "reflect" - "strings" - "sync" - "sync/atomic" - "testing" + "os/signal" + "strconv" + "syscall" "time" + // The runtime image is alpine without the tzdata package, so IANA names the + // dashboard sends would not resolve unless the binary carries the database. + _ "time/tzdata" + + "pentagi/migrations" + "pentagi/pkg/config" + "pentagi/pkg/controller" + "pentagi/pkg/database" + "pentagi/pkg/docker" + "pentagi/pkg/graph/subscriptions" + obs "pentagi/pkg/observability" + "pentagi/pkg/observability/profiling" + "pentagi/pkg/providers" + router "pentagi/pkg/server" + "pentagi/pkg/server/update" + "pentagi/pkg/version" + + "github.com/jackc/pgx/v5/pgxpool" + _ "github.com/lib/pq" + "github.com/pressly/goose/v3" + "github.com/sirupsen/logrus" + "go.opentelemetry.io/otel/attribute" ) +// HTTP shutdown and telemetry drain run on fresh contexts: the signal context +// is already cancelled by the time either starts. const ( - testWaitTimeout = 10 * time.Second - testShutdownTimeout = 15 * time.Second - testDrainTimeout = 5 * time.Second + httpShutdownTimeout = 15 * time.Second + telemetryDrainTimeout = 5 * time.Second ) -var errSimulatedServerFailure = errors.New("accept tcp: too many open files") - -type shutdownRecorder struct { - mu sync.Mutex - events []string -} - -func (r *shutdownRecorder) add(event string) { - r.mu.Lock() - defer r.mu.Unlock() - - r.events = append(r.events, event) +func main() { + // os.Exit skips deferred calls, so every deferred cleanup lives in + // runApplication and the exit status is only decided once it has returned. + os.Exit(runApplication()) } -func (r *shutdownRecorder) snapshot() []string { - r.mu.Lock() - defer r.mu.Unlock() +func runApplication() int { + // Setup graceful shutdown context with signal handling + ctx, cancelOnSignal := signal.NotifyContext( + context.Background(), + os.Interrupt, + syscall.SIGTERM, + syscall.SIGQUIT, + ) + defer cancelOnSignal() - return append([]string(nil), r.events...) -} - -func requireEvents(t *testing.T, got, want []string) { - t.Helper() - - if !reflect.DeepEqual(got, want) { - t.Fatalf("events = %v, want %v", got, want) + // Before anything is logged or started. `-info` describes this server and + // exits, and its output is a JSON document on stdout — a greeting printed + // first would be the first thing a reader has to strip back out. + if wantsInfo(os.Args[1:]) { + return runInfo(ctx) } -} -type ctxSnapshot struct { - err error - deadline time.Time - hasDeadline bool - taken time.Time -} - -// Snapshot at call time: production code cancels these contexts on return. -func snapshotCtx(ctx context.Context) ctxSnapshot { - deadline, ok := ctx.Deadline() - - return ctxSnapshot{err: ctx.Err(), deadline: deadline, hasDeadline: ok, taken: time.Now()} -} + logrus.Infof("Starting PentAGI %s", version.GetBinaryVersion()) -func (c ctxSnapshot) requireFresh(t *testing.T, name string, limit time.Duration) { - t.Helper() - - if c.err != nil { - t.Errorf("%s context was already done: %v", name, c.err) - } - if !c.hasDeadline { - t.Errorf("%s context has no deadline", name) - return - } - if remaining := c.deadline.Sub(c.taken); remaining <= 0 || remaining > limit { - t.Errorf("%s deadline is %v away, want within (0, %v]", name, remaining, limit) + cfg, err := config.NewConfig() + if err != nil { + // Fatal rather than Error: cfg is nil on failure, so continuing would + // panic anyway, and an invalid TENANT_ID must never boot into a namespace + // it did not intend to own. + logrus.WithError(err).Fatal("Unable to load config") + } + + // Configure logrus log level based on DEBUG env variable + if cfg.Debug { + logrus.SetLevel(logrus.DebugLevel) + logrus.Debug("Debug logging enabled") + } else { + logrus.SetLevel(logrus.InfoLevel) + } + + // Record the identity of this instance for the record: which tenant namespace + // it owns and which data directory it writes to. Both are operator-supplied + // and must differ between instances that share a host; having them in the + // startup log makes an accidental overlap obvious after the fact. + logrus.WithFields(logrus.Fields{ + "tenant_id": cfg.TenantID, + "data_dir": cfg.DataDir, + "schema": cfg.SchemaName(), + "installation_id": cfg.InstallationID, + }).Info("Instance identity") + + // Telemetry is optional — degrade to a no-op observer on init failure so an + // unreachable collector can't take the app down. + lfclient, err := obs.NewLangfuseClient(ctx, cfg) + if err != nil && !errors.Is(err, obs.ErrNotConfigured) { + logrus.WithError(err).Warn("langfuse telemetry disabled: client init failed") + lfclient = nil + } + + otelclient, err := obs.NewTelemetryClient(ctx, cfg) + if err != nil && !errors.Is(err, obs.ErrNotConfigured) { + logrus.WithError(err).Warn("opentelemetry disabled: client init failed") + otelclient = nil + } + + obs.InitObserver(ctx, lfclient, otelclient, []logrus.Level{ + logrus.DebugLevel, + logrus.InfoLevel, + logrus.WarnLevel, + logrus.ErrorLevel, + }) + + _ = obs.Observer.StartProcessMetricCollect(attribute.String("component", "server")) + _ = obs.Observer.StartGoRuntimeMetricCollect(attribute.String("component", "server")) + + // Create this tenant's schema and repoint DATABASE_URL at it before any + // consumer reads the DSN. No-op when TENANT_ID is empty. + if err := database.EnsureTenantSchema(ctx, cfg); err != nil { + logrus.WithError(err).Fatal("Tenant schema initialization failed") + } + + boundedURL, err := database.WithStatementTimeout(cfg.DatabaseURL, database.StatementTimeout) + if err != nil { + logrus.WithError(err).Fatal("Unable to bound database statements") } -} - -type stubHTTPServer struct { - shutdownFn func(ctx context.Context) error - closeFn func() error -} -func (s *stubHTTPServer) Shutdown(ctx context.Context) error { return s.shutdownFn(ctx) } -func (s *stubHTTPServer) Close() error { return s.closeFn() } - -func newGracefulStub(emit func(string)) (srv *stubHTTPServer, serve func() error, shutdownCtx *ctxSnapshot) { - stopped := make(chan struct{}) - var once sync.Once - stop := func() { once.Do(func() { close(stopped) }) } - - shutdownCtx = &ctxSnapshot{} - srv = &stubHTTPServer{ - shutdownFn: func(ctx context.Context) error { - *shutdownCtx = snapshotCtx(ctx) - emit("shutdown") - stop() - return nil - }, - closeFn: func() error { - emit("close") - stop() - return nil - }, - } - serve = func() error { - <-stopped - return http.ErrServerClosed + db, err := sql.Open("postgres", boundedURL) + if err != nil { + logrus.WithError(err).Fatal("Unable to open database") } - return srv, serve, shutdownCtx -} - -type spyHTTPServer struct { - *http.Server - - shutdownStarted chan struct{} - startOnce sync.Once - closeCalls atomic.Int32 -} - -func (s *spyHTTPServer) Shutdown(ctx context.Context) error { - s.startOnce.Do(func() { close(s.shutdownStarted) }) - return s.Server.Shutdown(ctx) -} - -func (s *spyHTTPServer) Close() error { - s.closeCalls.Add(1) - return s.Server.Close() -} - -type blockedHandler struct { - entered chan struct{} - release chan struct{} - enterOnce sync.Once - releaseOnce sync.Once -} - -func (h *blockedHandler) ServeHTTP(w http.ResponseWriter, _ *http.Request) { - h.enterOnce.Do(func() { close(h.entered) }) - <-h.release - _, _ = io.WriteString(w, "done") -} - -func (h *blockedHandler) unblock() { h.releaseOnce.Do(func() { close(h.release) }) } - -func startBlockedServer(t *testing.T) (spy *spyHTTPServer, serve func() error, h *blockedHandler, url string) { - t.Helper() - - h = &blockedHandler{entered: make(chan struct{}), release: make(chan struct{})} - ts := httptest.NewUnstartedServer(h) - - t.Cleanup(ts.Close) - t.Cleanup(h.unblock) - - spy = &spyHTTPServer{Server: ts.Config, shutdownStarted: make(chan struct{})} - serve = func() error { return ts.Config.Serve(ts.Listener) } - - return spy, serve, h, "http://" + ts.Listener.Addr().String() -} - -type requestResult struct { - status int - body string - err error -} - -func startRequest(url string) <-chan requestResult { - result := make(chan requestResult, 1) + db.SetMaxOpenConns(cfg.DBMaxOpenConns) + db.SetMaxIdleConns(cfg.DBMaxIdleConns) + db.SetConnMaxLifetime(time.Hour) - go func() { - client := &http.Client{ - Timeout: testWaitTimeout, - Transport: &http.Transport{DisableKeepAlives: true}, + // One pool shared by sqlc and GORM, closed here exactly once. Deferred, so + // it runs only after the HTTP server has stopped (and after the pgxpool + // defer below, which was registered later). + defer func() { + if err := db.Close(); err != nil { + logrus.WithError(err).Warn("failed to close the database connection pool") } - - resp, err := client.Get(url) - if err != nil { - result <- requestResult{err: err} - return - } - defer resp.Body.Close() - - body, err := io.ReadAll(resp.Body) - result <- requestResult{status: resp.StatusCode, body: string(body), err: err} }() - return result -} - -func waitClosed(t *testing.T, ch <-chan struct{}, what string) { - t.Helper() - - select { - case <-ch: - case <-time.After(testWaitTimeout): - t.Fatalf("timed out waiting for %s", what) - } -} - -func waitError(t *testing.T, ch <-chan error, what string) error { - t.Helper() - - select { - case err := <-ch: - return err - case <-time.After(testWaitTimeout): - t.Fatalf("timed out waiting for %s", what) - return nil + if err := database.VerifySearchPath(ctx, db, cfg); err != nil { + logrus.WithError(err).Fatal("Tenant schema verification failed") } -} -func waitResult(t *testing.T, ch <-chan requestResult, what string) requestResult { - t.Helper() + queries := database.New(db) - select { - case res := <-ch: - return res - case <-time.After(testWaitTimeout): - t.Fatalf("timed out waiting for %s", what) - return requestResult{} + // Pass the same *sql.DB to GORM so both sqlc and GORM share one connection + // pool. Previously each opened its own *sql.DB, consuming up to 40 + // Postgres connections. Now together they consume at most DBMaxOpenConns. + orm, err := database.NewGorm(db, cfg.Debug) + if err != nil { + logrus.WithError(err).Fatal("Unable to open database with gorm") } -} - -func TestMain_Timeouts_MatchContract(t *testing.T) { - shutdown, drain := httpShutdownTimeout, telemetryDrainTimeout - if shutdown != 15*time.Second { - t.Errorf("httpShutdownTimeout = %v, want 15s", shutdown) + // Create a shared pgxpool for all pgvector stores so that each executor + // reuses pooled connections instead of opening a dedicated pgx.Connect. + pgPoolConfig, err := pgxpool.ParseConfig(boundedURL) + if err != nil { + logrus.WithError(err).Fatal("Failed to parse pgxpool config") } - if drain != 5*time.Second { - t.Errorf("telemetryDrainTimeout = %v, want 5s", drain) + pgPoolConfig.MaxConns = int32(cfg.DBVectorMaxConns) + pgPool, err := pgxpool.NewWithConfig(ctx, pgPoolConfig) + if err != nil { + logrus.WithError(err).Fatal("Failed to create pgxpool") } -} - -func TestMain_SuperviseHTTPServer_ActiveRequestCompletes(t *testing.T) { - spy, serve, handler, url := startBlockedServer(t) - - ctx, sendSignal := context.WithCancel(context.Background()) - defer sendSignal() - - done := make(chan error, 1) - go func() { done <- superviseHTTPServer(ctx, spy, serve, testShutdownTimeout) }() + defer pgPool.Close() - response := startRequest(url) - waitClosed(t, handler.entered, "the request to reach the handler") + // Attach the live pool to the config so router and tool executors can use + // pgvector.WithConn(cfg.PgxPool) without any interface changes. + cfg.PgxPool = pgPool - sendSignal() - waitClosed(t, spy.shutdownStarted, "Shutdown to be called") + goose.SetBaseFS(migrations.EmbedMigrations) - select { - case err := <-done: - t.Fatalf("shutdown returned while a request was in flight: %v", err) - default: + if err := goose.SetDialect("postgres"); err != nil { + logrus.WithError(err).Fatal("Database dialect configuration failed") } - handler.unblock() + // goose's own queries are unqualified, so without this a fresh tenant + // schema silently inherits public's version table via search_path and + // skips its migrations. See pkg/database/tenant.go for the + // schema/search_path setup. + goose.SetTableName(cfg.SchemaName() + ".goose_db_version") - if err := waitError(t, done, "graceful shutdown to finish"); err != nil { - t.Fatalf("graceful shutdown failed: %v", err) + // Hold an advisory lock so simultaneous boots cannot execute the same + // migration set concurrently; the initial migration uses bare CREATE TABLE, + // so the loser would otherwise abort on "relation already exists". + // Migrations get a connection of their own, without the statement ceiling: + // one of them may legitimately run longer than any request would. + migrator, err := sql.Open("postgres", cfg.DatabaseURL) + if err != nil { + logrus.WithError(err).Fatal("Unable to open database for migrations") } - res := waitResult(t, response, "the in-flight response") - if res.err != nil || res.status != http.StatusOK || res.body != "done" { - t.Fatalf("in-flight request = (status %d, body %q, err %v), want it to complete normally", - res.status, res.body, res.err) - } + migrationErr := database.RunMigrations(ctx, migrator, cfg, func(db *sql.DB) error { + return goose.Up(db, "sql") + }) - if n := spy.closeCalls.Load(); n != 0 { - t.Errorf("Close called %d times after a successful graceful shutdown, want 0", n) + if err := migrator.Close(); err != nil { + logrus.WithError(err).Warn("failed to close the migration connection") } -} -func TestMain_SuperviseHTTPServer_DeadlineExpiryForcesClose(t *testing.T) { - spy, serve, handler, url := startBlockedServer(t) - - ctx, sendSignal := context.WithCancel(context.Background()) - defer sendSignal() - - done := make(chan error, 1) - go func() { done <- superviseHTTPServer(ctx, spy, serve, 100*time.Millisecond) }() - - response := startRequest(url) - waitClosed(t, handler.entered, "the request to reach the handler") - - sendSignal() - - err := waitError(t, done, "shutdown to give up") - if !errors.Is(err, context.DeadlineExceeded) { - t.Fatalf("err = %v, want it to report context.DeadlineExceeded", err) + if err := migrationErr; err != nil { + // Fatal: continuing on a half-migrated schema and then serving traffic is + // strictly worse than refusing to start. + logrus.WithError(err).Fatal("Schema migration execution failed") } - if n := spy.closeCalls.Load(); n != 1 { - t.Errorf("Close called %d times, want exactly 1 forced close", n) - } + logrus.Info("Database schema updated successfully") - if res := waitResult(t, response, "the stuck request to be cut"); res.err == nil { - t.Errorf("stuck request finished with status %d, want its connection closed", res.status) + // Keep the update service informed about this installation, periodically. + // After the migrations, deliberately: on a first boot the tables it reads do + // not exist until this point, and a summary of a schema that is not there yet + // describes an installation with nothing in it — which is also what an + // installation nobody uses looks like. + // + // A goroutine that owns nothing: it reads, it sends, and every failure along + // the way is a debug line. A server must not fail to start because something + // it only talks to is unreachable. What it learns is handed to the router, so + // the UI can show whether this build is current; nil means it never runs, and + // the UI says so. + var updates *update.Service + if built, err := update.New(cfg, queries); err != nil { + logrus.WithError(err).Debug("update service integration is not configured") + } else { + updates = built + go updates.Run(ctx) } -} -func TestMain_ShutdownHTTPServer_ForceCloseErrorIsReported(t *testing.T) { - closeErr := errors.New("close failed") - srv := &stubHTTPServer{ - shutdownFn: func(context.Context) error { return context.DeadlineExceeded }, - closeFn: func() error { return closeErr }, + if cfg.PprofAddr != "" { + go profiling.Start(cfg.PprofAddr) } - err := shutdownHTTPServer(srv, testShutdownTimeout) - - if !errors.Is(err, context.DeadlineExceeded) { - t.Errorf("err = %v, want it to keep the graceful-shutdown failure", err) - } - if !errors.Is(err, closeErr) { - t.Errorf("err = %v, want it to include the forced-close failure", err) + client, err := docker.NewDockerClient(ctx, queries, cfg) + if err != nil { + logrus.WithError(err).Fatal("Docker runtime client initialization failed") } -} - -func TestMain_SuperviseHTTPServer_ErrServerClosedAfterShutdownIsExpected(t *testing.T) { - rec := &shutdownRecorder{} - srv, serve, shutdownCtx := newGracefulStub(rec.add) - ctx, sendSignal := context.WithCancel(context.Background()) - - if err := superviseHTTPServer(ctx, srv, serve, testShutdownTimeout); err != nil { - t.Fatalf("http.ErrServerClosed after a requested shutdown must not be an error, got: %v", err) + providers, err := providers.NewProviderController(cfg, queries, client) + if err != nil { + logrus.WithError(err).Fatal("LLM provider controller initialization failed") } + subscriptions := subscriptions.NewSubscriptionsController() + controller := controller.NewFlowController(queries, cfg, client, providers, subscriptions) - requireEvents(t, rec.snapshot(), []string{"shutdown"}) - shutdownCtx.requireFresh(t, "Shutdown", testShutdownTimeout) -} - -func TestMain_SuperviseHTTPServer_ErrServerClosedWithoutShutdownIsUnexpected(t *testing.T) { - rec := &shutdownRecorder{} - srv, _, _ := newGracefulStub(rec.add) - serve := func() error { return http.ErrServerClosed } + if err := controller.LoadFlows(ctx); err != nil { + logrus.WithError(err).Fatal("Active flows restoration failed") + } - err := superviseHTTPServer(context.Background(), srv, serve, testShutdownTimeout) + r := router.NewRouter(queries, orm, cfg, providers, controller, subscriptions, client, updates) - if !errors.Is(err, http.ErrServerClosed) { - t.Fatalf("err = %v, want the unrequested closure reported", err) + listen := net.JoinHostPort(cfg.ServerHost, strconv.Itoa(cfg.ServerPort)) + srv := &http.Server{ + Addr: listen, + Handler: r.Handler(), + ReadHeaderTimeout: router.ReadHeaderTimeout, + IdleTimeout: router.IdleTimeout, } -} -func TestMain_SuperviseHTTPServer_ServerErrorRacingSignalIsNotLost(t *testing.T) { - stopped := make(chan struct{}) - srv := &stubHTTPServer{ - shutdownFn: func(context.Context) error { close(stopped); return nil }, - closeFn: func() error { return nil }, - } serve := func() error { - <-stopped - return errSimulatedServerFailure - } + logrus.Infof("API server listening on %s", listen) - ctx, sendSignal := context.WithCancel(context.Background()) - sendSignal() - - err := superviseHTTPServer(ctx, srv, serve, testShutdownTimeout) + if cfg.ServerUseSSL && cfg.ServerSSLCrt != "" && cfg.ServerSSLKey != "" { + logrus.Info("Starting server with TLS enabled") + return srv.ListenAndServeTLS(cfg.ServerSSLCrt, cfg.ServerSSLKey) + } - if !errors.Is(err, errSimulatedServerFailure) { - t.Fatalf("err = %v, want the genuine server error reported", err) + logrus.Info("Starting server without TLS (HTTP only)") + return srv.ListenAndServe() } -} -func TestMain_SuperviseHTTPServer_UnexpectedFailureIsReportedAndServerStopped(t *testing.T) { - rec := &shutdownRecorder{} - srv, _, shutdownCtx := newGracefulStub(rec.add) - serve := func() error { return errSimulatedServerFailure } + // Blocks until a shutdown signal or a server failure, then stops the HTTP + // server and drains telemetry. The exit code is returned rather than raised + // with Fatal so the deferred pool cleanup above still runs. + exitCode := serveUntilShutdown(ctx, srv, serve, httpShutdownTimeout, func(ctx context.Context) error { + return obs.Observer.Drain(ctx) + }) - err := superviseHTTPServer(context.Background(), srv, serve, testShutdownTimeout) + // On a server failure ctx is still live; cancel it so work bound to it (the + // update service, flows) stops before the deferred calls close the pools. + // Idempotent after a signal. + cancelOnSignal() - if !errors.Is(err, errSimulatedServerFailure) { - t.Fatalf("err = %v, want the server failure propagated", err) + if exitCode == 0 { + logrus.Info("Application shutdown completed successfully") } - requireEvents(t, rec.snapshot(), []string{"shutdown"}) - shutdownCtx.requireFresh(t, "Shutdown", testShutdownTimeout) + return exitCode } -func TestMain_ServeUntilShutdown_SignalShutsDownThenDrains(t *testing.T) { - rec := &shutdownRecorder{} - srv, serve, _ := newGracefulStub(rec.add) - - var drainCtx ctxSnapshot - drain := func(ctx context.Context) error { - drainCtx = snapshotCtx(ctx) - rec.add("drain") - return nil - } - - ctx, sendSignal := context.WithCancel(context.Background()) - sendSignal() - - if code := serveUntilShutdown(ctx, srv, serve, testShutdownTimeout, drain); code != 0 { - t.Fatalf("exit code = %d, want 0 after a clean signal shutdown", code) - } - - requireEvents(t, rec.snapshot(), []string{"shutdown", "drain"}) - drainCtx.requireFresh(t, "telemetry drain", testDrainTimeout) +// httpServer is the part of *http.Server the shutdown logic needs, so tests can +// drive the deadline and failure paths without waiting on real timers. +type httpServer interface { + Shutdown(ctx context.Context) error + Close() error } -func TestMain_ServeUntilShutdown_DrainFailureDoesNotChangeExitCode(t *testing.T) { - rec := &shutdownRecorder{} - srv, serve, _ := newGracefulStub(rec.add) - drain := func(context.Context) error { return errors.New("collector unreachable") } - - ctx, sendSignal := context.WithCancel(context.Background()) - sendSignal() - - if code := serveUntilShutdown(ctx, srv, serve, testShutdownTimeout, drain); code != 0 { - t.Fatalf("exit code = %d, want 0: an incomplete drain is only logged", code) - } -} +// serveUntilShutdown runs serve until ctx is cancelled or serve fails, stops the +// HTTP server, and only then drains telemetry, so requests that finish during +// graceful shutdown can still emit it. It returns the process exit code and +// never exits itself, which keeps the caller's deferred cleanup reachable. +func serveUntilShutdown( + ctx context.Context, + srv httpServer, + serve func() error, + shutdownTimeout time.Duration, + drain func(context.Context) error, +) int { + err := superviseHTTPServer(ctx, srv, serve, shutdownTimeout) -func failureScenario(emit func(string)) int { - defer emit("cleanup") + drainTelemetry(drain) - srv := &stubHTTPServer{ - shutdownFn: func(context.Context) error { emit("shutdown"); return nil }, - closeFn: func() error { emit("close"); return nil }, + if err != nil { + logrus.WithError(err).Error("HTTP server terminated with an error") + return 1 } - serve := func() error { return errSimulatedServerFailure } - drain := func(context.Context) error { emit("drain"); return nil } - return serveUntilShutdown(context.Background(), srv, serve, testShutdownTimeout, drain) + return 0 } -func TestMain_ServeUntilShutdown_UnexpectedFailureCleansUpAndReturnsFailure(t *testing.T) { - rec := &shutdownRecorder{} +// superviseHTTPServer runs serve in the background and returns once the server +// has been stopped. It returns nil only for a shutdown it requested itself and +// that completed gracefully. +func superviseHTTPServer( + ctx context.Context, + srv httpServer, + serve func() error, + shutdownTimeout time.Duration, +) error { + // Buffered so the goroutine can always deliver its result and exit, even + // when nobody is left to receive it. + serveErr := make(chan error, 1) + go func() { serveErr <- serve() }() - code := failureScenario(rec.add) + select { + case <-ctx.Done(): + logrus.Warn("Shutdown signal received, cleaning up resources...") + shutdownErr := shutdownHTTPServer(srv, shutdownTimeout) + + // Shutdown makes serve return http.ErrServerClosed. Wait for it so a real + // error that raced the signal is reported instead of lost. + var failure error + if err := <-serveErr; err != nil && !errors.Is(err, http.ErrServerClosed) { + failure = fmt.Errorf("API server failed: %w", err) + } - if code != 1 { - t.Fatalf("exit code = %d, want 1", code) - } + return errors.Join(failure, shutdownErr) - requireEvents(t, rec.snapshot(), []string{"shutdown", "drain", "cleanup"}) -} + case err := <-serveErr: + // Nobody asked the server to stop, so even http.ErrServerClosed or a nil + // return means it died on its own. + if err == nil { + err = errors.New("serve returned without an error") + } + failure := fmt.Errorf("API server failed: %w", err) -func TestMain_ServeUntilShutdown_ListenFailureOnRealServer(t *testing.T) { - occupied, err := net.Listen("tcp", "127.0.0.1:0") - if err != nil { - t.Fatalf("listen: %v", err) + // Connections accepted before the failure are still being served; give + // them the same bounded shutdown before their resources are released. + return errors.Join(failure, shutdownHTTPServer(srv, shutdownTimeout)) } - defer occupied.Close() - - srv := &http.Server{Addr: occupied.Addr().String(), Handler: http.NotFoundHandler()} +} - ctx, cancel := context.WithTimeout(context.Background(), testWaitTimeout) +// shutdownHTTPServer gracefully stops srv within timeout and force-closes what +// remains if that fails. Shutdown does not wait for hijacked connections +// (WebSockets), nor does it stop background services or running flows. +func shutdownHTTPServer(srv httpServer, timeout time.Duration) error { + ctx, cancel := context.WithTimeout(context.Background(), timeout) defer cancel() - drained := false - drain := func(context.Context) error { drained = true; return nil } - - code := serveUntilShutdown(ctx, srv, srv.ListenAndServe, testShutdownTimeout, drain) - - if code != 1 { - t.Errorf("exit code = %d, want 1 when the address is already in use", code) - } - if !drained { - t.Error("telemetry was not drained after the listener failed") + err := srv.Shutdown(ctx) + if err == nil { + return nil } -} -const failureHelperEnv = "PENTAGI_TEST_SERVER_FAILURE_HELPER" + logrus.WithError(err).Warn("Graceful HTTP shutdown failed, closing remaining connections") + shutdownErr := fmt.Errorf("graceful HTTP shutdown: %w", err) -// Child process of TestMain_ServeUntilShutdown_UnexpectedFailureExitStatus. -func TestMain_ServeUntilShutdown_FailureHelperProcess(t *testing.T) { - if os.Getenv(failureHelperEnv) != "1" { - t.Skip("helper process for TestMain_ServeUntilShutdown_UnexpectedFailureExitStatus") + if closeErr := srv.Close(); closeErr != nil { + logrus.WithError(closeErr).Error("Forced HTTP server close failed") + return errors.Join(shutdownErr, fmt.Errorf("force close HTTP server: %w", closeErr)) } - os.Exit(failureScenario(func(event string) { fmt.Println("EVENT:" + event) })) + return shutdownErr } -// runApplication itself needs PostgreSQL and Docker, so a stand-in with the same shape runs instead. -func TestMain_ServeUntilShutdown_UnexpectedFailureExitStatus(t *testing.T) { - cmd := exec.Command(os.Args[0], "-test.run=^TestMain_ServeUntilShutdown_FailureHelperProcess$") - cmd.Env = append(os.Environ(), failureHelperEnv+"=1") - - out, err := cmd.Output() - - var exitErr *exec.ExitError - if !errors.As(err, &exitErr) { - t.Fatalf("child err = %v, want a non-zero exit\nstdout:\n%s", err, out) - } - if got := exitErr.ExitCode(); got != 1 { - t.Fatalf("child exit status = %d, want 1\nstdout:\n%s\nstderr:\n%s", got, out, exitErr.Stderr) - } +// drainTelemetry flushes buffered telemetry on its own deadline. A failure is +// only logged: losing a final batch must not change the exit status. +func drainTelemetry(drain func(context.Context) error) { + ctx, cancel := context.WithTimeout(context.Background(), telemetryDrainTimeout) + defer cancel() - var events []string - for _, line := range strings.Split(string(out), "\n") { - line = strings.TrimSuffix(line, "\r") - if strings.HasPrefix(line, "EVENT:") { - events = append(events, strings.TrimPrefix(line, "EVENT:")) - } + if err := drain(ctx); err != nil { + logrus.WithError(err).Warn("Telemetry drain incomplete") } - - requireEvents(t, events, []string{"shutdown", "drain", "cleanup"}) } From a7ad77d74b6cb91611549e887fde73fab235c404 Mon Sep 17 00:00:00 2001 From: yash27git Date: Sat, 31 Oct 2026 10:29:39 +0530 Subject: [PATCH 5/5] ci: trigger CI workflow