From 5398dddb02a058ac0795cfd2037ff28980624747 Mon Sep 17 00:00:00 2001
From: Laurence Withers <l@lwithers.me.uk>
Date: Sat, 26 Nov 2022 10:39:47 +0000
Subject: [PATCH] Add --graceful-shutdown-delay

This option allows a configurable delay after receiving SIGTERM (or
SIGINT) but before the HTTP server stops accepting new connections. It
is quite useful for distributed systems where callers are only notified
asynchronously (e.g. via service discovery) that a service is being shut
down; it prevents the shut down from occurring prior to callers
processing the notification.

This required some minor refactoring to allow the Shutdown() method on
http.Server to be accessed.
---
 cmd/packserver/main.go | 65 ++++++++++++++++++++++++++++++++++++------
 1 file changed, 57 insertions(+), 8 deletions(-)

diff --git a/cmd/packserver/main.go b/cmd/packserver/main.go
index c18b90b..2000102 100644
--- a/cmd/packserver/main.go
+++ b/cmd/packserver/main.go
@@ -5,12 +5,17 @@ package main
 
 import (
 	"bufio"
+	"context"
 	"errors"
 	"fmt"
 	"net/http"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strings"
+	"sync/atomic"
+	"syscall"
+	"time"
 
 	"github.com/spf13/cobra"
 	"src.lwithers.me.uk/go/htpack"
@@ -53,6 +58,8 @@ func main() {
 		"Name of file to return if response would be 404 (spa.html or similar)")
 	rootCmd.Flags().String("frames", "sameorigin",
 		"Override X-Frame-Options header (can be sameorigin, deny, allow)")
+	rootCmd.Flags().Duration("graceful-shutdown-delay", 3*time.Second,
+		"Number of seconds to wait after receiving SIGTERM before initiating graceful shutdown")
 
 	if err := rootCmd.Execute(); err != nil {
 		fmt.Fprintln(os.Stderr, err)
@@ -151,6 +158,15 @@ func run(c *cobra.Command, args []string) error {
 		return err
 	}
 
+	// graceful shutdown delay must be > 0
+	gracefulShutdownDelay, err := c.Flags().GetDuration("graceful-shutdown-delay")
+	if err != nil {
+		return err
+	}
+	if gracefulShutdownDelay <= 0 {
+		return errors.New("graceful shutdown delay must be > 0s")
+	}
+
 	// verify .htpack specifications
 	if len(args) == 0 {
 		return errors.New("must specify one or more .htpack files")
@@ -176,6 +192,7 @@ func run(c *cobra.Command, args []string) error {
 	}
 
 	// load packfiles, registering handlers as we go
+	var handler http.Handler
 	for prefix, packfile := range packPaths {
 		packHandler, err := htpack.New(packfile)
 		if err != nil {
@@ -190,26 +207,58 @@ func run(c *cobra.Command, args []string) error {
 		}
 		packHandler.SetHeader("X-Frame-Options", framesHeader)
 
-		handler := &addHeaders{
+		handler = &addHeaders{
 			extraHeaders: extraHeaders,
 			handler:      packHandler,
 		}
 
 		if prefix != "/" {
-			http.Handle(prefix+"/",
-				http.StripPrefix(prefix, handler))
-		} else {
-			http.Handle("/", handler)
+			handler = http.StripPrefix(prefix, handler)
 		}
 	}
 
+	// HTTP server object setup
+	sv := &http.Server{
+		Addr:    bindAddr,
+		Handler: handler,
+	}
+
+	// register SIGINT, SIGTERM handler
+	sigch := make(chan os.Signal, 1)
+	signal.Notify(sigch, syscall.SIGINT, syscall.SIGTERM)
+	var (
+		// if we are shut down by a signal, then http.ListenAndServe()
+		// returns straight away, but we actually need to wait for
+		// Shutdown() to complete prior to returning / exiting
+		isSignalled atomic.Bool
+		signalDone  = make(chan struct{})
+	)
+	go func() {
+		<-sigch
+		time.Sleep(gracefulShutdownDelay)
+		isSignalled.Store(true)
+		shutctx, shutcancel := context.WithTimeout(context.Background(), gracefulShutdownDelay)
+		sv.Shutdown(shutctx)
+		shutcancel()
+		close(signalDone)
+	}()
+
 	// main server loop
 	if keyFile == "" {
-		err = http.ListenAndServe(bindAddr, nil)
+		err = sv.ListenAndServe()
 	} else {
-		err = http.ListenAndServeTLS(bindAddr, certFile, keyFile, nil)
+		err = sv.ListenAndServeTLS(certFile, keyFile)
 	}
-	if err != nil {
+
+	// if we were shut down by a signal, wait for Shutdown() to return
+	if isSignalled.Load() {
+		<-signalDone
+	}
+
+	switch err {
+	case nil, http.ErrServerClosed:
+		// OK
+	default:
 		fmt.Fprintln(os.Stderr, err)
 		os.Exit(1)
 	}