diff --git a/pkg/server/server.go b/pkg/server/server.go index 315da19a5..d491f3514 100644 --- a/pkg/server/server.go +++ b/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,7 +200,8 @@ 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{ @@ -208,6 +209,17 @@ func ServeConnect(ctx context.Context, logger *slog.Logger, cfg Config, deps api 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 { diff --git a/pkg/server/server_test.go b/pkg/server/server_test.go new file mode 100644 index 000000000..fc33ea814 --- /dev/null +++ b/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") + } + }) + } +}