Frontend health check waits until membership initialized (#3936)

Author: dnr

common/membership/interfaces.go

@@ -27,6 +27,7 @@
 package membership
 
 import (
+	"context"
 	"errors"
 
 	"go.temporal.io/api/serviceerror"
@@ -82,6 +83,11 @@ type (
 		// GetMemberCount returns the number of reachable members
 		// currently in this node's membership list for the given service
 		GetMemberCount(service primitives.ServiceName) (int, error)
+		// WaitUntilInitialized blocks until initialization is completed and returns the result
+		// of initialization. The current implementation does log.Fatal if it can't initialize,
+		// so currently this will never return non-nil, except for context cancel/timeout. A
+		// future implementation might return more errors.
+		WaitUntilInitialized(context.Context) error
 	}
 
 	// ServiceResolver provides membership information for a specific temporal service.

common/membership/interfaces_mock.go

@@ -36,6 +36,7 @@ import (
 	"time"
 
 	"go.temporal.io/server/common/convert"
+	"go.temporal.io/server/common/future"
 	"go.temporal.io/server/common/headers"
 	"go.temporal.io/server/common/primitives"
 
@@ -69,6 +70,7 @@ type ringpopMonitor struct {
 	metadataManager           persistence.ClusterMetadataManager
 	broadcastHostPortResolver func() (string, error)
 	hostID                    uuid.UUID
+	initialized               *future.FutureImpl[struct{}]
 }
 
 var _ Monitor = (*ringpopMonitor)(nil)
@@ -103,6 +105,7 @@ func NewRingpopMonitor(
 		metadataManager:           metadataManager,
 		broadcastHostPortResolver: broadcastHostPortResolver,
 		hostID:                    uuid.NewUUID(),
+		initialized:               future.NewFuture[struct{}](),
 	}
 	for service, port := range services {
 		rpo.rings[service] = newRingpopServiceResolver(service, port, rp, logger)
@@ -152,6 +155,13 @@ func (rpo *ringpopMonitor) Start() {
 	for _, ring := range rpo.rings {
 		ring.Start()
 	}
+
+	rpo.initialized.Set(struct{}{}, nil)
+}
+
+func (rpo *ringpopMonitor) WaitUntilInitialized(ctx context.Context) error {
+	_, err := rpo.initialized.Get(ctx)
+	return err
 }
 
 func ServiceNameToServiceTypeEnum(name primitives.ServiceName) (persistence.ServiceType, error) {

common/membership/rpMonitor.go

@@ -166,6 +166,7 @@ func NewTest(
 	membershipMonitor.EXPECT().GetResolver(primitives.MatchingService).Return(matchingServiceResolver, nil).AnyTimes()
 	membershipMonitor.EXPECT().GetResolver(primitives.HistoryService).Return(historyServiceResolver, nil).AnyTimes()
 	membershipMonitor.EXPECT().GetResolver(primitives.WorkerService).Return(workerServiceResolver, nil).AnyTimes()
+	membershipMonitor.EXPECT().WaitUntilInitialized(gomock.Any()).Return(nil).AnyTimes()
 
 	scope := tally.NewTestScope("test", nil)
 	serviceName, _ := metrics.MetricsServiceIdxToServiceName(serviceMetricsIndex)

common/resourcetest/resourceTest.go

@@ -555,6 +555,7 @@ func HandlerProvider(
 	clusterMetadata cluster.Metadata,
 	archivalMetadata archiver.ArchivalMetadata,
 	healthServer *health.Server,
+	membershipMonitor membership.Monitor,
 ) Handler {
 	wfHandler := NewWorkflowHandler(
 		serviceConfig,
@@ -576,6 +577,7 @@ func HandlerProvider(
 		archivalMetadata,
 		healthServer,
 		timeSource,
+		membershipMonitor,
 	)
 	return wfHandler
 }

service/frontend/fx.go

@@ -67,6 +67,7 @@ import (
 	"go.temporal.io/server/common/headers"
 	"go.temporal.io/server/common/log"
 	"go.temporal.io/server/common/log/tag"
+	"go.temporal.io/server/common/membership"
 	"go.temporal.io/server/common/metrics"
 	"go.temporal.io/server/common/namespace"
 	"go.temporal.io/server/common/payload"
@@ -123,6 +124,7 @@ type (
 		archivalMetadata                archiver.ArchivalMetadata
 		healthServer                    *health.Server
 		overrides                       *Overrides
+		membershipMonitor               membership.Monitor
 	}
 )
 
@@ -147,6 +149,7 @@ func NewWorkflowHandler(
 	archivalMetadata archiver.ArchivalMetadata,
 	healthServer *health.Server,
 	timeSource clock.TimeSource,
+	membershipMonitor membership.Monitor,
 ) *WorkflowHandler {
 
 	handler := &WorkflowHandler{
@@ -187,9 +190,10 @@ func NewWorkflowHandler(
 			visibilityMrg.GetIndexName(),
 			visibility.AllowListForValidation(visibilityMrg.GetName()),
 		),
-		archivalMetadata: archivalMetadata,
-		healthServer:     healthServer,
-		overrides:        NewOverrides(),
+		archivalMetadata:  archivalMetadata,
+		healthServer:      healthServer,
+		overrides:         NewOverrides(),
+		membershipMonitor: membershipMonitor,
 	}
 
 	return handler
@@ -202,7 +206,13 @@ func (wh *WorkflowHandler) Start() {
 		common.DaemonStatusInitialized,
 		common.DaemonStatusStarted,
 	) {
-		wh.healthServer.SetServingStatus(WorkflowServiceName, healthpb.HealthCheckResponse_SERVING)
+		// Start in NOT_SERVING state and switch to SERVING after membership is ready
+		wh.healthServer.SetServingStatus(WorkflowServiceName, healthpb.HealthCheckResponse_NOT_SERVING)
+		go func() {
+			_ = wh.membershipMonitor.WaitUntilInitialized(context.Background())
+			wh.healthServer.SetServingStatus(WorkflowServiceName, healthpb.HealthCheckResponse_SERVING)
+			wh.logger.Info("Frontend is now healthy")
+		}()
 	}
 }
 

service/frontend/workflow_handler.go

@@ -190,6 +190,7 @@ func (s *workflowHandlerSuite) getWorkflowHandler(config *Config) *WorkflowHandl
 		s.mockResource.GetArchivalMetadata(),
 		health.NewServer(),
 		clock.NewRealTimeSource(),
+		s.mockResource.GetMembershipMonitor(),
 	)
 }
 

service/frontend/workflow_handler_test.go

@@ -25,6 +25,7 @@
 package tests
 
 import (
+	"context"
 	"fmt"
 
 	"go.temporal.io/server/common/membership"
@@ -92,3 +93,7 @@ func (s *simpleMonitor) GetReachableMembers() ([]string, error) {
 func (s *simpleMonitor) GetMemberCount(service primitives.ServiceName) (int, error) {
 	return 0, nil
 }
+
+func (s *simpleMonitor) WaitUntilInitialized(_ context.Context) error {
+	return nil
+}