aws_dynamodb_batcher: encapsulate throttling logic in pinClock

internal/impl/aws/dynamodb/batcher.go

@@ -31,6 +31,37 @@ const (
 	throttlePinWarnInterval = 5 * time.Minute
 )
 
+// pinClock encapsulates continuous throttling duration tracking and rate-limited
+// warning log state.
+type pinClock struct {
+	// throttledSince records when reservations started continuously failing
+	// without an intervening successful reservation/check; zero while not
+	// throttled.
+	throttledSince time.Time
+	// lastPinWarn is the timestamp when the pinned-throttle warning was last
+	// logged, used to rate-limit warnings.
+	lastPinWarn time.Time
+}
+
+func (c *pinClock) clear() {
+	c.throttledSince = time.Time{}
+}
+
+// check records the start of a throttling event or checks if a warning is due.
+// Returns the total duration spent continuously throttled and whether a log warning should be emitted.
+func (c *pinClock) check(now time.Time) (since time.Duration, shouldWarn bool) {
+	if c.throttledSince.IsZero() {
+		c.throttledSince = now
+		return 0, false
+	}
+	since = now.Sub(c.throttledSince)
+	if since >= throttlePinWarnAfter && now.Sub(c.lastPinWarn) >= throttlePinWarnInterval {
+		c.lastPinWarn = now
+		return since, true
+	}
+	return since, false
+}
+
 // RecordBatcher tracks in-flight message batches and persists shard
 // checkpoints in stream order.
 //
@@ -84,13 +115,9 @@ type RecordBatcher struct {
 	// shard with nothing in flight always admits one batch regardless, so
 	// progress is guaranteed for any cap/batch-size combination.
 	perShardCap int
-	// throttledSince is when reservations last started failing on the global
-	// budget without one passing the global check in between (a reservation
-	// the shard's own cap then refuses still proves the global budget has
-	// room); zero while not throttled. lastPinWarn rate-limits the
-	// pinned-throttle warning.
-	throttledSince time.Time
-	lastPinWarn    time.Time
+	// pinClock tracks continuous throttling duration and rate-limits pinned
+	// warnings on the global budget.
+	pinClock pinClock
 }
 
 // trackedBatch is the settlement handle for one dispatched batch, returned by
@@ -142,15 +169,9 @@ type shardAckTracker struct {
 	// Together with the shard's entry in the batcher's reservedByShard it is
 	// bounded by perShardCap.
 	inflight int
-	// throttledSince is when reservations last started failing on this
-	// shard's in-flight cap without a successful reserve in between; zero
-	// while not throttled. A shard pinned at its cap never reaches the
-	// global-budget branch (fewer than four pinned shards leave the global
-	// budget under 100%), so the per-shard refusal needs its own pin clock
-	// or a wedged shard parks in silence. lastPinWarn rate-limits the
-	// warning.
-	throttledSince time.Time
-	lastPinWarn    time.Time
+	// pinClock tracks continuous throttling duration and rate-limits pinned
+	// warnings on this shard's in-flight cap.
+	pinClock pinClock
 }
 
 // NewRecordBatcher creates a new [RecordBatcher] for DynamoDB CDC.
