From 73662318a522b095991d3db0fc4970c6b658bea4 Mon Sep 17 00:00:00 2001 From: "pj.vauthier" Date: Mon, 31 Aug 2026 14:06:56 -0400 Subject: [PATCH 1/4] aws_dynamodb_cdc: simplify the counting of in-flight messages This makes it easier for a human to reason about the code. It is functionally identical. The code comment has been reworder to be easier to read and understand. --- internal/impl/aws/dynamodb/batcher.go | 37 +++++++++++++++++++--- internal/impl/aws/dynamodb/batcher_test.go | 30 ++++++++++++++++++ 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/internal/impl/aws/dynamodb/batcher.go b/internal/impl/aws/dynamodb/batcher.go index 33e35fa196..ec27adf003 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 maxium 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. + inFlight := b.inFlightCountLocked() + if inFlight > 0 && inFlight+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,30 @@ 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()) +} From ac04891e9e2be7aa257ae05c357d22613f40a317 Mon Sep 17 00:00:00 2001 From: "pj.vauthier" Date: Mon, 31 Aug 2026 14:25:25 -0400 Subject: [PATCH 2/4] aws_dynamodb_batcher: fix typo in code comment --- internal/impl/aws/dynamodb/batcher.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/internal/impl/aws/dynamodb/batcher.go b/internal/impl/aws/dynamodb/batcher.go index ec27adf003..7d722b4a63 100644 --- a/internal/impl/aws/dynamodb/batcher.go +++ b/internal/impl/aws/dynamodb/batcher.go @@ -188,7 +188,7 @@ func (b *RecordBatcher) TryReserve(shardID string, n int) bool { defer b.mu.Unlock() // If the batcher has zero messages in flight, always allow the first reservation - // to go through, even if it exceeds the maxium budget. + // 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 From ada60a3e9d925cb87415cbe7177d5b5e8373771f Mon Sep 17 00:00:00 2001 From: "pj.vauthier" Date: Mon, 31 Aug 2026 14:29:30 -0400 Subject: [PATCH 3/4] aws_dynamodb_batcher: reformat comment (line length limit) --- internal/impl/aws/dynamodb/batcher.go | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/internal/impl/aws/dynamodb/batcher.go b/internal/impl/aws/dynamodb/batcher.go index 7d722b4a63..6efe5baca9 100644 --- a/internal/impl/aws/dynamodb/batcher.go +++ b/internal/impl/aws/dynamodb/batcher.go @@ -566,10 +566,13 @@ func (b *RecordBatcher) PendingCount(shardID string) int { return 0 } -// inFlightCountLocked returns the total message load across both stages of ingestion: +// 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. +// 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. From 8729c909968983d13d8f7919a33f1d3ce88e3555 Mon Sep 17 00:00:00 2001 From: "pj.vauthier" Date: Mon, 31 Aug 2026 14:53:42 -0400 Subject: [PATCH 4/4] aws_dynamodb_batcher: rename variable to `globalInFlight` for clarity Improves readability by providing a more descriptive variable name. --- internal/impl/aws/dynamodb/batcher.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/internal/impl/aws/dynamodb/batcher.go b/internal/impl/aws/dynamodb/batcher.go index 6efe5baca9..540448caba 100644 --- a/internal/impl/aws/dynamodb/batcher.go +++ b/internal/impl/aws/dynamodb/batcher.go @@ -193,8 +193,8 @@ func (b *RecordBatcher) TryReserve(shardID string, n int) bool { // 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. - inFlight := b.inFlightCountLocked() - if inFlight > 0 && inFlight+n > b.maxTrackedMessages { + 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.