feat: Heartbeat shard statistics

service/sharddistributor/store/etcd/executorstore/etcdstore.go

@@ -8,6 +8,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"math"
 	"strconv"
 	"time"
 
@@ -153,9 +154,108 @@ func (s *executorStoreImpl) RecordHeartbeat(ctx context.Context, namespace, exec
 	if err != nil {
 		return fmt.Errorf("record heartbeat: %w", err)
 	}
+
+	s.recordShardStatistics(ctx, namespace, executorID, request.ReportedShards)
+
 	return nil
 }
 
+func (s *executorStoreImpl) recordShardStatistics(ctx context.Context, namespace, executorID string, reported map[string]*types.ShardStatusReport) {
+	if len(reported) == 0 {
+		return
+	}
+
+	now := s.timeSource.Now().Unix()
+
+	for shardID, report := range reported {
+		if report == nil {
+			s.logger.Warn("empty report; skipping EWMA update",
+				tag.ShardNamespace(namespace),
+				tag.ShardExecutor(executorID),
+				tag.ShardKey(shardID),
+			)
+			continue
+		}
+
+		load := report.ShardLoad
+		if math.IsNaN(load) || math.IsInf(load, 0) {
+			s.logger.Warn(
+				"invalid shard load reported; skipping EWMA update",
+				tag.ShardNamespace(namespace),
+				tag.ShardExecutor(executorID),
+				tag.ShardKey(shardID),
+			)
+			continue
+		}
+
+		shardStatsKey, err := etcdkeys.BuildShardKey(s.prefix, namespace, shardID, etcdkeys.ShardStatisticsKey)
+		if err != nil {
+			s.logger.Warn(
+				"failed to build shard statistics key from heartbeat",
+				tag.ShardNamespace(namespace),
+				tag.ShardExecutor(executorID),
+				tag.ShardKey(shardID),
+				tag.Error(err),
+			)
+			continue
+		}
+
+		statsResp, err := s.client.Get(ctx, shardStatsKey)
+		if err != nil {
+			s.logger.Warn(
+				"failed to read shard statistics for heartbeat update",
+				tag.ShardNamespace(namespace),
+				tag.ShardExecutor(executorID),
+				tag.ShardKey(shardID),
+				tag.Error(err),
+			)
+			continue
+		}
+
+		var stats store.ShardStatistics
+		if len(statsResp.Kvs) > 0 {
+			err := common.DecompressAndUnmarshal(statsResp.Kvs[0].Value, &stats)
+			if err != nil {
+				s.logger.Warn(
+					"failed to unmarshal shard statistics for heartbeat update",
+					tag.ShardNamespace(namespace),
+					tag.ShardExecutor(executorID),
+					tag.ShardKey(shardID),
+					tag.Error(err),
+				)
+				continue
+			}
+		}
+
+		// Update smoothed load via EWMA.
+		stats.SmoothedLoad = ewmaSmoothedLoad(stats.SmoothedLoad, load, stats.LastUpdateTime, now)
+		stats.LastUpdateTime = now
+
+		payload, err := json.Marshal(stats)
+		if err != nil {
+			s.logger.Warn(
+				"failed to marshal shard statistics after heartbeat",
+				tag.ShardNamespace(namespace),
+				tag.ShardExecutor(executorID),
+				tag.ShardKey(shardID),
+				tag.Error(err),
+			)
+			continue
+		}
+
+		_, err = s.client.Put(ctx, shardStatsKey, string(payload))
+		if err != nil {
+			s.logger.Warn(
+				"failed to persist shard statistics from heartbeat",
+				tag.ShardNamespace(namespace),
+				tag.ShardExecutor(executorID),
+				tag.ShardKey(shardID),
+				tag.Error(err),
+			)
+		}
+	}
+}
+
 // GetHeartbeat retrieves the last known heartbeat state for a single executor.
 func (s *executorStoreImpl) GetHeartbeat(ctx context.Context, namespace string, executorID string) (*store.HeartbeatState, *store.AssignedState, error) {
 	// The prefix for all keys related to a single executor.
@@ -717,3 +817,13 @@ func (s *executorStoreImpl) applyShardStatisticsUpdates(ctx context.Context, nam
 		}
 	}
 }
