fix(lifecycle): finish shutdown before the process exits

Makefile

@@ -26,9 +26,16 @@ install-tools:
 
 build:
 	go build -ldflags '$(LDFLAGS)' -o bin/gomodel ./cmd/gomodel
-# Run the application
+# Run the application.
+#
+# Built and exec'd rather than `go run`: `go run` exits 1 when it is
+# interrupted, even after the program it supervises shuts down cleanly, so
+# Ctrl+C always ended in a bogus "make: *** [run] Error 1". exec replaces the
+# recipe shell, which also puts the gateway directly under make's signal
+# handling instead of behind a supervisor.
 run:
-	LOG_LEVEL=$(LOG_LEVEL) SWAGGER_ENABLED=$(SWAGGER_ENABLED) go run -tags=swagger ./cmd/gomodel
+	go build -tags=swagger -ldflags '$(LDFLAGS)' -o bin/gomodel ./cmd/gomodel
+	LOG_LEVEL="$(LOG_LEVEL)" SWAGGER_ENABLED="$(SWAGGER_ENABLED)" exec ./bin/gomodel
 
 # Seed the local SQLite database and start GoModel with a populated dashboard.
 demo: seed-demo-data

internal/server/http.go

@@ -45,6 +45,20 @@ const (
 	inboundServerReadTimeout       = 30 * time.Second
 	inboundServerReadHeaderTimeout = 10 * time.Second
 	inboundServerWriteTimeout      = 30 * time.Second
+
+	// GracefulDrainTimeout bounds how long the HTTP server waits for in-flight
+	// requests to finish once shutdown begins. Streamed model responses run far
+	// longer than any drain window worth waiting for, so this is deliberately a
+	// cutoff rather than a promise: past it the remaining connections are cut
+	// and shutdown moves on to flushing usage and audit records. It is sized
+	// for the short requests that can actually finish — a dashboard fetch, a
+	// health check — not for model traffic.
+	//
+	// It must stay below run.shutdownTimeout, which covers this drain plus
+	// those flushes and the database close that follow it. Exported so that
+	// relationship is checked rather than merely asserted here — see
+	// TestGracefulDrainFitsInsideTheShutdownBudget in the run package.
+	GracefulDrainTimeout = 10 * time.Second
 )
 
 // Config holds server configuration options
@@ -556,6 +570,17 @@ func newGatewayStartConfig(addr string) echo.StartConfig {
 		BeforeServeFunc: func(server *http.Server) error {
 			return configureGatewayHTTPServer(server)
 		},
+		// Echo's own default is an implicit 10s that reports the cutoff as an
+		// unexplained error. Both are set here so the drain window is sized
+		// against the application's shutdown budget and so cutting a stream
+		// short on Ctrl+C reads as the routine event it is.
+		GracefulTimeout: GracefulDrainTimeout,
+		OnShutdownError: func(err error) {
+			slog.Warn("closing requests still in flight at the shutdown deadline",
+				"graceful_timeout", GracefulDrainTimeout,
+				"error", err,
+			)
+		},
 	}
 }
 

internal/server/http_start_test.go

@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"net/http"
 	"net/http/httptest"
 	"testing"
@@ -52,6 +53,25 @@ func TestNewGatewayStartConfig_AppliesTimeoutOverrides(t *testing.T) {
 	}
 }
 
