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"
 
@@ -208,6 +208,9 @@ func ServeConnect(ctx context.Context, logger *slog.Logger, cfg Config, deps api
 		Handler: handler,
 	}
 
+	// counts shutdown goroutines still draining their servers
+	var shutdownWG sync.WaitGroup
+
 	// start dedicated metrics server if configured
 	if cfg.MetricsPort > 0 {
 		metricsMux := http.NewServeMux()
@@ -227,22 +230,35 @@ func ServeConnect(ctx context.Context, logger *slog.Logger, cfg Config, deps api
 				logger.Error("metrics server failed", "err", err)
 			}
 		}()
+		shutdownWG.Add(1)
 		go func() {
+			defer shutdownWG.Done()
 			<-ctx.Done()
-			metricsServer.Shutdown(context.Background())
+
+			ctxShutdown, cancel := context.WithTimeout(context.Background(), connectServerGracePeriod)
+			defer cancel()
+
+			if err := metricsServer.Shutdown(ctxShutdown); err != nil {
+				logger.ErrorContext(ctxShutdown, "metrics server shutdown error", "err", err)
+				return
+			}
+			logger.Info("Graceful shutdown of metrics server complete")
 		}()
 	}
 
 	logger.Info("connect server starting", "port", cfg.Connect.Port)
 
+	shutdownWG.Add(1)
 	go func() {
+		defer shutdownWG.Done()
 		<-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.ErrorContext(ctxShutdown, "HTTP shutdown error", "err", err)
+			return
 		}
 
 		logger.Info("Graceful shutdown of connect server complete")
@@ -253,6 +269,11 @@ func ServeConnect(ctx context.Context, logger *slog.Logger, cfg Config, deps api
 		return fmt.Errorf("connect server failed: %w", err)
 	}
 
+	// Shutdown closes the listener first, which makes ListenAndServe return
+	// while in-flight requests are still draining. Wait for the drain to
+	// finish so callers don't tear down the database and other dependencies
+	// under requests that are still running.
+	shutdownWG.Wait()
 	return nil
 }
 

pkg/server/server_test.go

@@ -0,0 +1,87 @@
+package server
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"log/slog"
+	"net"
+	"net/http"
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/raystack/frontier/internal/api"
+)
+
+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")
+			}
+		})
+	}
+}