Go http graceful shutdown

main.go

package main

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

func main() {
	log.SetFlags(log.Flags() | log.Lshortfile)

	mux := http.NewServeMux()
	mux.HandleFunc("/hello", func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusOK)
	})

	ctx := InterruptContext()

	if err := ListenAndServe(ctx, ":8080", mux, 5*time.Second); err != nil {
		log.Printf("ListenAndServe error %s", err)
	}
}

// ListenAndServe adds graceful shutdown to http.ListenAndServer.
// When ctx context is canceled graceful shutdown with timeout will be started.
func ListenAndServe(ctx context.Context, addr string, handler http.Handler, shutdownTimeout time.Duration) error {
	// context which can be cancelled by parent or calling serveCancel
	serveCtx, serveCancel := context.WithCancel(ctx)

	server := &http.Server{
		Addr:    addr,
		Handler: handler,
	}

	var errServe error
	go func() {
		if err := server.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) {
			errServe = err
		}
		serveCancel()
	}()

	// wait for serveCancel or parent context to be canceled
	<-serveCtx.Done()
	if errServe != nil { // ListenAndServe failed
		return errServe
	}

	// parent context canceled, start graceful shutdown
	shutdownCtx, shutdownRelease := context.WithTimeout(context.Background(), shutdownTimeout)
	defer shutdownRelease()
	return server.Shutdown(shutdownCtx)
}

func WaitForInterupt() {
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, syscall.SIGINT, syscall.SIGTERM)
	<-quit
}

// InterruptContext returns context which will be closed on application interrupt
func InterruptContext() context.Context {
	ctx, cancel := context.WithCancel(context.Background())
	go func() {
		WaitForInterupt()
		cancel()
	}()
	return ctx
}