Merge pull request #4 from tangerineyu/feat/add_pprof

feat(observability): 为 API 和 worker 增加独立 pprof 诊断端口

backend/cmd/main.go

@@ -7,6 +7,7 @@ import (
 	apphttp "feedsystem_video_go/internal/http"
 	rabbitmq "feedsystem_video_go/internal/middleware/rabbitmq"
 	rediscache "feedsystem_video_go/internal/middleware/redis"
+	"feedsystem_video_go/internal/observability"
 	"log"
 	"strconv"
 	"time"
@@ -64,6 +65,16 @@ func main() {
 		defer rmq.Close()
 		log.Printf("RabbitMQ connected")
 	}
+	// Pprof
+	pprofServer, err := observability.NewPprofServer(
+		"API",
+		cfg.ObservabilityConfig.Pprof.Enabled,
+		cfg.ObservabilityConfig.Pprof.ApiAddr,
+	)
+	if err != nil {
+		log.Printf("Failed to start API pprof server: %v", err)
+	}
+	defer pprofServer.Close()
 
 	// 设置路由
 	r := apphttp.SetRouter(sqlDB, cache, rmq)

backend/cmd/worker/main.go

@@ -5,10 +5,12 @@ import (
 	"feedsystem_video_go/internal/config"
 	"feedsystem_video_go/internal/db"
 	rediscache "feedsystem_video_go/internal/middleware/redis"
+	"feedsystem_video_go/internal/observability"
 	"feedsystem_video_go/internal/social"
 	"feedsystem_video_go/internal/video"
 	"feedsystem_video_go/internal/worker"
 	"log"
+	
 	"os"
 	"os/signal"
 	"strconv"
@@ -120,6 +122,16 @@ func main() {
 	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
 	defer stop()
 
+	pprofServer, err := observability.NewPprofServer(
+		"Worker",
+		cfg.ObservabilityConfig.Pprof.Enabled,
+		cfg.ObservabilityConfig.Pprof.WorkerAddr,
+	)
+	if err != nil {
+		log.Printf("Failed to start worker pprof server: %v", err)
+	}
+	defer pprofServer.Close()
+
 	errCh := make(chan error, 4)
 	log.Printf("Worker started, consuming queue=%s", socialQueue)
 	go func() { errCh <- socialWorker.Run(ctx) }()

backend/configs/config.docker.yaml

@@ -20,3 +20,8 @@ rabbitmq:
   username: admin
   password: password123
 
+observability:
+  pprof:
+    enabled: false
+    api_addr: localhost:6060
+    worker_addr: localhost:6061

backend/configs/config.yaml

@@ -19,4 +19,9 @@ rabbitmq:
   port: 5672
   username: admin
   password: password123
-  
+  
+observability:
+  pprof:
+    enabled: true
+    api_addr: localhost:6060
+    worker_addr: localhost:6061

backend/internal/config/loadconfig.go

@@ -12,6 +12,7 @@ type Config struct {
 	Database DatabaseConfig `yaml:"database"`
 	Redis    RedisConfig    `yaml:"redis"`
 	RabbitMQ RabbitMQConfig `yaml:"rabbitmq"`
+	ObservabilityConfig ObservabilityConfig `yaml:"observability"`
 }
 
 type ServerConfig struct {
@@ -40,6 +41,14 @@ type RabbitMQConfig struct {
 	Password string `yaml:"password"`
 }
 
+type ObservabilityConfig struct {
+	Pprof PprofConfig `yaml:"pprof"`
+}
+type PprofConfig struct {
+	Enabled bool `yaml:"enabled"`
+	ApiAddr string `yaml:"api_addr"`
+	WorkerAddr string `yaml:"worker_addr"`
+}
 func Load(filename string) (Config, error) {
 	data, err := os.ReadFile(filename)
 	if err != nil {
@@ -90,5 +99,12 @@ func DefaultLocalConfig() Config {
 			Username: "admin",
 			Password: "password123",
 		},
+		ObservabilityConfig: ObservabilityConfig{
+			Pprof: PprofConfig{
+				Enabled: true,
+				ApiAddr: "localhost:6060",
+				WorkerAddr: "localhost:6061",
+			},
+		},
 	}
 }

backend/internal/observability/pprof.go

@@ -0,0 +1,75 @@
+package observability
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log"
+	"net"
+	"net/http"
+	"net/http/pprof"
+	"time"
+
+)
+
+type PprofServer struct {
+	name string
+	server *http.Server
+	shutdownTimeout time.Duration
+}
+func NewPprofMux() *http.ServeMux {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/debug/pprof/", pprof.Index)
+	mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+	mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+	mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+	mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+
+	return mux
+}
+
+func NewPprofServer(name string, enabled bool, addr string) (*PprofServer, error) {
+	pprofServer := &PprofServer{
+		name: name,
+		shutdownTimeout: 3 * time.Second,
+	}
+	if !enabled || addr == "" {
+		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)
+	}
+	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 := pprofServer.server.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
+			log.Printf("%s pprof server error: %v", name, err)
+		}
+	}()
+	return pprofServer, nil
+}
+
+func Shutdown(ctx context.Context, srv *http.Server) error{
+	if srv == nil {
+		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

@@ -0,0 +1,42 @@
+package observability
+
+import (
+	"net/http"
+	"net/http/httptest"
+	"testing"
+)
+func TestNewPprofMux(t *testing.T) {
+	t.Parallel()
+
+	req := httptest.NewRequest(http.MethodGet, "/debug/pprof/", nil)
+	rr := httptest.NewRecorder()
+
+	NewPprofMux().ServeHTTP(rr, req)
+
+	if rr.Code != http.StatusOK {
+		t.Errorf("Expected status code 200, got %d", rr.Code)
+	}
+}
+func TestNewPprofServerWithDisabled(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 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)
+	}
+}