cdp/ingestion/ingest/cmd/server/main.go

// Command server runs the CDP ingest HTTP API.
package main

import (
	"context"
	"errors"
	"log"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/redis/rueidis"
	"go.uber.org/zap"

	"github.com/dbiz/cdp/ingestion/ingest/internal/config"
	"github.com/dbiz/cdp/ingestion/ingest/internal/dedup"
	"github.com/dbiz/cdp/ingestion/ingest/internal/handler"
	"github.com/dbiz/cdp/ingestion/ingest/internal/kafka"
	"github.com/dbiz/cdp/ingestion/ingest/internal/live"
	mw "github.com/dbiz/cdp/ingestion/ingest/internal/middleware"
	"github.com/dbiz/cdp/ingestion/ingest/internal/ratelimit"
	"github.com/dbiz/cdp/ingestion/ingest/internal/repo"
	"github.com/dbiz/cdp/ingestion/ingest/internal/service"
)

func main() {
	if err := run(); err != nil {
		log.Fatal(err)
	}
}

func run() error {
	cfg, err := config.Load()
	if err != nil {
		return err
	}

	logger, err := newLogger(cfg.LogLevel)
	if err != nil {
		return err
	}
	defer func() { _ = logger.Sync() }()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()

	// ---- infra clients ----------------------------------------------------
	pg, err := repo.NewPool(ctx, cfg.PostgresDSN)
	if err != nil {
		return err
	}
	defer pg.Close()

	redisClient, err := rueidis.NewClient(rueidis.ClientOption{
		InitAddress: []string{cfg.RedisAddr},
	})
	if err != nil {
		return err
	}
	defer redisClient.Close()

	producer, err := kafka.NewProducer(cfg.KafkaBrokers,
		cfg.KafkaTopicIngest, cfg.KafkaTopicDLQ, cfg.KafkaTopicRetry, logger)
	if err != nil {
		return err
	}
	defer producer.Close()

	// ---- repos / services -------------------------------------------------
	writeKeyRepo := repo.NewWriteKeyRepo(pg)
	schemaRepo := repo.NewSchemaRepo(pg)

	authSvc := service.NewAuthService(writeKeyRepo, redisClient, cfg.WriteKeyCacheTTL, logger)
	ingestSvc := service.NewIngestService(service.IngestDeps{
		Producer:     producer,
		Limiter:      ratelimit.New(redisClient),
		Dedup:        dedup.New(redisClient, time.Duration(cfg.DedupTTLHours)*time.Hour),
		Schema:       schemaRepo,
		Log:          logger,
		LateAfter:    time.Duration(cfg.LateEventHours) * time.Hour,
		RateLimitRPS: cfg.RateLimitRPS,
	})
	evHandler := handler.NewEventHandler(ingestSvc, logger)

	liveStreamer := live.New(cfg.KafkaBrokers, cfg.KafkaTopicIngest, logger)
	liveHandler := handler.NewLiveHandler(liveStreamer, logger)

	// ---- HTTP router ------------------------------------------------------
	r := chi.NewRouter()
	r.Use(mw.RequestID)
	r.Use(mw.Recover(logger))
	r.Use(mw.Logger(logger))
	r.Use(mw.CORS)
	r.Use(mw.PayloadLimit(cfg.PayloadLimitKB))

	// public health endpoints
	r.Get("/health", evHandler.Health)
	r.Get("/ready", evHandler.Ready)

	// SSE stream of events flowing through Kafka. Intentionally outside the
	// auth group so the console can subscribe without forwarding the write
	// key; lock this down before production.
	r.Get("/live/events", liveHandler.Stream)

	// authenticated routes
	r.Group(func(rr chi.Router) {
		rr.Use(mw.Auth(authSvc))
		rr.Post("/track", evHandler.Track)
		rr.Post("/identify", evHandler.Identify)
		rr.Post("/page", evHandler.Page)
		rr.Post("/group", evHandler.Group)
		rr.Post("/alias", evHandler.Alias)
		rr.Post("/screen", evHandler.Screen)
		// batch has its own (larger) payload limit
		rr.With(mw.PayloadLimit(cfg.BatchLimitKB)).Post("/batch", evHandler.Batch)
		// Segment compatibility paths
		rr.With(mw.PayloadLimit(cfg.BatchLimitKB)).Post("/v1/batch", evHandler.Batch)
		rr.Post("/v1/track", evHandler.Track)
		rr.Post("/v1/identify", evHandler.Identify)
		rr.Post("/v1/page", evHandler.Page)
		rr.Post("/v1/group", evHandler.Group)
		rr.Post("/v1/alias", evHandler.Alias)
		rr.Post("/v1/screen", evHandler.Screen)
	})

	srv := &http.Server{
		Addr:              cfg.HTTPAddr,
		Handler:           r,
		ReadHeaderTimeout: 5 * time.Second,
		ReadTimeout:       30 * time.Second,
		// WriteTimeout intentionally 0 (no deadline) -- /live/events is a
		// long-lived SSE stream. Per-handler deadlines apply via ctx.
		WriteTimeout: 0,
		IdleTimeout:  120 * time.Second,
	}

	// ---- graceful shutdown ------------------------------------------------
	shutdownErr := make(chan error, 1)
	go func() {
		sigCh := make(chan os.Signal, 1)
		signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
		<-sigCh
		logger.Info("shutdown signal received; draining...")
		shutCtx, cancel := context.WithTimeout(context.Background(), cfg.ShutdownTimeout)
		defer cancel()
		shutdownErr <- srv.Shutdown(shutCtx)
	}()

	logger.Info("ingest listening", zap.String("addr", cfg.HTTPAddr))
	if err := srv.ListenAndServe(); err != nil && !errors.Is(err, http.ErrServerClosed) {
		return err
	}
	return <-shutdownErr
}

func newLogger(level string) (*zap.Logger, error) {
	lvl, err := zap.ParseAtomicLevel(level)
	if err != nil {
		lvl = zap.NewAtomicLevelAt(zap.InfoLevel)
	}
	cfg := zap.NewProductionConfig()
	cfg.Level = lvl
	cfg.EncoderConfig.TimeKey = "ts"
	cfg.EncoderConfig.MessageKey = "msg"
	return cfg.Build()
}