ING-1429: Ensure resources are release during shutdown

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
@@ -409,6 +410,9 @@ func (g *Gateway) Run(ctx context.Context) error {
 			Password:        config.Password,
 		})
 
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+
 		config.Logger.Info("initializing protostellar system")
 		gatewaySys, err := system.NewSystem(&system.SystemOptions{
 			Logger:      config.Logger.Named("gateway-system"),
@@ -430,8 +434,9 @@ func (g *Gateway) Run(ctx context.Context) error {
 				ClientCAs:  config.ClientCaCert,
 				ClientAuth: tls.VerifyClientCertIfGiven,
 			},
-			AlphaEndpoints: config.AlphaEndpoints,
-			Debug:          config.Debug,
+			ShutdownTimeout: config.ShutdownTimeout,
+			AlphaEndpoints:  config.AlphaEndpoints,
+			Debug:           config.Debug,
 		})
 		if err != nil {
 			config.Logger.Error("error creating legacy proxy")

gateway/system/system.go

@@ -46,6 +46,7 @@ import (
 )
 
 const maxMsgSize = 25 * 1024 * 1024 // 25MiB
+const defaultShutdownTimeout = time.Second * 30
 
 type SystemOptions struct {
 	Logger *zap.Logger
@@ -59,13 +60,17 @@ type SystemOptions struct {
 	DapiTlsConfig  *tls.Config
 	AlphaEndpoints bool
 	Debug          bool
+
+	ShutdownTimeout time.Duration
 }
 
 type System struct {
 	logger *zap.Logger
 
 	dataServer *grpc.Server
 	dapiServer *http.Server
+
+	shutdownTimeout time.Duration
 }
 
 func NewSystem(opts *SystemOptions) (*System, error) {
@@ -238,10 +243,16 @@ func NewSystem(opts *SystemOptions) (*System, error) {
 		TLSConfig:         opts.DapiTlsConfig,
 	}
 
+	if opts.ShutdownTimeout == 0 {
+		opts.Logger.Warn("no shutdown timeout configured using default", zap.Duration("shutdownTimeout", defaultShutdownTimeout))
+		opts.ShutdownTimeout = defaultShutdownTimeout
+	}
+
 	s := &System{
-		logger:     opts.Logger,
-		dataServer: dataSrv,
-		dapiServer: dapiSrv,
+		logger:          opts.Logger,
+		dataServer:      dataSrv,
+		dapiServer:      dapiSrv,
+		shutdownTimeout: opts.ShutdownTimeout,
 	}
 
 	return s, nil
@@ -289,7 +300,20 @@ func (s *System) Shutdown() {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			s.dataServer.GracefulStop()
+
+			// GracefulStop has no timeout mechanism so we need to take this approach
+			done := make(chan struct{})
+			go func() {
+				s.dataServer.GracefulStop()
+				close(done)
+			}()
+
+			select {
+			case <-done:
+			case <-time.After(s.shutdownTimeout):
+				s.logger.Warn("data server shutdown timed out, forcing stop")
+				s.dataServer.Stop()
+			}
 		}()
 	}
 
@@ -299,7 +323,13 @@ func (s *System) Shutdown() {
 			defer wg.Done()
 			s.dapiServer.SetKeepAlivesEnabled(false)
 			time.Sleep(time.Second * 5)
-			_ = s.dapiServer.Shutdown(context.Background())
+			ctx, cancel := context.WithTimeout(context.Background(), s.shutdownTimeout)
+			defer cancel()
+			err := s.dapiServer.Shutdown(ctx)
+			if err != nil {
+				s.logger.Warn("data api server shutdown failed", zap.Error(err))
+				_ = s.dapiServer.Close()
+			}
 		}()
 	}