Better lifecycle management in the shard executor library (#7095)

What changed?

Better shutdown of the executor process loop
Made the updateShardAssignment function use waitgroups, so it blocks until all the processors has started or stopped
Made sure we do not stop heartbeating if an assignment takes a long time.
If we are still running an old assignment when we get a new one, we will skip the new one.
Protected the state of the managed processor with an atomic int
Why?
Make sure the life cycle management of the executor library is robust

How did you test it?
Unit tests

Potential risks

Release notes

Documentation Changes

Author: jakobht

service/sharddistributor/executorclient/clientimpl.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/uber/cadence/client/sharddistributorexecutor"
@@ -14,7 +15,7 @@ import (
 	"github.com/uber/cadence/service/sharddistributor/executorclient/syncgeneric"
 )
 
-type processorState int
+type processorState int32
 
 const (
 	processorStateStarting processorState = iota
@@ -24,7 +25,25 @@ const (
 
 type managedProcessor[SP ShardProcessor] struct {
 	processor SP
-	state     processorState
+	state     atomic.Int32
+}
+
+func (mp *managedProcessor[SP]) setState(state processorState) {
+	mp.state.Store(int32(state))
+}
+
+func (mp *managedProcessor[SP]) getState() processorState {
+	return processorState(mp.state.Load())
+}
+
+func newManagedProcessor[SP ShardProcessor](processor SP, state processorState) *managedProcessor[SP] {
+	managed := &managedProcessor[SP]{
+		processor: processor,
+		state:     atomic.Int32{},
+	}
+
+	managed.setState(state)
+	return managed
 }
 
 type executorImpl[SP ShardProcessor] struct {
@@ -37,14 +56,21 @@ type executorImpl[SP ShardProcessor] struct {
 	managedProcessors      syncgeneric.Map[string, *managedProcessor[SP]]
 	executorID             string
 	timeSource             clock.TimeSource
+	processLoopWG          sync.WaitGroup
+	assignmentMutex        sync.Mutex
 }
 
 func (e *executorImpl[SP]) Start(ctx context.Context) {
-	go e.heartbeatloop(ctx)
+	e.processLoopWG.Add(1)
+	go func() {
+		defer e.processLoopWG.Done()
+		e.heartbeatloop(ctx)
+	}()
 }
 
 func (e *executorImpl[SP]) Stop() {
 	close(e.stopC)
+	e.processLoopWG.Wait()
 }
 
 func (e *executorImpl[SP]) GetShardProcess(shardID string) (SP, error) {
@@ -63,6 +89,9 @@ func (e *executorImpl[SP]) heartbeatloop(ctx context.Context) {
 	for {
 		select {
 		case <-ctx.Done():
+			e.logger.Info("shard distributorexecutor context done, stopping")
+			e.stopShardProcessors()
+			return
 		case <-e.stopC:
 			e.logger.Info("shard distributorexecutor stopped")
 			e.stopShardProcessors()
@@ -73,7 +102,14 @@ func (e *executorImpl[SP]) heartbeatloop(ctx context.Context) {
 				e.logger.Error("failed to heartbeat", tag.Error(err))
 				continue // TODO: should we stop the executor, and drop all the shards?
 			}
-			e.updateShardAssignment(ctx, shardAssignment)
+			if !e.assignmentMutex.TryLock() {
+				e.logger.Warn("already doing shard assignment, will skip this assignment")
+				continue
+			}
+			go func() {
+				defer e.assignmentMutex.Unlock()
+				e.updateShardAssignment(ctx, shardAssignment)
+			}()
 		}
 	}
 }
@@ -82,7 +118,7 @@ func (e *executorImpl[SP]) heartbeat(ctx context.Context) (shardAssignments map[
 	// Fill in the shard status reports
 	shardStatusReports := make(map[string]*types.ShardStatusReport)
 	e.managedProcessors.Range(func(shardID string, managedProcessor *managedProcessor[SP]) bool {
-		if managedProcessor.state == processorStateStarted {
+		if managedProcessor.getState() == processorStateStarted {
 			shardStatusReports[shardID] = &types.ShardStatusReport{
 				ShardLoad: managedProcessor.processor.GetShardLoad(),
 				Status:    types.ShardStatusREADY,
@@ -109,11 +145,15 @@ func (e *executorImpl[SP]) heartbeat(ctx context.Context) (shardAssignments map[
 }
 
 func (e *executorImpl[SP]) updateShardAssignment(ctx context.Context, shardAssignments map[string]*types.ShardAssignment) {
+	wg := sync.WaitGroup{}
+
 	// 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 {
+			wg.Add(1)
 			go func() {
-				managedProcessor.state = processorStateStopping
+				defer wg.Done()
+				managedProcessor.setState(processorStateStopping)
 				managedProcessor.processor.Stop()
 				e.managedProcessors.Delete(shardID)
 			}()
@@ -125,37 +165,42 @@ func (e *executorImpl[SP]) updateShardAssignment(ctx context.Context, shardAssig
 	for shardID, assignment := range shardAssignments {
 		if assignment.Status == types.AssignmentStatusREADY {
 			if _, ok := e.managedProcessors.Load(shardID); !ok {
+				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))
 						return
 					}
+					managedProcessor := newManagedProcessor(processor, processorStateStarting)
+					e.managedProcessors.Store(shardID, managedProcessor)
+
 					processor.Start(ctx)
-					e.managedProcessors.Store(shardID, &managedProcessor[SP]{
-						processor: processor,
-						state:     processorStateStarted,
-					})
+
+					managedProcessor.setState(processorStateStarted)
 				}()
 			}
 		}
 	}
+
+	wg.Wait()
 }
 
 func (e *executorImpl[SP]) stopShardProcessors() {
 	wg := sync.WaitGroup{}
 
 	e.managedProcessors.Range(func(shardID string, managedProcessor *managedProcessor[SP]) bool {
 		// If the processor is already stopping, skip it
-		if managedProcessor.state == processorStateStopping {
+		if managedProcessor.getState() == processorStateStopping {
 			return true
 		}
 
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
 
-			managedProcessor.state = processorStateStopping
+			managedProcessor.setState(processorStateStopping)
 			managedProcessor.processor.Stop()
 			e.managedProcessors.Delete(shardID)
 		}()

service/sharddistributor/executorclient/clientimpl_test.go

@@ -129,14 +129,8 @@ func TestHeartbeat(t *testing.T) {
 		executorID:             "test-executor-id",
 	}
 
-	executor.managedProcessors.Store("test-shard-id1", &managedProcessor[*MockShardProcessor]{
-		processor: shardProcessorMock1,
-		state:     processorStateStarted,
-	})
-	executor.managedProcessors.Store("test-shard-id2", &managedProcessor[*MockShardProcessor]{
-		processor: shardProcessorMock2,
-		state:     processorStateStarted,
-	})
+	executor.managedProcessors.Store("test-shard-id1", newManagedProcessor(shardProcessorMock1, processorStateStarted))
+	executor.managedProcessors.Store("test-shard-id2", newManagedProcessor(shardProcessorMock2, processorStateStarted))
 
 	// Do the call to heartbeat
 	shardAssignments, err := executor.heartbeat(context.Background())
@@ -166,14 +160,8 @@ func TestHeartBeartLoop_ShardAssignmentChange(t *testing.T) {
 		shardProcessorFactory: shardProcessorFactory,
 	}
 
-	executor.managedProcessors.Store("test-shard-id1", &managedProcessor[*MockShardProcessor]{
-		processor: shardProcessorMock1,
-		state:     processorStateStarted,
-	})
-	executor.managedProcessors.Store("test-shard-id2", &managedProcessor[*MockShardProcessor]{
-		processor: shardProcessorMock2,
-		state:     processorStateStarted,
-	})
+	executor.managedProcessors.Store("test-shard-id1", newManagedProcessor(shardProcessorMock1, processorStateStarted))
+	executor.managedProcessors.Store("test-shard-id2", newManagedProcessor(shardProcessorMock2, processorStateStarted))
 
 	// We expect to get a new assignment with shards 2 and 3 assigned to it
 	newAssignment := map[string]*types.ShardAssignment{