From 9a0ba45a51f589ddc7c2bb056c2407db2ea4f4a0 Mon Sep 17 00:00:00 2001 From: tangerineyu Date: Wed, 25 Mar 2026 19:27:33 +0800 Subject: [PATCH 1/2] =?UTF-8?q?feat(observability):=20=E4=B8=BA=20API=20?= =?UTF-8?q?=E5=92=8C=20worker=20=E5=A2=9E=E5=8A=A0=E7=8B=AC=E7=AB=8B=20ppr?= =?UTF-8?q?of=20=E8=AF=8A=E6=96=AD=E7=AB=AF=E5=8F=A3?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/cmd/main.go | 17 +++++++ backend/cmd/worker/main.go | 18 +++++++ backend/configs/config.docker.yaml | 5 ++ backend/configs/config.yaml | 7 ++- backend/internal/config/loadconfig.go | 16 ++++++ backend/internal/observability/pprof.go | 51 ++++++++++++++++++++ backend/internal/observability/pprof_test.go | 31 ++++++++++++ 7 files changed, 144 insertions(+), 1 deletion(-) create mode 100644 backend/internal/observability/pprof.go create mode 100644 backend/internal/observability/pprof_test.go 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") + } +} From 145bf3fa8698d37e20dd544689b9204477c2983b Mon Sep 17 00:00:00 2001 From: tangerineyu Date: Thu, 26 Mar 2026 12:17:05 +0800 Subject: [PATCH 2/2] =?UTF-8?q?refactor:=E7=AE=80=E5=8C=96main=E5=87=BD?= =?UTF-8?q?=E6=95=B0=E4=B8=ADpprof=E7=9A=84=E5=85=B3=E9=97=AD=E8=B0=83?= =?UTF-8?q?=E7=94=A8?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- backend/cmd/main.go | 10 ++---- backend/cmd/worker/main.go | 10 ++---- backend/internal/observability/pprof.go | 36 ++++++++++++++++---- backend/internal/observability/pprof_test.go | 23 +++++++++---- 4 files changed, 51 insertions(+), 28 deletions(-) diff --git a/backend/cmd/main.go b/backend/cmd/main.go index 9d677f8..e302e66 100644 --- a/backend/cmd/main.go +++ b/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() { - 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) - } - }() + defer pprofServer.Close() // 设置路由 r := apphttp.SetRouter(sqlDB, cache, rmq) diff --git a/backend/cmd/worker/main.go b/backend/cmd/worker/main.go index eaab8d0..5c74bb1 100644 --- a/backend/cmd/worker/main.go +++ b/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() { - 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) - } - }() + defer pprofServer.Close() errCh := make(chan error, 4) log.Printf("Worker started, consuming queue=%s", socialQueue) diff --git a/backend/internal/observability/pprof.go b/backend/internal/observability/pprof.go index 7b3c2e2..d4d470f 100644 --- a/backend/internal/observability/pprof.go +++ b/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 } \ No newline at end of file diff --git a/backend/internal/observability/pprof_test.go b/backend/internal/observability/pprof_test.go index 168f8cd..3c45df6 100644 --- a/backend/internal/observability/pprof_test.go +++ b/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) + } +} \ No newline at end of file