chore(push-gateway): ai02 module updates - config, hub, ws, kafka, health, docs
This commit is contained in:
SpecialX
@@ -1,8 +1,33 @@
|
||||
// Command push-gateway is the Edu platform's real-time WebSocket push gateway.
|
||||
//
|
||||
// It maintains per-user WebSocket connections, delivers directed and broadcast
|
||||
// messages from msg (via HTTP /internal/* APIs and Kafka), and synchronizes
|
||||
// online presence across instances via Redis Pub/Sub.
|
||||
//
|
||||
// Startup sequence:
|
||||
// 1. config.Load (env vars, DevMode detection)
|
||||
// 2. observability.InitLogger (slog JSON/text)
|
||||
// 3. observability.InitTracer (shared-go/tracer → OTLP)
|
||||
// 4. observability.NewMetrics (Prometheus)
|
||||
// 5. hub.NewHub (in-memory connection registry)
|
||||
// 6. redisclient.New + SetHub + RebuildPresenceOnStartup (ISSUE-058)
|
||||
// 7. redisclient.SubscribeAll (cross-instance Pub/Sub fanout)
|
||||
// 8. kafkaconsumer.New + goroutine Run (edu.notification.requested)
|
||||
// 9. ws.NewHandler (JWT RS256, X-Internal-Token, WebSocket upgrade)
|
||||
// 10. health.NewReadyzer (/readyz soft-failure probe)
|
||||
// 11. gin router + http.Server
|
||||
//
|
||||
// Graceful shutdown (SIGINT/SIGTERM):
|
||||
// 1. Hub.CloseAll (send close frame 1001 to every live connection)
|
||||
// 2. http.Server.Shutdown (stop accepting new HTTP/WebSocket requests)
|
||||
// 3. Kafka consumer goroutine cancel + reader.Close
|
||||
// 4. Redis Pub/Sub cancel + client.Close
|
||||
// 5. Tracer shutdown (flush pending spans)
|
||||
package main
|
||||
|
||||
import (
|
||||
"context"
|
||||
"log"
|
||||
"errors"
|
||||
"net/http"
|
||||
"os"
|
||||
"os/signal"
|
||||
@@ -10,76 +35,195 @@ import (
|
||||
"time"
|
||||
|
||||
"github.com/edu-cloud/push-gateway/internal/config"
|
||||
"github.com/edu-cloud/push-gateway/internal/health"
|
||||
"github.com/edu-cloud/push-gateway/internal/hub"
|
||||
"github.com/edu-cloud/push-gateway/internal/kafkaconsumer"
|
||||
"github.com/edu-cloud/push-gateway/internal/observability"
|
||||
"github.com/edu-cloud/push-gateway/internal/redisclient"
|
||||
"github.com/edu-cloud/push-gateway/internal/ws"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/prometheus/client_golang/prometheus/promhttp"
|
||||
"go.opentelemetry.io/contrib/instrumentation/github.com/gin-gonic/gin/otelgin"
|
||||
)
|
||||
|
||||
const serviceName = "push-gateway"
|
||||
|
||||
// shutdownTimeout bounds the graceful shutdown of the HTTP server. After this
|
||||
// window in-flight requests are forcibly closed. WebSocket long-poll
|
||||
// connections are drained by Hub.CloseAll before this timer starts.
|
||||
const shutdownTimeout = 10 * time.Second
|
||||
|
||||
func main() {
|
||||
cfg := config.Load()
|
||||
// 初始化 OpenTelemetry tracer（endpoint 为空时自动跳过）
|
||||
tracerShutdown := observability.InitTracer("push-gateway", cfg.OTLPEndpoint)
|
||||
|
||||
// 1. Structured logger (slog).
|
||||
logger := observability.InitLogger(cfg.DevMode)
|
||||
logger.Info("push-gateway starting",
|
||||
"dev_mode", cfg.DevMode,
|
||||
"instance_id", cfg.InstanceID,
|
||||
"port", cfg.Port,
|
||||
"max_conns_per_user", cfg.MaxConnsPerUser,
|
||||
"jwks_url", cfg.JWKSURL,
|
||||
"redis_url", cfg.RedisURL,
|
||||
"kafka_brokers", cfg.KafkaBrokers,
|
||||
"kafka_topic", cfg.KafkaNotificationTopic,
|
||||
)
|
||||
|
||||
// 2. OpenTelemetry tracer (best-effort).
|
||||
tracerShutdown := observability.InitTracer(serviceName, cfg.OTLPEndpoint)
|
||||
defer tracerShutdown()
|
||||
|
||||
// 3. Prometheus metrics.
|
||||
metrics := observability.NewMetrics()
|
||||
|
||||
// 4. Hub: in-memory connection registry.
|
||||
h := hub.NewHub(cfg.MaxConnsPerUser)
|
||||
|
||||
// 5. Redis client (optional in DevMode). When present, wire presence
|
||||
// hooks to the Hub and rebuild the online SET (ISSUE-058).
|
||||
var redisClient *redisclient.Client
|
||||
if cfg.RedisURL != "" {
|
||||
rc, err := redisclient.New(cfg.RedisURL, cfg.InstanceID, metrics)
|
||||
if err != nil {
|
||||
logger.Error("redis connect failed; running without cross-instance fanout",
|
||||
"err", err, "url", cfg.RedisURL)
|
||||
} else {
|
||||
redisClient = rc
|
||||
redisClient.SetHub(h)
|
||||
rebuildCtx, rebuildCancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
if err := redisClient.RebuildPresenceOnStartup(rebuildCtx); err != nil {
|
||||
logger.Warn("redis presence rebuild failed (continuing)", "err", err)
|
||||
}
|
||||
rebuildCancel()
|
||||
logger.Info("redis presence synced", "url", cfg.RedisURL)
|
||||
}
|
||||
} else {
|
||||
logger.Warn("REDIS_URL empty; cross-instance fanout disabled")
|
||||
}
|
||||
|
||||
// 6. Redis Pub/Sub subscriber (cross-instance message fanout).
|
||||
var pubsubCancel func() error
|
||||
if redisClient != nil {
|
||||
subCtx, subCancel := context.WithCancel(context.Background())
|
||||
defer subCancel()
|
||||
var err error
|
||||
pubsubCancel, err = redisClient.SubscribeAll(subCtx)
|
||||
if err != nil {
|
||||
logger.Error("redis SubscribeAll failed; cross-instance fanout disabled",
|
||||
"err", err)
|
||||
} else {
|
||||
logger.Info("redis pubsub subscriber active")
|
||||
}
|
||||
}
|
||||
|
||||
// 7. Kafka consumer (edu.notification.requested). Started only when at
|
||||
// least one broker is configured; otherwise the HTTP /internal/push API
|
||||
// is the only delivery channel.
|
||||
var kafkaConsumer *kafkaconsumer.Consumer
|
||||
var kafkaCancel context.CancelFunc
|
||||
if len(cfg.KafkaBrokers) > 0 && cfg.KafkaNotificationTopic != "" {
|
||||
kafkaConsumer = kafkaconsumer.New(kafkaconsumer.Config{
|
||||
Brokers: cfg.KafkaBrokers,
|
||||
Topic: cfg.KafkaNotificationTopic,
|
||||
GroupID: cfg.KafkaConsumerGroup,
|
||||
}, h, redisClient, metrics)
|
||||
kafkaCtx, cancel := context.WithCancel(context.Background())
|
||||
kafkaCancel = cancel
|
||||
go func() {
|
||||
logger.Info("kafka consumer starting",
|
||||
"topic", cfg.KafkaNotificationTopic, "group", cfg.KafkaConsumerGroup)
|
||||
if err := kafkaConsumer.Run(kafkaCtx); err != nil &&
|
||||
!errors.Is(err, context.Canceled) {
|
||||
logger.Error("kafka consumer exited with error", "err", err)
|
||||
}
|
||||
}()
|
||||
} else {
|
||||
logger.Warn("KAFKA_BROKERS empty; notification consumption disabled")
|
||||
}
|
||||
|
||||
// 8. WebSocket + internal HTTP handlers.
|
||||
wsHandler := ws.NewHandler(h, cfg, redisClient, metrics)
|
||||
|
||||
// 9. /readyz probe (soft failure per ARB-015 §17.4).
|
||||
readyzer := health.NewReadyzer(h, redisClient, kafkaConsumer, serviceName, cfg.InstanceID)
|
||||
|
||||
// 10. Gin router.
|
||||
gin.SetMode(gin.ReleaseMode)
|
||||
|
||||
h := hub.NewHub()
|
||||
wsHandler := ws.NewHandler(h, cfg.JWTSecret, cfg.DevMode, cfg.InternalAPIKey)
|
||||
|
||||
r := gin.New()
|
||||
r.Use(gin.Recovery())
|
||||
// OpenTelemetry 自动埋点（HTTP 请求/响应 span）
|
||||
r.Use(otelgin.Middleware("push-gateway"))
|
||||
r.Use(otelgin.Middleware(serviceName))
|
||||
|
||||
r.GET("/healthz", func(c *gin.Context) {
|
||||
c.JSON(200, gin.H{"status": "ok", "service": "push-gateway"})
|
||||
})
|
||||
|
||||
// 就绪探针：检查 WebSocket hub 状态
|
||||
r.GET("/readyz", func(c *gin.Context) {
|
||||
c.JSON(200, gin.H{
|
||||
"status": "ok",
|
||||
"service": "push-gateway",
|
||||
"connections": h.ClientCount(),
|
||||
})
|
||||
})
|
||||
|
||||
// Prometheus 指标端点
|
||||
// Liveness (no auth, no dependency checks).
|
||||
r.GET("/healthz", health.Healthz(serviceName))
|
||||
// Readiness (soft failure on Redis/Kafka).
|
||||
r.GET("/readyz", readyzer.Handler)
|
||||
// Prometheus metrics.
|
||||
r.GET("/metrics", gin.WrapH(promhttp.Handler()))
|
||||
|
||||
// WebSocket 升级端点
|
||||
// WebSocket upgrade (JWT auth via query ?token= or Authorization header).
|
||||
r.GET("/ws", wsHandler.HandleWebSocket)
|
||||
|
||||
// 内部推送 API（Msg 服务调用）
|
||||
api := r.Group("/internal")
|
||||
api.POST("/push", wsHandler.PushHandler)
|
||||
api.POST("/broadcast", wsHandler.BroadcastHandler)
|
||||
// Internal HTTP APIs consumed by msg (X-Internal-Token auth).
|
||||
internal := r.Group("/internal")
|
||||
internal.POST("/push", wsHandler.PushHandler)
|
||||
internal.POST("/broadcast", wsHandler.BroadcastHandler)
|
||||
internal.GET("/online/:userID", wsHandler.OnlineHandler)
|
||||
|
||||
srv := &http.Server{
|
||||
Addr: ":" + cfg.Port,
|
||||
Handler: r,
|
||||
ReadTimeout: 10 * time.Second,
|
||||
WriteTimeout: 60 * time.Second, // WebSocket 长连接
|
||||
WriteTimeout: 0, // WebSocket connections are long-lived; no write timeout.
|
||||
}
|
||||
|
||||
// 11. Start HTTP server.
|
||||
go func() {
|
||||
log.Printf("Push Gateway listening on :%s", cfg.Port)
|
||||
if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
|
||||
log.Fatalf("listen: %s\n", err)
|
||||
logger.Info("http server listening", "addr", srv.Addr)
|
||||
if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
|
||||
logger.Error("http server fatal error", "err", err)
|
||||
os.Exit(1)
|
||||
}
|
||||
}()
|
||||
|
||||
// 12. Wait for termination signal.
|
||||
quit := make(chan os.Signal, 1)
|
||||
signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
|
||||
<-quit
|
||||
log.Println("Shutting down Push Gateway...")
|
||||
sig := <-quit
|
||||
logger.Info("shutdown signal received", "signal", sig.String())
|
||||
|
||||
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
||||
defer cancel()
|
||||
if err := srv.Shutdown(ctx); err != nil {
|
||||
log.Fatal("Server forced to shutdown:", err)
|
||||
// 13. Graceful shutdown sequence (see package doc).
|
||||
shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), shutdownTimeout)
|
||||
defer shutdownCancel()
|
||||
|
||||
// 13a. Hub.CloseAll sends close frame 1001 to every live connection.
|
||||
h.CloseAll()
|
||||
logger.Info("hub closeAll complete; draining connections",
|
||||
"active_connections", h.ActiveConnections())
|
||||
|
||||
// 13b. Stop accepting new HTTP requests; drain in-flight.
|
||||
if err := srv.Shutdown(shutdownCtx); err != nil {
|
||||
logger.Error("http server shutdown error", "err", err)
|
||||
}
|
||||
|
||||
// 13c. Stop Kafka consumer (cancel context + close reader).
|
||||
if kafkaCancel != nil {
|
||||
kafkaCancel()
|
||||
}
|
||||
if kafkaConsumer != nil {
|
||||
if err := kafkaConsumer.Close(); err != nil {
|
||||
logger.Warn("kafka consumer close error", "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
// 13d. Stop Redis Pub/Sub subscriber + close client.
|
||||
if pubsubCancel != nil {
|
||||
if err := pubsubCancel(); err != nil {
|
||||
logger.Warn("redis pubsub cancel error", "err", err)
|
||||
}
|
||||
}
|
||||
if redisClient != nil {
|
||||
if err := redisClient.Close(); err != nil {
|
||||
logger.Warn("redis close error", "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
logger.Info("push-gateway shutdown complete")
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user