fix: unify shutdown resource cleanup

db/clickhouse.go

@@ -209,13 +209,23 @@ func (p *Pool) Get() clickhouse.Conn {
 	return p.conn
 }
 
-// Close 释放底层连接。
+// Close starts closing the underlying connection and waits within ctx.
-func (p *Pool) Close() {
+func (p *Pool) Close(ctx context.Context) error {
 	p.mu.Lock()
-	defer p.mu.Unlock()
+	conn := p.conn
-	if p.conn != nil {
-		_ = p.conn.Close()
 	p.conn = nil
-	}
 	p.healthy = false
+	p.mu.Unlock()
+	if conn == nil {
+		return nil
+	}
+
+	done := make(chan error, 1)
+	go func() { done <- conn.Close() }()
+	select {
+	case err := <-done:
+		return err
+	case <-ctx.Done():
+		return ctx.Err()
+	}
 }

main.go

@@ -25,6 +25,36 @@ func main() {
 	}
 }
 
+type shutdownStep struct {
+	name string
+	run  func(context.Context) error
+}
+
+func waitForShutdown(stop <-chan os.Signal, serverErr <-chan error) error {
+	select {
+	case <-stop:
+		log.Printf("[main] shutdown signal received")
+		return nil
+	case err := <-serverErr:
+		if errors.Is(err, http.ErrServerClosed) {
+			return nil
+		}
+		return fmt.Errorf("server: %w", err)
+	}
+}
+
+func shutdownAll(ctx context.Context, cancelRoot context.CancelFunc, steps ...shutdownStep) error {
+	defer cancelRoot()
+	var shutdownErr error
+	for _, step := range steps {
+		if err := step.run(ctx); err != nil {
+			log.Printf("[main] %s shutdown: %v", step.name, err)
+			shutdownErr = errors.Join(shutdownErr, fmt.Errorf("%s shutdown: %w", step.name, err))
+		}
+	}
+	return shutdownErr
+}
+
 func run() error {
 	cfg, err := config.Load()
 	if err != nil {
@@ -73,37 +103,18 @@ func run() error {
 
 	stop := make(chan os.Signal, 1)
 	signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM)
-	var runErr error
+	runErr := waitForShutdown(stop, serverErr)
-	select {
-	case <-stop:
-		log.Printf("[main] shutdown signal received")
-	case err := <-serverErr:
-		if !errors.Is(err, http.ErrServerClosed) {
-			runErr = fmt.Errorf("server: %w", err)
-		}
-	}
 	signal.Stop(stop)
 
-	// HTTP、升级连接和日志排空共享同一个关闭总预算。
+	// HTTP、升级连接、日志排空和数据库共享同一个关闭总预算。
 	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second)
 	defer shutdownCancel()
-	if err := srv.Shutdown(shutdownCtx); err != nil {
+	shutdownErr := shutdownAll(shutdownCtx, cancel,
-		log.Printf("[main] http shutdown: %v", err)
+		shutdownStep{name: "http", run: srv.Shutdown},
-	}
+		shutdownStep{name: "upgraded connection", run: h.Shutdown},
-	if err := h.Shutdown(shutdownCtx); err != nil {
+		shutdownStep{name: "queue", run: func(ctx context.Context) error { return queue.Stop(ctx) }},
-		log.Printf("[main] upgraded connection shutdown: %v", err)
+		shutdownStep{name: "db", run: pool.Close},
-	}
+	)
-	queueStopped := true
-	if err := queue.Stop(shutdownCtx); err != nil {
-		log.Printf("[main] queue shutdown: %v", err)
-		queueStopped = false
-	}
-	cancel()
-	if queueStopped {
-		pool.Close()
-	} else {
-		log.Printf("[main] skip db close while queue workers are still exiting")
-	}
 	log.Printf("[main] bye")
-	return runErr
+	return errors.Join(runErr, shutdownErr)
 }

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
+}