feat(shard distributor): retain shard stats while shards are within heartbeat TTL

Signed-off-by: Andreas Holt <[email protected]>

Author: AndreasHolt

service/sharddistributor/leader/process/processor.go

@@ -297,10 +297,19 @@ func (p *namespaceProcessor) cleanupStaleShardStats(ctx context.Context) {
 
 	// append all shard stats that are not in the active shards set
 	var staleShardStats []string
-	for shardID := range namespaceState.ShardStats {
-		if _, ok := activeShards[shardID]; !ok {
-			staleShardStats = append(staleShardStats, shardID)
+	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 {

service/sharddistributor/leader/process/processor_test.go

@@ -215,7 +215,7 @@ func TestCleanupStaleShardStats(t *testing.T) {
 		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.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{
@@ -232,6 +232,30 @@ func TestCleanupStaleShardStats(t *testing.T) {
 		processor.cleanupStaleShardStats(context.Background())
 	})
 
+	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()
+		state := &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()},
+			},
+		}
+
+		mocks.store.EXPECT().GetState(gomock.Any(), mocks.cfg.Name).Return(state, nil)
+		processor.cleanupStaleShardStats(context.Background())
+
+		// No delete expected since stats are recent.
+	})
+
 }
 
 func TestRebalance_StoreErrors(t *testing.T) {