diff --git a/cmd/root.go b/cmd/root.go index 9dbf0b963..4db15cf3b 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -463,6 +463,8 @@ the Proxy will then pick-up automatically.`) "Enable debug logging") localFlags.Uint64Var(&c.conf.MaxConnections, "max-connections", 0, "Limit the number of connections. Default is no limit.") + localFlags.DurationVar(&c.conf.WaitBeforeClose, "min-sigterm-delay", 0, + "The number of seconds to accept new connections after receiving a TERM signal.") localFlags.DurationVar(&c.conf.WaitOnClose, "max-sigterm-delay", 0, "Maximum number of seconds to wait for connections to close after receiving a TERM signal.") localFlags.StringVar(&c.conf.TelemetryProject, "telemetry-project", "", @@ -1144,12 +1146,16 @@ func runSignalWrapper(cmd *Command) (err error) { switch { case errors.Is(err, errSigInt): cmd.logger.Infof("SIGINT signal received. Shutting down...") + time.Sleep(cmd.conf.WaitBeforeClose) case errors.Is(err, errSigTerm): cmd.logger.Infof("SIGTERM signal received. Shutting down...") + time.Sleep(cmd.conf.WaitBeforeClose) case errors.Is(err, errSigTermZero): cmd.logger.Infof("SIGTERM signal received. Shutting down...") + time.Sleep(cmd.conf.WaitBeforeClose) case errors.Is(err, errQuitQuitQuit): cmd.logger.Infof("/quitquitquit received request. Shutting down...") + time.Sleep(cmd.conf.WaitBeforeClose) default: cmd.logger.Errorf("The proxy has encountered a terminal error: %v", err) } diff --git a/cmd/root_test.go b/cmd/root_test.go index 027c2efff..4a4be698c 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -340,6 +340,13 @@ func TestNewCommandArguments(t *testing.T) { MaxConnections: 1, }), }, + { + desc: "using min-sigterm-delay flag", + args: []string{"--min-sigterm-delay", "10s", "proj:region:inst"}, + want: withDefaults(&proxy.Config{ + WaitBeforeClose: 10 * time.Second, + }), + }, { desc: "using wait after signterm flag", args: []string{"--max-sigterm-delay", "10s", "proj:region:inst"}, diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index 26d32af74..d0eb3da82 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -170,6 +170,11 @@ type Config struct { // connections. A zero-value indicates no limit. MaxConnections uint64 + // WaitBeforeClose sets the duration to wait after receiving a shutdown signal + // but before closing the process. Not setting this field means to initiate + // the shutdown process immediately. + WaitBeforeClose time.Duration + // WaitOnClose sets the duration to wait for connections to close before // shutting down. Not setting this field means to close immediately // regardless of any open connections.