feat(shard distributor): Persist Shard-Level Statistics for Load Balancing, and Add Cleanup Function

common/metrics/defs.go

@@ -1475,6 +1475,7 @@ const (
 	ShardDistributorStoreAssignShardScope
 	ShardDistributorStoreAssignShardsScope
 	ShardDistributorStoreDeleteExecutorsScope
+	ShardDistributorStoreDeleteShardStatsScope
 	ShardDistributorStoreGetHeartbeatScope
 	ShardDistributorStoreGetStateScope
 	ShardDistributorStoreRecordHeartbeatScope
@@ -2151,18 +2152,19 @@ var ScopeDefs = map[ServiceIdx]map[ScopeIdx]scopeDefinition{
 		DiagnosticsWorkflowScope:               {operation: "DiagnosticsWorkflow"},
 	},
 	ShardDistributor: {
-		ShardDistributorGetShardOwnerScope:        {operation: "GetShardOwner"},
-		ShardDistributorHeartbeatScope:            {operation: "ExecutorHeartbeat"},
-		ShardDistributorAssignLoopScope:           {operation: "ShardAssignLoop"},
-		ShardDistributorExecutorScope:             {operation: "Executor"},
-		ShardDistributorStoreGetShardOwnerScope:   {operation: "StoreGetShardOwner"},
-		ShardDistributorStoreAssignShardScope:     {operation: "StoreAssignShard"},
-		ShardDistributorStoreAssignShardsScope:    {operation: "StoreAssignShards"},
-		ShardDistributorStoreDeleteExecutorsScope: {operation: "StoreDeleteExecutors"},
-		ShardDistributorStoreGetHeartbeatScope:    {operation: "StoreGetHeartbeat"},
-		ShardDistributorStoreGetStateScope:        {operation: "StoreGetState"},
-		ShardDistributorStoreRecordHeartbeatScope: {operation: "StoreRecordHeartbeat"},
-		ShardDistributorStoreSubscribeScope:       {operation: "StoreSubscribe"},
+		ShardDistributorGetShardOwnerScope:         {operation: "GetShardOwner"},
+		ShardDistributorHeartbeatScope:             {operation: "ExecutorHeartbeat"},
+		ShardDistributorAssignLoopScope:            {operation: "ShardAssignLoop"},
+		ShardDistributorExecutorScope:              {operation: "Executor"},
+		ShardDistributorStoreGetShardOwnerScope:    {operation: "StoreGetShardOwner"},
+		ShardDistributorStoreAssignShardScope:      {operation: "StoreAssignShard"},
+		ShardDistributorStoreAssignShardsScope:     {operation: "StoreAssignShards"},
+		ShardDistributorStoreDeleteExecutorsScope:  {operation: "StoreDeleteExecutors"},
+		ShardDistributorStoreDeleteShardStatsScope: {operation: "StoreDeleteShardStats"},
+		ShardDistributorStoreGetHeartbeatScope:     {operation: "StoreGetHeartbeat"},
+		ShardDistributorStoreGetStateScope:         {operation: "StoreGetState"},
+		ShardDistributorStoreRecordHeartbeatScope:  {operation: "StoreRecordHeartbeat"},
+		ShardDistributorStoreSubscribeScope:        {operation: "StoreSubscribe"},
 	},
 }
 

service/sharddistributor/leader/process/processor.go

@@ -230,16 +230,21 @@ func (p *namespaceProcessor) runCleanupLoop(ctx context.Context) {
 			return
 		case <-ticker.Chan():
 			p.logger.Info("Periodic heartbeat cleanup triggered.")
-			p.cleanupStaleExecutors(ctx)
+			namespaceState, err := p.shardStore.GetState(ctx, p.namespaceCfg.Name)
+			if err != nil {
+				p.logger.Error("Failed to get state for cleanup", tag.Error(err))
+				continue
+			}
+			p.cleanupStaleExecutors(ctx, namespaceState)
+			p.cleanupStaleShardStats(ctx, namespaceState)
 		}
 	}
 }
 
 // cleanupStaleExecutors removes executors who have not reported a heartbeat recently.
