diff --git a/internal/impl/aws/dynamodb/batcher.go b/internal/impl/aws/dynamodb/batcher.go index 33e35fa196..540448caba 100644 --- a/internal/impl/aws/dynamodb/batcher.go +++ b/internal/impl/aws/dynamodb/batcher.go @@ -187,11 +187,14 @@ func (b *RecordBatcher) TryReserve(shardID string, n int) bool { b.mu.Lock() defer b.mu.Unlock() - // A completely empty batcher admits one batch even above the global - // budget (mirroring the per-shard cap's idle escape): config validation - // bounds batch_size below the budget's floor, but a reservation that - // could never fit must hang the read loop under no circumstances. - if b.trackedMessages+b.reserved > 0 && b.trackedMessages+b.reserved+n > b.maxTrackedMessages { + // If the batcher has zero messages in flight, always allow the first reservation + // to go through, even if it exceeds the maximum budget. + // + // Why? Even though configuration validation normally ensures batch sizes are smaller + // than the budget, a batch that is larger than the limit must never cause the reader + // loop to hang forever waiting for space that can never clear. + globalInFlight := b.inFlightCountLocked() + if globalInFlight > 0 && globalInFlight+n > b.maxTrackedMessages { // Surface a continuously pinned budget: every shard reader parks on // this check, so if in-flight messages never settle the input is // stalled with no other signal. @@ -563,6 +566,33 @@ func (b *RecordBatcher) PendingCount(shardID string) int { return 0 } +// inFlightCountLocked returns the total message load across both stages of +// ingestion: +// +// 1. Reserved (b.reserved): Messages currently being fetched over the +// network (GetRecords in flight). +// 2. Tracked (b.trackedMessages): Messages received into memory and +// awaiting downstream processing/acknowledgement. +// +// Counting pending reads as in-flight ensures concurrent shard readers cannot +// flood the pipeline before the first batch of records even arrives. +// +// Callers must hold b.mu. +func (b *RecordBatcher) inFlightCountLocked() int { + return b.trackedMessages + b.reserved +} + +// InFlightCount returns the total number of in-flight (tracked and reserved) +// messages. Exported for testing and diagnostics. +func (b *RecordBatcher) InFlightCount() int { + if b == nil { + return 0 + } + b.mu.Lock() + defer b.mu.Unlock() + return b.inFlightCountLocked() +} + // 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 1974b0d0c1..1d13582fe1 100644 --- a/internal/impl/aws/dynamodb/batcher_test.go +++ b/internal/impl/aws/dynamodb/batcher_test.go @@ -999,3 +999,33 @@ func TestBatcherSettleIsIdempotent(t *testing.T) { batcher.AddMessages(other, "shard-001") assert.Equal(t, 2, batcher.TrackedMessageCount()) } + +func TestBatcherInFlightCount(t *testing.T) { + var nilBatcher *RecordBatcher + assert.Equal(t, 0, nilBatcher.InFlightCount()) + + batcher := NewRecordBatcher(100, 100, service.MockResources().Logger()) + assert.Equal(t, 0, batcher.InFlightCount()) + assert.Equal(t, 0, batcher.TrackedMessageCount()) + + // Reserve 10 messages + require.True(t, batcher.TryReserve("shard-001", 10)) + assert.Equal(t, 10, batcher.InFlightCount()) + 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, 6, batcher.TrackedMessageCount()) + + // Release remaining 4 reserved messages + batcher.Release("shard-001", 4) + assert.Equal(t, 6, batcher.InFlightCount()) + assert.Equal(t, 6, batcher.TrackedMessageCount()) + + // Settle batch + batcher.RemoveBatch(tb) + assert.Equal(t, 0, batcher.InFlightCount()) + assert.Equal(t, 0, batcher.TrackedMessageCount()) +}