mux/serve.go

package mux

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

type ServeCB func(srv *http.Server) error

const (
	shutdownDelay     = time.Second * 10
	shutdownHardDelay = time.Second * 5
	drainDelay        = time.Second
)

// Serve with graceful shutdown
func (m *Mux) Serve(cb ServeCB) {
	rootCtx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
	defer stop()

	srvCtx, cancelSrvCtx := context.WithCancel(context.Background())
	srv := &http.Server{
		Handler: m,
		BaseContext: func(_ net.Listener) context.Context {
			return srvCtx
		},
	}

	go func() {
		if err := cb(srv); !errors.Is(err, http.ErrServerClosed) {
			log.Fatalf("server error: %v", err)
		}
	}()

	// Wait for interrupt signal
	<-rootCtx.Done()

	stop()
	m.IsShuttingDown.Store(true)
	slog.Info("received interrupt signal, shutting down")
	time.Sleep(drainDelay)
	slog.Info("readiness check propagated, now waiting for ongoing requests to finish.")

	shutdownCtx, cancel := context.WithTimeout(context.Background(), shutdownDelay)
	defer cancel()

	err := srv.Shutdown(shutdownCtx)
	cancelSrvCtx()
	if err != nil {
		log.Println("failed to wait for ongoing requests to finish, waiting for forced cancellation")
		time.Sleep(shutdownHardDelay)
	}

	slog.Info("server shut down gracefully")
}