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