fix: flyteadmin doesn't shutdown servers gracefully

Signed-off-by: Ryan Lo <wenchih@apache.org>

flyteadmin/pkg/server/service.go

@@ -6,7 +6,10 @@ import (
 	"fmt"
 	"net"
 	"net/http"
+	"os"
+	"os/signal"
 	"strings"
+	"syscall"
 	"time"
 
 	"github.com/gorilla/handlers"
@@ -81,8 +84,8 @@ func SetMetricKeys(appConfig *runtimeIfaces.ApplicationConfig) {
 
 // Creates a new gRPC Server with all the configuration
 func newGRPCServer(ctx context.Context, pluginRegistry *plugins.Registry, cfg *config.ServerConfig,
-	storageCfg *storage.Config, authCtx interfaces.AuthenticationContext,
-	scope promutils.Scope, sm core.SecretManager, opts ...grpc.ServerOption) (*grpc.Server, error) {
+ storageCfg *storage.Config, authCtx interfaces.AuthenticationContext,
+ scope promutils.Scope, sm core.SecretManager, opts ...grpc.ServerOption) (*grpc.Server, error) {
 
 	logger.Infof(ctx, "Registering default middleware with blanket auth validation")
 	pluginRegistry.RegisterDefault(plugins.PluginIDUnaryServiceMiddleware, grpcmiddleware.ChainUnaryServer(
@@ -200,8 +203,8 @@ func healthCheckFunc(w http.ResponseWriter, _ *http.Request) {
 }
 
 func newHTTPServer(ctx context.Context, pluginRegistry *plugins.Registry, cfg *config.ServerConfig, _ *authConfig.Config, authCtx interfaces.AuthenticationContext,
-	additionalHandlers map[string]func(http.ResponseWriter, *http.Request),
-	grpcAddress string, grpcConnectionOpts ...grpc.DialOption) (*http.ServeMux, error) {
+ additionalHandlers map[string]func(http.ResponseWriter, *http.Request),
+ grpcAddress string, grpcConnectionOpts ...grpc.DialOption) (*http.ServeMux, error) {
 
 	// Register the server that will serve HTTP/REST Traffic
 	mux := http.NewServeMux()
@@ -329,9 +332,9 @@ func generateRequestID() string {
 }
 
 func serveGatewayInsecure(ctx context.Context, pluginRegistry *plugins.Registry, cfg *config.ServerConfig,
-	authCfg *authConfig.Config, storageConfig *storage.Config,
-	additionalHandlers map[string]func(http.ResponseWriter, *http.Request),
-	scope promutils.Scope) error {
+ authCfg *authConfig.Config, storageConfig *storage.Config,
+ additionalHandlers map[string]func(http.ResponseWriter, *http.Request),
+ scope promutils.Scope) error {
 	logger.Infof(ctx, "Serving Flyte Admin Insecure")
 
 	// This will parse configuration and create the necessary objects for dealing with auth
@@ -422,11 +425,29 @@ 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 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(), 10*time.Second)
+	defer cancel()
+
+	if err := server.Shutdown(shutdownCtx); err != nil {
+		logger.Errorf(ctx, "Failed to shut down HTTP server: %v", err)
 	}
 
+	grpcServer.GracefulStop()
+
+	logger.Infof(ctx, "Servers gracefully stopped")
 	return nil
 }
 
@@ -445,9 +466,9 @@ func grpcHandlerFunc(grpcServer *grpc.Server, otherHandler http.Handler) http.Ha
 }
 
 func serveGatewaySecure(ctx context.Context, pluginRegistry *plugins.Registry, cfg *config.ServerConfig, authCfg *authConfig.Config,
-	storageCfg *storage.Config,
-	additionalHandlers map[string]func(http.ResponseWriter, *http.Request),
-	scope promutils.Scope) error {
+ storageCfg *storage.Config,
+ additionalHandlers map[string]func(http.ResponseWriter, *http.Request),
+ scope promutils.Scope) error {
 	certPool, cert, err := GetSslCredentials(ctx, cfg.Security.Ssl.CertificateFile, cfg.Security.Ssl.KeyFile)
 	sm := secretmanager.NewFileEnvSecretManager(secretmanager.GetConfig())