+// Leaving GracefulTimeout unset takes Echo's implicit 10s default and reports
+// the cutoff through Echo's own logger as a bare "context deadline exceeded",
+// which is what an operator saw on Ctrl+C while a stream was open. The drain
+// window has to be the gateway's own decision, and sized against the
+// application shutdown budget that has to contain it.
+func TestNewGatewayStartConfig_ConfiguresGracefulDrain(t *testing.T) {
+	cfg := newGatewayStartConfig(":0")
+
+	if cfg.GracefulTimeout != GracefulDrainTimeout {
+		t.Fatalf("GracefulTimeout = %v, want %v", cfg.GracefulTimeout, GracefulDrainTimeout)
+	}
+	if cfg.OnShutdownError == nil {
+		t.Fatal("OnShutdownError = nil, want the drain cutoff reported by the gateway")
+	}
+	// A nil handler is Echo's signal to log it itself; ours must absorb the
+	// error without panicking on the deadline it will actually be handed.
+	cfg.OnShutdownError(context.DeadlineExceeded)
+}
+
 func TestModelInteractionWriteDeadlineMiddleware_ClearsDeadlineForModelRoutes(t *testing.T) {
 	e := echo.New()
 	writer := &deadlineTrackingWriter{ResponseRecorder: httptest.NewRecorder()}

run/lifecycle_test.go

@@ -5,11 +5,13 @@ import (
 	"errors"
 	"slices"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
 	"github.com/enterpilot/gomodel/config"
 	"github.com/enterpilot/gomodel/internal/providers"
+	"github.com/enterpilot/gomodel/internal/server"
 )
 
 type stubLifecycleApp struct {
@@ -58,11 +60,14 @@ func (s *stubLifecycleApp) capturedShutdownContext() context.Context {
 	return s.shutdownCtx
 }
 
-func TestStartApplication_ShutsDownOnStartFailure(t *testing.T) {
+// A server that never came up still holds a database handle and whatever the
+// loggers buffered while it was being built, so it gets torn down — once, on
+// one shutdownTimeout budget, from the same place every other exit uses.
+func TestServeUntilShutdown_TearsDownOnceAfterAFailedStart(t *testing.T) {
 	startErr := errors.New("listen tcp :8080: bind: address already in use")
 	app := &stubLifecycleApp{startErr: startErr}
 
-	err := startApplication(app, ":8080")
+	err := serveUntilShutdown(context.Background(), app, ":8080")
 	if !errors.Is(err, startErr) {
 		t.Fatalf("error = %v, want start error %v", err, startErr)
 	}
@@ -85,41 +90,24 @@ func TestStartApplication_ShutsDownOnStartFailure(t *testing.T) {
 	}
 }
 
-func TestStartApplication_ReportsShutdownFailure(t *testing.T) {
+// The start error is what the operator needs to see and what sets the exit
+// code, so a teardown that also fails is logged rather than wrapped around it.
+func TestServeUntilShutdown_ShutdownFailureDoesNotMaskTheStartError(t *testing.T) {
 	startErr := errors.New("listen failed")
-	shutdownErr := errors.New("close failed")
-	app := &stubLifecycleApp{
-		startErr:    startErr,
-		shutdownErr: shutdownErr,
-	}
+	app := &stubLifecycleApp{startErr: startErr, shutdownErr: errors.New("close failed")}
 
-	err := startApplication(app, ":8080")
+	err := serveUntilShutdown(context.Background(), app, ":8080")
 	if !errors.Is(err, startErr) {
 		t.Fatalf("error = %v, want start error %v", err, startErr)
 	}
-	if !errors.Is(err, shutdownErr) {
-		t.Fatalf("error = %v, want shutdown error %v", err, shutdownErr)
-	}
 	if calls := app.shutdownCallCount(); calls != 1 {
 		t.Fatalf("shutdownCalls = %d, want 1", calls)
 	}
 }
 
-func TestStartApplication_DoesNotShutdownOnSuccess(t *testing.T) {
-	app := &stubLifecycleApp{}
-
-	if err := startApplication(app, ":8080"); err != nil {
-		t.Fatalf("startApplication() error = %v, want nil", err)
-	}
-	if calls := app.startCallCount(); calls != 1 {
-		t.Fatalf("startCalls = %d, want 1", calls)
-	}
-	if calls := app.shutdownCallCount(); calls != 0 {
-		t.Fatalf("shutdownCalls = %d, want 0", calls)
-	}
-}
-
-func TestStartApplication_StopsWaitingWhenShutdownTimesOut(t *testing.T) {
+// A teardown that wedges must not wedge the process with it: the wait is
+// bounded by shutdownTimeout and serveUntilShutdown returns regardless.
+func TestServeUntilShutdown_StopsWaitingWhenShutdownTimesOut(t *testing.T) {
 	previousTimeout := shutdownTimeout
 	shutdownTimeout = 10 * time.Millisecond
 	defer func() {
@@ -130,23 +118,127 @@ func TestStartApplication_StopsWaitingWhenShutdownTimesOut(t *testing.T) {
 	shutdownBlock := make(chan struct{})
 	defer close(shutdownBlock)
 
-	app := &stubLifecycleApp{
-		startErr:      startErr,
-		shutdownBlock: shutdownBlock,
+	app := &stubLifecycleApp{startErr: startErr, shutdownBlock: shutdownBlock}
+
+	done := make(chan error, 1)
+	go func() {
+		done <- serveUntilShutdown(context.Background(), app, ":8080")
+	}()
+
+	select {
+	case err := <-done:
+		if !errors.Is(err, startErr) {
+			t.Fatalf("error = %v, want start error %v", err, startErr)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("serveUntilShutdown blocked on a shutdown that never returned")
+	}
+	if calls := app.shutdownCallCount(); calls != 1 {
+		t.Fatalf("shutdownCalls = %d, want 1", calls)
 	}
+}
 
-	err := startApplication(app, ":8080")
-	if !errors.Is(err, startErr) {
-		t.Fatalf("error = %v, want start error %v", err, startErr)
+// The drain window and the shutdown budget live in different packages, so the
+// comment tying them together is only as good as this check: the budget has to
+// cover the drain plus the usage and audit flushes that follow it.
+func TestGracefulDrainFitsInsideTheShutdownBudget(t *testing.T) {
+	if server.GracefulDrainTimeout >= shutdownTimeout {
+		t.Fatalf("GracefulDrainTimeout = %v must be shorter than shutdownTimeout = %v",
+			server.GracefulDrainTimeout, shutdownTimeout)
+	}
+	if headroom := shutdownTimeout - server.GracefulDrainTimeout; headroom < 5*time.Second {
+		t.Fatalf("only %v left for flushing after the drain; widen shutdownTimeout or shorten the drain", headroom)
+	}
+}
+
+// servingApp mirrors the ordering that matters in the real App: Start blocks
+// until Shutdown stops the server, and Shutdown keeps working afterwards —
+// flushing buffered usage and audit records, closing the database — before it
+// returns.
+type servingApp struct {
+	serverStopped chan struct{} // closed by Shutdown, releases Start
+	flushing      chan struct{} // closed by the test, releases Shutdown
+	shutdownDone  atomic.Bool
+}
+
+func newServingApp() *servingApp {
+	return &servingApp{
+		serverStopped: make(chan struct{}),
+		flushing:      make(chan struct{}),
 	}
-	if !errors.Is(err, context.DeadlineExceeded) {
-		t.Fatalf("error = %v, want context deadline exceeded", err)
+}
+
+func (a *servingApp) Start(context.Context, string) error {
+	<-a.serverStopped
+	return nil
+}
+
+func (a *servingApp) Shutdown(context.Context) error {
+	close(a.serverStopped)
+	<-a.flushing
+	a.shutdownDone.Store(true)
+	return nil
+}
+
+// Run returns straight into process exit, so returning while Shutdown is still
+// flushing loses whatever it had not written yet. That is what happened on
+// every Ctrl+C: the server stopped, Start returned, the process left, and
+// "application shutdown complete" was never reached.
+func TestServeUntilShutdown_WaitsForTeardownToFinish(t *testing.T) {
+	app := newServingApp()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	returned := make(chan error, 1)
+	go func() {
+		returned <- serveUntilShutdown(ctx, app, ":0")
+	}()
+
+	cancel() // the SIGINT equivalent
+
+	// Start has returned by now; Shutdown is still flushing.
+	select {
+	case err := <-returned:
+		t.Fatalf("serveUntilShutdown returned mid-teardown (error = %v)", err)
+	case <-time.After(100 * time.Millisecond):
+	}
+
+	close(app.flushing)
+	select {
+	case err := <-returned:
+		if err != nil {
+			t.Fatalf("serveUntilShutdown() error = %v, want nil", err)
+		}
+	case <-time.After(5 * time.Second):
+		t.Fatal("serveUntilShutdown did not return after teardown finished")
+	}
+	if !app.shutdownDone.Load() {
+		t.Fatal("teardown did not run to completion")
+	}
+}
+
+// A server that stops without a signal still owns a database handle and
+// buffered records, so it gets the same teardown.
+func TestServeUntilShutdown_TearsDownWhenServerStopsOnItsOwn(t *testing.T) {
+	app := &stubLifecycleApp{}
+
+	if err := serveUntilShutdown(context.Background(), app, ":0"); err != nil {
+		t.Fatalf("serveUntilShutdown() error = %v, want nil", err)
 	}
 	if calls := app.shutdownCallCount(); calls != 1 {
 		t.Fatalf("shutdownCalls = %d, want 1", calls)
 	}
 }
 
+func TestServeUntilShutdown_ReturnsStartFailure(t *testing.T) {
+	startErr := errors.New("listen tcp :8080: bind: address already in use")
+	app := &stubLifecycleApp{startErr: startErr}
+
+	if err := serveUntilShutdown(context.Background(), app, ":8080"); !errors.Is(err, startErr) {
+		t.Fatalf("serveUntilShutdown() error = %v, want start error %v", err, startErr)
+	}
+}
+
 func TestMain_KimicodeProviderRegistration(t *testing.T) {
 	factory := defaultProviderFactory(&config.Config{})
 

run/run.go

@@ -189,18 +189,9 @@ func Run(ctx context.Context, opts Options) error {
 
 	signalCtx, stop := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
 	defer stop()
-	go func() {
-		<-signalCtx.Done()
-		shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
-		defer cancel()
-
-		if err := shutdownApplication(application, shutdownCtx); err != nil {
-			slog.Error("application shutdown error", "error", err)
-		}
-	}()
 
 	addr := ":" + result.Config.Server.Port
-	if err := startApplication(application, addr); err != nil {
+	if err := serveUntilShutdown(signalCtx, application, addr); err != nil {
 		slog.Error("application failed", "error", err)
 		return err
 	}
@@ -217,6 +208,43 @@ type lifecycleApp interface {
 	Shutdown(ctx context.Context) error
 }
 
+// serveUntilShutdown starts the application and returns only once the server
+// has stopped *and* the teardown that stopped it has finished.
+//
+// The teardown has to run on its own goroutine because Start blocks until the
+// server stops and Shutdown is what stops it. Waiting for that goroutine here
+// is the load-bearing part: the process exits the moment Run returns, so
+// anything Shutdown had not reached yet — the buffered usage and audit
+// records, the database handle — would be dropped on every Ctrl+C.
+//
+// This is the only caller of shutdownApplication, so teardown runs exactly
+// once per exit and on a single shutdownTimeout budget, whichever way the
+// server ended: a signal, a stop of its own accord, or a Start that never got
+// off the ground all converge here. Routing the failed-start path through the
+// same place is what removes the second teardown that used to run alongside
+// it, and with it any reliance on Shutdown being idempotent.
+func serveUntilShutdown(ctx context.Context, application lifecycleApp, addr string) error {
+	serverReturned := make(chan struct{})
+	shutdownDone := make(chan error, 1)
+	go func() {
+		select {
+		case <-ctx.Done():
+		case <-serverReturned:
+		}
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
+		defer cancel()
+		shutdownDone <- shutdownApplication(application, shutdownCtx)
+	}()
+
+	startErr := application.Start(context.Background(), addr)
+	close(serverReturned)
+
+	if err := <-shutdownDone; err != nil {
+		slog.Error("application shutdown error", "error", err)
+	}
+	return startErr
+}
+
 func shutdownApplication(application lifecycleApp, ctx context.Context) error {
 	done := make(chan error, 1)
 	go func() {
@@ -230,19 +258,3 @@ func shutdownApplication(application lifecycleApp, ctx context.Context) error {
 		return ctx.Err()
 	}
 }
-
-// startApplication calls lifecycleApp.Start and, if Start fails, attempts a
-// graceful shutdown via shutdownApplication using shutdownTimeout before
-// returning the original start error or a combined start/shutdown error.
-func startApplication(application lifecycleApp, addr string) error {
-	if err := application.Start(context.Background(), addr); err != nil {
-		shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownTimeout)
-		defer cancel()
-
-		if shutdownErr := shutdownApplication(application, shutdownCtx); shutdownErr != nil {
-			return fmt.Errorf("server failed to start: %w", errors.Join(err, fmt.Errorf("shutdown after start failure: %w", shutdownErr)))
-		}
-		return err
-	}
-	return nil
-}