Edu/services/push-gateway/main.go

// 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.notify.notification.sent, ARB-013)
//  9. ws.NewHandler (JWT RS256, X-Internal-Key, 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"
	"errors"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"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()

	// 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.notify.notification.sent, ARB-013). 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)
	r := gin.New()
	r.Use(gin.Recovery())
	r.Use(otelgin.Middleware(serviceName))

	// 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 upgrade (JWT auth via query ?token= or Authorization header).
	r.GET("/ws", wsHandler.HandleWebSocket)
	// Internal HTTP APIs consumed by msg (X-Internal-Key auth, ARB-013).
	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: 0, // WebSocket connections are long-lived; no write timeout.
	}

	// 11. Start HTTP server.
	go func() {
		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)
	sig := <-quit
	logger.Info("shutdown signal received", "signal", sig.String())

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