diff --git a/main.go b/main.go index 8c45cdc..9961a60 100644 --- a/main.go +++ b/main.go @@ -83,6 +83,10 @@ func run() error { } log.Printf("[main] filter mode=%s patterns=%d", filter.Mode, len(filter.Patterns)) + stop := make(chan os.Signal, 1) + signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM) + defer signal.Stop(stop) + rootCtx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -114,10 +118,7 @@ func run() error { serverErr <- srv.ListenAndServe() }() - stop := make(chan os.Signal, 1) - signal.Notify(stop, syscall.SIGINT, syscall.SIGTERM) runErr := waitForShutdown(stop, serverErr) - signal.Stop(stop) // HTTP、升级连接、日志排空和数据库共享同一个关闭总预算。 shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 30*time.Second) diff --git a/main_test.go b/main_test.go index ee483e1..4fc13bb 100644 --- a/main_test.go +++ b/main_test.go @@ -3,6 +3,9 @@ package main import ( "context" "errors" + "go/ast" + "go/parser" + "go/token" "net/http" "os" "reflect" @@ -11,6 +14,63 @@ import ( "time" ) +func TestRunRegistersSignalsBeforeStartingResources(t *testing.T) { + fset := token.NewFileSet() + file, err := parser.ParseFile(fset, "main.go", nil, 0) + if err != nil { + t.Fatalf("parse main.go: %v", err) + } + + var runBody *ast.BlockStmt + for _, declaration := range file.Decls { + function, ok := declaration.(*ast.FuncDecl) + if ok && function.Name.Name == "run" { + runBody = function.Body + break + } + } + if runBody == nil { + t.Fatal("main.go does not define run") + } + + positions := make(map[string]token.Pos) + ast.Inspect(runBody, func(node ast.Node) bool { + call, ok := node.(*ast.CallExpr) + if !ok { + return true + } + selector, ok := call.Fun.(*ast.SelectorExpr) + if !ok { + return true + } + owner, ok := selector.X.(*ast.Ident) + if !ok { + return true + } + name := owner.Name + "." + selector.Sel.Name + switch name { + case "signal.Notify", "db.NewPool", "queue.Start", "srv.ListenAndServe": + positions[name] = call.Pos() + } + return true + }) + + notifyPos, ok := positions["signal.Notify"] + if !ok { + t.Fatal("run does not register for shutdown signals") + } + for _, start := range []string{"db.NewPool", "queue.Start", "srv.ListenAndServe"} { + startPos, ok := positions[start] + if !ok { + t.Fatalf("run does not call %s", start) + } + if notifyPos >= startPos { + t.Errorf("signal.Notify at line %d must precede %s at line %d", + fset.Position(notifyPos).Line, start, fset.Position(startPos).Line) + } + } +} + func TestWaitForShutdownReturnsListenerError(t *testing.T) { listenErr := errors.New("listen failed") serverErr := make(chan error, 1)