Restore ReadinessProbe for ML Storage sidecar, rename shutdown -> controller service

Also small clean-up of duplicate Ter function

Author: nikita-vanyasin

CHANGELOG.md

@@ -9,6 +9,7 @@
 - (Documentation) Use relative links for generated docs
 - (Improvement) Change default logging level to info. Add --log.sampling (default true). Adjust log levels.
 - (Maintenance) Bump Go to 1.21.6
+- (Feature) (ML) Restore ReadinessProbe for ML Storage sidecar
 
 ## [1.2.36](https://github.com/arangodb/kube-arangodb/tree/1.2.36) (2024-01-08)
 - (Documentation) Improvements and fixes for rendered documentation (GH pages)

cmd/ml_storage.go

@@ -25,8 +25,11 @@ import (
 
 	"github.com/rs/zerolog/log"
 	"github.com/spf13/cobra"
+	"google.golang.org/grpc"
 
+	pbShutdown "github.com/arangodb/kube-arangodb/pkg/api/shutdown/v1"
 	"github.com/arangodb/kube-arangodb/pkg/ml/storage"
+	"github.com/arangodb/kube-arangodb/pkg/util/probe"
 	"github.com/arangodb/kube-arangodb/pkg/util/shutdown"
 	"github.com/arangodb/kube-arangodb/pkg/util/svc"
 )
@@ -49,8 +52,8 @@ var (
 		storage.ServiceConfig
 	}
 
-	cmdMLShutdownOptions struct {
-		shutdown.ServiceConfig
+	cmdMLStorageControllerOptions struct {
+		svc.GRPCConfig
 	}
 )
 
