agnhost/net: support --delay-shutdown flag

Signed-off-by: Andrew Sy Kim <andrewsy@google.com>
This commit is contained in:
Andrew Sy Kim 2022-03-16 16:32:38 -04:00
parent d981aa118f
commit ff997ae5a0

View File

@ -24,7 +24,10 @@ import (
"log"
"net/http"
"os"
"os/signal"
"strings"
"syscall"
"time"
"github.com/spf13/cobra"
@ -38,9 +41,10 @@ var (
// flags for the command line. See usage args below for
// descriptions.
flags struct {
Serve string
Runner string
Options string
Serve string
Runner string
Options string
DelayShutdown int
}
// runners is a map from runner name to runner instance.
runners = makeRunnerMap()
@ -79,6 +83,7 @@ func init() {
"HTTP requests.")
CmdNet.Flags().StringVar(&flags.Runner, "runner", "", "Runner to execute (available:"+legalRunners+")")
CmdNet.Flags().StringVar(&flags.Options, "options", "", "JSON options to the Runner")
CmdNet.Flags().IntVar(&flags.DelayShutdown, "delay-shutdown", 0, "Number of seconds to delay shutdown when receiving SIGTERM.")
}
func main(cmd *cobra.Command, args []string) {
@ -86,6 +91,17 @@ func main(cmd *cobra.Command, args []string) {
log.Fatalf("Must set either --runner or --serve, see --help")
}
if flags.DelayShutdown > 0 {
termCh := make(chan os.Signal, 1)
signal.Notify(termCh, syscall.SIGTERM)
go func() {
<-termCh
log.Printf("Sleeping %d seconds before terminating...", flags.DelayShutdown)
time.Sleep(time.Duration(flags.DelayShutdown) * time.Second)
os.Exit(0)
}()
}
log.SetFlags(log.Flags() | log.Lshortfile)
if flags.Serve == "" {