server.go

package httpx

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

// Graceful shutdown implementation inspired by:
// https://victoriametrics.com/blog/go-graceful-shutdown/

const (
	// shutdownPeriod is the maximum time to wait for ongoing requests to complete
	// during graceful shutdown before forcing termination.
	shutdownPeriod = 15 * time.Second

	// shutdownHardPeriod is the additional time to wait after a failed graceful
	// shutdown before returning an error.
	shutdownHardPeriod = 3 * time.Second

	// readinessDrainDelay is the time to wait after marking the server as shutting
	// down to allow readiness checks to propagate before stopping new connections.
	readinessDrainDelay = 5 * time.Second
)

var (
	// IsShuttingDown is an atomic boolean flag indicating whether the server
	// is currently in the shutdown process. This can be used by readiness
	// checks to fail health checks during graceful shutdown.
	IsShuttingDown atomic.Bool

	// ErrServerHasStopped is returned when the server has initiated shutdown
	// and is no longer accepting new requests.
	ErrServerHasStopped = errors.New("server has stopped")
)

// ServerConfFn is a function type used to configure an http.Server instance
// before starting it. This allows callers to customize server settings such
// as timeouts, TLS configuration, and other http.Server fields.
type ServerConfFn func(*http.Server)

// WithHTTPMiddleware returns a ServerConfFn that wraps the server's handler
// with the given HTTP middlewares. Middlewares are applied so that they execute
// in the order they are provided (first middleware runs first).
// This is useful for functionality like CORS, request ID generation, or metrics
// that should apply to all requests including those that don't match any route.
//
// Example:
//
//	httpx.ListenAndServe(ctx, router,
//	    httpx.WithHTTPMiddleware(CORSMiddleware, RequestIDMiddleware),
//	)
func WithHTTPMiddleware(middlewares ...HTTPMiddleware) ServerConfFn {
	return func(s *http.Server) {
		handler := s.Handler
		// Apply in reverse order so they execute in the order they were provided
		for i := len(middlewares) - 1; i >= 0; i-- {
			handler = middlewares[i](handler)
		}
		s.Handler = handler
	}
}

// ListenAndServe starts an HTTP server with graceful shutdown capabilities.
// It listens for SIGINT and SIGTERM signals to initiate a graceful shutdown sequence.
//
// The shutdown sequence follows these steps:
//  1. Upon receiving a shutdown signal or when the given context is done, sets IsShuttingDown to true
//  2. Waits readinessDrainDelay for readiness checks to propagate
//  3. Stops accepting new connections and waits up to shutdownPeriod for ongoing requests
//  4. If graceful shutdown fails, waits an additional shutdownHardPeriod before returning
//
// Returns an error if the server fails to start or encounters an error during shutdown.
func ListenAndServe(ctx context.Context, router http.Handler, serverConfs ...ServerConfFn) error {
	rootCtx, stopRootCtx := signal.NotifyContext(ctx, syscall.SIGINT, syscall.SIGTERM)
	defer stopRootCtx()

	handlerCtx, stopHandlerCtx := context.WithCancelCause(context.Background())

	server := &http.Server{
		Addr:    ":8080",
		Handler: router,
	}
	for _, conf := range serverConfs {
		conf(server)
	}
	server.BaseContext = func(_ net.Listener) context.Context {
		return handlerCtx
	}

	// Channel to signal server startup or error
	startupErr := make(chan error, 1)

	go func() {
		log.Println("Server starting on " + server.Addr)
		if err := server.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
			startupErr <- err
		}
	}()

	// Wait a moment for server to start or fail
	select {
	case err := <-startupErr:
		stopHandlerCtx(err)
		return err
	case <-time.After(100 * time.Millisecond):
		// Server started successfully
	}

	// Wait for signal
	<-rootCtx.Done()
	stopRootCtx()
	log.Println("Received shutdown signal, shutting down.")

	// Give time for readiness check to propagate
	IsShuttingDown.Store(true)
	time.Sleep(readinessDrainDelay)
	log.Println("Readiness check propagated, now waiting for ongoing requests to finish.")

	stopHandlerCtx(ErrServerHasStopped)
	shutdownCtx, stopShutdownCtx := context.WithTimeout(context.Background(), shutdownPeriod)
	defer stopShutdownCtx()
	if err := server.Shutdown(shutdownCtx); err != nil {
		log.Println("Failed to wait for ongoing requests to finish, waiting for forced cancellation.")
		time.Sleep(shutdownHardPeriod)
		return err
	}

	log.Println("Server shut down gracefully.")

	return nil
}