We now use the interface correctly to create and ping shards in one go

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

Author: jakobht

service/sharddistributor/canary/pinger/pingAndLog.go

@@ -0,0 +1,37 @@
+package pinger
+
+import (
+	"context"
+	"time"
+
+	sharddistributorv1 "github.com/uber/cadence/.gen/proto/sharddistributor/v1"
+	"github.com/uber/cadence/service/sharddistributor/client/spectatorclient"
+	"go.uber.org/yarpc"
+	"go.uber.org/zap"
+)
+
+const (
+	pingTimeout = 5 * time.Second
+)
+
+func PingShard(ctx context.Context, canaryClient sharddistributorv1.ShardDistributorExecutorCanaryAPIYARPCClient, logger *zap.Logger, namespace, shardKey string) {
+	request := &sharddistributorv1.PingRequest{
+		ShardKey:  shardKey,
+		Namespace: namespace,
+	}
+
+	ctx, cancel := context.WithTimeout(ctx, pingTimeout)
+	defer cancel()
+
+	response, err := canaryClient.Ping(ctx, request, yarpc.WithShardKey(shardKey), yarpc.WithHeader(spectatorclient.NamespaceHeader, namespace))
+	if err != nil {
+		logger.Error("Failed to ping shard", zap.String("namespace", namespace), zap.String("shard_key", shardKey), zap.Error(err))
+	}
+
+	// Verify response
+	if !response.GetOwnsShard() {
+		logger.Warn("Executor does not own shard", zap.String("namespace", namespace), zap.String("shard_key", shardKey), zap.String("executor_id", response.GetExecutorId()))
+	}
+
+	logger.Info("Successfully pinged shard owner", zap.String("namespace", namespace), zap.String("shard_key", shardKey), zap.String("executor_id", response.GetExecutorId()))
+}

service/sharddistributor/canary/pinger/pinger.go

@@ -8,21 +8,18 @@ import (
 	"time"
 
 	"go.uber.org/fx"
-	"go.uber.org/yarpc"
 	"go.uber.org/zap"
 
 	sharddistributorv1 "github.com/uber/cadence/.gen/proto/sharddistributor/v1"
 	"github.com/uber/cadence/common/backoff"
 	"github.com/uber/cadence/common/clock"
-	"github.com/uber/cadence/service/sharddistributor/client/spectatorclient"
 )
 
 //go:generate mockgen -package $GOPACKAGE -destination canary_client_mock.go github.com/uber/cadence/.gen/proto/sharddistributor/v1 ShardDistributorExecutorCanaryAPIYARPCClient
 
 const (
 	pingInterval    = 1 * time.Second
 	pingJitterCoeff = 0.1 // 10% jitter
-	pingTimeout     = 5 * time.Second
 )
 
 // Pinger periodically pings shard owners in the fixed namespace
@@ -96,23 +93,5 @@ func (p *Pinger) pingRandomShard() {
 	shardNum := rand.Intn(p.numShards)
 	shardKey := fmt.Sprintf("%d", shardNum)
 
-	request := &sharddistributorv1.PingRequest{
-		ShardKey:  shardKey,
-		Namespace: p.namespace,
-	}
-
-	ctx, cancel := context.WithTimeout(p.ctx, pingTimeout)
-	defer cancel()
-
-	response, err := p.canaryClient.Ping(ctx, request, yarpc.WithShardKey(shardKey), yarpc.WithHeader(spectatorclient.NamespaceHeader, p.namespace))
-	if err != nil {
-		p.logger.Error("Failed to ping shard", zap.String("namespace", p.namespace), zap.String("shard_key", shardKey), zap.Error(err))
-	}
-
-	// Verify response
-	if !response.GetOwnsShard() {
-		p.logger.Warn("Executor does not own shard", zap.String("namespace", p.namespace), zap.String("shard_key", shardKey), zap.String("executor_id", response.GetExecutorId()))
-	}
-
-	p.logger.Info("Successfully pinged shard owner", zap.String("namespace", p.namespace), zap.String("shard_key", shardKey), zap.String("executor_id", response.GetExecutorId()))
+	PingShard(p.ctx, p.canaryClient, p.logger, p.namespace, shardKey)
 }

service/sharddistributor/canary/processorephemeral/shardcreator.go

