From 9b42e883ebf62b65032bbfe6abf2303bd9de3702 Mon Sep 17 00:00:00 2001 From: "pj.vauthier" Date: Mon, 31 Aug 2026 14:24:18 -0400 Subject: [PATCH] aws_dynamodb_cdc: extract per shard in-flight message counting logic Extracted new functions (`shardInFlightLocked` and `ShardInFlightCount`) to make the main `TryReserve` logic easier to reason about. This also allows us to more easily unit test this counting logic in isolation. --- internal/impl/aws/dynamodb/batcher.go | 41 ++++++++++++++-------- internal/impl/aws/dynamodb/batcher_test.go | 9 ++++- 2 files changed, 35 insertions(+), 15 deletions(-) diff --git a/internal/impl/aws/dynamodb/batcher.go b/internal/impl/aws/dynamodb/batcher.go index 540448caba..9f17b05fd4 100644 --- a/internal/impl/aws/dynamodb/batcher.go +++ b/internal/impl/aws/dynamodb/batcher.go @@ -213,16 +213,8 @@ func (b *RecordBatcher) TryReserve(shardID string, n int) bool { // global exhaustion would report a pin spanning the drained interval. b.throttledSince = time.Time{} - // The shard-tracker entry may not exist yet - it is only materialised by - // AddMessages, so shards that never yield records don't count against - // maxTrackedShards - in which case the shard has nothing in flight. - st := b.shards[shardID] - inflight := 0 - if st != nil { - inflight = st.inflight - } - resv := b.reservedByShard[shardID] - if inflight+resv > 0 && inflight+resv+n > b.perShardCap && b.otherShardActiveLocked(shardID) { + shardInFlight := b.shardInFlightLocked(shardID) + if shardInFlight > 0 && shardInFlight+n > b.perShardCap && b.otherShardActiveLocked(shardID) { // The shard is pinned by its own unsettled messages; park it alone // without touching the global throttle clock, but surface a // continuous pin on the shard's own clock - in a topology with fewer @@ -231,23 +223,23 @@ func (b *RecordBatcher) TryReserve(shardID string, n int) bool { // shard always has tracked messages (its single reader never holds a // reservation while reserving again), so st is non-nil here; the // guard is belt and braces. - if st != nil { + 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 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), inflight+resv, b.perShardCap) + shardID, since.Round(time.Second), shardInFlight, b.perShardCap) } } return false } - if st != nil { + if st := b.shards[shardID]; st != nil { st.throttledSince = time.Time{} } - b.reservedByShard[shardID] = resv + n + b.reservedByShard[shardID] += n b.reserved += n return true } @@ -593,6 +585,27 @@ func (b *RecordBatcher) InFlightCount() int { return b.inFlightCountLocked() } +// shardInFlightLocked returns the in-flight (tracked and reserved) message count +// for a specific shard. Callers must hold b.mu. +func (b *RecordBatcher) shardInFlightLocked(shardID string) int { + inflight := 0 + if st := b.shards[shardID]; st != nil { + inflight = st.inflight + } + return inflight + b.reservedByShard[shardID] +} + +// ShardInFlightCount returns the in-flight (tracked and reserved) message count +// for a specific shard. Exported for testing and diagnostics. +func (b *RecordBatcher) ShardInFlightCount(shardID string) int { + if b == nil { + return 0 + } + b.mu.Lock() + defer b.mu.Unlock() + return b.shardInFlightLocked(shardID) +} + // TrackedMessageCount returns the number of tracked messages. Exported for testing. func (b *RecordBatcher) TrackedMessageCount() int { b.mu.Lock() diff --git a/internal/impl/aws/dynamodb/batcher_test.go b/internal/impl/aws/dynamodb/batcher_test.go index 1d13582fe1..397516f3f6 100644 --- a/internal/impl/aws/dynamodb/batcher_test.go +++ b/internal/impl/aws/dynamodb/batcher_test.go @@ -1003,29 +1003,36 @@ func TestBatcherSettleIsIdempotent(t *testing.T) { func TestBatcherInFlightCount(t *testing.T) { var nilBatcher *RecordBatcher assert.Equal(t, 0, nilBatcher.InFlightCount()) + assert.Equal(t, 0, nilBatcher.ShardInFlightCount("shard-001")) batcher := NewRecordBatcher(100, 100, service.MockResources().Logger()) assert.Equal(t, 0, batcher.InFlightCount()) + assert.Equal(t, 0, batcher.ShardInFlightCount("shard-001")) assert.Equal(t, 0, batcher.TrackedMessageCount()) - // Reserve 10 messages + // Reserve 10 messages on shard-001 require.True(t, batcher.TryReserve("shard-001", 10)) assert.Equal(t, 10, batcher.InFlightCount()) + assert.Equal(t, 10, batcher.ShardInFlightCount("shard-001")) + assert.Equal(t, 0, batcher.ShardInFlightCount("shard-002")) assert.Equal(t, 0, batcher.TrackedMessageCount()) // Add 6 messages (consumes 6 of the reservation) batch := createTestMessages(6, "shard-001", 1) tb := batcher.AddMessages(batch, "shard-001") assert.Equal(t, 10, batcher.InFlightCount()) // 4 reserved + 6 tracked + assert.Equal(t, 10, batcher.ShardInFlightCount("shard-001")) assert.Equal(t, 6, batcher.TrackedMessageCount()) // Release remaining 4 reserved messages batcher.Release("shard-001", 4) assert.Equal(t, 6, batcher.InFlightCount()) + assert.Equal(t, 6, batcher.ShardInFlightCount("shard-001")) assert.Equal(t, 6, batcher.TrackedMessageCount()) // Settle batch batcher.RemoveBatch(tb) assert.Equal(t, 0, batcher.InFlightCount()) + assert.Equal(t, 0, batcher.ShardInFlightCount("shard-001")) assert.Equal(t, 0, batcher.TrackedMessageCount()) }