diff --git a/backend/cmd/main.go b/backend/cmd/main.go index e853ed5..e302e66 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,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) diff --git a/backend/cmd/worker/main.go b/backend/cmd/worker/main.go index 1618b3a..5c74bb1 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,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) }() 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..d4d470f --- /dev/null +++ b/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 +} \ 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..3c45df6 --- /dev/null +++ b/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) + } +} \ No newline at end of file