fix: wait for HTTP server drain during graceful shutdown

pkg/server/server.go

@@ -4,17 +4,17 @@ import (
 	"context"
 	"encoding/json"
 	"fmt"
+	"log/slog"
 	"net/http"
 	"net/http/httputil"
 	"net/url"
 	"strings"
+	"sync"
 	"time"
 
 	"golang.org/x/net/http2"
 	"golang.org/x/net/http2/h2c"
 
-	"log/slog"
-
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/raystack/salt/server/spa"
 
@@ -200,14 +200,26 @@ func ServeConnect(ctx context.Context, logger *slog.Logger, cfg Config, deps api
 	})
 
 	// Configure and create the server
-	handler := h2c.NewHandler(mux, &http2.Server{})
+	h2s := &http2.Server{}
+	handler := h2c.NewHandler(mux, h2s)
 	handler = connectinterceptors.WithConnectCORS(handler, cfg.ConnectCors)
 
 	server := &http.Server{
 		Addr:    fmt.Sprintf(":%d", cfg.Connect.Port),
 		Handler: handler,
 	}
 
+	// Shutdown cannot reach h2c connections on its own: h2c hijacks them out
+	// of the server's connection tracking. This hook makes Shutdown send
+	// GOAWAY on them so clients finish up and reconnect elsewhere
+	// (golang/go#26682).
+	if err := http2.ConfigureServer(server, h2s); err != nil {
+		return fmt.Errorf("configure http2 server: %w", err)
+	}
+
+	// counts shutdown goroutines still draining their servers
+	var shutdownWG sync.WaitGroup
+
 	// start dedicated metrics server if configured
 	if cfg.MetricsPort > 0 {
 		metricsMux := http.NewServeMux()
@@ -221,41 +233,64 @@ func ServeConnect(ctx context.Context, logger *slog.Logger, cfg Config, deps api
 			Addr:    fmt.Sprintf(":%d", cfg.MetricsPort),
 			Handler: metricsMux,
 		}
+		metricsFailed := make(chan struct{})
 		go func() {
 			logger.Info("metrics server starting", "port", cfg.MetricsPort)
 			if err := metricsServer.ListenAndServe(); err != nil && err != http.ErrServerClosed {
 				logger.Error("metrics server failed", "err", err)
+				close(metricsFailed)
 			}
 		}()
-		go func() {
-			<-ctx.Done()
-			metricsServer.Shutdown(context.Background())
-		}()
+		shutdownWG.Add(1)
+		go gracefulShutdown(ctx, logger, &shutdownWG, metricsServer, "metrics server", metricsFailed)
 	}
 
 	logger.Info("connect server starting", "port", cfg.Connect.Port)
 
-	go func() {
-		<-ctx.Done()
-
-		ctxShutdown, cancel := context.WithTimeout(context.Background(), connectServerGracePeriod)
-		defer cancel()
-
-		if err := server.Shutdown(ctxShutdown); err != nil {
-			logger.ErrorContext(ctxShutdown, "HTTP shutdown error", "error", err)
-		}
-
-		logger.Info("Graceful shutdown of connect server complete")
-	}()
+	serveFailed := make(chan struct{})
+	shutdownWG.Add(1)
+	go gracefulShutdown(ctx, logger, &shutdownWG, server, "connect server", serveFailed)
 
 	// Start server
 	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+		close(serveFailed)
+		// No shutdownWG.Wait here: the metrics watcher only unblocks when ctx
+		// is cancelled, which never happens on this path; it dies with the
+		// process right after.
 		return fmt.Errorf("connect server failed: %w", err)
 	}
 
+	// Wait for Shutdown to finish draining HTTP/1.1 requests before returning,
+	// so callers don't tear down the database under them. Hijacked h2c
+	// connections are not tracked by Shutdown and are not waited on; they get
+	// GOAWAY through the http2.ConfigureServer hook instead.
+	shutdownWG.Wait()
 	return nil
 }
 
