Add ping verification to ephemeral shard creator

After creating a shard, the ephemeral shard creator now pings the owner
to verify that:
1. The shard was created successfully
2. The owner can be reached via gRPC
3. The owner actually owns the shard

This provides end-to-end validation of the shard creation and routing
mechanisms in the canary test environment.

Changes:
- Switch from using ShardDistributor client to using Spectators
- Add pingShardOwner() method that sends canary ping after creation
- Verify executor ID and ownership in ping response
- Log warnings if executor ID mismatches or ownership is incorrect
- Add test coverage for successful ping flow (78.5% total coverage)

Signed-off-by: Jakob Haahr Taankvist <[email protected]>

Author: jakobht

service/sharddistributor/canary/processorephemeral/canary_client_mock_test.go

@@ -2,51 +2,58 @@ package processorephemeral
 
 import (
 	"context"
+	"fmt"
 	"sync"
 	"time"
 
 	"github.com/google/uuid"
 	"go.uber.org/fx"
+	"go.uber.org/yarpc"
 	"go.uber.org/zap"
 
-	"github.com/uber/cadence/client/sharddistributor"
+	sharddistributorv1 "github.com/uber/cadence/.gen/proto/sharddistributor/v1"
 	"github.com/uber/cadence/common/clock"
-	"github.com/uber/cadence/common/types"
+	"github.com/uber/cadence/service/sharddistributor/client/spectatorclient"
 )
 
+//go:generate mockgen -package $GOPACKAGE -destination canary_client_mock_test.go github.com/uber/cadence/.gen/proto/sharddistributor/v1 ShardDistributorExecutorCanaryAPIYARPCClient
+
 const (
 	shardCreationInterval = 1 * time.Second
 )
 
 // ShardCreator creates shards at regular intervals for ephemeral canary testing
 type ShardCreator struct {
-	logger           *zap.Logger
-	timeSource       clock.TimeSource
-	shardDistributor sharddistributor.Client
+	logger       *zap.Logger
+	timeSource   clock.TimeSource
+	spectators   map[string]spectatorclient.Spectator // namespace -> spectator
+	canaryClient sharddistributorv1.ShardDistributorExecutorCanaryAPIYARPCClient
+	namespaces   []string
 
 	stopChan    chan struct{}
 	goRoutineWg sync.WaitGroup
-	namespaces  []string
 }
 
 // ShardCreatorParams contains the dependencies needed to create a ShardCreator
 type ShardCreatorParams struct {
 	fx.In
 
-	Logger           *zap.Logger
-	TimeSource       clock.TimeSource
-	ShardDistributor sharddistributor.Client
+	Logger       *zap.Logger
+	TimeSource   clock.TimeSource
+	Spectators   map[string]spectatorclient.Spectator
+	CanaryClient sharddistributorv1.ShardDistributorExecutorCanaryAPIYARPCClient
 }
 
 // NewShardCreator creates a new ShardCreator instance with the given parameters and namespace
 func NewShardCreator(params ShardCreatorParams, namespaces []string) *ShardCreator {
 	return &ShardCreator{
-		logger:           params.Logger,
-		timeSource:       params.TimeSource,
-		shardDistributor: params.ShardDistributor,
-		stopChan:         make(chan struct{}),
-		goRoutineWg:      sync.WaitGroup{},
-		namespaces:       namespaces,
+		logger:       params.Logger,
+		timeSource:   params.TimeSource,
+		spectators:   params.Spectators,
+		canaryClient: params.CanaryClient,
+		stopChan:     make(chan struct{}),
+		goRoutineWg:  sync.WaitGroup{},
+		namespaces:   namespaces,
 	}
 }
 
