diff --git a/backend/cmd/main.go b/backend/cmd/main.go index 9d677f8..e302e66 100644 --- a/backend/cmd/main.go +++ b/backend/cmd/main.go @@ -66,7 +66,7 @@ func main() { log.Printf("RabbitMQ connected") } // Pprof - pprofServer, err := observability.StartPprofServer( + pprofServer, err := observability.NewPprofServer( "API", cfg.ObservabilityConfig.Pprof.Enabled, cfg.ObservabilityConfig.Pprof.ApiAddr, @@ -74,13 +74,7 @@ func main() { if err != nil { log.Printf("Failed to start API pprof server: %v", err) } - defer func() { - shutdownCtx, cancel := context.WithTimeout(context.Background(), 3 * time.Second) - defer cancel() - if err := observability.Shutdown(shutdownCtx, pprofServer); err != nil { - log.Printf("Failed to shutdown API pprof server: %v", err) - } - }() + defer pprofServer.Close() // 设置路由 r := apphttp.SetRouter(sqlDB, cache, rmq) diff --git a/backend/cmd/worker/main.go b/backend/cmd/worker/main.go index eaab8d0..5c74bb1 100644 --- a/backend/cmd/worker/main.go +++ b/backend/cmd/worker/main.go @@ -122,7 +122,7 @@ func main() { ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() - pprofServer, err := observability.StartPprofServer( + pprofServer, err := observability.NewPprofServer( "Worker", cfg.ObservabilityConfig.Pprof.Enabled, cfg.ObservabilityConfig.Pprof.WorkerAddr, @@ -130,13 +130,7 @@ func main() { if err != nil { log.Printf("Failed to start worker pprof server: %v", err) } - defer func() { - shutdownCtx, cancel := context.WithTimeout(context.Background(), 3 * time.Second) - defer cancel() - if err := observability.Shutdown(shutdownCtx, pprofServer); err != nil { - log.Printf("Failed to shutdown worker pprof server: %v", err) - } - }() + defer pprofServer.Close() errCh := make(chan error, 4) log.Printf("Worker started, consuming queue=%s", socialQueue) diff --git a/backend/internal/observability/pprof.go b/backend/internal/observability/pprof.go index 7b3c2e2..d4d470f 100644 --- a/backend/internal/observability/pprof.go +++ b/backend/internal/observability/pprof.go @@ -1,6 +1,7 @@ package observability import ( + "context" "errors" "fmt" "log" @@ -8,8 +9,14 @@ import ( "net/http" "net/http/pprof" "time" - "context" + ) + +type PprofServer struct { + name string + server *http.Server + shutdownTimeout time.Duration +} func NewPprofMux() *http.ServeMux { mux := http.NewServeMux() mux.HandleFunc("/debug/pprof/", pprof.Index) @@ -21,26 +28,30 @@ func NewPprofMux() *http.ServeMux { return mux } -func StartPprofServer(name string, enabled bool, addr string) (*http.Server, error) { +func NewPprofServer(name string, enabled bool, addr string) (*PprofServer, error) { + pprofServer := &PprofServer{ + name: name, + shutdownTimeout: 3 * time.Second, + } if !enabled || addr == "" { - return nil, nil + return pprofServer, nil } ln, err := net.Listen("tcp", addr) if err != nil { return nil, fmt.Errorf("failed to start %s pprof server on %s: %w", name, addr, err) } - server := &http.Server{ + pprofServer.server = &http.Server{ Addr: addr, Handler: NewPprofMux(), ReadHeaderTimeout: 5 * time.Second, } go func() { log.Printf("%s pprof listening on %s", name, addr) - if err := server.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) { + if err := pprofServer.server.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) { log.Printf("%s pprof server error: %v", name, err) } }() - return server, nil + return pprofServer, nil } func Shutdown(ctx context.Context, srv *http.Server) error{ @@ -48,4 +59,17 @@ func Shutdown(ctx context.Context, srv *http.Server) error{ return nil } return srv.Shutdown(ctx) +} + +func (s *PprofServer) Close() error { + if s == nil { + return nil + } + shutdownCtx, cancel := context.WithTimeout(context.Background(), 3 * time.Second) + defer cancel() + if err := Shutdown(shutdownCtx, s.server); err != nil { + log.Printf("Failed to shutdown %s pprof server: %v", s.name, err) + return err + } + return nil } \ No newline at end of file diff --git a/backend/internal/observability/pprof_test.go b/backend/internal/observability/pprof_test.go index 168f8cd..3c45df6 100644 --- a/backend/internal/observability/pprof_test.go +++ b/backend/internal/observability/pprof_test.go @@ -17,15 +17,26 @@ func TestNewPprofMux(t *testing.T) { t.Errorf("Expected status code 200, got %d", rr.Code) } } - -func TestStartPprofServerWithDisabled(t *testing.T) { +func TestNewPprofServerWithDisabled(t *testing.T) { t.Parallel() - server, err := StartPprofServer("api", false, "localhost:6060") + pprofServer, err := NewPprofServer("api", false, "localhost:6060") if err != nil { - t.Fatalf("Failed to start pprof server: %v", err) + t.Fatalf("Failed to create pprof server: %v", err) } - if server != nil { - t.Fatalf("Expected nil server when pprof is disabled, got non-nil") + if pprofServer != nil { + t.Fatalf("Expected nil pprof server when disabled, got non-nil") } } + +func TestPprofServerCloseWithDisabledServer(t *testing.T) { + t.Parallel() + + pprofServer, err := NewPprofServer("api", false, "localhost:6060") + if err != nil { + t.Fatalf("Failed to create pprof server: %v", err) + } + if err := pprofServer.Close(); err != nil { + t.Fatalf("Expected no error when closing disabled pprof server, got: %v", err) + } +} \ No newline at end of file