dance-lessons-coach/pkg/server/server.go

package server

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

	"DanceLessonsCoach/pkg/config"
	"DanceLessonsCoach/pkg/greet"
	"DanceLessonsCoach/pkg/telemetry"
	"go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp"
	sdktrace "go.opentelemetry.io/otel/sdk/trace"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
	"github.com/rs/zerolog/log"
)

type Server struct {
	router      *chi.Mux
	readyCtx    context.Context
	withOTEL    bool
	config      *config.Config
	tracerProvider *sdktrace.TracerProvider
}

func NewServer(cfg *config.Config, readyCtx context.Context) *Server {
	s := &Server{
		router:   chi.NewRouter(),
		readyCtx: readyCtx,
		withOTEL: cfg.GetTelemetryEnabled(),
		config:   cfg,
	}
	s.setupRoutes()
	return s
}

func (s *Server) setupRoutes() {
	// Use Zerolog middleware instead of Chi's default logger
	s.router.Use(middleware.RequestLogger(&middleware.DefaultLogFormatter{
		Logger:  &log.Logger,
		NoColor: false,
	}))

	// Health endpoint at root level
	s.router.Get("/api/health", s.handleHealth)

	// Readiness endpoint at root level
	s.router.Get("/api/ready", s.handleReadiness)

	// API routes
	s.router.Route("/api/v1", func(r chi.Router) {
		r.Use(s.getAllMiddlewares()...)
		s.registerApiV1Routes(r)
	})
}

func (s *Server) registerApiV1Routes(r chi.Router) {
	greetService := greet.NewService()
	greetHandler := greet.NewApiV1GreetHandler(greetService)
	r.Route("/greet", func(r chi.Router) {
		greetHandler.RegisterRoutes(r)
	})
}

// getAllMiddlewares returns all middleware including OpenTelemetry if enabled
func (s *Server) getAllMiddlewares() []func(http.Handler) http.Handler {
	middlewares := []func(http.Handler) http.Handler{
		middleware.StripSlashes,
		middleware.Recoverer,
	}

	if s.withOTEL {
		middlewares = append(middlewares, func(next http.Handler) http.Handler {
			return otelhttp.NewHandler(next, "")
		})
	}

	return middlewares
}

func (s *Server) handleHealth(w http.ResponseWriter, r *http.Request) {
	log.Info().Msg("Health check requested")
	w.Write([]byte(`{"status":"healthy"}`))
}

func (s *Server) handleReadiness(w http.ResponseWriter, r *http.Request) {
	log.Info().Msg("Readiness check requested")

	select {
	case <-s.readyCtx.Done():
		log.Info().Msg("Readiness check: not ready (shutting down)")
		w.WriteHeader(http.StatusServiceUnavailable)
		w.Write([]byte(`{"ready":false}`))
	default:
		log.Info().Msg("Readiness check: ready")
		w.Write([]byte(`{"ready":true}`))
	}
}

func (s *Server) Router() http.Handler {
	return s.router
}

// Run starts the HTTP server and handles graceful shutdown
func (s *Server) Run() error {
	// Initialize OpenTelemetry if enabled
	var err error
	if s.withOTEL {
		log.Info().Msg("Initializing OpenTelemetry tracing")

		telemetrySetup := &telemetry.Setup{
			ServiceName:  s.config.GetServiceName(),
			OTLPEndpoint: s.config.GetOTLPEndpoint(),
			Insecure:     s.config.GetTelemetryInsecure(),
			SamplerType:  s.config.GetSamplerType(),
			SamplerRatio: s.config.GetSamplerRatio(),
		}

		if s.tracerProvider, err = telemetrySetup.InitializeTracing(context.Background()); err != nil {
			log.Error().Err(err).Msg("Failed to initialize OpenTelemetry, continuing without tracing")
			s.withOTEL = false
		} else {
			log.Info().Msg("OpenTelemetry tracing initialized successfully")
		}
	}

	// Setup signal context for graceful shutdown
	rootCtx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	// Create ongoing context for active requests
	ongoingCtx, stopOngoingGracefully := context.WithCancel(context.Background())
	defer stopOngoingGracefully()

	// Create HTTP server
	log.Info().Str("address", s.config.GetServerAddress()).Msg("Server running")

	srv := &http.Server{
		Addr:    s.config.GetServerAddress(),
		Handler: s.router,
		BaseContext: func(_ net.Listener) context.Context {
			return ongoingCtx
		},
	}

	// Start the HTTP server in a separate goroutine
	serverErrors := make(chan error, 1)
	go func() {
		if err := srv.ListenAndServe(); err != nil && err != http.ErrServerClosed {
			serverErrors <- err
		}
		close(serverErrors)
	}()

	// Wait for signal
	<-rootCtx.Done()
	stop()
	log.Info().Msg("Shutdown signal received")

	// Cancel readiness context to stop accepting new requests
	if cancelReady, ok := s.readyCtx.(interface{ Cancel() }); ok {
		cancelReady.Cancel()
	}
	log.Info().Msg("Readiness set to false, no longer accepting new requests")

	// Give time for readiness check to propagate
	time.Sleep(1 * time.Second)
	log.Info().Msg("Readiness check propagated, now waiting for ongoing requests to finish.")

	// Create shutdown context with timeout from config
	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), s.config.Shutdown.Timeout)
	defer shutdownCancel()

	if err := srv.Shutdown(shutdownCtx); err != nil {
		log.Error().Err(err).Msg("Server shutdown failed")
		return err
	}
	log.Info().Msg("Server shutdown complete")

	// Shutdown OpenTelemetry tracer provider
	if s.tracerProvider != nil {
		if err := telemetry.Shutdown(context.Background(), s.tracerProvider); err != nil {
			log.Error().Err(err).Msg("Failed to shutdown OpenTelemetry tracer provider")
		} else {
			log.Info().Msg("OpenTelemetry tracer provider shutdown complete")
		}
	}

	// Return any server errors
	if err, ok := <-serverErrors; ok {
		return err
	}

	return nil
}