@@ -92,16 +99,83 @@ func (s *ShardCreator) process(ctx context.Context) {
 			for _, namespace := range s.namespaces {
 				shardKey := uuid.New().String()
 				s.logger.Info("Creating shard", zap.String("shardKey", shardKey), zap.String("namespace", namespace))
-				response, err := s.shardDistributor.GetShardOwner(ctx, &types.GetShardOwnerRequest{
-					ShardKey:  shardKey,
-					Namespace: namespace,
-				})
+
+				// Get spectator for this namespace
+				spectator, ok := s.spectators[namespace]
+				if !ok {
+					s.logger.Warn("No spectator for namespace, skipping shard creation",
+						zap.String("namespace", namespace),
+						zap.String("shardKey", shardKey))
+					continue
+				}
+
+				// Create shard and get owner via spectator (spectator will create if not exists)
+				owner, err := spectator.GetShardOwner(ctx, shardKey)
 				if err != nil {
-					s.logger.Error("create shard failed", zap.Error(err))
+					s.logger.Error("Failed to get/create shard owner",
+						zap.Error(err),
+						zap.String("namespace", namespace),
+						zap.String("shardKey", shardKey))
 					continue
 				}
-				s.logger.Info("shard created", zap.String("shardKey", shardKey), zap.String("shardOwner", response.Owner), zap.String("namespace", response.Namespace))
+
+				s.logger.Info("Shard created, got owner from spectator",
+					zap.String("shardKey", shardKey),
+					zap.String("namespace", namespace),
+					zap.String("executor_id", owner.ExecutorID))
+
+				// Now ping the owner to verify
+				if err := s.pingShardOwner(ctx, owner, namespace, shardKey); err != nil {
+					s.logger.Error("Failed to ping shard owner",
+						zap.Error(err),
+						zap.String("namespace", namespace),
+						zap.String("shardKey", shardKey))
+				}
 			}
 		}
 	}
 }
+
+func (s *ShardCreator) pingShardOwner(ctx context.Context, owner *spectatorclient.ShardOwner, namespace, shardKey string) error {
+	s.logger.Debug("Pinging shard owner after creation",
+		zap.String("namespace", namespace),
+		zap.String("shardKey", shardKey),
+		zap.String("expected_executor_id", owner.ExecutorID))
+
+	// Create ping request
+	request := &sharddistributorv1.PingRequest{
+		ShardKey:  shardKey,
+		Namespace: namespace,
+	}
+
+	// SIMPLE CANARY CODE: Just pass the shard key and namespace!
+	// The SpectatorPeerChooser library code handles routing to the right executor
+	response, err := s.canaryClient.Ping(ctx, request, yarpc.WithShardKey(shardKey), yarpc.WithHeader(spectatorclient.NamespaceHeader, namespace))
+	if err != nil {
+		return fmt.Errorf("ping rpc failed: %w", err)
+	}
+
+	// Verify response matches the owner we got from spectator
+	if response.GetExecutorId() != owner.ExecutorID {
+		s.logger.Warn("Executor ID mismatch",
+			zap.String("namespace", namespace),
+			zap.String("shardKey", shardKey),
+			zap.String("expected", owner.ExecutorID),
+			zap.String("actual", response.GetExecutorId()))
+	}
+
+	if !response.GetOwnsShard() {
+		s.logger.Warn("Executor does not own shard",
+			zap.String("namespace", namespace),
+			zap.String("shardKey", shardKey),
+			zap.String("executor_id", response.GetExecutorId()))
+		return fmt.Errorf("executor %s does not own shard %s", response.GetExecutorId(), shardKey)
+	}
+
+	s.logger.Info("Successfully verified shard owner after creation",
+		zap.String("namespace", namespace),
+		zap.String("shardKey", shardKey),
+		zap.String("executor_id", response.GetExecutorId()))
+
+	return nil
+}

service/sharddistributor/canary/processorephemeral/shardcreator.go

