feat(executor-client): split function for local shard assignment (#7446)

<!-- Describe what has changed in this PR -->
**What changed?**
Instead of having a function that replace the assignment generated with
the local logic we have two different a functions, one to assign a new
shard and the other to remove the assignment. Also these function are
used in the canary to simulate the shard assignment.
The changes to the canary are made to make more similar to the matching
service.


<!-- Tell your future self why have you made these changes -->
**Why?**
we need to onboard matching and to model the local logic with assignment
replacement is adding too much complexity.

<!-- How have you verified this change? Tested locally? Added a unit
test? Checked in staging env? -->
**How did you test it?**
local testing and unit testing

<!-- Assuming the worst case, what can be broken when deploying this
change to production? -->
**Potential risks**
the change is only used in canary service which has been tested

<!-- Is it notable for release? e.g. schema updates, configuration or
data migration required? If so, please mention it, and also update
CHANGELOG.md -->
**Release notes**

<!-- Is there any documentation updates should be made for config,
https://cadenceworkflow.io/docs/operation-guide/setup/ ? If so, please
open an PR in https://github.com/cadence-workflow/cadence-docs -->
**Documentation Changes**

Signed-off-by: edigregorio <[email protected]>

Author: eleonoradgr

service/sharddistributor/canary/executors/executors.go

@@ -85,12 +85,6 @@ func NewExecutorsModule(params ExecutorsParams) {
 func Module(fixedNamespace, ephemeralNamespace, externalAssignmentNamespace string) fx.Option {
 	return fx.Module(
 		"Executors",
-		// Executors that are used for testing namespaces with the different modes of the migration
-		fx.Provide(
-			NewExecutorLocalPassthroughNamespace,
-			NewExecutorLocalPassthroughShadowNamespace,
-			NewExecutorDistributedPassthroughNamespace,
-		),
 		// Executor that is used for testing a namespace with fixed shards
 		fx.Provide(
 			func(params executorclient.Params[*processor.ShardProcessor]) (ExecutorResult, error) {

service/sharddistributor/canary/externalshardassignment/shardassigner.go

@@ -2,6 +2,7 @@ package externalshardassignment
 
 import (
 	"context"
+	"math/rand"
 	"strconv"
 	"sync"
 	"time"
@@ -19,17 +20,19 @@ import (
 )
 
 const (
-	shardAssignmentInterval = 1 * time.Second
+	shardAssignmentInterval = 4 * time.Second
+	minimumAssignedShards   = 3
 )
 
 // ShardAssigneer assigns shards to the executor for canary testing
 type ShardAssigner struct {
-	logger         log.Logger
-	timeSource     clock.TimeSource
-	executorclient executorclient.Executor[*processorephemeral.ShardProcessor]
-	stopChan       chan struct{}
-	goRoutineWg    sync.WaitGroup
-	namespace      string
+	logger          log.Logger
+	timeSource      clock.TimeSource
+	shardProcessors map[string]*processorephemeral.ShardProcessor
+	executorclient  executorclient.Executor[*processorephemeral.ShardProcessor]
+	stopChan        chan struct{}
+	goRoutineWg     sync.WaitGroup
+	namespace       string
 }
 
 // ShardAssignerParams contains the dependencies needed to create a ShardParam
@@ -42,13 +45,15 @@ type ShardAssignerParams struct {
 
 // NewShardCreator creates a new ShardCreator instance with the given parameters and namespace
 func NewShardAssigner(params ShardAssignerParams, namespace string) *ShardAssigner {
+	sp := make(map[string]*processorephemeral.ShardProcessor)
 	return &ShardAssigner{
-		logger:         params.Logger,
-		timeSource:     params.TimeSource,
-		executorclient: params.Executorclient,
-		stopChan:       make(chan struct{}),
-		goRoutineWg:    sync.WaitGroup{},
-		namespace:      namespace,
+		logger:          params.Logger,
+		timeSource:      params.TimeSource,
+		shardProcessors: sp,
+		executorclient:  params.Executorclient,
+		stopChan:        make(chan struct{}),
+		goRoutineWg:     sync.WaitGroup{},
+		namespace:       namespace,
 	}
 }
 
@@ -92,21 +97,46 @@ func (s *ShardAssigner) process(ctx context.Context) {
 		case <-s.stopChan:
 			return
 		case <-ticker.Chan():
+			if len(s.shardProcessors) > minimumAssignedShards {
+				var shardToRemove string
+				shardToRemoveIndex := rand.Intn(len(s.shardProcessors))
+				for shardID := range s.shardProcessors {
+					if shardToRemoveIndex == 0 {
+						shardToRemove = shardID
+						break
+					}
+					shardToRemoveIndex--
+				}
+				err := s.executorclient.RemoveShardsFromLocalLogic([]string{shardToRemove})
+				if err != nil {
+					s.logger.Error("Failed to remove shards", tag.Error(err))
+					continue
+				}
+				delete(s.shardProcessors, shardToRemove)
+				s.logger.Info("Removed a shard from external source", tag.ShardKey(shardToRemove))
+
+			}
+
+			// Simulate the assignment of new shards
 			newAssignedShard := uuid.New().String()
 			s.logger.Info("Assign a new shard from external source", tag.ShardKey(newAssignedShard))
 			shardAssignment := map[string]*types.ShardAssignment{
 				newAssignedShard: {
 					Status: types.AssignmentStatusREADY,
 				},
 			}
-			s.executorclient.AssignShardsFromLocalLogic(context.Background(), shardAssignment)
+			err := s.executorclient.AssignShardsFromLocalLogic(context.Background(), shardAssignment)
+			if err != nil {
+				s.logger.Error("Failed to assign shard from external source", tag.Error(err))
+				continue
+			}
 			sp, err := s.executorclient.GetShardProcess(ctx, newAssignedShard)
 			if err != nil {
 				s.logger.Error("failed to get shard assigned", tag.ShardKey(newAssignedShard), tag.Error(err))
 			} else {
 				s.logger.Info("shard assigned", tag.ShardStatus(string(sp.GetShardReport().Status)), tag.ShardLoad(strconv.FormatFloat(sp.GetShardReport().ShardLoad, 'f', -1, 64)))
 			}
-
+			s.shardProcessors[newAssignedShard] = sp
 		}
 	}
 }

service/sharddistributor/canary/module.go

@@ -48,6 +48,6 @@ func opts(names NamespacesNames) fx.Option {
 		// Instantiate executors for multiple namespaces
 		executors.Module(names.FixedNamespace, names.EphemeralNamespace, names.ExternalAssignmentNamespace),
 
-		processorephemeral.ShardCreatorModule([]string{names.EphemeralNamespace, names.ExternalAssignmentNamespace, "test-local-passthrough-shadow"}),
+		processorephemeral.ShardCreatorModule([]string{names.EphemeralNamespace}),
 	)
 }

service/sharddistributor/client/executorclient/client.go

@@ -50,8 +50,10 @@ type Executor[SP ShardProcessor] interface {
 	// Get the current metadata of the executor
 	GetMetadata() map[string]string
 
-	// Used during the migration during local-passthrough and local-passthrough-shadow
-	AssignShardsFromLocalLogic(ctx context.Context, shardAssignment map[string]*types.ShardAssignment)
+	// AssignShardsFromLocalLogic is used for the migration during local-passthrough, local-passthrough-shadow, distributed-passthrough
+	AssignShardsFromLocalLogic(ctx context.Context, shardAssignment map[string]*types.ShardAssignment) error
+	// RemoveShardsFromLocalLogic is used for the migration during local-passthrough, local-passthrough-shadow, distributed-passthrough
+	RemoveShardsFromLocalLogic(shardIDs []string) error
 }
 
 type Params[SP ShardProcessor] struct {

service/sharddistributor/client/executorclient/clientimpl.go

@@ -156,12 +156,26 @@ func (e *executorImpl[SP]) GetShardProcess(ctx context.Context, shardID string)
 	return shardProcess.processor, nil
 }
 
-func (e *executorImpl[SP]) AssignShardsFromLocalLogic(ctx context.Context, shardAssignment map[string]*types.ShardAssignment) {
+func (e *executorImpl[SP]) AssignShardsFromLocalLogic(ctx context.Context, shardAssignment map[string]*types.ShardAssignment) error {
 	e.assignmentMutex.Lock()
 	defer e.assignmentMutex.Unlock()
-
+	if e.getMigrationMode() == types.MigrationModeONBOARDED {
+		return fmt.Errorf("migration mode is onborded, no local assignemnt allowed")
+	}
 	e.logger.Info("Executing external shard assignment")
-	e.updateShardAssignment(ctx, shardAssignment)
+	e.addNewShards(ctx, shardAssignment)
+	return nil
+}
+
+func (e *executorImpl[SP]) RemoveShardsFromLocalLogic(shardIDs []string) error {
+	e.assignmentMutex.Lock()
+	defer e.assignmentMutex.Unlock()
+	if e.getMigrationMode() == types.MigrationModeONBOARDED {
+		return fmt.Errorf("migration mode is onborded, no local assignemnt allowed")
+	}
+	e.logger.Info("Executing external shard deletion assignment")
+	e.deleteShards(shardIDs)
+	return nil
 }
 
 func (e *executorImpl[SP]) heartbeatloop(ctx context.Context) {
@@ -291,7 +305,9 @@ func (e *executorImpl[SP]) heartbeat(ctx context.Context) (shardAssignments map[
 	if previousMode != currentMode {
 		e.logger.Info("migration mode transition",
 			tag.Dynamic("previous", previousMode),
-			tag.Dynamic("current", currentMode))
+			tag.Dynamic("current", currentMode),
+			tag.ShardNamespace(e.namespace),
+			tag.ShardExecutor(e.executorID))
 		e.setMigrationMode(currentMode)
 	}
 
@@ -304,71 +320,102 @@ func (e *executorImpl[SP]) updateShardAssignment(ctx context.Context, shardAssig
 	// Stop shard processing for shards not assigned to this executor
 	e.managedProcessors.Range(func(shardID string, managedProcessor *managedProcessor[SP]) bool {
 		if assignment, ok := shardAssignments[shardID]; !ok || assignment.Status != types.AssignmentStatusREADY {
-			e.metrics.Counter(metricsconstants.ShardDistributorExecutorShardsStopped).Inc(1)
-
 			wg.Add(1)
-			go func() {
+			go func(shardID string) {
 				defer wg.Done()
-				managedProcessor.setState(processorStateStopping)
-				managedProcessor.processor.Stop()
-				e.managedProcessors.Delete(shardID)
-			}()
+				e.stopManagerProcessor(shardID)
+			}(shardID)
 		}
 		return true
 	})
 
 	// Start shard processing for shards assigned to this executor
 	for shardID, assignment := range shardAssignments {
 		if assignment.Status == types.AssignmentStatusREADY {
-			if _, ok := e.managedProcessors.Load(shardID); !ok {
-				e.metrics.Counter(metricsconstants.ShardDistributorExecutorShardsStarted).Inc(1)
-
-				wg.Add(1)
-				go func() {
-					defer wg.Done()
-					processor, err := e.shardProcessorFactory.NewShardProcessor(shardID)
-					if err != nil {
-						e.logger.Error("failed to create shard processor", tag.Error(err))
-						e.metrics.Counter(metricsconstants.ShardDistributorExecutorProcessorCreationFailures).Inc(1)
-						return
-					}
-					managedProcessor := newManagedProcessor(processor, processorStateStarting)
-					e.managedProcessors.Store(shardID, managedProcessor)
-
-					processor.Start(ctx)
-
-					managedProcessor.setState(processorStateStarted)
-				}()
-			}
+			wg.Add(1)
+			go func(shardID string) {
+				defer wg.Done()
+				e.addManagerProcessor(ctx, shardID)
+			}(shardID)
 		}
 	}
 
 	wg.Wait()
 }
 
-func (e *executorImpl[SP]) stopShardProcessors() {
+func (e *executorImpl[SP]) addNewShards(ctx context.Context, shardAssignments map[string]*types.ShardAssignment) {
 	wg := sync.WaitGroup{}
 
-	e.managedProcessors.Range(func(shardID string, managedProcessor *managedProcessor[SP]) bool {
-		// If the processor is already stopping, skip it
-		if managedProcessor.getState() == processorStateStopping {
-			return true
+	for shardID, assignment := range shardAssignments {
+		if assignment.Status == types.AssignmentStatusREADY {
+			wg.Add(1)
+			go func(shardID string) {
+				defer wg.Done()
+				e.addManagerProcessor(ctx, shardID)
+			}(shardID)
 		}
+	}
+
+	wg.Wait()
+}
 
+func (e *executorImpl[SP]) deleteShards(shardIDs []string) {
+	wg := sync.WaitGroup{}
+	for _, shardID := range shardIDs {
 		wg.Add(1)
-		go func() {
+		go func(shardID string) {
 			defer wg.Done()
+			e.stopManagerProcessor(shardID)
+		}(shardID)
+	}
+	wg.Wait()
+}
 
-			managedProcessor.setState(processorStateStopping)
-			managedProcessor.processor.Stop()
-			e.managedProcessors.Delete(shardID)
-		}()
+func (e *executorImpl[SP]) stopShardProcessors() {
+	wg := sync.WaitGroup{}
+
+	e.managedProcessors.Range(func(shardID string, managedProcessor *managedProcessor[SP]) bool {
+		wg.Add(1)
+		go func(shardID string) {
+			defer wg.Done()
+			e.stopManagerProcessor(shardID)
+		}(shardID)
 		return true
 	})
 
 	wg.Wait()
 }
 
+func (e *executorImpl[SP]) addManagerProcessor(ctx context.Context, shardID string) {
+	if _, ok := e.managedProcessors.Load(shardID); !ok {
+		e.metrics.Counter(metricsconstants.ShardDistributorExecutorShardsStarted).Inc(1)
+		processor, err := e.shardProcessorFactory.NewShardProcessor(shardID)
+		if err != nil {
+			e.logger.Error("failed to create shard processor", tag.Error(err))
+			e.metrics.Counter(metricsconstants.ShardDistributorExecutorProcessorCreationFailures).Inc(1)
+			return
+		}
+		managedProcessor := newManagedProcessor(processor, processorStateStarting)
+		e.managedProcessors.Store(shardID, managedProcessor)
+
+		processor.Start(ctx)
+
+		managedProcessor.setState(processorStateStarted)
+
+	}
+}
+func (e *executorImpl[SP]) stopManagerProcessor(shardID string) {
+	managedProcessor, ok := e.managedProcessors.Load(shardID)
+	// If the processor do not exist for the shard, or it is already stopping, skip it
+	if !ok || managedProcessor.getState() == processorStateStopping {
+		return
+	}
+	e.metrics.Counter(metricsconstants.ShardDistributorExecutorShardsStopped).Inc(1)
+	managedProcessor.setState(processorStateStopping)
+	managedProcessor.processor.Stop()
+	e.managedProcessors.Delete(shardID)
+}
+
 // compareAssignments compares the local assignments with the heartbeat response assignments
 // and emits convergence or divergence metrics
 func (e *executorImpl[SP]) compareAssignments(heartbeatAssignments map[string]*types.ShardAssignment) {