diff --git a/test/images/agnhost/net/main.go b/test/images/agnhost/net/main.go index 6e238765c49..71423e90aee 100644 --- a/test/images/agnhost/net/main.go +++ b/test/images/agnhost/net/main.go @@ -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 == "" {