From f94d704d03a76a1d20c4bdebe50a3a0a7fe9933c Mon Sep 17 00:00:00 2001
From: George <me@georgemac.com>
Date: Wed, 5 Oct 2022 14:36:35 +0100
Subject: [PATCH] fix: initialize shutdown context timeout after interrupt
 (#1057)

* fix: initialize shutdown context timeout after interrupt

* Move shutdown log msgs

Co-authored-by: Mark Phelps <209477+markphelps@users.noreply.github.com>
---
 cmd/flipt/main.go | 29 +++++++++++++++++++----------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/cmd/flipt/main.go b/cmd/flipt/main.go
index a186f435a2..e93f0112c8 100644
--- a/cmd/flipt/main.go
+++ b/cmd/flipt/main.go
@@ -250,9 +250,6 @@ func run(ctx context.Context, logger *zap.Logger) error {
 	signal.Notify(interrupt, os.Interrupt, syscall.SIGTERM)
 	defer signal.Stop(interrupt)
 
-	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
-	defer shutdownCancel()
-
 	var (
 		isRelease = isRelease()
 		isConsole = cfg.Log.Encoding == config.LogEncodingConsole
@@ -384,6 +381,8 @@ func run(ctx context.Context, logger *zap.Logger) error {
 	var (
 		grpcServer *grpc.Server
 		httpServer *http.Server
+
+		shutdownFuncs = []func(context.Context){}
 	)
 
 	// starts grpc server
@@ -496,7 +495,7 @@ func run(ctx context.Context, logger *zap.Logger) error {
 					DB:       cfg.Cache.Redis.DB,
 				})
 
-				defer rdb.Shutdown(shutdownCtx)
+				shutdownFuncs = append(shutdownFuncs, func(ctx context.Context) { _ = rdb.Shutdown(ctx) })
 
 				status := rdb.Ping(ctx)
 				if status == nil {
@@ -533,6 +532,12 @@ func run(ctx context.Context, logger *zap.Logger) error {
 		// initialize grpc server
 		grpcServer = grpc.NewServer(grpcOpts...)
 
+		// register grpcServer graceful stop on shutdown
+		shutdownFuncs = append(shutdownFuncs, func(context.Context) {
+			grpcServer.GracefulStop()
+			logger.Info("grpc server shutdown gracefully")
+		})
+
 		pb.RegisterFliptServer(grpcServer, srv)
 		grpc_prometheus.EnableHandlingTimeHistogram()
 		grpc_prometheus.Register(grpcServer)
@@ -657,6 +662,12 @@ func run(ctx context.Context, logger *zap.Logger) error {
 			MaxHeaderBytes: 1 << 20,
 		}
 
+		// register httpServer graceful stop on shutdown
+		shutdownFuncs = append(shutdownFuncs, func(ctx context.Context) {
+			_ = httpServer.Shutdown(ctx)
+			logger.Info("http server shutdown gracefully")
+		})
+
 		logger.Debug("starting http server")
 
 		var (
@@ -703,7 +714,6 @@ func run(ctx context.Context, logger *zap.Logger) error {
 			return fmt.Errorf("http server: %w", err)
 		}
 
-		logger.Info("server shutdown gracefully")
 		return nil
 	})
 
@@ -718,12 +728,11 @@ func run(ctx context.Context, logger *zap.Logger) error {
 
 	cancel()
 
-	if httpServer != nil {
-		_ = httpServer.Shutdown(shutdownCtx)
-	}
+	shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer shutdownCancel()
 
-	if grpcServer != nil {
-		grpcServer.GracefulStop()
+	for _, shutdown := range shutdownFuncs {
+		shutdown(shutdownCtx)
 	}
 
 	return g.Wait()