-func (p *namespaceProcessor) cleanupStaleExecutors(ctx context.Context) {
-	namespaceState, err := p.shardStore.GetState(ctx, p.namespaceCfg.Name)
-	if err != nil {
-		p.logger.Error("Failed to get state for heartbeat cleanup", tag.Error(err))
+func (p *namespaceProcessor) cleanupStaleExecutors(ctx context.Context, namespaceState *store.NamespaceState) {
+	if namespaceState == nil {
+		p.logger.Error("Namespace state missing for heartbeat cleanup")
 		return
 	}
 
@@ -264,6 +269,73 @@ func (p *namespaceProcessor) cleanupStaleExecutors(ctx context.Context) {
 	}
 }
 
+func (p *namespaceProcessor) cleanupStaleShardStats(ctx context.Context, namespaceState *store.NamespaceState) {
+	if namespaceState == nil {
+		p.logger.Error("Namespace state missing for shard stats cleanup")
+		return
+	}
+
+	activeShards := make(map[string]struct{})
+	now := p.timeSource.Now().Unix()
+	shardStatsTTL := int64(p.cfg.HeartbeatTTL.Seconds())
+
+	// 1. build set of active executors
+
+	// add all assigned shards from executors that are ACTIVE and not stale
+	for executorID, assignedState := range namespaceState.ShardAssignments {
+		executor, exists := namespaceState.Executors[executorID]
+		if !exists {
+			continue
+		}
+
+		isActive := executor.Status == types.ExecutorStatusACTIVE
+		isNotStale := (now - executor.LastHeartbeat) <= shardStatsTTL
+		if isActive && isNotStale {
+			for shardID := range assignedState.AssignedShards {
+				activeShards[shardID] = struct{}{}
+			}
+		}
+	}
+
+	// add all shards in ReportedShards where the status is not DONE
+	for _, heartbeatState := range namespaceState.Executors {
+		for shardID, shardStatusReport := range heartbeatState.ReportedShards {
+			if shardStatusReport.Status != types.ShardStatusDONE {
+				activeShards[shardID] = struct{}{}
+			}
+		}
+	}
+
+	// 2. build set of stale shard stats
+
+	// append all shard stats that are not in the active shards set
+	var staleShardStats []string
+	for shardID, stats := range namespaceState.ShardStats {
+		if _, ok := activeShards[shardID]; ok {
+			continue
+		}
+		recentUpdate := stats.LastUpdateTime > 0 && (now-stats.LastUpdateTime) <= shardStatsTTL
+		recentMove := stats.LastMoveTime > 0 && (now-stats.LastMoveTime) <= shardStatsTTL
+		if recentUpdate || recentMove {
+			// Preserve stats that have been updated recently to allow cooldown/load history to
+			// survive executor churn. These shards are likely awaiting reassignment,
+			// so we don't want to delete them.
+			continue
+		}
+		staleShardStats = append(staleShardStats, shardID)
+	}
+
+	if len(staleShardStats) == 0 {
+		return
+	}
+
+	p.logger.Info("Removing stale shard stats")
+	// Use the leader guard for the delete operation.
+	if err := p.shardStore.DeleteShardStats(ctx, p.namespaceCfg.Name, staleShardStats, p.election.Guard()); err != nil {
+		p.logger.Error("Failed to delete stale shard stats", tag.Error(err))
+	}
+}
+
 // rebalanceShards is the core logic for distributing shards among active executors.
 func (p *namespaceProcessor) rebalanceShards(ctx context.Context) (err error) {
 	metricsLoopScope := p.metricsClient.Scope(metrics.ShardDistributorAssignLoopScope)

service/sharddistributor/leader/process/processor_test.go

@@ -178,11 +178,80 @@ func TestCleanupStaleExecutors(t *testing.T) {
 		"exec-stale":  {LastHeartbeat: now.Add(-2 * time.Second).Unix()},
 	}
 
-	mocks.store.EXPECT().GetState(gomock.Any(), mocks.cfg.Name).Return(&store.NamespaceState{Executors: heartbeats}, nil)
+	namespaceState := &store.NamespaceState{Executors: heartbeats}
 	mocks.election.EXPECT().Guard().Return(store.NopGuard())
 	mocks.store.EXPECT().DeleteExecutors(gomock.Any(), mocks.cfg.Name, []string{"exec-stale"}, gomock.Any()).Return(nil)
 
-	processor.cleanupStaleExecutors(context.Background())
+	processor.cleanupStaleExecutors(context.Background(), namespaceState)
+}
+
+func TestCleanupStaleShardStats(t *testing.T) {
+	t.Run("stale shard stats are deleted", func(t *testing.T) {
+		mocks := setupProcessorTest(t, config.NamespaceTypeFixed)
+		defer mocks.ctrl.Finish()
+		processor := mocks.factory.CreateProcessor(mocks.cfg, mocks.store, mocks.election).(*namespaceProcessor)
+
+		now := mocks.timeSource.Now()
+
+		heartbeats := map[string]store.HeartbeatState{
+			"exec-active": {LastHeartbeat: now.Unix(), Status: types.ExecutorStatusACTIVE},
+			"exec-stale":  {LastHeartbeat: now.Add(-2 * time.Second).Unix()},
+		}
+
+		assignments := map[string]store.AssignedState{
+			"exec-active": {
+				AssignedShards: map[string]*types.ShardAssignment{
+					"shard-1": {Status: types.AssignmentStatusREADY},
+					"shard-2": {Status: types.AssignmentStatusREADY},
+				},
+			},
+			"exec-stale": {
+				AssignedShards: map[string]*types.ShardAssignment{
+					"shard-3": {Status: types.AssignmentStatusREADY},
+				},
+			},
+		}
+
+		shardStats := map[string]store.ShardStatistics{
+			"shard-1": {SmoothedLoad: 1.0, LastUpdateTime: now.Unix(), LastMoveTime: now.Unix()},
+			"shard-2": {SmoothedLoad: 2.0, LastUpdateTime: now.Unix(), LastMoveTime: now.Unix()},
+			"shard-3": {SmoothedLoad: 3.0, LastUpdateTime: now.Add(-2 * time.Second).Unix(), LastMoveTime: now.Add(-2 * time.Second).Unix()},
+		}
+
+		namespaceState := &store.NamespaceState{
+			Executors:        heartbeats,
+			ShardAssignments: assignments,
+			ShardStats:       shardStats,
+		}
+
+		mocks.election.EXPECT().Guard().Return(store.NopGuard())
+		mocks.store.EXPECT().DeleteShardStats(gomock.Any(), mocks.cfg.Name, []string{"shard-3"}, gomock.Any()).Return(nil)
+		processor.cleanupStaleShardStats(context.Background(), namespaceState)
+	})
+
+	t.Run("recent shard stats are preserved", func(t *testing.T) {
+		mocks := setupProcessorTest(t, config.NamespaceTypeFixed)
+		defer mocks.ctrl.Finish()
+		processor := mocks.factory.CreateProcessor(mocks.cfg, mocks.store, mocks.election).(*namespaceProcessor)
+
+		now := mocks.timeSource.Now()
+
+		expiredExecutor := now.Add(-2 * time.Second).Unix()
+		namespaceState := &store.NamespaceState{
+			Executors: map[string]store.HeartbeatState{
+				"exec-stale": {LastHeartbeat: expiredExecutor},
+			},
+			ShardAssignments: map[string]store.AssignedState{},
+			ShardStats: map[string]store.ShardStatistics{
+				"shard-1": {SmoothedLoad: 5.0, LastUpdateTime: now.Unix(), LastMoveTime: now.Unix()},
+			},
+		}
+
+		processor.cleanupStaleShardStats(context.Background(), namespaceState)
+
+		// No delete expected since stats are recent.
+	})
+
 }
 
 func TestRebalance_StoreErrors(t *testing.T) {
@@ -213,16 +282,13 @@ func TestCleanup_StoreErrors(t *testing.T) {
 	processor := mocks.factory.CreateProcessor(mocks.cfg, mocks.store, mocks.election).(*namespaceProcessor)
 	expectedErr := errors.New("store is down")
 
-	mocks.store.EXPECT().GetState(gomock.Any(), mocks.cfg.Name).Return(nil, expectedErr)
-	processor.cleanupStaleExecutors(context.Background())
-
-	mocks.store.EXPECT().GetState(gomock.Any(), mocks.cfg.Name).Return(&store.NamespaceState{
+	namespaceState := &store.NamespaceState{
 		Executors:      map[string]store.HeartbeatState{"stale": {LastHeartbeat: 0}},
 		GlobalRevision: 1,
-	}, nil)
+	}
 	mocks.election.EXPECT().Guard().Return(store.NopGuard())
 	mocks.store.EXPECT().DeleteExecutors(gomock.Any(), mocks.cfg.Name, gomock.Any(), gomock.Any()).Return(expectedErr)
-	processor.cleanupStaleExecutors(context.Background())
+	processor.cleanupStaleExecutors(context.Background(), namespaceState)
 }
 
 func TestRunLoop_SubscriptionError(t *testing.T) {

service/sharddistributor/store/etcd/etcdkeys/etcdkeys.go

@@ -11,6 +11,7 @@ const (
 	ExecutorReportedShardsKey = "reported_shards"
 	ExecutorAssignedStateKey  = "assigned_state"
 	ShardAssignedKey          = "assigned"
+	ShardStatisticsKey        = "statistics"
 	ExecutorMetadataKey       = "metadata"
 )
 
@@ -70,6 +71,30 @@ func ParseExecutorKey(prefix string, namespace, key string) (executorID, keyType
 	return parts[0], parts[1], nil
 }
 
+func BuildShardPrefix(prefix string, namespace string) string {
+	return fmt.Sprintf("%s/shards/", BuildNamespacePrefix(prefix, namespace))
+}
+
+func BuildShardKey(prefix string, namespace, shardID, keyType string) (string, error) {
+	if keyType != ShardStatisticsKey {
+		return "", fmt.Errorf("invalid shard key type: %s", keyType)
+	}
+	return fmt.Sprintf("%s%s/%s", BuildShardPrefix(prefix, namespace), shardID, keyType), nil
+}
+
+func ParseShardKey(prefix string, namespace, key string) (shardID, keyType string, err error) {
+	prefix = BuildShardPrefix(prefix, namespace)
+	if !strings.HasPrefix(key, prefix) {
+		return "", "", fmt.Errorf("key '%s' does not have expected prefix '%s'", key, prefix)
+	}
+	remainder := strings.TrimPrefix(key, prefix)
+	parts := strings.Split(remainder, "/")
+	if len(parts) != 2 {
+		return "", "", fmt.Errorf("unexpected shard key format: %s", key)
+	}
+	return parts[0], parts[1], nil
+}
+
 func BuildMetadataKey(prefix string, namespace, executorID, metadataKey string) string {
 	metadataKeyPrefix, err := BuildExecutorKey(prefix, namespace, executorID, ExecutorMetadataKey)
 	if err != nil {

service/sharddistributor/store/etcd/etcdkeys/etcdkeys_test.go

@@ -16,6 +16,11 @@ func TestBuildExecutorPrefix(t *testing.T) {
 	assert.Equal(t, "/cadence/test-ns/executors/", got)
 }
 
+func TestBuildShardPrefix(t *testing.T) {
+	got := BuildShardPrefix("/cadence", "test-ns")
+	assert.Equal(t, "/cadence/test-ns/shards/", got)
+}
+
 func TestBuildExecutorKey(t *testing.T) {
 	got, err := BuildExecutorKey("/cadence", "test-ns", "exec-1", "heartbeat")
 	assert.NoError(t, err)
@@ -27,6 +32,17 @@ func TestBuildExecutorKeyFail(t *testing.T) {
 	assert.ErrorContains(t, err, "invalid key type: invalid")
 }
 
+func TestBuildShardKey(t *testing.T) {
+	got, err := BuildShardKey("/cadence", "test-ns", "shard-1", "statistics")
+	assert.NoError(t, err)
+	assert.Equal(t, "/cadence/test-ns/shards/shard-1/statistics", got)
+}
+
+func TestBuildShardKeyFail(t *testing.T) {
+	_, err := BuildShardKey("/cadence", "test-ns", "shard-1", "invalid")
+	assert.ErrorContains(t, err, "invalid shard key type: invalid")
+}
+
 func TestParseExecutorKey(t *testing.T) {
 	// Valid key
 	executorID, keyType, err := ParseExecutorKey("/cadence", "test-ns", "/cadence/test-ns/executors/exec-1/heartbeat")
@@ -43,6 +59,22 @@ func TestParseExecutorKey(t *testing.T) {
 	assert.ErrorContains(t, err, "unexpected key format: /cadence/test-ns/executors/exec-1/heartbeat/extra")
 }
 
+func TestParseShardKey(t *testing.T) {
+	// Valid key
+	shardID, keyType, err := ParseShardKey("/cadence", "test-ns", "/cadence/test-ns/shards/shard-1/statistics")
+	assert.NoError(t, err)
+	assert.Equal(t, "shard-1", shardID)
+	assert.Equal(t, "statistics", keyType)
+
+	// Prefix missing
+	_, _, err = ParseShardKey("/cadence", "test-ns", "/cadence/other/shards/shard-1/statistics")
+	assert.ErrorContains(t, err, "key '/cadence/other/shards/shard-1/statistics' does not have expected prefix '/cadence/test-ns/shards/'")
+
+	// Unexpected format
+	_, _, err = ParseShardKey("/cadence", "test-ns", "/cadence/test-ns/shards/shard-1/statistics/extra")
+	assert.ErrorContains(t, err, "unexpected shard key format: /cadence/test-ns/shards/shard-1/statistics/extra")
+}
+
 func TestBuildMetadataKey(t *testing.T) {
 	got := BuildMetadataKey("/cadence", "test-ns", "exec-1", "my-metadata-key")
 	assert.Equal(t, "/cadence/test-ns/executors/exec-1/metadata/my-metadata-key", got)