fix(shard-distributor): send initial state to new subscribers (#7499)

<!-- Describe what has changed in this PR -->
**What changed?**
Modified the `Subscribe` method in `namespaceShardToExecutor` to send
the initial executor state to new subscribers immediately upon
subscription, and refactored `getExecutorState` into a separate method
to safely retrieve a copy of the current state.

Removed logic for sending initial state from the handler. The logic
didn't send the metadata for the executors.

Improved test infrastructure by updating `setupExecutorWithShards` to
use atomic etcd transactions and to properly delete old testdata.

<!-- Tell your future self why have you made these changes -->
**Why?**
New subscribers were not receiving the initial state when they
subscribed to the executor state pub/sub, causing them to wait for the
next update before getting data. This lead to delays in subscribers
getting the current shard distribution.

The test improvements ensure executor setup happens atomically and makes
tests more maintainable by ensuring deletion of stale state.

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

<!-- Assuming the worst case, what can be broken when deploying this
change to production? -->
**Potential risks**

<!-- 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: Jakob Haahr Taankvist <jht@uber.com>

Author: jakobht

service/sharddistributor/handler/handler.go

@@ -157,16 +157,6 @@ func (h *handlerImpl) WatchNamespaceState(request *types.WatchNamespaceStateRequ
 		return &types.InternalServiceError{Message: fmt.Sprintf("failed to subscribe to namespace state: %v", err)}
 	}
 
-	// Send initial state immediately so client doesn't have to wait for first update
-	state, err := h.storage.GetState(server.Context(), request.Namespace)
-	if err != nil {
-		return &types.InternalServiceError{Message: fmt.Sprintf("failed to get namespace state: %v", err)}
-	}
-	response := toWatchNamespaceStateResponse(state)
-	if err := server.Send(response); err != nil {
-		return fmt.Errorf("send initial state: %w", err)
-	}
-
 	// Stream subsequent updates
 	for {
 		select {
@@ -177,7 +167,7 @@ func (h *handlerImpl) WatchNamespaceState(request *types.WatchNamespaceStateRequ
 				return fmt.Errorf("unexpected close of updates channel")
 			}
 			response := &types.WatchNamespaceStateResponse{
-				Executors: make([]*types.ExecutorShardAssignment, 0, len(state.ShardAssignments)),
+				Executors: make([]*types.ExecutorShardAssignment, 0, len(assignmentChanges)),
 			}
 			for executor, shardIDs := range assignmentChanges {
 				response.Executors = append(response.Executors, &types.ExecutorShardAssignment{
@@ -195,27 +185,6 @@ func (h *handlerImpl) WatchNamespaceState(request *types.WatchNamespaceStateRequ
 	}
 }
 
-func toWatchNamespaceStateResponse(state *store.NamespaceState) *types.WatchNamespaceStateResponse {
-	response := &types.WatchNamespaceStateResponse{
-		Executors: make([]*types.ExecutorShardAssignment, 0, len(state.ShardAssignments)),
-	}
-
-	for executorID, assignment := range state.ShardAssignments {
-		// Extract shard IDs from the assigned shards map
-		shardIDs := make([]string, 0, len(assignment.AssignedShards))
-		for shardID := range assignment.AssignedShards {
-			shardIDs = append(shardIDs, shardID)
-		}
-
-		response.Executors = append(response.Executors, &types.ExecutorShardAssignment{
-			ExecutorID:     executorID,
-			AssignedShards: WrapShards(shardIDs),
-			Metadata:       state.Executors[executorID].Metadata,
-		})
-	}
-	return response
-}
-
 func WrapShards(shardIDs []string) []*types.Shard {
 	shards := make([]*types.Shard, 0, len(shardIDs))
 	for _, shardID := range shardIDs {

service/sharddistributor/handler/handler_test.go

@@ -243,62 +243,31 @@ func TestWatchNamespaceState(t *testing.T) {
 		startWG:              sync.WaitGroup{},
 	}
 
-	t.Run("successful streaming", func(t *testing.T) {
-		ctx, cancel := context.WithCancel(context.Background())
+	ctx, cancel := context.WithCancel(context.Background())
 
-		initialState := &store.NamespaceState{
-			ShardAssignments: map[string]store.AssignedState{
-				"executor-1": {
-					AssignedShards: map[string]*types.ShardAssignment{
-						"shard-1": {},
-					},
-				},
-			},
-		}
-
-		updatesChan := make(chan map[*store.ShardOwner][]string, 1)
-		unsubscribe := func() { close(updatesChan) }
-
-		mockServer.EXPECT().Context().Return(ctx).AnyTimes()
-		mockStorage.EXPECT().GetState(gomock.Any(), "test-ns").Return(initialState, nil)
-		mockStorage.EXPECT().SubscribeToAssignmentChanges(gomock.Any(), "test-ns").Return(updatesChan, unsubscribe, nil)
-
-		// Expect initial state send
-		mockServer.EXPECT().Send(gomock.Any()).DoAndReturn(func(resp *types.WatchNamespaceStateResponse) error {
-			require.Len(t, resp.Executors, 1)
-			require.Equal(t, "executor-1", resp.Executors[0].ExecutorID)
-			return nil
-		})
-
-		// Expect update send
-		mockServer.EXPECT().Send(gomock.Any()).DoAndReturn(func(resp *types.WatchNamespaceStateResponse) error {
-			require.Len(t, resp.Executors, 1)
-			require.Equal(t, "executor-2", resp.Executors[0].ExecutorID)
-			return nil
-		})
+	updatesChan := make(chan map[*store.ShardOwner][]string, 1)
+	unsubscribe := func() { close(updatesChan) }
 
-		// Send update, then cancel
-		go func() {
-			time.Sleep(10 * time.Millisecond)
-			updatesChan <- map[*store.ShardOwner][]string{
-				{ExecutorID: "executor-2", Metadata: map[string]string{}}: {"shard-2"},
-			}
-			cancel()
-		}()
+	mockServer.EXPECT().Context().Return(ctx).AnyTimes()
+	mockStorage.EXPECT().SubscribeToAssignmentChanges(gomock.Any(), "test-ns").Return(updatesChan, unsubscribe, nil)
 
-		err := handler.WatchNamespaceState(&types.WatchNamespaceStateRequest{Namespace: "test-ns"}, mockServer)
-		require.Error(t, err)
-		require.ErrorIs(t, err, context.Canceled)
+	// Expect update send
+	mockServer.EXPECT().Send(gomock.Any()).DoAndReturn(func(resp *types.WatchNamespaceStateResponse) error {
+		require.Len(t, resp.Executors, 1)
+		require.Equal(t, "executor-1", resp.Executors[0].ExecutorID)
+		return nil
 	})
 
-	t.Run("storage error on initial state", func(t *testing.T) {
-		ctx := context.Background()
-		mockServer.EXPECT().Context().Return(ctx).AnyTimes()
-		mockStorage.EXPECT().GetState(gomock.Any(), "test-ns").Return(nil, errors.New("storage error"))
-		mockStorage.EXPECT().SubscribeToAssignmentChanges(gomock.Any(), "test-ns").Return(make(chan map[*store.ShardOwner][]string), func() {}, nil)
+	// Send update, then cancel
+	go func() {
+		time.Sleep(10 * time.Millisecond)
+		updatesChan <- map[*store.ShardOwner][]string{
+			{ExecutorID: "executor-1", Metadata: map[string]string{}}: {"shard-1"},
+		}
+		cancel()
+	}()
 
-		err := handler.WatchNamespaceState(&types.WatchNamespaceStateRequest{Namespace: "test-ns"}, mockServer)
-		require.Error(t, err)
-		require.Contains(t, err.Error(), "failed to get namespace state: storage error")
-	})
+	err := handler.WatchNamespaceState(&types.WatchNamespaceStateRequest{Namespace: "test-ns"}, mockServer)
+	require.Error(t, err)
+	require.ErrorIs(t, err, context.Canceled)
 }

service/sharddistributor/store/etcd/executorstore/shardcache/namespaceshardcache.go

@@ -97,7 +97,22 @@ func (n *namespaceShardToExecutor) GetExecutorModRevisionCmp() ([]clientv3.Cmp,
 }
 
 func (n *namespaceShardToExecutor) Subscribe(ctx context.Context) (<-chan map[*store.ShardOwner][]string, func()) {
-	return n.pubSub.subscribe(ctx)
+	subCh, unSub := n.pubSub.subscribe(ctx)
+
+	// The go routine sends the initial state to the subscriber.
+	go func() {
+		initialState := n.getExecutorState()
+
+		select {
+		case <-ctx.Done():
+			n.logger.Warn("context finnished before initial state was sent", tag.ShardNamespace(n.namespace))
+		case subCh <- initialState:
+			n.logger.Info("initial state sent to subscriber", tag.ShardNamespace(n.namespace), tag.Value(initialState))
+		}
+
+	}()
+
+	return subCh, unSub
 }
 
 func (n *namespaceShardToExecutor) nameSpaceRefreashLoop() {
@@ -135,16 +150,20 @@ func (n *namespaceShardToExecutor) refresh(ctx context.Context) error {
 		return fmt.Errorf("refresh executor state: %w", err)
 	}
 
+	n.pubSub.publish(n.getExecutorState())
+	return nil
+}
+
+func (n *namespaceShardToExecutor) getExecutorState() map[*store.ShardOwner][]string {
 	n.RLock()
+	defer n.RUnlock()
 	executorState := make(map[*store.ShardOwner][]string)
 	for executor, shardIDs := range n.executorState {
 		executorState[executor] = make([]string, len(shardIDs))
 		copy(executorState[executor], shardIDs)
 	}
-	n.RUnlock()
 
-	n.pubSub.publish(executorState)
-	return nil
+	return executorState
 }
 
 func (n *namespaceShardToExecutor) refreshExecutorState(ctx context.Context) error {

service/sharddistributor/store/etcd/executorstore/shardcache/namespaceshardcache_test.go

@@ -9,69 +9,30 @@ import (
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	clientv3 "go.etcd.io/etcd/client/v3"
 
 	"github.com/uber/cadence/common/log/testlogger"
 	"github.com/uber/cadence/common/types"
+	"github.com/uber/cadence/service/sharddistributor/store"
 	"github.com/uber/cadence/service/sharddistributor/store/etcd/etcdkeys"
 	"github.com/uber/cadence/service/sharddistributor/store/etcd/etcdtypes"
 	"github.com/uber/cadence/service/sharddistributor/store/etcd/testhelper"
 )
 
-// setupExecutorWithShards creates an executor in etcd with assigned shards and metadata
-func setupExecutorWithShards(t *testing.T, testCluster *testhelper.StoreTestCluster, namespace, executorID string, shards []string, metadata map[string]string) {
-	// Create assigned state
-	assignedState := &etcdtypes.AssignedState{
-		AssignedShards: make(map[string]*types.ShardAssignment),
-	}
-	for _, shardID := range shards {
-		assignedState.AssignedShards[shardID] = &types.ShardAssignment{Status: types.AssignmentStatusREADY}
-	}
-	assignedStateJSON, err := json.Marshal(assignedState)
-	require.NoError(t, err)
-
-	executorAssignedStateKey := etcdkeys.BuildExecutorKey(testCluster.EtcdPrefix, namespace, executorID, etcdkeys.ExecutorAssignedStateKey)
-	testCluster.Client.Put(context.Background(), executorAssignedStateKey, string(assignedStateJSON))
-
-	// Add metadata
-	for key, value := range metadata {
-		metadataKey := etcdkeys.BuildMetadataKey(testCluster.EtcdPrefix, namespace, executorID, key)
-		testCluster.Client.Put(context.Background(), metadataKey, value)
-	}
-}
-
-// verifyShardOwner checks that a shard has the expected owner and metadata
-func verifyShardOwner(t *testing.T, cache *namespaceShardToExecutor, shardID, expectedExecutorID string, expectedMetadata map[string]string) {
-	owner, err := cache.GetShardOwner(context.Background(), shardID)
-	require.NoError(t, err)
-	require.NotNil(t, owner)
-	assert.Equal(t, expectedExecutorID, owner.ExecutorID)
-	for key, expectedValue := range expectedMetadata {
-		assert.Equal(t, expectedValue, owner.Metadata[key])
-	}
-
-	executor, err := cache.GetExecutor(context.Background(), expectedExecutorID)
-	require.NoError(t, err)
-	require.NotNil(t, executor)
-	assert.Equal(t, expectedExecutorID, executor.ExecutorID)
-	for key, expectedValue := range expectedMetadata {
-		assert.Equal(t, expectedValue, executor.Metadata[key])
-	}
-}
-
 func TestNamespaceShardToExecutor_Lifecycle(t *testing.T) {
 	testCluster := testhelper.SetupStoreTestCluster(t)
 	logger := testlogger.New(t)
 	stopCh := make(chan struct{})
 	defer close(stopCh)
 
 	// Setup: Create executor-1 with shard-1
-	setupExecutorWithShards(t, testCluster, "test-ns", "executor-1", []string{"shard-1"}, map[string]string{
+	setupExecutorWithShards(t, testCluster, "executor-1", []string{"shard-1"}, map[string]string{
 		"hostname": "executor-1-host",
 		"version":  "v1.0.0",
 	})
 
 	// Start the cache
-	namespaceShardToExecutor, err := newNamespaceShardToExecutor(testCluster.EtcdPrefix, "test-ns", testCluster.Client, stopCh, logger)
+	namespaceShardToExecutor, err := newNamespaceShardToExecutor(testCluster.EtcdPrefix, testCluster.Namespace, testCluster.Client, stopCh, logger)
 	assert.NoError(t, err)
 	namespaceShardToExecutor.Start(&sync.WaitGroup{})
 	time.Sleep(50 * time.Millisecond)
@@ -90,7 +51,7 @@ func TestNamespaceShardToExecutor_Lifecycle(t *testing.T) {
 	namespaceShardToExecutor.RUnlock()
 
 	// Add executor-2 with shard-2 to trigger watch update
-	setupExecutorWithShards(t, testCluster, "test-ns", "executor-2", []string{"shard-2"}, map[string]string{
+	setupExecutorWithShards(t, testCluster, "executor-2", []string{"shard-2"}, map[string]string{
 		"hostname": "executor-2-host",
 		"region":   "us-west",
 	})
@@ -109,3 +70,125 @@ func TestNamespaceShardToExecutor_Lifecycle(t *testing.T) {
 		"region":   "us-west",
 	})
 }
+
+func TestNamespaceShardToExecutor_Subscribe(t *testing.T) {
+	testCluster := testhelper.SetupStoreTestCluster(t)
+	logger := testlogger.New(t)
+	stopCh := make(chan struct{})
+	defer close(stopCh)
+
+	// Setup: Create executor-1 with shard-1
+	setupExecutorWithShards(t, testCluster, "executor-1", []string{"shard-1"}, map[string]string{
+		"hostname": "executor-1-host",
+		"version":  "v1.0.0",
+	})
+
+	// Start the cache
+	namespaceShardToExecutor, err := newNamespaceShardToExecutor(testCluster.EtcdPrefix, testCluster.Namespace, testCluster.Client, stopCh, logger)
+	assert.NoError(t, err)
+	namespaceShardToExecutor.Start(&sync.WaitGroup{})
+
+	// Refresh the cache to get the initial state
+	namespaceShardToExecutor.refresh(context.Background())
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+	subCh, unSub := namespaceShardToExecutor.Subscribe(ctx)
+	defer unSub()
+
+	var wg sync.WaitGroup
+	wg.Add(1)
+
+	// start listener
+	go func() {
+		defer wg.Done()
+		// Check that we get the initial state
+		state := <-subCh
+		assert.Len(t, state, 1)
+		verifyExecutorInState(t, state, "executor-1", []string{"shard-1"}, map[string]string{
+			"hostname": "executor-1-host",
+			"version":  "v1.0.0",
+		})
+
+		// Check that we get the updated state
+		state = <-subCh
+		assert.Len(t, state, 2)
+		verifyExecutorInState(t, state, "executor-1", []string{"shard-1"}, map[string]string{
+			"hostname": "executor-1-host",
+			"version":  "v1.0.0",
+		})
+		verifyExecutorInState(t, state, "executor-2", []string{"shard-2"}, map[string]string{
+			"hostname": "executor-2-host",
+			"region":   "us-west",
+		})
+	}()
+	time.Sleep(10 * time.Millisecond)
+
+	// Add executor-2 with shard-2 to trigger new subscription update
+	setupExecutorWithShards(t, testCluster, "executor-2", []string{"shard-2"}, map[string]string{
+		"hostname": "executor-2-host",
+		"region":   "us-west",
+	})
+
+	wg.Wait()
+}
+
+// setupExecutorWithShards creates an executor in etcd with assigned shards and metadata
+func setupExecutorWithShards(t *testing.T, testCluster *testhelper.StoreTestCluster, executorID string, shards []string, metadata map[string]string) {
+	// Create assigned state
+	assignedState := &etcdtypes.AssignedState{
+		AssignedShards: make(map[string]*types.ShardAssignment),
+	}
+	for _, shardID := range shards {
+		assignedState.AssignedShards[shardID] = &types.ShardAssignment{Status: types.AssignmentStatusREADY}
+	}
+	assignedStateJSON, err := json.Marshal(assignedState)
+	require.NoError(t, err)
+
+	var operations []clientv3.Op
+
+	executorAssignedStateKey := etcdkeys.BuildExecutorKey(testCluster.EtcdPrefix, testCluster.Namespace, executorID, etcdkeys.ExecutorAssignedStateKey)
+	operations = append(operations, clientv3.OpPut(executorAssignedStateKey, string(assignedStateJSON)))
+
+	// Add metadata
+	for key, value := range metadata {
+		metadataKey := etcdkeys.BuildMetadataKey(testCluster.EtcdPrefix, testCluster.Namespace, executorID, key)
+		operations = append(operations, clientv3.OpPut(metadataKey, value))
+	}
+
+	txnResp, err := testCluster.Client.Txn(context.Background()).Then(operations...).Commit()
+	require.NoError(t, err)
+	require.True(t, txnResp.Succeeded)
+}
+
+func verifyExecutorInState(t *testing.T, state map[*store.ShardOwner][]string, executorID string, shards []string, metadata map[string]string) {
+	executorInState := false
+	for executor, executorShards := range state {
+		if executor.ExecutorID == executorID {
+			assert.Equal(t, shards, executorShards)
+			assert.Equal(t, metadata, executor.Metadata)
+			executorInState = true
+			break
+		}
+	}
+	assert.True(t, executorInState)
+}
+
+// verifyShardOwner checks that a shard has the expected owner and metadata
+func verifyShardOwner(t *testing.T, cache *namespaceShardToExecutor, shardID, expectedExecutorID string, expectedMetadata map[string]string) {
+	owner, err := cache.GetShardOwner(context.Background(), shardID)
+	require.NoError(t, err)
+	require.NotNil(t, owner)
+	assert.Equal(t, expectedExecutorID, owner.ExecutorID)
+	for key, expectedValue := range expectedMetadata {
+		assert.Equal(t, expectedValue, owner.Metadata[key])
+	}
+
+	executor, err := cache.GetExecutor(context.Background(), expectedExecutorID)
+	require.NoError(t, err)
+	require.NotNil(t, executor)
+	assert.Equal(t, expectedExecutorID, executor.ExecutorID)
+	for key, expectedValue := range expectedMetadata {
+		assert.Equal(t, expectedValue, executor.Metadata[key])
+	}
+}