[Feature] [ML] Shutdown Handler (#1529)

CHANGELOG.md

@@ -30,6 +30,7 @@
 - (Improvement) (ML) Switch to fsnotify for file watching for MacOS support
 - (Feature) (ML) Unify Images, Resources and Lifecycle
 - (Improvement) (ML) CronJob status update
+- (Improvement) (ML) Job Sidecar Shutdown
 
 ## [1.2.35](https://github.com/arangodb/kube-arangodb/tree/1.2.35) (2023-11-06)
 - (Maintenance) Update go-driver to v1.6.0, update IsNotFound() checks

cmd/ml_storage.go

@@ -21,14 +21,14 @@
 package cmd
 
 import (
-	"context"
 	"os"
 
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
 
 	"github.com/arangodb/kube-arangodb/pkg/ml/storage"
-	"github.com/arangodb/kube-arangodb/pkg/util"
+	"github.com/arangodb/kube-arangodb/pkg/util/shutdown"
+	"github.com/arangodb/kube-arangodb/pkg/util/svc"
 )
 
 var (
@@ -48,13 +48,18 @@ var (
 	cmdMLStorageS3Options struct {
 		storage.ServiceConfig
 	}
+
+	cmdMLShutdownOptions struct {
+		shutdown.ServiceConfig
+	}
 )
 
 func init() {
 	cmdML.AddCommand(cmdMLStorage)
 	cmdMLStorage.AddCommand(cmdMLStorageS3)
 
 	f := cmdMLStorageS3.PersistentFlags()
+	f.StringVar(&cmdMLShutdownOptions.ListenAddress, "shutdown.address", "", "Address the GRPC shutdown service will listen on (IP:port)")
 	f.StringVar(&cmdMLStorageS3Options.ListenAddress, "server.address", "", "Address the GRPC service will listen on (IP:port)")
 
 	f.StringVar(&cmdMLStorageS3Options.S3.Endpoint, "s3.endpoint", "", "Endpoint of S3 API implementation")
@@ -76,12 +81,10 @@ func cmdMLStorageS3Run(cmd *cobra.Command, _ []string) {
 }
 
 func cmdMLStorageS3RunE(_ *cobra.Command) error {
-	ctx := util.CreateSignalContext(context.Background())
-
-	svc, err := storage.NewService(ctx, storage.StorageTypeS3Proxy, cmdMLStorageS3Options.ServiceConfig)
+	service, err := storage.NewService(shutdown.Context(), storage.StorageTypeS3Proxy, cmdMLStorageS3Options.ServiceConfig)
 	if err != nil {
 		return err
 	}
 
-	return svc.Run(ctx)
+	return svc.RunServices(shutdown.Context(), service, shutdown.ServiceCentral(cmdMLShutdownOptions.ServiceConfig))
 }

docs/api/ArangoMLStorage.V1Alpha1.md

@@ -143,6 +143,16 @@ Resources holds resource requests & limits for container
 Links:
 * [Documentation of core.ResourceRequirements](https://kubernetes.io/docs/reference/generated/kubernetes-api/v1.26/#resourcerequirements-v1-core)
 
+***
+
+### .spec.mode.sidecar.shutdownListenPort
+
+Type: `integer` <sup>[\[ref\]](https://github.com/arangodb/kube-arangodb/blob/1.2.35/pkg/apis/ml/v1alpha1/storage_spec_mode_sidecar.go#L36)</sup>
+
+ShutdownListenPort defines on which port the sidecar container will be listening for shutdown connections
+
+Default Value: `9202`
+
 ## Status
 
 ### .status.conditions

pkg/apis/ml/v1alpha1/extension_spec_deployment.go

@@ -107,9 +107,7 @@ func (s *ArangoMLExtensionSpecDeployment) Validate() error {
 	}
 
 	errs := []error{
-		shared.PrefixResourceErrors("service", shared.ValidateOptional(s.GetService(), func(service ArangoMLExtensionSpecDeploymentService) error {
-			return service.Validate()
-		})),
+		shared.PrefixResourceErrors("service", shared.ValidateOptional(s.GetService(), func(s ArangoMLExtensionSpecDeploymentService) error { return s.Validate() })),
 	}
 
 	if s.GetReplicas() < 0 || s.GetReplicas() > 10 {

pkg/apis/ml/v1alpha1/storage_spec_mode_sidecar.go

@@ -31,6 +31,10 @@ type ArangoMLStorageSpecModeSidecar struct {
 	// +doc/default: 9201
 	ListenPort *uint16 `json:"listenPort,omitempty"`
 
+	// ShutdownListenPort defines on which port the sidecar container will be listening for shutdown connections
+	// +doc/default: 9202
+	ShutdownListenPort *uint16 `json:"shutdownListenPort,omitempty"`
+
 	// Image define default image used for the extension
 	*sharedApi.Image `json:",inline"`
 
@@ -65,6 +69,10 @@ func (s *ArangoMLStorageSpecModeSidecar) Validate() error {
 		err = append(err, shared.PrefixResourceErrors("listenPort", errors.Newf("must be positive")))
 	}
 
+	if s.GetShutdownListenPort() < 1 {
+		err = append(err, shared.PrefixResourceErrors("shutdownListenPort", errors.Newf("must be positive")))
+	}
+
 	err = append(err, s.GetResources().Validate())
 
 	return shared.WithErrors(err...)
@@ -76,3 +84,10 @@ func (s *ArangoMLStorageSpecModeSidecar) GetListenPort() uint16 {
 	}
 	return *s.ListenPort
 }
+
+func (s *ArangoMLStorageSpecModeSidecar) GetShutdownListenPort() uint16 {
+	if s == nil || s.ShutdownListenPort == nil {
+		return 9202
+	}
+	return *s.ShutdownListenPort
+}

pkg/apis/shared/v1/image.go

@@ -38,6 +38,14 @@ type Image struct {
 	PullSecrets []string `json:"pullSecrets,omitempty"`
 }
 
+func (i *Image) GetImage() string {
+	if i == nil || i.Image == nil {
+		return ""
+	}
+
+	return *i.Image
+}
+
 func (i *Image) Validate() error {
 	if i == nil {
 		return nil