@@ -199,10 +220,7 @@ func (b *RecordBatcher) TryReserve(shardID string, n int) bool {
 		// this check, so if in-flight messages never settle the input is
 		// stalled with no other signal.
 		now := time.Now()
-		if b.throttledSince.IsZero() {
-			b.throttledSince = now
-		} else if since := now.Sub(b.throttledSince); since >= throttlePinWarnAfter && now.Sub(b.lastPinWarn) >= throttlePinWarnInterval {
-			b.lastPinWarn = now
+		if since, shouldWarn := b.pinClock.check(now); shouldWarn {
 			b.log.Warnf("Shard readers throttled for %v: %d/%d in-flight messages are still awaiting downstream acknowledgement (top shards: %s); no records are being read while this persists",
 				since.Round(time.Second), b.trackedMessages, b.maxTrackedMessages, b.topInflightShardsLocked(3))
 		}
@@ -211,7 +229,7 @@ func (b *RecordBatcher) TryReserve(shardID string, n int) bool {
 	// The global budget can admit this reservation, so it is not pinned:
 	// clear its clock even if the shard's own cap refuses below, or the next
 	// global exhaustion would report a pin spanning the drained interval.
-	b.throttledSince = time.Time{}
+	b.pinClock.clear()
 
 	shardInFlight := b.shardInFlightLocked(shardID)
 	if shardInFlight > 0 && shardInFlight+n > b.perShardCap && b.otherShardActiveLocked(shardID) {
@@ -225,10 +243,7 @@ func (b *RecordBatcher) TryReserve(shardID string, n int) bool {
 		// guard is belt and braces.
 		if st := b.shards[shardID]; st != nil {
 			now := time.Now()
-			if st.throttledSince.IsZero() {
-				st.throttledSince = now
-			} else if since := now.Sub(st.throttledSince); since >= throttlePinWarnAfter && now.Sub(st.lastPinWarn) >= throttlePinWarnInterval {
-				st.lastPinWarn = now
+			if since, shouldWarn := st.pinClock.check(now); shouldWarn {
 				b.log.Warnf("Shard %s reader throttled for %v: %d/%d in-flight messages on this shard are still awaiting downstream acknowledgement; no records are being read from it while this persists",
 					shardID, since.Round(time.Second), shardInFlight, b.perShardCap)
 			}
@@ -237,7 +252,7 @@ func (b *RecordBatcher) TryReserve(shardID string, n int) bool {
 	}
 
 	if st := b.shards[shardID]; st != nil {
-		st.throttledSince = time.Time{}
+		st.pinClock.clear()
 	}
 	b.reservedByShard[shardID] += n
 	b.reserved += n

internal/impl/aws/dynamodb/batcher_test.go

@@ -593,6 +593,67 @@ func TestBatcherAckOverShardCapStillDrainsTracker(t *testing.T) {
 	assert.Equal(t, "00002", cp.get("shard-002"))
 }
 
+func TestPinClock(t *testing.T) {
+	now := time.Now()
+	var clock pinClock
+
+	tests := []struct {
+		name      string
+		at        time.Time
+		wantSince time.Duration
+		wantWarn  bool
+	}{
+		{
+			name:      "first check starts throttling",
+			at:        now,
+			wantSince: 0,
+			wantWarn:  false,
+		},
+		{
+			name:      "check before warn threshold returns duration without warning",
+			at:        now.Add(throttlePinWarnAfter - time.Second),
+			wantSince: throttlePinWarnAfter - time.Second,
+			wantWarn:  false,
+		},
+		{
+			name:      "check at warn threshold triggers warning",
+			at:        now.Add(throttlePinWarnAfter),
+			wantSince: throttlePinWarnAfter,
+			wantWarn:  true,
+		},
+		{
+			name:      "immediate check within interval is rate-limited",
+			at:        now.Add(throttlePinWarnAfter + time.Second),
+			wantSince: throttlePinWarnAfter + time.Second,
+			wantWarn:  false,
+		},
+		{
+			name:      "check after interval triggers next warning",
+			at:        now.Add(throttlePinWarnAfter + throttlePinWarnInterval),
+			wantSince: throttlePinWarnAfter + throttlePinWarnInterval,
+			wantWarn:  true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			since, shouldWarn := clock.check(tc.at)
+			assert.Equal(t, tc.wantSince, since)
+			assert.Equal(t, tc.wantWarn, shouldWarn)
+		})
+	}
+
+	t.Run("clear resets throttling", func(t *testing.T) {
+		clock.clear()
+
+		// Next check starts fresh at t2 instead of measuring from now
+		t2 := now.Add(1 * time.Hour)
+		since, warn := clock.check(t2)
+		assert.Equal(t, time.Duration(0), since)
+		assert.False(t, warn)
+	})
+}
+
 // TestBatcherWarnsWhenThrottlePinned: a throttle that never releases is an
 // otherwise-silent stall, so ShouldThrottle must surface it once the tracker
 // has been pinned past the warn threshold, and dropping below the threshold
@@ -613,12 +674,12 @@ func TestBatcherWarnsWhenThrottlePinned(t *testing.T) {
 
 	// Backdate the pin start beyond the warn threshold.
 	batcher.mu.Lock()
-	batcher.throttledSince = time.Now().Add(-2 * throttlePinWarnAfter)
+	batcher.pinClock.throttledSince = time.Now().Add(-2 * throttlePinWarnAfter)
 	batcher.mu.Unlock()
 
 	require.False(t, batcher.TryReserve("shard-003", 100))
 	batcher.mu.Lock()
-	warned := !batcher.lastPinWarn.IsZero()
+	warned := !batcher.pinClock.lastPinWarn.IsZero()
 	batcher.mu.Unlock()
 	assert.True(t, warned, "a continuously pinned budget must emit the stall warning")
 
@@ -627,7 +688,7 @@ func TestBatcherWarnsWhenThrottlePinned(t *testing.T) {
 	batcher.RemoveBatch(tbBatchb)
 	require.True(t, batcher.TryReserve("shard-003", 100))
 	batcher.mu.Lock()
-	reset := batcher.throttledSince.IsZero()
+	reset := batcher.pinClock.throttledSince.IsZero()
 	batcher.mu.Unlock()
 	assert.True(t, reset, "a successful reservation must reset the pin clock")
 }
@@ -653,13 +714,13 @@ func TestBatcherWarnsWhenShardPinned(t *testing.T) {
 	// Backdate the shard's pin start beyond the warn threshold.
 	batcher.mu.Lock()
 	st := batcher.shards["shard-001"]
-	st.throttledSince = time.Now().Add(-2 * throttlePinWarnAfter)
+	st.pinClock.throttledSince = time.Now().Add(-2 * throttlePinWarnAfter)
 	batcher.mu.Unlock()
 
 	require.False(t, batcher.TryReserve("shard-001", 100))
 	batcher.mu.Lock()
-	warned := !st.lastPinWarn.IsZero()
-	globalUntouched := batcher.throttledSince.IsZero()
+	warned := !st.pinClock.lastPinWarn.IsZero()
+	globalUntouched := batcher.pinClock.throttledSince.IsZero()
 	batcher.mu.Unlock()
 	assert.True(t, warned, "a continuously pinned shard must emit the stall warning")
 	assert.True(t, globalUntouched, "a per-shard refusal must not start the global pin clock")
@@ -669,7 +730,7 @@ func TestBatcherWarnsWhenShardPinned(t *testing.T) {
 	batcher.RemoveBatch(tb)
 	require.True(t, batcher.TryReserve("shard-001", 100))
 	batcher.mu.Lock()
-	reset := st.throttledSince.IsZero()
+	reset := st.pinClock.throttledSince.IsZero()
 	batcher.mu.Unlock()
 	assert.True(t, reset, "a successful reservation must reset the shard's pin clock")
 }
@@ -692,7 +753,7 @@ func TestBatcherPerShardRefusalClearsGlobalClock(t *testing.T) {
 	require.False(t, batcher.TryReserve("shard-005", 100),
 		"the global budget is exhausted")
 	batcher.mu.Lock()
-	pinned := !batcher.throttledSince.IsZero()
+	pinned := !batcher.pinClock.throttledSince.IsZero()
 	batcher.mu.Unlock()
 	require.True(t, pinned)
 
@@ -702,7 +763,7 @@ func TestBatcherPerShardRefusalClearsGlobalClock(t *testing.T) {
 	require.False(t, batcher.TryReserve("shard-001", 1000),
 		"shard-001 is still at its per-shard cap")
 	batcher.mu.Lock()
-	cleared := batcher.throttledSince.IsZero()
+	cleared := batcher.pinClock.throttledSince.IsZero()
 	batcher.mu.Unlock()
 	assert.True(t, cleared,
 		"passing the global check must clear the global pin clock even when the per-shard cap refuses")