+
+func ewmaSmoothedLoad(prev, current float64, lastUpdate, now int64) float64 {
+	const tauSeconds = 30.0 // smaller = more responsive, larger = smoother
+	if lastUpdate <= 0 || tauSeconds <= 0 {
+		return current
+	}
+	dt := max(now-lastUpdate, 0)
+	alpha := 1 - math.Exp(-float64(dt)/tauSeconds)
+	return (1-alpha)*prev + alpha*current
+}

service/sharddistributor/store/etcd/executorstore/etcdstore_test.go

@@ -91,6 +91,93 @@ func TestRecordHeartbeat(t *testing.T) {
 	assert.Equal(t, "value-2", string(resp.Kvs[0].Value))
 }
 
+func TestRecordHeartbeatUpdatesShardStatistics(t *testing.T) {
+	tc := testhelper.SetupStoreTestCluster(t)
+	executorStore := createStore(t, tc)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	executorID := "executor-shard-stats"
+	shardID := "shard-with-load"
+
+	initialStats := store.ShardStatistics{
+		SmoothedLoad:   1.23,
+		LastUpdateTime: 10,
+		LastMoveTime:   123,
+	}
+
+	shardStatsKey, err := etcdkeys.BuildShardKey(tc.EtcdPrefix, tc.Namespace, shardID, etcdkeys.ShardStatisticsKey)
+	require.NoError(t, err)
+	payload, err := json.Marshal(initialStats)
+	require.NoError(t, err)
+	_, err = tc.Client.Put(ctx, shardStatsKey, string(payload))
+	require.NoError(t, err)
+
+	nowTS := time.Now().Unix()
+
+	req := store.HeartbeatState{
+		LastHeartbeat: nowTS,
+		Status:        types.ExecutorStatusACTIVE,
+		ReportedShards: map[string]*types.ShardStatusReport{
+			shardID: {
+				Status:    types.ShardStatusREADY,
+				ShardLoad: 45.6,
+			},
+		},
+	}
+
+	require.NoError(t, executorStore.RecordHeartbeat(ctx, tc.Namespace, executorID, req))
+
+	nsState, err := executorStore.GetState(ctx, tc.Namespace)
+	require.NoError(t, err)
+
+	require.Contains(t, nsState.ShardStats, shardID)
+	updated := nsState.ShardStats[shardID]
+
+	assert.InDelta(t, 45.6, updated.SmoothedLoad, 1e-9)
+	assert.GreaterOrEqual(t, updated.LastUpdateTime, nowTS)
+	assert.Equal(t, initialStats.LastMoveTime, updated.LastMoveTime)
+}
+
+func TestRecordHeartbeatSkipsShardStatisticsWithNilReport(t *testing.T) {
+	tc := testhelper.SetupStoreTestCluster(t)
+	executorStore := createStore(t, tc)
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	executorID := "executor-missing-load"
+	validShardID := "shard-with-valid-load"
+	skippedShardID := "shard-missing-load"
+
+	nowTS := time.Now().Unix()
+
+	req := store.HeartbeatState{
+		LastHeartbeat: nowTS,
+		Status:        types.ExecutorStatusACTIVE,
+		ReportedShards: map[string]*types.ShardStatusReport{
+			validShardID: {
+				Status:    types.ShardStatusREADY,
+				ShardLoad: 3.21,
+			},
+			skippedShardID: nil,
+		},
+	}
+
+	require.NoError(t, executorStore.RecordHeartbeat(ctx, tc.Namespace, executorID, req))
+
+	nsState, err := executorStore.GetState(ctx, tc.Namespace)
+	require.NoError(t, err)
+
+	require.Contains(t, nsState.ShardStats, validShardID)
+	validStats := nsState.ShardStats[validShardID]
+	assert.InDelta(t, 3.21, validStats.SmoothedLoad, 1e-9)
+	assert.Greater(t, validStats.LastUpdateTime, int64(0))
+
+	assert.NotContains(t, nsState.ShardStats, skippedShardID)
+}
+
 func TestGetHeartbeat(t *testing.T) {
 	tc := testhelper.SetupStoreTestCluster(t)
 	executorStore := createStore(t, tc)