fix: unify shutdown resource cleanup

main_test.go

@@ -0,0 +1,89 @@
+package main
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"os"
+	"reflect"
+	"strings"
+	"testing"
+	"time"
+)
+
+func TestWaitForShutdownReturnsListenerError(t *testing.T) {
+	listenErr := errors.New("listen failed")
+	serverErr := make(chan error, 1)
+	serverErr <- listenErr
+
+	err := waitForShutdown(make(chan os.Signal), serverErr)
+	if !errors.Is(err, listenErr) || !strings.Contains(err.Error(), "server") {
+		t.Fatalf("waitForShutdown() error = %v, want wrapped listener error", err)
+	}
+}
+
+func TestWaitForShutdownAcceptsSignalAndServerClosed(t *testing.T) {
+	t.Run("signal", func(t *testing.T) {
+		stop := make(chan os.Signal, 1)
+		stop <- os.Interrupt
+		if err := waitForShutdown(stop, make(chan error)); err != nil {
+			t.Fatalf("waitForShutdown() error = %v, want nil", err)
+		}
+	})
+
+	t.Run("server closed", func(t *testing.T) {
+		serverErr := make(chan error, 1)
+		serverErr <- http.ErrServerClosed
+		if err := waitForShutdown(make(chan os.Signal), serverErr); err != nil {
+			t.Fatalf("waitForShutdown() error = %v, want nil", err)
+		}
+	})
+}
+
+func TestShutdownAllContinuesAfterTimeoutWithSharedDeadline(t *testing.T) {
+	ctx, cancel := context.WithTimeout(context.Background(), time.Minute)
+	defer cancel()
+	wantDeadline, _ := ctx.Deadline()
+
+	var order []string
+	var contexts []context.Context
+	rootCtx, rootCancel := context.WithCancel(context.Background())
+	steps := []shutdownStep{
+		{name: "http", run: recordShutdown(&order, &contexts, "http", nil)},
+		{name: "upgraded", run: recordShutdown(&order, &contexts, "upgraded", nil)},
+		{name: "queue", run: recordShutdown(&order, &contexts, "queue", context.DeadlineExceeded)},
+		{name: "db", run: recordShutdown(&order, &contexts, "db", nil)},
+	}
+
+	err := shutdownAll(ctx, rootCancel, steps...)
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatalf("shutdownAll() error = %v, want deadline exceeded", err)
+	}
+	if want := []string{"http", "upgraded", "queue", "db", "root"}; !reflect.DeepEqual(orderWithRoot(order, rootCtx), want) {
+		t.Fatalf("shutdown order = %v, want %v", orderWithRoot(order, rootCtx), want)
+	}
+	for i, gotCtx := range contexts {
+		gotDeadline, ok := gotCtx.Deadline()
+		if !ok || !gotDeadline.Equal(wantDeadline) {
+			t.Errorf("step %d deadline = %v, %v; want %v, true", i, gotDeadline, ok, wantDeadline)
+		}
+	}
+}
+
+func recordShutdown(order *[]string, contexts *[]context.Context, name string, err error) func(context.Context) error {
+	return func(ctx context.Context) error {
+		*order = append(*order, name)
+		*contexts = append(*contexts, ctx)
+		return err
+	}
+}
+
+func orderWithRoot(order []string, rootCtx context.Context) []string {
+	got := append([]string(nil), order...)
+	select {
+	case <-rootCtx.Done():
+		got = append(got, "root")
+	default:
+	}
+	return got
+}