feat: make the shutdown grace period configurable

config/sample.config.yaml

@@ -56,6 +56,8 @@ app:
     port: 8002
   # port for application metrics
   metrics_port: 9000
+  # how long each server waits for in-flight requests to finish during shutdown
+  shutdown_grace_period: 10s
   # enable pprof endpoints for cpu/mem/mutex profiling
   profiler: false
   host: "127.0.0.1"

pkg/server/config.go

@@ -1,6 +1,8 @@
 package server
 
 import (
+	"time"
+
 	"github.com/raystack/frontier/core/userpat"
 	"github.com/raystack/frontier/core/webhook"
 
@@ -54,6 +56,10 @@ type Config struct {
 	// metrics port
 	MetricsPort int `yaml:"metrics_port" mapstructure:"metrics_port" default:"9000"`
 
+	// ShutdownGracePeriod is how long each server waits for in-flight
+	// requests to finish during shutdown before cutting them off
+	ShutdownGracePeriod time.Duration `yaml:"shutdown_grace_period" mapstructure:"shutdown_grace_period" default:"10s"`
+
 	// Profiler enables /debug/pprof under metrics port
 	Profiler bool `yaml:"profiler" mapstructure:"profiler" default:"false"`
 

pkg/server/serve_ui_test.go

@@ -35,7 +35,7 @@ func TestServeUIReturnsAfterShutdownOnContextCancel(t *testing.T) {
 	done := make(chan struct{})
 	go func() {
 		defer close(done)
-		ServeUI(ctx, logger, UIConfig{Port: port}, Config{})
+		ServeUI(ctx, logger, UIConfig{Port: port}, Config{ShutdownGracePeriod: 5 * time.Second})
 	}()
 
 	url := fmt.Sprintf("http://127.0.0.1:%d/configs", port)

pkg/server/server.go

@@ -37,10 +37,6 @@ import (
 	frontierv1beta1connect "github.com/raystack/frontier/proto/v1beta1/frontierv1beta1connect"
 )
 
-const (
-	connectServerGracePeriod = 10 * time.Second
-)
-
 type WebhooksConfigApiResponse struct {
 	EnableDelete bool `json:"enable_delete"`
 }
@@ -123,7 +119,7 @@ func ServeUI(ctx context.Context, logger *slog.Logger, uiConfig UIConfig, apiSer
 	case <-ctx.Done():
 	}
 
-	ctxShutdown, cancel := context.WithTimeout(context.Background(), connectServerGracePeriod)
+	ctxShutdown, cancel := context.WithTimeout(context.Background(), apiServerConfig.ShutdownGracePeriod)
 	defer cancel()
 
 	if err := server.Shutdown(ctxShutdown); err != nil {
@@ -266,14 +262,14 @@ func ServeConnect(ctx context.Context, logger *slog.Logger, cfg Config, deps api
 			}
 		}()
 		shutdownWG.Add(1)
-		go gracefulShutdown(ctx, logger, &shutdownWG, metricsServer, "metrics server", metricsFailed)
+		go gracefulShutdown(ctx, logger, &shutdownWG, metricsServer, "metrics server", cfg.ShutdownGracePeriod, metricsFailed)
 	}
 
 	logger.Info("connect server starting", "port", cfg.Connect.Port)
 
 	serveFailed := make(chan struct{})
 	shutdownWG.Add(1)
-	go gracefulShutdown(ctx, logger, &shutdownWG, server, "connect server", serveFailed)
+	go gracefulShutdown(ctx, logger, &shutdownWG, server, "connect server", cfg.ShutdownGracePeriod, serveFailed)
 
 	// Start server
 	if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
@@ -295,7 +291,7 @@ func ServeConnect(ctx context.Context, logger *slog.Logger, cfg Config, deps api
 // 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{}) {
+func gracefulShutdown(ctx context.Context, logger *slog.Logger, wg *sync.WaitGroup, srv *http.Server, name string, gracePeriod time.Duration, failed <-chan struct{}) {
 	defer wg.Done()
 
 	select {
@@ -304,7 +300,7 @@ func gracefulShutdown(ctx context.Context, logger *slog.Logger, wg *sync.WaitGro
 		return
 	}
 
-	ctxShutdown, cancel := context.WithTimeout(context.Background(), connectServerGracePeriod)
+	ctxShutdown, cancel := context.WithTimeout(context.Background(), gracePeriod)
 	defer cancel()
 
 	if err := srv.Shutdown(ctxShutdown); err != nil {

pkg/server/server_test.go

@@ -39,7 +39,7 @@ func TestGracefulShutdownDrainsInflightRequests(t *testing.T) {
 
 	var wg sync.WaitGroup
 	wg.Add(1)
-	go gracefulShutdown(ctx, logger, &wg, srv, "test server", make(chan struct{}))
+	go gracefulShutdown(ctx, logger, &wg, srv, "test server", 5*time.Second, make(chan struct{}))
 
 	serveReturned := make(chan struct{})
 	go func() {
@@ -77,6 +77,59 @@ func TestGracefulShutdownDrainsInflightRequests(t *testing.T) {
 	}
 }
 
+func TestGracefulShutdownCutsOffRequestsAfterGracePeriod(t *testing.T) {
+	logger := slog.New(slog.NewTextHandler(io.Discard, nil))
+
+	entered := make(chan struct{})
+	handlerDone := make(chan struct{})
+	mux := http.NewServeMux()
+	mux.HandleFunc("/stuck", func(w http.ResponseWriter, r *http.Request) {
+		close(entered)
+		time.Sleep(2 * time.Second)
+		close(handlerDone)
+	})
+
+	l, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("listen: %v", err)
+	}
+	srv := &http.Server{Handler: mux}
+	t.Cleanup(func() { srv.Close() })
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+	go gracefulShutdown(ctx, logger, &wg, srv, "test server", 200*time.Millisecond, make(chan struct{}))
+
+	go func() {
+		_ = srv.Serve(l)
+	}()
+
+	go func() {
+		resp, err := http.Get(fmt.Sprintf("http://%s/stuck", l.Addr()))
+		if err == nil {
+			resp.Body.Close()
+		}
+	}()
+
+	<-entered
+	cancel()
+
+	waitStart := time.Now()
+	wg.Wait()
+
+	if elapsed := time.Since(waitStart); elapsed >= 2*time.Second {
+		t.Fatalf("grace period did not cut off the drain, wait took %v", elapsed)
+	}
+	select {
+	case <-handlerDone:
+		t.Fatal("handler finished before the wait released, cutoff never happened")
+	default:
+	}
+}
+
 func freePort(t *testing.T) int {
 	t.Helper()
 	l, err := net.Listen("tcp", "127.0.0.1:0")
@@ -118,6 +171,7 @@ func TestServeConnectReturnsAfterShutdownOnContextCancel(t *testing.T) {
 
 			var cfg Config
 			cfg.Connect.Port = freePort(t)
+			cfg.ShutdownGracePeriod = 5 * time.Second
 			if tc.withMetrics {
 				cfg.MetricsPort = freePort(t)
 			}