Graceful Shutdown Behind a Load Balancer

package main

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

func main() {
	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	var ready atomic.Bool
	ready.Store(true)

	started := make(chan struct{}, 1)
	mux := http.NewServeMux()
	mux.HandleFunc("/readyz", func(w http.ResponseWriter, r *http.Request) {
		if !ready.Load() {
			http.Error(w, "draining", http.StatusServiceUnavailable)
			return
		}
		fmt.Fprintln(w, "ready")
	})
	mux.HandleFunc("/work", func(w http.ResponseWriter, r *http.Request) {
		started <- struct{}{}
		select {
		case <-time.After(300 * time.Millisecond):
			fmt.Fprintln(w, "finished")
		case <-r.Context().Done():
		}
	})

	listener, err := net.Listen("tcp", "127.0.0.1:0")
	if err != nil {
		fmt.Println("listen:", err)
		return
	}
	server := &http.Server{Handler: mux, ReadHeaderTimeout: 2 * time.Second}
	go func() {
		serveErr := server.Serve(listener)
		if serveErr != nil && !errors.Is(serveErr, http.ErrServerClosed) {
			fmt.Println("serve:", serveErr)
		}
	}()

	go func() {
		<-started
		time.Sleep(50 * time.Millisecond)
		stop() // simulate SIGTERM while a request is still running
	}()
	go func() {
		time.Sleep(10 * time.Millisecond)
		response, requestErr := http.Get("http://" + listener.Addr().String() + "/work")
		if requestErr != nil {
			fmt.Println("request:", requestErr)
			return
		}
		defer response.Body.Close()
		fmt.Println("client status:", response.Status)
	}()

	<-ctx.Done()
	ready.Store(false)
	fmt.Println("readiness probe now fails; waiting for load balancers to stop routing")
	time.Sleep(100 * time.Millisecond)

	shutdownCtx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel()
	if err := server.Shutdown(shutdownCtx); err != nil {
		fmt.Println("shutdown:", err)
		return
	}
	fmt.Println("server stopped after draining in-flight requests")
}