feat(shard distributor): throttle shard-stat writes

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

Author: AndreasHolt

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

@@ -36,8 +36,15 @@ type executorStoreImpl struct {
 	logger     log.Logger
 	shardCache *shardcache.ShardToExecutorCache
 	timeSource clock.TimeSource
+	// Max interval (seconds) before we force a shard-stat persist.
+	maxStatsPersistIntervalSeconds int64
 }
 
+// Constants for gating shard statistics writes to reduce etcd load.
+const (
+	shardStatsEpsilon = 0.05
+)
+
 // shardStatisticsUpdate holds the staged statistics for a shard so we can write them
 // to etcd after the main AssignShards transaction commits.
 type shardStatisticsUpdate struct {
@@ -90,11 +97,12 @@ func NewStore(p ExecutorStoreParams) (store.Store, error) {
 	}
 
 	store := &executorStoreImpl{
-		client:     etcdClient,
-		prefix:     etcdCfg.Prefix,
-		logger:     p.Logger,
-		shardCache: shardCache,
-		timeSource: timeSource,
+		client:                         etcdClient,
+		prefix:                         etcdCfg.Prefix,
+		logger:                         p.Logger,
+		shardCache:                     shardCache,
+		timeSource:                     timeSource,
+		maxStatsPersistIntervalSeconds: deriveStatsPersistInterval(p.Cfg.Process.HeartbeatTTL),
 	}
 
 	p.Lifecycle.Append(fx.StartStopHook(store.Start, store.Stop))
@@ -160,6 +168,15 @@ func (s *executorStoreImpl) RecordHeartbeat(ctx context.Context, namespace, exec
 	return nil
 }
 
+func deriveStatsPersistInterval(heartbeatTTL time.Duration) int64 {
+	ttlSeconds := int64(heartbeatTTL.Seconds())
+	interval := ttlSeconds - 1
+	if interval < 1 {
+		interval = 1
+	}
+	return interval
+}
+
 func (s *executorStoreImpl) recordShardStatistics(ctx context.Context, namespace, executorID string, reported map[string]*types.ShardStatusReport) {
 	if len(reported) == 0 {
 		return
@@ -228,7 +245,28 @@ func (s *executorStoreImpl) recordShardStatistics(ctx context.Context, namespace
 		}
 
 		// Update smoothed load via EWMA.
-		stats.SmoothedLoad = ewmaSmoothedLoad(stats.SmoothedLoad, load, stats.LastUpdateTime, now)
+		prevSmoothed := stats.SmoothedLoad
+		prevUpdate := stats.LastUpdateTime
+		newSmoothed := ewmaSmoothedLoad(prevSmoothed, load, prevUpdate, now)
+
+		// Decide whether to persist this update. We always persist if this is the
+		// first observation (prevUpdate == 0). Otherwise, if the change is small
+		// and the previous persist is recent, skip the write to reduce etcd load.
+		shouldPersist := true
+		if prevUpdate > 0 {
+			age := now - prevUpdate
+			delta := math.Abs(newSmoothed - prevSmoothed)
+			if delta < shardStatsEpsilon && age < s.maxStatsPersistIntervalSeconds {
+				shouldPersist = false
+			}
+		}
+
+		if !shouldPersist {
+			// Skip persisting, proceed to next shard.
+			continue
+		}
+
+		stats.SmoothedLoad = newSmoothed
 		stats.LastUpdateTime = now
 
 		payload, err := json.Marshal(stats)

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

@@ -11,6 +11,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/fx/fxtest"
 
+	"github.com/uber/cadence/common/clock"
 	"github.com/uber/cadence/common/log/testlogger"
 	"github.com/uber/cadence/common/types"
 	"github.com/uber/cadence/service/sharddistributor/store"
@@ -178,6 +179,66 @@ func TestRecordHeartbeatSkipsShardStatisticsWithNilReport(t *testing.T) {
 	assert.NotContains(t, nsState.ShardStats, skippedShardID)
 }
 
+func TestRecordHeartbeatShardStatisticsThrottlesWrites(t *testing.T) {
+	tc := testhelper.SetupStoreTestCluster(t)
+	tc.LeaderCfg.Process.HeartbeatTTL = 10 * time.Second
+	mockTS := clock.NewMockedTimeSourceAt(time.Unix(1000, 0))
+	executorStore := createStoreWithTimeSource(t, tc, mockTS)
+	esImpl, ok := executorStore.(*executorStoreImpl)
+	require.True(t, ok, "unexpected store implementation")
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	executorID := "executor-shard-stats-throttle"
+	shardID := "shard-stats-throttle"
+
+	baseLoad := 0.40
+	smallDelta := shardStatsEpsilon / 2
+	intervalSeconds := esImpl.maxStatsPersistIntervalSeconds
+	halfIntervalSeconds := intervalSeconds / 2
+	if halfIntervalSeconds == 0 {
+		halfIntervalSeconds = 1
+	}
+
+	// First heartbeat should always persist stats.
+	require.NoError(t, executorStore.RecordHeartbeat(ctx, tc.Namespace, executorID, store.HeartbeatState{
+		LastHeartbeat: mockTS.Now().Unix(),
+		Status:        types.ExecutorStatusACTIVE,
+		ReportedShards: map[string]*types.ShardStatusReport{
+			shardID: {Status: types.ShardStatusREADY, ShardLoad: baseLoad},
+		},
+	}))
+	statsAfterFirst := getShardStats(t, executorStore, ctx, tc.Namespace, shardID)
+	require.NotNil(t, statsAfterFirst)
+
+	// Advance time by less than the persist interval and provide a small delta: should skip the write.
+	mockTS.Advance(time.Duration(halfIntervalSeconds) * time.Second)
+	require.NoError(t, executorStore.RecordHeartbeat(ctx, tc.Namespace, executorID, store.HeartbeatState{
+		LastHeartbeat: mockTS.Now().Unix(),
+		Status:        types.ExecutorStatusACTIVE,
+		ReportedShards: map[string]*types.ShardStatusReport{
+			shardID: {Status: types.ShardStatusREADY, ShardLoad: baseLoad + smallDelta},
+		},
+	}))
+	statsAfterSkip := getShardStats(t, executorStore, ctx, tc.Namespace, shardID)
+	require.NotNil(t, statsAfterSkip)
+	assert.Equal(t, statsAfterFirst.LastUpdateTime, statsAfterSkip.LastUpdateTime, "small recent deltas should not trigger a persist")
+
+	// Advance time beyond the max persist interval, even small deltas should now persist.
+	mockTS.Advance(time.Duration(intervalSeconds) * time.Second)
+	require.NoError(t, executorStore.RecordHeartbeat(ctx, tc.Namespace, executorID, store.HeartbeatState{
+		LastHeartbeat: mockTS.Now().Unix(),
+		Status:        types.ExecutorStatusACTIVE,
+		ReportedShards: map[string]*types.ShardStatusReport{
+			shardID: {Status: types.ShardStatusREADY, ShardLoad: baseLoad + smallDelta/2},
+		},
+	}))
+	statsAfterForce := getShardStats(t, executorStore, ctx, tc.Namespace, shardID)
+	require.NotNil(t, statsAfterForce)
+	assert.Greater(t, statsAfterForce.LastUpdateTime, statsAfterSkip.LastUpdateTime, "stale stats must be refreshed even if delta is small")
+}
+
 func TestGetHeartbeat(t *testing.T) {
 	tc := testhelper.SetupStoreTestCluster(t)
 	executorStore := createStore(t, tc)
@@ -695,3 +756,27 @@ func createStore(t *testing.T, tc *testhelper.StoreTestCluster) store.Store {
 	require.NoError(t, err)
 	return store
 }
+
+func createStoreWithTimeSource(t *testing.T, tc *testhelper.StoreTestCluster, ts clock.TimeSource) store.Store {
+	t.Helper()
+	store, err := NewStore(ExecutorStoreParams{
+		Client:     tc.Client,
+		Cfg:        tc.LeaderCfg,
+		Lifecycle:  fxtest.NewLifecycle(t),
+		Logger:     testlogger.New(t),
+		TimeSource: ts,
+	})
+	require.NoError(t, err)
+	return store
+}
+
+func getShardStats(t *testing.T, s store.Store, ctx context.Context, namespace, shardID string) *store.ShardStatistics {
+	t.Helper()
+	nsState, err := s.GetState(ctx, namespace)
+	require.NoError(t, err)
+	stats, ok := nsState.ShardStats[shardID]
+	if !ok {
+		return nil
+	}
+	return &stats
+}