+// gracefulShutdown drains srv within the grace period once ctx is cancelled.
+// It returns without logging when the server already failed, so a server
+// that never started is not reported as gracefully shut down.
+func gracefulShutdown(ctx context.Context, logger *slog.Logger, wg *sync.WaitGroup, srv *http.Server, name string, failed <-chan struct{}) {
+	defer wg.Done()
+
+	select {
+	case <-ctx.Done():
+	case <-failed:
+		return
+	}
+
+	ctxShutdown, cancel := context.WithTimeout(context.Background(), connectServerGracePeriod)
+	defer cancel()
+
+	if err := srv.Shutdown(ctxShutdown); err != nil {
+		logger.ErrorContext(ctxShutdown, name+" shutdown error", "err", err)
+		return
+	}
+
+	logger.Info("Graceful shutdown of " + name + " complete")
+}
+
 func getSessionCookieCutter(blockSecretKey string, hashSecretKey string, logger *slog.Logger) securecookie.Codec {
 	var sessionCookieCutter securecookie.Codec
 	if len(hashSecretKey) != 32 || len(blockSecretKey) != 32 {

pkg/server/server_test.go

@@ -0,0 +1,150 @@
+package server
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+	"net"
+	"net/http"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/raystack/frontier/internal/api"
+)
+
+func TestGracefulShutdownDrainsInflightRequests(t *testing.T) {
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+
+	entered := make(chan struct{})
+	requestDone := make(chan struct{})
+	mux := http.NewServeMux()
+	mux.HandleFunc("/slow", func(w http.ResponseWriter, r *http.Request) {
+		close(entered)
+		time.Sleep(300 * time.Millisecond)
+		close(requestDone)
+		w.WriteHeader(http.StatusOK)
+	})
+
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	srv := &http.Server{Handler: mux}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go gracefulShutdown(ctx, logger, &wg, srv, "test server", make(chan struct{}))
+
+	serveReturned := make(chan struct{})
+	go func() {
+		defer close(serveReturned)
+		if err := srv.Serve(l); err != nil && err != http.ErrServerClosed {
+			t.Errorf("serve failed: %v", err)
+		}
+	}()
+
+	requestErr := make(chan error, 1)
+	go func() {
+		resp, err := http.Get(fmt.Sprintf("http://%s/slow", l.Addr()))
+		if err == nil {
+			resp.Body.Close()
+			if resp.StatusCode != http.StatusOK {
+				err = fmt.Errorf("unexpected status %d", resp.StatusCode)
+			}
+		}
+		requestErr <- err
+	}()
+
+	<-entered
+	cancel()
+
+	<-serveReturned
+	wg.Wait()
+
+	select {
+	case <-requestDone:
+	default:
+		t.Fatal("shutdown wait released before the in-flight request finished")
+	}
+	if err := <-requestErr; err != nil {
+		t.Fatalf("in-flight request failed: %v", err)
+	}
+}
+
+func freePort(t *testing.T) int {
+	t.Helper()
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("could not find a free port: %v", err)
+	}
+	defer l.Close()
+	return l.Addr().(*net.TCPAddr).Port
+}
+
+func waitForHTTP(t *testing.T, url string) {
+	t.Helper()
+	deadline := time.Now().Add(10 * time.Second)
+	for time.Now().Before(deadline) {
+		resp, err := http.Get(url)
+		if err == nil {
+			resp.Body.Close()
+			if resp.StatusCode == http.StatusOK {
+				return
+			}
+		}
+		time.Sleep(50 * time.Millisecond)
+	}
+	t.Fatalf("server did not respond at %s in time", url)
+}
+
+func TestServeConnectReturnsAfterShutdownOnContextCancel(t *testing.T) {
+	tests := []struct {
+		name        string
+		withMetrics bool
+	}{
+		{name: "connect server only", withMetrics: false},
+		{name: "connect and metrics servers", withMetrics: true},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+
+			var cfg Config
+			cfg.Connect.Port = freePort(t)
+			if tc.withMetrics {
+				cfg.MetricsPort = freePort(t)
+			}
+
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			served := make(chan error, 1)
+			go func() {
+				served <- ServeConnect(ctx, logger, cfg, api.Deps{}, prometheus.NewRegistry())
+			}()
+
+			waitForHTTP(t, fmt.Sprintf("http://127.0.0.1:%d/ping", cfg.Connect.Port))
+			if tc.withMetrics {
+				waitForHTTP(t, fmt.Sprintf("http://127.0.0.1:%d/metrics", cfg.MetricsPort))
+			}
+
+			cancel()
+
+			select {
+			case err := <-served:
+				if err != nil {
+					t.Fatalf("ServeConnect returned error: %v", err)
+				}
+			case <-time.After(15 * time.Second):
+				t.Fatal("ServeConnect did not return after context cancel")
+			}
+		})
+	}
+}