feat(shard distributor): clean up the shard statistics

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

Author: AndreasHolt

common/metrics/defs.go

@@ -1469,6 +1469,7 @@ const (
 	ShardDistributorStoreAssignShardScope
 	ShardDistributorStoreAssignShardsScope
 	ShardDistributorStoreDeleteExecutorsScope
+	ShardDistributorStoreDeleteShardStatsScope
 	ShardDistributorStoreGetHeartbeatScope
 	ShardDistributorStoreGetStateScope
 	ShardDistributorStoreRecordHeartbeatScope

service/sharddistributor/leader/process/processor.go

@@ -221,6 +221,7 @@ func (p *namespaceProcessor) runCleanupLoop(ctx context.Context) {
 		case <-ticker.Chan():
 			p.logger.Info("Periodic heartbeat cleanup triggered.")
 			p.cleanupStaleExecutors(ctx)
+			p.cleanupStaleShardStats(ctx)
 		}
 	}
 }
@@ -254,6 +255,65 @@ func (p *namespaceProcessor) cleanupStaleExecutors(ctx context.Context) {
 	}
 }
 
+func (p *namespaceProcessor) cleanupStaleShardStats(ctx context.Context) {
+	namespaceState, err := p.shardStore.GetState(ctx, p.namespaceCfg.Name)
+	if err != nil {
+		p.logger.Error("Failed to get state for shard stats cleanup", tag.Error(err))
+		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 := range namespaceState.ShardStats {
+		if _, ok := activeShards[shardID]; !ok {
+			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/store/etcd/executorstore/etcdstore.go

@@ -580,6 +580,41 @@ func (s *executorStoreImpl) DeleteExecutors(ctx context.Context, namespace strin
 	return nil
 }
 
+func (s *executorStoreImpl) DeleteShardStats(ctx context.Context, namespace string, shardIDs []string, guard store.GuardFunc) error {
+	if len(shardIDs) == 0 {
+		return nil
+	}
+	var ops []clientv3.Op
+	for _, shardID := range shardIDs {
+		shardStatsKey, err := etcdkeys.BuildShardKey(s.prefix, namespace, shardID, etcdkeys.ShardStatisticsKey)
+		if err != nil {
+			return fmt.Errorf("build shard statistics key: %w", err)
+		}
+		ops = append(ops, clientv3.OpDelete(shardStatsKey))
+	}
+
+	nativeTxn := s.client.Txn(ctx)
+	guardedTxn, err := guard(nativeTxn)
+
+	if err != nil {
+		return fmt.Errorf("apply transaction guard: %w", err)
+	}
+	etcdGuardedTxn, ok := guardedTxn.(clientv3.Txn)
+	if !ok {
+		return fmt.Errorf("guard function returned invalid transaction type")
+	}
+
+	etcdGuardedTxn = etcdGuardedTxn.Then(ops...)
+	resp, err := etcdGuardedTxn.Commit()
+	if err != nil {
+		return fmt.Errorf("commit shard statistics deletion: %w", err)
+	}
+	if !resp.Succeeded {
+		return fmt.Errorf("transaction failed, leadership may have changed")
+	}
+	return nil
+}
+
 func (s *executorStoreImpl) GetShardOwner(ctx context.Context, namespace, shardID string) (string, error) {
 	return s.shardCache.GetShardOwner(ctx, namespace, shardID)
 }

service/sharddistributor/store/store.go

@@ -60,6 +60,7 @@ type Store interface {
 	AssignShards(ctx context.Context, namespace string, request AssignShardsRequest, guard GuardFunc) error
 	Subscribe(ctx context.Context, namespace string) (<-chan int64, error)
 	DeleteExecutors(ctx context.Context, namespace string, executorIDs []string, guard GuardFunc) error
+	DeleteShardStats(ctx context.Context, namespace string, shardIDs []string, guard GuardFunc) error
 
 	GetShardOwner(ctx context.Context, namespace, shardID string) (string, error)
 	AssignShard(ctx context.Context, namespace, shardID, executorID string) error