@@ -2,18 +2,16 @@ package processorephemeral
 
 import (
 	"context"
-	"fmt"
 	"sync"
 	"time"
 
 	"github.com/google/uuid"
 	"go.uber.org/fx"
-	"go.uber.org/yarpc"
 	"go.uber.org/zap"
 
 	sharddistributorv1 "github.com/uber/cadence/.gen/proto/sharddistributor/v1"
 	"github.com/uber/cadence/common/clock"
-	"github.com/uber/cadence/service/sharddistributor/client/spectatorclient"
+	"github.com/uber/cadence/service/sharddistributor/canary/pinger"
 )
 
 //go:generate mockgen -package $GOPACKAGE -destination canary_client_mock_test.go github.com/uber/cadence/.gen/proto/sharddistributor/v1 ShardDistributorExecutorCanaryAPIYARPCClient
@@ -26,7 +24,6 @@ const (
 type ShardCreator struct {
 	logger       *zap.Logger
 	timeSource   clock.TimeSource
-	spectators   *spectatorclient.Spectators
 	canaryClient sharddistributorv1.ShardDistributorExecutorCanaryAPIYARPCClient
 	namespaces   []string
 
@@ -40,7 +37,6 @@ type ShardCreatorParams struct {
 
 	Logger       *zap.Logger
 	TimeSource   clock.TimeSource
-	Spectators   *spectatorclient.Spectators
 	CanaryClient sharddistributorv1.ShardDistributorExecutorCanaryAPIYARPCClient
 }
 
@@ -49,7 +45,6 @@ func NewShardCreator(params ShardCreatorParams, namespaces []string) *ShardCreat
 	return &ShardCreator{
 		logger:       params.Logger,
 		timeSource:   params.TimeSource,
-		spectators:   params.Spectators,
 		canaryClient: params.CanaryClient,
 		stopChan:     make(chan struct{}),
 		goRoutineWg:  sync.WaitGroup{},
@@ -100,83 +95,8 @@ func (s *ShardCreator) process(ctx context.Context) {
 				shardKey := uuid.New().String()
 				s.logger.Info("Creating shard", zap.String("shardKey", shardKey), zap.String("namespace", namespace))
 
-				// Get spectator for this namespace
-				spectator, err := s.spectators.ForNamespace(namespace)
-				if err != nil {
-					s.logger.Warn("No spectator for namespace, skipping shard creation",
-						zap.String("namespace", namespace),
-						zap.String("shardKey", shardKey),
-						zap.Error(err))
-					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("Failed to get/create shard owner",
-						zap.Error(err),
-						zap.String("namespace", namespace),
-						zap.String("shardKey", shardKey))
-					continue
-				}
-
-				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))
-				}
+				pinger.PingShard(ctx, s.canaryClient, s.logger, namespace, 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_test.go

@@ -11,40 +11,18 @@ import (
 
 	sharddistributorv1 "github.com/uber/cadence/.gen/proto/sharddistributor/v1"
 	"github.com/uber/cadence/common/clock"
-	"github.com/uber/cadence/service/sharddistributor/client/spectatorclient"
 )
 
-func TestShardCreator_Lifecycle(t *testing.T) {
+func TestShardCreator_PingsShards(t *testing.T) {
 	goleak.VerifyNone(t)
 
 	logger := zaptest.NewLogger(t)
 	timeSource := clock.NewMockedTimeSource()
 	ctrl := gomock.NewController(t)
 
 	namespace := "test-namespace"
-
-	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{}, shardKey string) (*spectatorclient.ShardOwner, error) {
-			assert.NotEmpty(t, shardKey)
-			return nil, assert.AnError
-		})
-
-	// 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
-		}).
-		After(firstCall)
-
 	// Ping happens after successful GetShardOwner
 	mockCanaryClient.EXPECT().
 		Ping(gomock.Any(), gomock.Any(), gomock.Any()).
@@ -57,27 +35,17 @@ func TestShardCreator_Lifecycle(t *testing.T) {
 			}, nil
 		})
 
-	spectators := map[string]spectatorclient.Spectator{
-		namespace: mockSpectator,
-	}
-
 	params := ShardCreatorParams{
 		Logger:       logger,
 		TimeSource:   timeSource,
-		Spectators:   spectators,
 		CanaryClient: mockCanaryClient,
 	}
 
 	creator := NewShardCreator(params, []string{namespace})
 	creator.Start()
 
+	// Wait for the goroutine to start and do it's ping
 	timeSource.BlockUntil(1)
-
-	// First cycle - GetShardOwner fails, no ping
-	timeSource.Advance(shardCreationInterval + 100*time.Millisecond)
-	time.Sleep(10 * time.Millisecond)
-
-	// Second cycle - GetShardOwner succeeds, ping happens
 	timeSource.Advance(shardCreationInterval + 100*time.Millisecond)
 	time.Sleep(10 * time.Millisecond)