diff --git a/backend/cmd/main.go b/backend/cmd/main.go index e853ed5..9d677f8 100644 --- a/backend/cmd/main.go +++ b/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,22 @@ func main() { defer rmq.Close() log.Printf("RabbitMQ connected") } + // Pprof + pprofServer, err := observability.StartPprofServer( + "API", + cfg.ObservabilityConfig.Pprof.Enabled, + cfg.ObservabilityConfig.Pprof.ApiAddr, + ) + 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) + } + }() // 设置路由 r := apphttp.SetRouter(sqlDB, cache, rmq) diff --git a/backend/cmd/worker/main.go b/backend/cmd/worker/main.go index 1618b3a..eaab8d0 100644 --- a/backend/cmd/worker/main.go +++ b/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,22 @@ func main() { ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM) defer stop() + pprofServer, err := observability.StartPprofServer( + "Worker", + cfg.ObservabilityConfig.Pprof.Enabled, + cfg.ObservabilityConfig.Pprof.WorkerAddr, + ) + 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) + } + }() + errCh := make(chan error, 4) log.Printf("Worker started, consuming queue=%s", socialQueue) go func() { errCh <- socialWorker.Run(ctx) }() diff --git a/backend/configs/config.docker.yaml b/backend/configs/config.docker.yaml index 2ee2181..5712403 100644 --- a/backend/configs/config.docker.yaml +++ b/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 \ No newline at end of file diff --git a/backend/configs/config.yaml b/backend/configs/config.yaml index 223f2c5..e1fa2cc 100644 --- a/backend/configs/config.yaml +++ b/backend/configs/config.yaml @@ -19,4 +19,9 @@ rabbitmq: port: 5672 username: admin password: password123 - \ No newline at end of file + +observability: + pprof: + enabled: true + api_addr: localhost:6060 + worker_addr: localhost:6061 \ No newline at end of file diff --git a/backend/internal/config/loadconfig.go b/backend/internal/config/loadconfig.go index 0cd3523..817ba98 100644 --- a/backend/internal/config/loadconfig.go +++ b/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", + }, + }, } } \ No newline at end of file diff --git a/backend/internal/observability/pprof.go b/backend/internal/observability/pprof.go new file mode 100644 index 0000000..7b3c2e2 --- /dev/null +++ b/backend/internal/observability/pprof.go @@ -0,0 +1,51 @@ +package observability + +import ( + "errors" + "fmt" + "log" + "net" + "net/http" + "net/http/pprof" + "time" + "context" +) +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 StartPprofServer(name string, enabled bool, addr string) (*http.Server, error) { + if !enabled || addr == "" { + return nil, 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{ + 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) { + log.Printf("%s pprof server error: %v", name, err) + } + }() + return server, nil +} + +func Shutdown(ctx context.Context, srv *http.Server) error{ + if srv == nil { + return nil + } + return srv.Shutdown(ctx) +} \ No newline at end of file diff --git a/backend/internal/observability/pprof_test.go b/backend/internal/observability/pprof_test.go new file mode 100644 index 0000000..168f8cd --- /dev/null +++ b/backend/internal/observability/pprof_test.go @@ -0,0 +1,31 @@ +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 TestStartPprofServerWithDisabled(t *testing.T) { + t.Parallel() + + server, err := StartPprofServer("api", false, "localhost:6060") + if err != nil { + t.Fatalf("Failed to start pprof server: %v", err) + } + if server != nil { + t.Fatalf("Expected nil server when pprof is disabled, got non-nil") + } +}