@@ -59,7 +62,7 @@ func init() {
 	cmdMLStorage.AddCommand(cmdMLStorageS3)
 
 	f := cmdMLStorageS3.PersistentFlags()
-	f.StringVar(&cmdMLShutdownOptions.ListenAddress, "shutdown.address", "", "Address the GRPC shutdown service will listen on (IP:port)")
+	f.StringVar(&cmdMLStorageControllerOptions.ListenAddress, "controller.address", "", "Address the GRPC controller 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")
@@ -81,10 +84,16 @@ func cmdMLStorageS3Run(cmd *cobra.Command, _ []string) {
 }
 
 func cmdMLStorageS3RunE(_ *cobra.Command) error {
-	service, err := storage.NewService(shutdown.Context(), storage.StorageTypeS3Proxy, cmdMLStorageS3Options.ServiceConfig)
+	storageService, err := storage.NewService(shutdown.Context(), storage.StorageTypeS3Proxy, cmdMLStorageS3Options.ServiceConfig)
 	if err != nil {
 		return err
 	}
 
-	return svc.RunServices(shutdown.Context(), service, shutdown.ServiceCentral(cmdMLShutdownOptions.ServiceConfig))
+	healthService := probe.NewHealthService()
+
+	controllerService := svc.NewGRPC(cmdMLStorageControllerOptions.GRPCConfig, func(server *grpc.Server) {
+		pbShutdown.RegisterShutdownServer(server, shutdown.NewShutdownableShutdownServer())
+		healthService.Register(server)
+	})
+	return svc.RunServices(shutdown.Context(), healthService, storageService, controllerService)
 }

pkg/apis/ml/v1alpha1/storage_spec_mode_sidecar.go

@@ -31,9 +31,9 @@ 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
+	// ControllerListenPort defines on which port the sidecar container will be listening for controller requests
 	// +doc/default: 9202
-	ShutdownListenPort *uint16 `json:"shutdownListenPort,omitempty"`
+	ControllerListenPort *uint16 `json:"controllerListenPort,omitempty"`
 
 	// ContainerTemplate Keeps the information about Container configuration
 	*sharedApi.ContainerTemplate `json:",inline"`
@@ -58,8 +58,8 @@ 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")))
+	if s.GetControllerListenPort() < 1 {
+		err = append(err, shared.PrefixResourceErrors("controllerListenPort", errors.Newf("must be positive")))
 	}
 
 	err = append(err, s.GetContainerTemplate().Validate())
@@ -74,9 +74,9 @@ func (s *ArangoMLStorageSpecModeSidecar) GetListenPort() uint16 {
 	return *s.ListenPort
 }
 
-func (s *ArangoMLStorageSpecModeSidecar) GetShutdownListenPort() uint16 {
-	if s == nil || s.ShutdownListenPort == nil {
+func (s *ArangoMLStorageSpecModeSidecar) GetControllerListenPort() uint16 {
+	if s == nil || s.ControllerListenPort == nil {
 		return 9202
 	}
-	return *s.ShutdownListenPort
+	return *s.ControllerListenPort
 }

pkg/apis/ml/v1alpha1/zz_generated.deepcopy.go

@@ -0,0 +1,57 @@
+//
+// DISCLAIMER
+//
+// Copyright 2023 ArangoDB GmbH, Cologne, Germany
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// Copyright holder is ArangoDB GmbH, Cologne, Germany
+//
+
+package probe
+
+import (
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/health"
+	pbHealth "google.golang.org/grpc/health/grpc_health_v1"
+)
+
+type HealthService interface {
+	Register(server *grpc.Server)
+	SetServing()
+	Shutdown()
+}
+
+func NewHealthService() HealthService {
+	return &grpcHealthService{
+		hs: health.NewServer(),
+	}
+}
+
+type grpcHealthService struct {
+	hs *health.Server
+}
+
+func (s *grpcHealthService) Register(server *grpc.Server) {
+	pbHealth.RegisterHealthServer(server, s.hs)
+}
+
+// SetServing marks the health response as Serving for all services
+func (s *grpcHealthService) SetServing() {
+	s.hs.SetServingStatus("", pbHealth.HealthCheckResponse_SERVING)
+}
+
+// Shutdown marks as not serving and forbids further changes in health status
+func (s *grpcHealthService) Shutdown() {
+	s.hs.Shutdown()
+}

pkg/util/probe/grpc.go

@@ -24,26 +24,12 @@ import (
 	"context"
 	"time"
 
-	"google.golang.org/grpc"
-
 	"github.com/arangodb/kube-arangodb/pkg/api/server"
 	pbShutdown "github.com/arangodb/kube-arangodb/pkg/api/shutdown/v1"
 )
 
-func RegisterCentral(pb grpc.ServiceRegistrar) {
-	pbShutdown.RegisterShutdownServer(pb, NewShutdownableShutdownCentralServer())
-}
-
-func Register(pb grpc.ServiceRegistrar, closer context.CancelFunc) {
-	pbShutdown.RegisterShutdownServer(pb, NewShutdownableShutdownServer(closer))
-}
-
-func NewShutdownableShutdownCentralServer() ShutdownableShutdownServer {
-	return NewShutdownableShutdownServer(stop)
-}
-
-func NewShutdownableShutdownServer(closer context.CancelFunc) ShutdownableShutdownServer {
-	return &impl{closer: closer}
+func NewShutdownableShutdownServer() ShutdownableShutdownServer {
+	return &impl{closer: stop}
 }
 
 type ShutdownableShutdownServer interface {

pkg/util/shutdown/grpc.go

@@ -18,36 +18,27 @@
 // Copyright holder is ArangoDB GmbH, Cologne, Germany
 //
 
-package shutdown
+package svc
 
 import (
 	"context"
 	"net"
 
 	"google.golang.org/grpc"
-
-	"github.com/arangodb/kube-arangodb/pkg/util/svc"
 )
 
-type ServiceConfig struct {
+type GRPCConfig struct {
 	ListenAddress string
 }
 
-func ServiceCentral(config ServiceConfig) svc.Service {
-	server := grpc.NewServer( /* currently no auth parameters required */ )
-
-	RegisterCentral(server)
-
-	return &service{
-		cfg:        config,
-		grpcServer: server,
-	}
-}
+type RegisterServerFunc func(server *grpc.Server)
 
-func Service(config ServiceConfig, closer context.CancelFunc) svc.Service {
+func NewGRPC(config GRPCConfig, registerFuncs ...RegisterServerFunc) Service {
 	server := grpc.NewServer( /* currently no auth parameters required */ )
 
-	Register(server, closer)
+	for _, fn := range registerFuncs {
+		fn(server)
+	}
 
 	return &service{
 		cfg:        config,
@@ -57,7 +48,7 @@ func Service(config ServiceConfig, closer context.CancelFunc) svc.Service {
 
 type service struct {
 	grpcServer *grpc.Server
-	cfg        ServiceConfig
+	cfg        GRPCConfig
 }
 
 func (s *service) Run(ctx context.Context) error {

pkg/util/shutdown/server.go renamed to pkg/util/svc/grpc.go

@@ -25,13 +25,14 @@ import (
 	"sync"
 
 	shared "github.com/arangodb/kube-arangodb/pkg/apis/shared"
+	"github.com/arangodb/kube-arangodb/pkg/util/probe"
 )
 
 type Service interface {
 	Run(ctx context.Context) error
 }
 
-func RunServices(ctx context.Context, services ...Service) error {
+func RunServices(ctx context.Context, healthService probe.HealthService, services ...Service) error {
 	if len(services) == 0 {
 		<-ctx.Done()
 		return nil
@@ -48,9 +49,12 @@ func RunServices(ctx context.Context, services ...Service) error {
 			defer wg.Done()
 
 			errors[id] = services[id].Run(ctx)
+
+			healthService.Shutdown()
 		}(id)
 	}
 
+	healthService.SetServing()
 	wg.Wait()
 
 	return shared.WithErrors(errors...)