@@ -9,9 +9,9 @@ import (
 	"go.uber.org/mock/gomock"
 	"go.uber.org/zap/zaptest"
 
-	"github.com/uber/cadence/client/sharddistributor"
+	sharddistributorv1 "github.com/uber/cadence/.gen/proto/sharddistributor/v1"
 	"github.com/uber/cadence/common/clock"
-	"github.com/uber/cadence/common/types"
+	"github.com/uber/cadence/service/sharddistributor/client/spectatorclient"
 )
 
 func TestShardCreator_Lifecycle(t *testing.T) {
@@ -21,40 +21,65 @@ func TestShardCreator_Lifecycle(t *testing.T) {
 	timeSource := clock.NewMockedTimeSource()
 	ctrl := gomock.NewController(t)
 
-	mockShardDistributor := sharddistributor.NewMockClient(ctrl)
 	namespace := "test-namespace"
 
-	// Set up expectation for GetShardOwner calls that return errors
-	mockShardDistributor.EXPECT().
+	mockSpectator := spectatorclient.NewMockSpectator(ctrl)
+	mockCanaryClient := NewMockShardDistributorExecutorCanaryAPIYARPCClient(ctrl)
+
+	// First call fails - no ping should happen
+	firstCall := mockSpectator.EXPECT().
 		GetShardOwner(gomock.Any(), gomock.Any()).
-		DoAndReturn(func(ctx interface{}, req *types.GetShardOwnerRequest, opts ...interface{}) (*types.GetShardOwnerResponse, error) {
-			// Verify the request contains the correct namespace even on error
-			assert.Equal(t, namespace, req.Namespace)
-			assert.NotEmpty(t, req.ShardKey)
+		DoAndReturn(func(ctx interface{}, shardKey string) (*spectatorclient.ShardOwner, error) {
+			assert.NotEmpty(t, shardKey)
+			return nil, assert.AnError
+		})
 
-			return nil, assert.AnError // Using testify's AnError for consistency
+	// Second call succeeds - ping should happen
+	mockSpectator.EXPECT().
+		GetShardOwner(gomock.Any(), gomock.Any()).
+		DoAndReturn(func(ctx interface{}, shardKey string) (*spectatorclient.ShardOwner, error) {
+			assert.NotEmpty(t, shardKey)
+			return &spectatorclient.ShardOwner{
+				ExecutorID: "executor-1",
+			}, nil
 		}).
-		Times(2)
+		After(firstCall)
+
+	// Ping happens after successful GetShardOwner
+	mockCanaryClient.EXPECT().
+		Ping(gomock.Any(), gomock.Any(), gomock.Any()).
+		DoAndReturn(func(ctx interface{}, req *sharddistributorv1.PingRequest, opts ...interface{}) (*sharddistributorv1.PingResponse, error) {
+			assert.NotEmpty(t, req.ShardKey)
+			assert.Equal(t, namespace, req.Namespace)
+			return &sharddistributorv1.PingResponse{
+				OwnsShard:  true,
+				ExecutorId: "executor-1",
+			}, nil
+		})
+
+	spectators := map[string]spectatorclient.Spectator{
+		namespace: mockSpectator,
+	}
 
 	params := ShardCreatorParams{
-		Logger:           logger,
-		TimeSource:       timeSource,
-		ShardDistributor: mockShardDistributor,
+		Logger:       logger,
+		TimeSource:   timeSource,
+		Spectators:   spectators,
+		CanaryClient: mockCanaryClient,
 	}
 
 	creator := NewShardCreator(params, []string{namespace})
 	creator.Start()
 
-	// Wait for the goroutine to start
 	timeSource.BlockUntil(1)
 
-	// Trigger shard creation that will fail
+	// First cycle - GetShardOwner fails, no ping
 	timeSource.Advance(shardCreationInterval + 100*time.Millisecond)
-	time.Sleep(10 * time.Millisecond) // Allow processing
+	time.Sleep(10 * time.Millisecond)
 
-	// Trigger another shard creation to ensure processing continues after error
+	// Second cycle - GetShardOwner succeeds, ping happens
 	timeSource.Advance(shardCreationInterval + 100*time.Millisecond)
-	time.Sleep(10 * time.Millisecond) // Allow processing
+	time.Sleep(10 * time.Millisecond)
 
 	creator.Stop()
 }