diff --git a/db/clickhouse.go b/db/clickhouse.go index 6638aed..dac9589 100644 --- a/db/clickhouse.go +++ b/db/clickhouse.go @@ -209,13 +209,23 @@ func (p *Pool) Get() clickhouse.Conn { return p.conn } -// Close 释放底层连接。 -func (p *Pool) Close() { +// Close starts closing the underlying connection and waits within ctx. +func (p *Pool) Close(ctx context.Context) error { p.mu.Lock() - defer p.mu.Unlock() - if p.conn != nil { - _ = p.conn.Close() - p.conn = nil - } + conn := p.conn + 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() + } } diff --git a/main.go b/main.go index ebe8e95..09e226b 100644 --- a/main.go +++ b/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 - select { - case <-stop: - log.Printf("[main] shutdown signal received") - case err := <-serverErr: - if !errors.Is(err, http.ErrServerClosed) { - runErr = fmt.Errorf("server: %w", err) - } - } + runErr := waitForShutdown(stop, serverErr) signal.Stop(stop) - // HTTP、升级连接和日志排空共享同一个关闭总预算。 + // HTTP、升级连接、日志排空和数据库共享同一个关闭总预算。 shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) defer shutdownCancel() - if err := srv.Shutdown(shutdownCtx); err != nil { - log.Printf("[main] http shutdown: %v", err) - } - if err := h.Shutdown(shutdownCtx); err != nil { - log.Printf("[main] upgraded connection shutdown: %v", err) - } - 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") - } + shutdownErr := shutdownAll(shutdownCtx, cancel, + shutdownStep{name: "http", run: srv.Shutdown}, + shutdownStep{name: "upgraded connection", run: h.Shutdown}, + shutdownStep{name: "queue", run: func(ctx context.Context) error { return queue.Stop(ctx) }}, + shutdownStep{name: "db", run: pool.Close}, + ) log.Printf("[main] bye") - return runErr + return errors.Join(runErr, shutdownErr) } diff --git a/main_test.go b/main_test.go new file mode 100644 index 0000000..5902be7 --- /dev/null +++ b/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 +}