Repository: flyteorg/flyte (go)

Issue:
[BUG] Flyte Admin Should Terminate Gracefully
### Describe the bug

Flyte Admin does not appear to respond to SIGTERM and as such does not gracefully terminate when Kubernetes is doing a rolling restart or upgrade. Servers should terminate gracefully so that they can drain buffers, connections, RPCs, etc and minimize errors or loss of data. 

### Expected behavior

When receiving a SIGTERM, I expect Flyte Admin to run a shutdown sequence where the RPC servers are shut down and any buffers are drained. 

### Additional context to reproduce

_No response_

### Screenshots

_No response_

### Are you sure this issue hasn't been raised already?

- [x] Yes

### Have you read the Code of Conduct?

- [x] Yes

Pull request:
fix: flyteadmin doesn't shutdown servers gracefully
## Tracking issue
Closes #6274 

## Why are the changes needed?
To handle SIGTERM properly, ensure cleanup of resources and provide smooth user experience
<!--
Please clarify why the changes are needed. For instance,
1. If you propose a new API, clarify the use case for a new API.
2. If you fix a bug, you can clarify why it is a bug.
-->

## What changes were proposed in this pull request?
Implement graceful shutdown for flyteadmin servers

<!--
Please clarify what changes you are proposing. The purpose of this section is to outline the changes and how this PR fixes the issue.
If possible, please consider writing useful notes for better and faster reviews in your PR. See the examples below.
1. If you refactor some codes with changing classes, showing the class hierarchy will help reviewers.
2. If there is design documentation, please add the link.
-->

## How was this patch tested?

<!--
If tests were added, say they were added here. Please make sure to add some test cases that check the changes thoroughly including negative and positive cases if possible.
If it was tested in a way different from regular unit tests, please clarify how you tested step by step, ideally copy and paste-able, so that other reviewers can test and check, and descendants can verify in the future.
If tests were not added, please describe why they were not added and/or why it was difficult to add.
-->
### Labels

Please add one or more of the following labels to categorize your PR:
- **added**: For new features.
- **changed**: For changes in existing functionality.
- **deprecated**: For soon-to-be-removed features.
- **removed**: For features being removed.
- **fixed**: For any bug fixed.
- **security**: In case of vulnerabilities

This is important to improve the readability of release notes.

### Setup process

### Screenshots

## Check all the applicable boxes <!-- Follow the above conventions to check the box -->

- [ ] I updated the documentation accordingly.
- [ ] All new and existing tests passed.
- [x] All commits are signed-off.

## Related PRs

<!-- Add related pull requests for reviewers to check -->

## Docs link

<!-- Add documentation link built by CI jobs here, and specify the changed place --> 
 <div id='description'>
    <a href="https://bito.ai#summarystart"></a>
<h3>Summary by Bito</h3>
This pull request enhances the flyteadmin servers by implementing a graceful shutdown mechanism with a configurable timeout for resource cleanup. It also improves error handling during gRPC server startup, enhancing overall service reliability.
<!-- Disabling unit_tests and post_effort_to_review
<br>
<br>
<b>Unit tests added</b>: False
<br>
<br>
<b>Estimated effort to review (1-5, lower is better)</b>: 2
-->
</div>

