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/go.mod

@@ -1,8 +1,15 @@
 module src.lwithers.me.uk/go/htpack/cmd/packserver
 
-go 1.13
+go 1.19
 
 require (
-	github.com/spf13/cobra v1.5.0
-	src.lwithers.me.uk/go/htpack v1.3.1
+	github.com/spf13/cobra v1.6.1
+	src.lwithers.me.uk/go/htpack v1.3.2
+)
+
+require (
+	github.com/gogo/protobuf v1.3.2 // indirect
+	github.com/inconshreveable/mousetrap v1.0.1 // indirect
+	github.com/spf13/pflag v1.0.5 // indirect
+	golang.org/x/sys v0.2.0 // indirect
 )

cmd/packserver/go.sum

@@ -1,13 +1,13 @@
 github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
 github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
 github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
-github.com/inconshreveable/mousetrap v1.0.0 h1:Z8tu5sraLXCXIcARxBp/8cbvlwVa7Z1NHg9XEKhtSvM=
-github.com/inconshreveable/mousetrap v1.0.0/go.mod h1:PxqpIevigyE2G7u3NXJIT2ANytuPF1OarO4DADm73n8=
+github.com/inconshreveable/mousetrap v1.0.1 h1:U3uMjPSQEBMNp1lFxmllqCPM6P5u/Xq7Pgzkat/bFNc=
+github.com/inconshreveable/mousetrap v1.0.1/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
 github.com/kisielk/errcheck v1.5.0/go.mod h1:pFxgyoBC7bSaBwPgfKdkLd5X25qrDl4LWUI2bnpBCr8=
 github.com/kisielk/gotool v1.0.0/go.mod h1:XhKaO+MFFWcvkIS/tQcRk01m1F5IRFswLeQ+oQHNcck=
 github.com/russross/blackfriday/v2 v2.1.0/go.mod h1:+Rmxgy9KzJVeS9/2gXHxylqXiyQDYRxCVz55jmeOWTM=
-github.com/spf13/cobra v1.5.0 h1:X+jTBEBqF0bHN+9cSMgmfuvv2VHJ9ezmFNf9Y/XstYU=
-github.com/spf13/cobra v1.5.0/go.mod h1:dWXEIy2H428czQCjInthrTRUg7yKbok+2Qi/yBIJoUM=
+github.com/spf13/cobra v1.6.1 h1:o94oiPyS4KD1mPy2fmcYYHHfCxLqYjJOhGsCHFZtEzA=
+github.com/spf13/cobra v1.6.1/go.mod h1:IOw/AERYS7UzyrGinqmz6HLUo219MORXGxhbaJUqzrY=
 github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
 github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
 github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
@@ -27,8 +27,8 @@ golang.org/x/sync v0.0.0-20201020160332-67f06af15bc9/go.mod h1:RxMgew5VJxzue5/jJ
 golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5hshtkjS+E42TnysNCUPdjciGhY=
 golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
 golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
-golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e h1:CsOuNlbOuf0mzxJIefr6Q4uAUetRUwZE4qt7VfzP+xo=
-golang.org/x/sys v0.0.0-20220704084225-05e143d24a9e/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.2.0 h1:ljd4t30dBnAvMZaQCevtY0xLLD0A+bRZXbgLMLU1F/A=
+golang.org/x/sys v0.2.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
@@ -40,6 +40,6 @@ golang.org/x/xerrors v0.0.0-20191011141410-1b5146add898/go.mod h1:I/5z698sn9Ka8T
 golang.org/x/xerrors v0.0.0-20191204190536-9bdfabe68543/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 golang.org/x/xerrors v0.0.0-20200804184101-5ec99f83aff1/go.mod h1:I/5z698sn9Ka8TeJc9MKroUUfqBBauWjQqLJ2OPfmY0=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
-gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
-src.lwithers.me.uk/go/htpack v1.3.1 h1:JvCeyVnrOjUOL//D9N/ctTOwobf60AuNU828G6WXf2M=
-src.lwithers.me.uk/go/htpack v1.3.1/go.mod h1:wpfOXweUh8oRPxYnVjUWDRNNwIXlAlsOrmGWc0fOls4=
+gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
+src.lwithers.me.uk/go/htpack v1.3.2 h1:WyJs0F1lHKMoDTY+HgDUS0mTcdFMGbq2CykJwgm5rkI=
+src.lwithers.me.uk/go/htpack v1.3.2/go.mod h1:rslOJzSAJTKf0fz+iULoYF+jXrffRJK5PwUNu1wADOk=

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)
 	}