From bfc319edb30fec2628c33f01458f877af49421f1 Mon Sep 17 00:00:00 2001
From: Fabio Bonelli <fb@fabiobonelli.it>
Date: Thu, 12 Nov 2026 01:44:00 +0000
Subject: [PATCH] fix: drain active requests on shutdown

---
 main.go          |  18 ++++++-
 shutdown_test.go | 122 +++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 139 insertions(+), 1 deletion(-)
 create mode 100644 shutdown_test.go

diff --git a/main.go b/main.go
index 828d044..456deac 100644
--- a/main.go
+++ b/main.go
@@ -1,7 +1,12 @@
 package main
 
 import (
+	"context"
+	"fmt"
 	"log"
+	"os"
+	"os/signal"
+	"syscall"
 	"time"
 
 	"github.com/gofiber/fiber/v3"
@@ -19,11 +24,22 @@ var (
 
 func main() {
 	app := Setup()
-	if err := app.Listen(":3000"); err != nil {
+	if err := listen(app, ":3000"); err != nil {
 		log.Fatal(err)
 	}
 }
 
+func listen(app *fiber.App, address string) error {
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
+	if err := app.Listen(address, fiber.ListenConfig{GracefulContext: ctx}); err != nil {
+		return fmt.Errorf("listen: %w", err)
+	}
+
+	return nil
+}
+
 func Setup() *fiber.App {
 	methods := append(fiber.DefaultMethods, "QUERY") //nolint:gocritic // We want a new slice here
 
diff --git a/shutdown_test.go b/shutdown_test.go
new file mode 100644
index 0000000..ea60818
--- /dev/null
+++ b/shutdown_test.go
@@ -0,0 +1,122 @@
+package main
+
+import (
+	"bufio"
+	"context"
+	"fmt"
+	"io"
+	"net/http"
+	"os"
+	"os/exec"
+	"strings"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/gofiber/fiber/v3"
+	"github.com/stretchr/testify/require"
+)
+
+func TestShutdownProcess(t *testing.T) {
+	if os.Getenv("TEST_SHUTDOWN_PROCESS") != "1" {
+		return
+	}
+
+	server := Setup()
+	server.Hooks().OnListen(func(data fiber.ListenData) error {
+		fmt.Println("test:address:http://" + data.Host + ":" + data.Port)
+		return nil
+	})
+	server.Hooks().OnPreShutdown(func() error {
+		fmt.Println("test:shutdown")
+		return nil
+	})
+	server.Get("/slow", func(ctx fiber.Ctx) error {
+		fmt.Println("test:started")
+		var release [1]byte
+		if _, err := os.Stdin.Read(release[:]); err != nil {
+			return err
+		}
+		return ctx.SendString("completed")
+	})
+	require.NoError(t, listen(server, "127.0.0.1:0"))
+}
+
+func TestShutdownFinishesRequest(t *testing.T) {
+	for name, sig := range map[string]os.Signal{"SIGTERM": syscall.SIGTERM, "SIGINT": os.Interrupt} {
+		t.Run(name, func(t *testing.T) {
+			ctx, cancel := context.WithTimeout(t.Context(), 15*time.Second)
+			defer cancel()
+			cmd := exec.CommandContext(ctx, os.Args[0], "-test.run=^TestShutdownProcess$")
+			cmd.Env = append(os.Environ(), "TEST_SHUTDOWN_PROCESS=1")
+			cmd.Stderr = os.Stderr
+			stdout, err := cmd.StdoutPipe()
+			require.NoError(t, err)
+			stdin, err := cmd.StdinPipe()
+			require.NoError(t, err)
+			require.NoError(t, cmd.Start())
+			done := make(chan struct{})
+			var exitErr error
+			t.Cleanup(func() {
+				_ = stdin.Close()
+				_ = cmd.Process.Kill()
+				<-done
+			})
+
+			events := make(chan string, 8)
+			go func() {
+				defer close(events)
+				scanner := bufio.NewScanner(stdout)
+				for scanner.Scan() {
+					if strings.HasPrefix(scanner.Text(), "test:") {
+						events <- scanner.Text()
+					}
+				}
+			}()
+			go func() {
+				exitErr = cmd.Wait()
+				close(done)
+			}()
+			nextEvent := func() string {
+				t.Helper()
+				select {
+				case event, ok := <-events:
+					require.True(t, ok, "server exited before the expected event")
+					return event
+				case <-ctx.Done():
+					t.Fatal("timed out waiting for the server")
+					return ""
+				}
+			}
+
+			address := nextEvent()
+			require.True(t, strings.HasPrefix(address, "test:address:"))
+			request, err := http.NewRequestWithContext(ctx, http.MethodGet,
+				strings.TrimPrefix(address, "test:address:")+"/slow", nil)
+			require.NoError(t, err)
+			response := make(chan *http.Response, 1)
+			requestErr := make(chan error, 1)
+			go func() {
+				res, err := http.DefaultClient.Do(request)
+				response <- res
+				requestErr <- err
+			}()
+
+			require.Equal(t, "test:started", nextEvent())
+			require.NoError(t, cmd.Process.Signal(sig))
+			require.Equal(t, "test:shutdown", nextEvent())
+			// Release the handler only after the real signal starts shutdown.
+			_, err = stdin.Write([]byte{1})
+			require.NoError(t, err)
+			require.NoError(t, <-requestErr)
+			res := <-response
+			defer res.Body.Close()
+			body, err := io.ReadAll(res.Body)
+			require.NoError(t, err)
+			require.Equal(t, http.StatusOK, res.StatusCode)
+			require.Equal(t, "completed", string(body))
+			<-done
+			require.NoError(t, exitErr, "shutdown must exit successfully")
+		})
+	}
+}