Patch:
diff --git a/flyteadmin/pkg/config/config.go b/flyteadmin/pkg/config/config.go
index 5b9e49460..e132ed1ef 100644
--- a/flyteadmin/pkg/config/config.go
+++ b/flyteadmin/pkg/config/config.go
@@ -25,9 +25,10 @@ type ServerConfig struct {
 	// Deprecated: please use auth.AppAuth.ThirdPartyConfig instead.
 	DeprecatedThirdPartyConfig authConfig.ThirdPartyConfigOptions `json:"thirdPartyConfig" pflag:",Deprecated please use auth.appAuth.thirdPartyConfig instead."`
 
-	DataProxy                DataProxyConfig  `json:"dataProxy" pflag:",Defines data proxy configuration."`
-	ReadHeaderTimeoutSeconds int              `json:"readHeaderTimeoutSeconds" pflag:",The amount of time allowed to read request headers."`
-	KubeClientConfig         KubeClientConfig `json:"kubeClientConfig" pflag:",Configuration to control the Kubernetes client"`
+	DataProxy                      DataProxyConfig  `json:"dataProxy" pflag:",Defines data proxy configuration."`
+	ReadHeaderTimeoutSeconds       int              `json:"readHeaderTimeoutSeconds" pflag:",The amount of time allowed to read request headers."`
+	KubeClientConfig               KubeClientConfig `json:"kubeClientConfig" pflag:",Configuration to control the Kubernetes client"`
+	GracefulShutdownTimeoutSeconds int              `json:"gracefulShutdownTimeoutSeconds" pflag:",Number of seconds to wait for graceful shutdown before forcefully terminating the server"`
 }
 
 type DataProxyConfig struct {
@@ -120,6 +121,7 @@ var defaultServerConfig = &ServerConfig{
 		Burst:   25,
 		Timeout: config.Duration{Duration: 30 * time.Second},
 	},
+	GracefulShutdownTimeoutSeconds: 10,
 }
 var serverConfig = config.MustRegisterSection(SectionKey, defaultServerConfig)
 
diff --git a/flyteadmin/pkg/config/serverconfig_flags.go b/flyteadmin/pkg/config/serverconfig_flags.go
index d60aeb353..315973818 100755
--- a/flyteadmin/pkg/config/serverconfig_flags.go
+++ b/flyteadmin/pkg/config/serverconfig_flags.go
@@ -82,5 +82,6 @@ func (cfg ServerConfig) GetPFlagSet(prefix string) *pflag.FlagSet {
 	cmdFlags.Int32(fmt.Sprintf("%v%v", prefix, "kubeClientConfig.qps"), defaultServerConfig.KubeClientConfig.QPS, "Max QPS to the master for requests to KubeAPI. 0 defaults to 5.")
 	cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "kubeClientConfig.burst"), defaultServerConfig.KubeClientConfig.Burst, "Max burst rate for throttle. 0 defaults to 10")
 	cmdFlags.String(fmt.Sprintf("%v%v", prefix, "kubeClientConfig.timeout"), defaultServerConfig.KubeClientConfig.Timeout.String(), "Max duration allowed for every request to KubeAPI before giving up. 0 implies no timeout.")
+	cmdFlags.Int(fmt.Sprintf("%v%v", prefix, "gracefulShutdownTimeoutSeconds"), defaultServerConfig.GracefulShutdownTimeoutSeconds, "Number of seconds to wait for graceful shutdown before forcefully terminating the server")
 	return cmdFlags
 }
diff --git a/flyteadmin/pkg/server/service.go b/flyteadmin/pkg/server/service.go
index 0575abc56..ec776f896 100644
--- a/flyteadmin/pkg/server/service.go
+++ b/flyteadmin/pkg/server/service.go
@@ -6,7 +6,11 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"os"
+	"os/signal"
 	"strings"
+	"sync"
+	"syscall"
 	"time"
 
 	"github.com/gorilla/handlers"
@@ -389,8 +393,9 @@ func serveGatewayInsecure(ctx context.Context, pluginRegistry *plugins.Registry,
 	}
 
 	go func() {
-		err := grpcServer.Serve(lis)
-		logger.Fatalf(ctx, "Failed to create GRPC Server, Err: ", err)
+		if err := grpcServer.Serve(lis); err != nil {
+			logger.Fatalf(ctx, "Failed to create GRPC Server, Err: %v", err)
+		}
 	}()
 
 	logger.Infof(ctx, "Starting HTTP/1 Gateway server on %s", cfg.GetHostAddress())
@@ -425,11 +430,45 @@ func serveGatewayInsecure(ctx context.Context, pluginRegistry *plugins.Registry,
 		ReadHeaderTimeout: time.Duration(cfg.ReadHeaderTimeoutSeconds) * time.Second,
 	}
 
