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

Author: ajanikow

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/api/shutdown/v1/operator.pb.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/api/shutdown/v1/operator.proto

@@ -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/api/shutdown/v1/operator_grpc.pb.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