refactor:简化main函数中pprof的关闭调用

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() {
+	defer pprofServer.Close()
-		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)
-		}
-	}()
 
 	// 设置路由
 	r := apphttp.SetRouter(sqlDB, cache, rmq)

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() {
+	defer pprofServer.Close()
-		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)
-		}
-	}()
 
 	errCh := make(chan error, 4)
 	log.Printf("Worker started, consuming queue=%s", socialQueue)

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
 }

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 TestNewPprofServerWithDisabled(t *testing.T) {
-func TestStartPprofServerWithDisabled(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 {
+	if pprofServer != nil {
-		t.Fatalf("Expected nil server when pprof is disabled, got non-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)
+	}
+}