-	err = server.ListenAndServe()
-	if err != nil {
-		return errors.Wrapf(err, "failed to Start HTTP Server")
-	}
+	go func() {
+		err = server.ListenAndServe()
+		if err != nil && err != http.ErrServerClosed {
+			logger.Fatalf(ctx, "Failed to start HTTP Server: %v", err)
+		}
+	}()
 
+	// Gracefully shut down the servers
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+	<-sigCh
+
+	// Forced shutdown because of timeout
+	logger.Infof(ctx, "Shutting down server... timeout: %d seconds", cfg.GracefulShutdownTimeoutSeconds)
+	shutdownTimeout := cfg.GracefulShutdownTimeoutSeconds
+	timer := time.AfterFunc(time.Duration(shutdownTimeout)*time.Second, func() {
+		logger.Infof(ctx, "Server couldn't stop gracefully in time. Doing force stop.")
+		server.Close()
+		grpcServer.Stop()
+	})
+	defer timer.Stop()
+
+	var wg sync.WaitGroup
+	wg.Add(2)
+
+	go func() {
+		defer wg.Done()
+		if err := server.Shutdown(ctx); err != nil {
+			logger.Errorf(ctx, "Failed to gracefully shutdown HTTP server: %v", err)
+		}
+	}()
+
+	go func() {
+		defer wg.Done()
+		grpcServer.GracefulStop()
+	}()
+
+	wg.Wait()
+	logger.Infof(ctx, "Servers gracefully stopped")
 	return nil
 }
 
@@ -537,10 +576,26 @@ func serveGatewaySecure(ctx context.Context, pluginRegistry *plugins.Registry, c
 		ReadHeaderTimeout: time.Duration(cfg.ReadHeaderTimeoutSeconds) * time.Second,
 	}
 
-	err = srv.Serve(tls.NewListener(conn, srv.TLSConfig))
+	go func() {
+		err = srv.Serve(tls.NewListener(conn, srv.TLSConfig))
+		if err != nil && err != http.ErrServerClosed {
+			logger.Errorf(ctx, "Failed to start HTTP/2 Server: %v", err)
+		}
+	}()
 
-	if err != nil {
-		return errors.Wrapf(err, "failed to Start HTTP/2 Server")
+	// Gracefully shutdown the servers
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, syscall.SIGINT, syscall.SIGTERM)
+	<-sigCh
+
+	// Create a context with timeout for the shutdown process
+	shutdownCtx, cancel := context.WithTimeout(context.Background(), time.Duration(cfg.GracefulShutdownTimeoutSeconds)*time.Second)
+	defer cancel()
+
+	if err := srv.Shutdown(shutdownCtx); err != nil {
+		logger.Errorf(ctx, "Failed to shutdown HTTP server: %v", err)
 	}
+
+	logger.Infof(ctx, "Servers gracefully stopped")
 	return nil
 }


Test patch:
diff --git a/flyteadmin/pkg/config/serverconfig_flags_test.go b/flyteadmin/pkg/config/serverconfig_flags_test.go
index f3535ae7f..77075d9f8 100755
--- a/flyteadmin/pkg/config/serverconfig_flags_test.go
+++ b/flyteadmin/pkg/config/serverconfig_flags_test.go
@@ -547,4 +547,18 @@ func TestServerConfig_SetFlags(t *testing.T) {
 			}
 		})
 	})
+	t.Run("Test_gracefulShutdownTimeoutSeconds", func(t *testing.T) {
+
+		t.Run("Override", func(t *testing.T) {
+			testValue := "1"
+
+			cmdFlags.Set("gracefulShutdownTimeoutSeconds", testValue)
+			if vInt, err := cmdFlags.GetInt("gracefulShutdownTimeoutSeconds"); err == nil {
+				testDecodeJson_ServerConfig(t, fmt.Sprintf("%v", vInt), &actual.GracefulShutdownTimeoutSeconds)
+
+			} else {
+				assert.FailNow(t, err.Error())
+			}
+		})
+	})
 }