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

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