httpx/server/server.go

package server

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

// Server is a production-ready HTTP server with graceful shutdown,
// middleware support, and signal handling.
type Server struct {
	httpServer      *http.Server
	listener        net.Listener
	addr            atomic.Value
	logger          *slog.Logger
	shutdownTimeout time.Duration
	onShutdown      []func()
	listenAddr      string
}

// New creates a new Server that will serve the given handler with the
// provided options. Middleware from options is applied to the handler.
func New(handler http.Handler, opts ...Option) *Server {
	o := &serverOptions{
		addr:            ":8080",
		shutdownTimeout: 15 * time.Second,
	}
	for _, opt := range opts {
		opt(o)
	}

	// Apply middleware chain to the handler.
	if len(o.middlewares) > 0 {
		handler = Chain(o.middlewares...)(handler)
	}

	srv := &Server{
		httpServer: &http.Server{
			Handler:           handler,
			ReadTimeout:       o.readTimeout,
			ReadHeaderTimeout: o.readHeaderTimeout,
			WriteTimeout:      o.writeTimeout,
			IdleTimeout:       o.idleTimeout,
		},
		logger:          o.logger,
		shutdownTimeout: o.shutdownTimeout,
		onShutdown:      o.onShutdown,
		listenAddr:      o.addr,
	}

	return srv
}

// ListenAndServe starts the server and blocks until a SIGINT or SIGTERM
// signal is received. It then performs a graceful shutdown within the
// configured shutdown timeout.
//
// Returns nil on clean shutdown or an error if listen/shutdown fails.
func (s *Server) ListenAndServe() error {
	ln, err := net.Listen("tcp", s.listenAddr)
	if err != nil {
		return err
	}
	s.listener = ln
	s.addr.Store(ln.Addr().String())

	s.log("server started", slog.String("addr", ln.Addr().String()))

	// Wait for signal in context.
	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	errCh := make(chan error, 1)
	go func() {
		if err := s.httpServer.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
			errCh <- err
		}
		close(errCh)
	}()

	select {
	case err := <-errCh:
		return err
	case <-ctx.Done():
		stop()
		return s.shutdown()
	}
}

// ListenAndServeTLS starts the server with TLS and blocks until a signal
// is received.
func (s *Server) ListenAndServeTLS(certFile, keyFile string) error {
	ln, err := net.Listen("tcp", s.listenAddr)
	if err != nil {
		return err
	}
	s.listener = ln
	s.addr.Store(ln.Addr().String())

	s.log("server started (TLS)", slog.String("addr", ln.Addr().String()))

	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	errCh := make(chan error, 1)
	go func() {
		if err := s.httpServer.ServeTLS(ln, certFile, keyFile); err != nil && !errors.Is(err, http.ErrServerClosed) {
			errCh <- err
		}
		close(errCh)
	}()

	select {
	case err := <-errCh:
		return err
	case <-ctx.Done():
		stop()
		return s.shutdown()
	}
}

// Shutdown gracefully shuts down the server. It calls any registered
// onShutdown hooks, then waits for active connections to drain within
// the shutdown timeout.
func (s *Server) Shutdown(ctx context.Context) error {
	s.runOnShutdown()
	return s.httpServer.Shutdown(ctx)
}

// Addr returns the listener address after the server has started.
// Returns an empty string if the server has not started yet.
func (s *Server) Addr() string {
	v := s.addr.Load()
	if v == nil {
		return ""
	}
	return v.(string)
}

func (s *Server) shutdown() error {
	s.log("shutting down")

	s.runOnShutdown()

	ctx, cancel := context.WithTimeout(context.Background(), s.shutdownTimeout)
	defer cancel()

	if err := s.httpServer.Shutdown(ctx); err != nil {
		s.log("shutdown error", slog.String("error", err.Error()))
		return err
	}

	s.log("server stopped")
	return nil
}

func (s *Server) runOnShutdown() {
	for _, fn := range s.onShutdown {
		fn()
	}
}

func (s *Server) log(msg string, attrs ...slog.Attr) {
	if s.logger != nil {
		s.logger.LogAttrs(context.Background(), slog.LevelInfo, msg, attrs...)
	}
}