Add flag to allow setting a shutdown wait period

Fixes open-policy-agent#2764

Signed-off-by: Björn Carlsson <bjorn.carlsson@bisnode.com>

cmd/run.go

@@ -181,6 +181,7 @@ To skip bundle verification, use the --skip-verify flag.
 	runCommand.Flags().VarP(cmdParams.logLevel, "log-level", "l", "set log level")
 	runCommand.Flags().VarP(cmdParams.logFormat, "log-format", "", "set log format")
 	runCommand.Flags().IntVar(&cmdParams.rt.GracefulShutdownPeriod, "shutdown-grace-period", 10, "set the time (in seconds) that the server will wait to gracefully shut down")
+	runCommand.Flags().IntVar(&cmdParams.rt.ShutdownWaitPeriod, "shutdown-wait-period", 0, "set the time (in seconds) that the server will wait before initiating shutdown")
 	addConfigOverrides(runCommand.Flags(), &cmdParams.rt.ConfigOverrides)
 	addConfigOverrideFiles(runCommand.Flags(), &cmdParams.rt.ConfigOverrideFiles)
 	addBundleModeFlag(runCommand.Flags(), &cmdParams.rt.BundleMode, false)

runtime/runtime.go

@@ -161,6 +161,9 @@ type Params struct {
 	// server to shutdown gracefully.
 	GracefulShutdownPeriod int
 
+	// ShutdownWaitPeriod is the time (in seconds) to wait before initiating shutdown.
+	ShutdownWaitPeriod int
+
 	// EnableVersionCheck flag controls whether OPA will report its version to an external service.
 	// If this flag is true, OPA will report its version to the external service
 	EnableVersionCheck bool
@@ -614,6 +617,11 @@ func (rt *Runtime) getBanner() string {
 }
 
 func (rt *Runtime) gracefulServerShutdown(s *server.Server) error {
+	if rt.Params.ShutdownWaitPeriod > 0 {
+		logrus.Infof("Waiting %vs before initiating shutdown...", rt.Params.ShutdownWaitPeriod)
+		time.Sleep(time.Duration(rt.Params.ShutdownWaitPeriod) * time.Second)
+	}
+
 	logrus.Info("Shutting down...")
 	ctx, cancel := context.WithTimeout(context.Background(), time.Duration(rt.Params.GracefulShutdownPeriod)*time.Second)
 	defer cancel()

test/e2e/shutdown/shutdown_test.go

@@ -0,0 +1,51 @@
+// Copyright 2020 The OPA Authors.  All rights reserved.
+// Use of this source code is governed by an Apache2
+// license that can be found in the LICENSE file.
+package shutdown
+
+import (
+	"flag"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/open-policy-agent/opa/test/e2e"
+)
+
+var testRuntime *e2e.TestRuntime
+
+func TestMain(m *testing.M) {
+	flag.Parse()
+	testServerParams := e2e.NewAPIServerTestParams()
+
+	testServerParams.GracefulShutdownPeriod = 1
+	testServerParams.ShutdownWaitPeriod = 2
+
+	var err error
+	testRuntime, err = e2e.NewTestRuntime(testServerParams)
+	if err != nil {
+		os.Exit(1)
+	}
+
+	os.Exit(testRuntime.RunTests(m))
+}
+
+func TestShutdownWaitPeriod(t *testing.T) {
+	proc, err := os.FindProcess(os.Getpid())
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	err = proc.Signal(os.Interrupt)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	time.Sleep(1500 * time.Millisecond)
+
+	// Ensure that OPA i still running
+	err = testRuntime.HealthCheck(testRuntime.URL())
+	if err != nil {
+		t.Fatalf("Expected health endpoint to be up but got:\n\n%v", err)
+	}
+}