diff --git a/cmd/root.go b/cmd/root.go index b252dd6c..d964346f 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -566,6 +566,9 @@ the Proxy will then pick-up automatically.`) localFlags.Uint64Var(&c.conf.MaxConnections, "max-connections", 0, `Limits the number of connections by refusing any additional connections. When this flag is not set, there 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. Defaults to 0s.`) localFlags.DurationVar(&c.conf.WaitOnClose, "max-sigterm-delay", 0, `Maximum amount of time to wait after for any open connections to close after receiving a TERM signal. The proxy will shut @@ -1179,8 +1182,10 @@ 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) 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 0a64dba4..00cf87d6 100644 --- a/cmd/root_test.go +++ b/cmd/root_test.go @@ -412,6 +412,13 @@ func TestNewCommandArguments(t *testing.T) { MaxConnections: 1, }), }, + { + desc: "using min-sigterm-delay flag", + args: []string{"--min-sigterm-delay", "10s", "projects/proj/locations/region/clusters/clust/instances/inst"}, + want: withDefaults(&proxy.Config{ + WaitBeforeClose: 10 * time.Second, + }), + }, { desc: "using wait after signterm flag", args: []string{"--max-sigterm-delay", "10s", "projects/proj/locations/region/clusters/clust/instances/inst"}, diff --git a/internal/proxy/proxy.go b/internal/proxy/proxy.go index c1e3f56f..877532b0 100644 --- a/internal/proxy/proxy.go +++ b/internal/proxy/proxy.go @@ -142,6 +142,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.