Added readiness support in server health-check

common/container/container.go

@@ -26,6 +26,8 @@ import (
 
 const (
 	maxGrpcFrameSize = 256 * 1024 * 1024
+
+	ReadinessProbeService = "oxia-readiness"
 )
 
 type GrpcServer interface {

deploy/charts/oxia-cluster/templates/_helpers.tpl

@@ -58,3 +58,14 @@ exec:
 initialDelaySeconds: 10
 timeoutSeconds: 10
 {{- end }}
+
+
+{{/*
+Probe
+*/}}
+{{- define "oxia-cluster.readiness-probe" -}}
+exec:
+  command: ["oxia", "health", "--port={{ . }}", "--service=oxia-readiness"]
+initialDelaySeconds: 10
+timeoutSeconds: 10
+{{- end }}

deploy/charts/oxia-cluster/templates/server-statefulset.yaml

@@ -64,7 +64,7 @@ spec:
           livenessProbe:
             {{- include "oxia-cluster.probe" .Values.server.ports.internal | nindent 12 }}
           readinessProbe:
-            {{- include "oxia-cluster.probe" .Values.server.ports.internal | nindent 12 }}
+            {{- include "oxia-cluster.readiness-probe" .Values.server.ports.internal | nindent 12 }}
   volumeClaimTemplates:
     - metadata:
         name: data

server/assignment_dispatcher.go

@@ -19,12 +19,15 @@ import (
 	"github.com/rs/zerolog"
 	"github.com/rs/zerolog/log"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/health"
+	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/metadata"
 	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
 	pb "google.golang.org/protobuf/proto"
 	"io"
 	"oxia/common"
+	"oxia/common/container"
 	"oxia/common/metrics"
 	"oxia/proto"
 	"oxia/server/util"
@@ -50,6 +53,7 @@ type shardAssignmentDispatcher struct {
 	clients      map[int64]chan *proto.ShardAssignments
 	nextClientId int64
 	standalone   bool
+	healthServer *health.Server
 
 	ctx    context.Context
 	cancel context.CancelFunc
@@ -209,6 +213,10 @@ func (s *shardAssignmentDispatcher) PushShardAssignments(stream proto.OxiaCoordi
 }
 
 func (s *shardAssignmentDispatcher) updateShardAssignment(assignments *proto.ShardAssignments) error {
+	// Once we receive the first update of the shards mapping, this service can be
+	// considered "ready" and it will be able to respond to service discovery requests
+	s.healthServer.SetServingStatus(container.ReadinessProbeService, grpc_health_v1.HealthCheckResponse_SERVING)
+
 	s.Lock()
 	defer s.Unlock()
 
@@ -230,10 +238,11 @@ func (s *shardAssignmentDispatcher) updateShardAssignment(assignments *proto.Sha
 	return nil
 }
 
-func NewShardAssignmentDispatcher() ShardAssignmentsDispatcher {
+func NewShardAssignmentDispatcher(healthServer *health.Server) ShardAssignmentsDispatcher {
 	s := &shardAssignmentDispatcher{
-		assignments: nil,
-		clients:     make(map[int64]chan *proto.ShardAssignments),
+		assignments:  nil,
+		healthServer: healthServer,
+		clients:      make(map[int64]chan *proto.ShardAssignments),
 		log: log.With().
 			Str("component", "shard-assignment-dispatcher").
 			Logger(),
@@ -254,7 +263,7 @@ func NewShardAssignmentDispatcher() ShardAssignmentsDispatcher {
 }
 
 func NewStandaloneShardAssignmentDispatcher(numShards uint32) ShardAssignmentsDispatcher {
-	assignmentDispatcher := NewShardAssignmentDispatcher().(*shardAssignmentDispatcher)
+	assignmentDispatcher := NewShardAssignmentDispatcher(health.NewServer()).(*shardAssignmentDispatcher)
 	assignmentDispatcher.standalone = true
 	res := &proto.ShardAssignments{
 		Namespaces: map[string]*proto.NamespaceShardsAssignment{

server/assignment_dispatcher_test.go

@@ -15,17 +15,23 @@
 package server
 
 import (
+	"context"
 	"github.com/stretchr/testify/assert"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/health"
+	"google.golang.org/grpc/health/grpc_health_v1"
+	"google.golang.org/grpc/status"
 	"math"
 	"oxia/common"
+	"oxia/common/container"
 	"oxia/proto"
 	"sync"
 	"testing"
 	"time"
 )
 
 func TestUninitializedAssignmentDispatcher(t *testing.T) {
-	dispatcher := NewShardAssignmentDispatcher()
+	dispatcher := NewShardAssignmentDispatcher(health.NewServer())
 	mockClient := newMockShardAssignmentClientStream()
 	assert.False(t, dispatcher.Initialized())
 	req := &proto.ShardAssignmentsRequest{Namespace: common.DefaultNamespace}
@@ -35,7 +41,7 @@ func TestUninitializedAssignmentDispatcher(t *testing.T) {
 }
 
 func TestShardAssignmentDispatcher_Initialized(t *testing.T) {
-	dispatcher := NewShardAssignmentDispatcher()
+	dispatcher := NewShardAssignmentDispatcher(health.NewServer())
 	coordinatorStream := newMockShardAssignmentControllerStream()
 	go func() {
 		err := dispatcher.PushShardAssignments(coordinatorStream)
@@ -76,12 +82,56 @@ func TestShardAssignmentDispatcher_Initialized(t *testing.T) {
 
 }
 
+func TestShardAssignmentDispatcher_ReadinessProbe(t *testing.T) {
+	healthServer := health.NewServer()
+	dispatcher := NewShardAssignmentDispatcher(healthServer)
+	coordinatorStream := newMockShardAssignmentControllerStream()
+	go func() {
+		err := dispatcher.PushShardAssignments(coordinatorStream)
+		assert.NoError(t, err)
+	}()
+
+	assert.False(t, dispatcher.Initialized())
+
+	// Readiness probe should fail while not initialized
+	resp, err := healthServer.Check(context.Background(), &grpc_health_v1.HealthCheckRequest{
+		Service: container.ReadinessProbeService,
+	})
+
+	assert.Equal(t, codes.NotFound, status.Code(err))
+	assert.Nil(t, resp)
+
+	coordinatorStream.AddRequest(&proto.ShardAssignments{
+		Namespaces: map[string]*proto.NamespaceShardsAssignment{
+			common.DefaultNamespace: {
+				Assignments: []*proto.ShardAssignment{
+					newShardAssignment(0, "server1", 0, 100),
+					newShardAssignment(1, "server2", 100, math.MaxUint32),
+				},
+				ShardKeyRouter: proto.ShardKeyRouter_XXHASH3,
+			},
+		},
+	})
+	assert.Eventually(t, func() bool {
+		return dispatcher.Initialized()
+	}, 10*time.Second, 10*time.Millisecond)
+
+	resp, err = healthServer.Check(context.Background(), &grpc_health_v1.HealthCheckRequest{
+		Service: container.ReadinessProbeService,
+	})
+	assert.NoError(t, err)
+	assert.Equal(t, grpc_health_v1.HealthCheckResponse_SERVING, resp.Status)
+
+	assert.NoError(t, dispatcher.Close())
+
+}
+
 func TestShardAssignmentDispatcher_AddClient(t *testing.T) {
 	shard0InitialAssignment := newShardAssignment(0, "server1", 0, 100)
 	shard1InitialAssignment := newShardAssignment(1, "server2", 100, math.MaxUint32)
 	shard1UpdatedAssignment := newShardAssignment(1, "server3", 100, math.MaxUint32)
 
-	dispatcher := NewShardAssignmentDispatcher()
+	dispatcher := NewShardAssignmentDispatcher(health.NewServer())
 
 	coordinatorStream := newMockShardAssignmentControllerStream()
 	go func() {
@@ -163,7 +213,7 @@ func TestShardAssignmentDispatcher_AddClient(t *testing.T) {
 }
 
 func TestShardAssignmentDispatcher_MultipleNamespaces(t *testing.T) {
-	dispatcher := NewShardAssignmentDispatcher()
+	dispatcher := NewShardAssignmentDispatcher(health.NewServer())
 
 	coordinatorStream := newMockShardAssignmentControllerStream()
 	go func() {

server/internal_rpc_server.go

@@ -43,11 +43,12 @@ type internalRpcServer struct {
 	log                  zerolog.Logger
 }
 
-func newInternalRpcServer(grpcProvider container.GrpcProvider, bindAddress string, shardsDirector ShardsDirector, assignmentDispatcher ShardAssignmentsDispatcher) (*internalRpcServer, error) {
+func newInternalRpcServer(grpcProvider container.GrpcProvider, bindAddress string, shardsDirector ShardsDirector,
+	assignmentDispatcher ShardAssignmentsDispatcher, healthServer *health.Server) (*internalRpcServer, error) {
 	server := &internalRpcServer{
 		shardsDirector:       shardsDirector,
 		assignmentDispatcher: assignmentDispatcher,
-		healthServer:         health.NewServer(),
+		healthServer:         healthServer,
 		log: log.With().
 			Str("component", "internal-rpc-server").
 			Logger(),
@@ -67,7 +68,6 @@ func newInternalRpcServer(grpcProvider container.GrpcProvider, bindAddress strin
 }
 
 func (s *internalRpcServer) Close() error {
-	s.healthServer.Shutdown()
 	return s.grpcServer.Close()
 }
 
@@ -82,6 +82,7 @@ func (s *internalRpcServer) PushShardAssignments(srv proto.OxiaCoordination_Push
 			Str("peer", common.GetPeer(srv.Context())).
 			Msg("Failed to provide shards assignments updates")
 	}
+
 	return err
 }
 

server/internal_rpc_server_test.go

@@ -20,13 +20,16 @@ import (
 	"github.com/stretchr/testify/assert"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/health"
 	"google.golang.org/grpc/health/grpc_health_v1"
 	"oxia/common/container"
 	"testing"
 )
 
 func TestInternalHealthCheck(t *testing.T) {
-	server, err := newInternalRpcServer(container.Default, "localhost:0", nil, NewShardAssignmentDispatcher())
+	healthServer := health.NewServer()
+	server, err := newInternalRpcServer(container.Default, "localhost:0", nil,
+		NewShardAssignmentDispatcher(healthServer), healthServer)
 	assert.NoError(t, err)
 
 	target := fmt.Sprintf("localhost:%d", server.grpcServer.Port())

server/server.go

@@ -17,6 +17,7 @@ package server
 import (
 	"github.com/rs/zerolog/log"
 	"go.uber.org/multierr"
+	"google.golang.org/grpc/health"
 	"oxia/common/container"
 	"oxia/common/metrics"
 	"oxia/server/kv"
@@ -45,6 +46,8 @@ type Server struct {
 	metrics                   *metrics.PrometheusMetrics
 	walFactory                wal.WalFactory
 	kvFactory                 kv.KVFactory
+
+	healthServer *health.Server
 }
 
 func New(config Config) (*Server, error) {
@@ -69,14 +72,15 @@ func NewWithGrpcProvider(config Config, provider container.GrpcProvider, replica
 		walFactory: wal.NewWalFactory(&wal.WalFactoryOptions{
 			LogDir: config.WalDir,
 		}),
-		kvFactory: kvFactory,
+		kvFactory:    kvFactory,
+		healthServer: health.NewServer(),
 	}
 
 	s.shardsDirector = NewShardsDirector(config, s.walFactory, s.kvFactory, replicationRpcProvider)
-	s.shardAssignmentDispatcher = NewShardAssignmentDispatcher()
+	s.shardAssignmentDispatcher = NewShardAssignmentDispatcher(s.healthServer)
 
 	s.internalRpcServer, err = newInternalRpcServer(provider, config.InternalServiceAddr,
-		s.shardsDirector, s.shardAssignmentDispatcher)
+		s.shardsDirector, s.shardAssignmentDispatcher, s.healthServer)
 	if err != nil {
 		return nil, err
 	}
@@ -105,6 +109,8 @@ func (s *Server) InternalPort() int {
 }
 
 func (s *Server) Close() error {
+	s.healthServer.Shutdown()
+
 	err := multierr.Combine(
 		s.shardAssignmentDispatcher.Close(),
 		s.shardsDirector.Close(),