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

This commit is contained in:
tangerineyu
2026-03-25 19:27:33 +08:00
parent 21efce6049
commit 9a0ba45a51
7 changed files with 144 additions and 1 deletions

View File

@@ -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)

View File

@@ -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) }()

View File

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

View File

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

View File

@@ -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",
},
},
}
}

View File

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

View File

@@ -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")
}
}