ING-1399: Make shutdown timeout configurable through flags

Author: Westwooo

cmd/gateway/main.go

@@ -94,6 +94,7 @@ func init() {
 	configFlags.String("dapi-cert", "", "path to data api tls cert for Data API")
 	configFlags.String("dapi-key", "", "path to data api private tls key for Data API")
 	configFlags.Int("rate-limit", 0, "specifies the maximum requests per second to allow")
+	configFlags.Duration("shutdown-timeout", 30*time.Second, "the graceful shutdown timeout")
 	configFlags.String("otlp-endpoint", "", "opentelemetry endpoint to send telemetry to")
 	configFlags.Bool("disable-traces", false, "disable tracing")
 	configFlags.Bool("disable-metrics", false, "disable metrics")
@@ -269,6 +270,7 @@ type config struct {
 	clusterCaCertPath     string
 	clientCaCertPath      string
 	rateLimit             int
+	shutdownTimeout       time.Duration
 	otlpEndpoint          string
 	disableTraces         bool
 	disableMetrics        bool
@@ -312,6 +314,7 @@ func readConfig(logger *zap.Logger) *config {
 		clusterCaCertPath:     viper.GetString("cluster-cert"),
 		clientCaCertPath:      viper.GetString("client-ca-cert"),
 		rateLimit:             viper.GetInt("rate-limit"),
+		shutdownTimeout:       viper.GetDuration("shutdown-timeout"),
 		otlpEndpoint:          viper.GetString("otlp-endpoint"),
 		disableTraces:         viper.GetBool("disable-traces"),
 		disableMetrics:        viper.GetBool("disable-metrics"),
@@ -354,6 +357,7 @@ func readConfig(logger *zap.Logger) *config {
 		zap.String("dapiKeyPath", config.dapiKeyPath),
 		zap.String("clusterCaCertPath", config.clusterCaCertPath),
 		zap.Int("rateLimit", config.rateLimit),
+		zap.Duration("shutdownTimeout", config.shutdownTimeout),
 		zap.String("otlpEndpoint", config.otlpEndpoint),
 		zap.Bool("disableTraces", config.disableTraces),
 		zap.Bool("disableMetrics", config.disableMetrics),
@@ -645,6 +649,7 @@ func startGateway() {
 		BindDapiPort:        config.dapiPort,
 		BindAddress:         config.bindAddress,
 		RateLimit:           config.rateLimit,
+		ShutdownTimeout:     config.shutdownTimeout,
 		GrpcCertificate:     grpcCertificate,
 		DapiCertificate:     dapiCertificate,
 		ClusterCaCert:       caCertPool,
@@ -684,8 +689,9 @@ func startGateway() {
 
 		if newConfig.bindAddress != config.bindAddress ||
 			newConfig.dataPort != config.dataPort ||
-			newConfig.dapiPort != config.dapiPort {
-			logger.Warn("config changes for bindAddress, dataPort or dapiPort require a restart")
+			newConfig.dapiPort != config.dapiPort ||
+			newConfig.shutdownTimeout != config.shutdownTimeout {
+			logger.Warn("config changes for bindAddress, dataPort, dapiPort or shutdownTimeout require a restart")
 		}
 
 		if newConfig.selfSign != config.selfSign {

gateway/gateway.go

@@ -66,7 +66,8 @@ type Config struct {
 	AdvertiseAddress string
 	AdvertisePorts   ServicePorts
 
-	RateLimit int
+	RateLimit       int
+	ShutdownTimeout time.Duration
 
 	GrpcCertificate tls.Certificate
 	DapiCertificate tls.Certificate
@@ -432,9 +433,10 @@ func (g *Gateway) Run(ctx context.Context) error {
 				ClientCAs:  config.ClientCaCert,
 				ClientAuth: tls.VerifyClientCertIfGiven,
 			},
-			AlphaEndpoints: config.AlphaEndpoints,
-			Debug:          config.Debug,
-			Cancel:         cancel,
+			ShutdownTimeout: config.ShutdownTimeout,
+			AlphaEndpoints:  config.AlphaEndpoints,
+			Debug:           config.Debug,
+			Cancel:          cancel,
 		})
 		if err != nil {
 			config.Logger.Error("error creating legacy proxy")

gateway/system/system.go

@@ -54,11 +54,12 @@ type SystemOptions struct {
 	DapiImpl *dapiimpl.Servers
 	Metrics  *metrics.SnMetrics
 
-	RateLimiter    ratelimiting.RateLimiter
-	GrpcTlsConfig  *tls.Config
-	DapiTlsConfig  *tls.Config
-	AlphaEndpoints bool
-	Debug          bool
+	RateLimiter     ratelimiting.RateLimiter
+	GrpcTlsConfig   *tls.Config
+	DapiTlsConfig   *tls.Config
+	ShutdownTimeout time.Duration
+	AlphaEndpoints  bool
+	Debug           bool
 
 	Cancel context.CancelFunc
 }
@@ -69,7 +70,8 @@ type System struct {
 	dataServer *grpc.Server
 	dapiServer *http.Server
 
-	cancelFunc context.CancelFunc
+	shutdownTimeout time.Duration
+	cancelFunc      context.CancelFunc
 }
 
 func NewSystem(opts *SystemOptions) (*System, error) {
@@ -243,10 +245,11 @@ func NewSystem(opts *SystemOptions) (*System, error) {
 	}
 
 	s := &System{
-		logger:     opts.Logger,
-		dataServer: dataSrv,
-		dapiServer: dapiSrv,
-		cancelFunc: opts.Cancel,
+		logger:          opts.Logger,
+		dataServer:      dataSrv,
+		dapiServer:      dapiSrv,
+		shutdownTimeout: opts.ShutdownTimeout,
+		cancelFunc:      opts.Cancel,
 	}
 
 	return s, nil
@@ -304,7 +307,7 @@ func (s *System) Shutdown() {
 			defer wg.Done()
 			s.dapiServer.SetKeepAlivesEnabled(false)
 			time.Sleep(time.Second * 5)
-			ctx, cancel := context.WithTimeout(context.Background(), time.Second*30)
+			ctx, cancel := context.WithTimeout(context.Background(), s.shutdownTimeout)
 			defer cancel()
 			_ = s.dapiServer.Shutdown(ctx)
 		}()

gateway/test/dapi_graceful_shutdown_test.go

@@ -49,6 +49,7 @@ func (s *GatewayOpsTestSuite) TestGracefulShutdown() {
 		StartupCallback: func(m *gateway.StartupInfo) {
 			gwStartInfoCh <- m
 		},
+		ShutdownTimeout: time.Second * 10,
 	})
 	if err != nil {
 		s.T().Fatalf("failed to initialize graceful-shutdown-gateway: %s", err)