From ff997ae5a0ee34059483dd15d3c692210620bee7 Mon Sep 17 00:00:00 2001
From: Andrew Sy Kim <andrewsy@google.com>
Date: Wed, 16 Mar 2022 16:32:38 -0400
Subject: [PATCH] agnhost/net: support --delay-shutdown flag

Signed-off-by: Andrew Sy Kim <andrewsy@google.com>
---
 test/images/agnhost/net/main.go | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

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 == "" {