Fix DynamoDB Streams checkpointing - preserve iterator position across shard re-discoveries

Track shard iterators by shard ID in a dictionary instead of a positional list.
On shard re-discovery, only create new iterators for newly discovered shards -
existing shards keep their current NextShardIterator position. This prevents
re-processing of already-invoked records when shards are re-discovered after
a shard closes or after empty poll thresholds.

Author: normj

Tools/LambdaTestTool-v2/src/Amazon.Lambda.TestTool/Processes/DynamoDBStreamsEventSource/DynamoDBStreamsEventSourceBackgroundService.cs

@@ -100,31 +100,32 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken)
 
     private async Task PollStream(string streamArn, CancellationToken stoppingToken)
     {
-        var shardIterators = await GetShardIterators(streamArn, stoppingToken);
+        // Track iterators by shard ID to preserve position across re-discoveries
+        var shardIterators = await GetShardIterators(streamArn, null, stoppingToken);
         var emptyPollCount = 0;
 
         while (!stoppingToken.IsCancellationRequested)
         {
             // Poll all shards concurrently
-            var tasks = new List<Task<(int Index, GetRecordsResponse? Response)>>();
-            for (int i = 0; i < shardIterators.Count; i++)
+            var shardIds = shardIterators.Keys.ToList();
+            var tasks = new List<Task<(string ShardId, GetRecordsResponse? Response)>>();
+            foreach (var shardId in shardIds)
             {
-                if (shardIterators[i] == null)
+                var iterator = shardIterators[shardId];
+                if (iterator == null)
                     continue;
 
-                var index = i;
-                var iterator = shardIterators[i]!;
-                tasks.Add(PollShard(index, iterator, stoppingToken));
+                tasks.Add(PollShard(shardId, iterator, stoppingToken));
             }
 
             var activeCount = tasks.Count;
             _logger.LogInformation("Polling {activeShardCount} active shard(s) out of {totalCount} total", activeCount, shardIterators.Count);
 
             if (activeCount == 0)
             {
-                // All iterators exhausted — re-discover shards
+                // All iterators exhausted — re-discover shards but preserve existing positions
                 _logger.LogInformation("All shard iterators exhausted, re-discovering shards");
-                shardIterators = await GetShardIterators(streamArn, stoppingToken);
+                shardIterators = await GetShardIterators(streamArn, shardIterators, stoppingToken);
                 emptyPollCount = 0;
                 if (shardIterators.Count == 0)
                 {
@@ -137,27 +138,27 @@ private async Task PollStream(string streamArn, CancellationToken stoppingToken)
             var results = await Task.WhenAll(tasks);
 
             var hasRecords = false;
-            var exhaustedInThisPoll = 0;
-            foreach (var (index, response) in results)
+            foreach (var (shardId, response) in results)
             {
                 if (response == null)
                     continue;
 
-                // Log when a shard iterator becomes null (shard closed)
                 if (response.NextShardIterator == null)
                 {
-                    exhaustedInThisPoll++;
-                    _logger.LogInformation("Shard at index {index} has been closed (NextShardIterator is null), records in final batch: {count}",
-                        index, response.Records.Count);
+                    _logger.LogInformation("Shard {shardId} has been closed (NextShardIterator is null), records in final batch: {count}",
+                        shardId, response.Records.Count);
+                    shardIterators.Remove(shardId);
+                }
+                else
+                {
+                    shardIterators[shardId] = response.NextShardIterator;
                 }
-
-                shardIterators[index] = response.NextShardIterator;
 
                 if (response.Records.Count == 0)
                     continue;
 
                 hasRecords = true;
-                _logger.LogInformation("Retrieved {recordCount} record(s) from shard index {index}", response.Records.Count, index);
+                _logger.LogInformation("Retrieved {recordCount} record(s) from shard {shardId}", response.Records.Count, shardId);
                 var lambdaRecords = ConvertToLambdaRecords(response.Records, streamArn);
 
                 var lambdaPayload = new DynamoDBEvent
@@ -184,14 +185,11 @@ private async Task PollStream(string streamArn, CancellationToken stoppingToken)
                 }
             }
 
-            // Remove exhausted shards (null iterators) and re-discover to pick up child shards
-            var exhaustedCount = shardIterators.Count(s => s == null);
-            if (exhaustedCount > 0)
+            // If any shards were removed (closed), re-discover to pick up child shards
+            if (results.Any(r => r.Response?.NextShardIterator == null))
             {
-                _logger.LogInformation("Removing {exhaustedCount} exhausted shard(s), re-discovering to find child shards", exhaustedCount);
-                // Re-discover shards immediately when any shard closes, since new records
-                // will be on child shards that we don't have iterators for yet.
-                shardIterators = await GetShardIterators(streamArn, stoppingToken);
+                _logger.LogInformation("Closed shard(s) detected, re-discovering to find child shards");
+                shardIterators = await GetShardIterators(streamArn, shardIterators, stoppingToken);
                 emptyPollCount = 0;
                 continue;
             }
@@ -203,11 +201,10 @@ private async Task PollStream(string streamArn, CancellationToken stoppingToken)
             else
             {
                 emptyPollCount++;
-                // After many empty polls, re-discover shards in case the stream topology changed
                 if (emptyPollCount >= 30)
                 {
                     _logger.LogInformation("No records after {count} consecutive polls, re-discovering shards", emptyPollCount);
-                    shardIterators = await GetShardIterators(streamArn, stoppingToken);
+                    shardIterators = await GetShardIterators(streamArn, shardIterators, stoppingToken);
                     emptyPollCount = 0;
                     continue;
                 }
@@ -218,24 +215,27 @@ private async Task PollStream(string streamArn, CancellationToken stoppingToken)
         }
     }
 
-    private async Task<(int Index, GetRecordsResponse? Response)> PollShard(int index, string iterator, CancellationToken stoppingToken)
+    private async Task<(string ShardId, GetRecordsResponse? Response)> PollShard(string shardId, string iterator, CancellationToken stoppingToken)
     {
         var response = await _streamsClient.GetRecordsAsync(new GetRecordsRequest
         {
             ShardIterator = iterator,
             Limit = _config.BatchSize
         }, stoppingToken);
 
-        return (index, response);
+        return (shardId, response);
     }
 
-    private async Task<List<string?>> GetShardIterators(string streamArn, CancellationToken stoppingToken)
+    /// <summary>
+    /// Discover open shards and get iterators. If existingIterators is provided, only creates new iterators
+    /// for shards not already tracked — preserving the stream position for known shards.
+    /// </summary>
+    private async Task<Dictionary<string, string?>> GetShardIterators(string streamArn, Dictionary<string, string?>? existingIterators, CancellationToken stoppingToken)
     {
         _logger.LogInformation("Discovering shards for stream {streamArn}", streamArn);
         var shards = new List<Shard>();
         string? lastEvaluatedShardId = null;
 
-        // Paginate through all shards
         do
         {
             var describeResponse = await _streamsClient.DescribeStreamAsync(new DescribeStreamRequest
@@ -250,24 +250,29 @@ private async Task PollStream(string streamArn, CancellationToken stoppingToken)
 
         _logger.LogInformation("Discovered {shardCount} total shard(s) for stream", shards.Count);
 
-        // Only get iterators for open (leaf) shards — shards without an EndingSequenceNumber.
-        // Closed shards with LATEST iterator type will never return new records.
         var openShards = shards.Where(s => s.SequenceNumberRange?.EndingSequenceNumber == null).ToList();
         _logger.LogInformation("Filtered to {openCount} open (leaf) shard(s) out of {totalCount} total", openShards.Count, shards.Count);
 
-        var iterators = new List<string?>();
+        var iterators = new Dictionary<string, string?>();
 
         foreach (var shard in openShards)
         {
+            // Preserve existing iterator position for shards we're already tracking
+            if (existingIterators != null && existingIterators.TryGetValue(shard.ShardId, out var existingIterator) && existingIterator != null)
+            {
+                iterators[shard.ShardId] = existingIterator;
+                continue;
+            }
+
             var iteratorResponse = await _streamsClient.GetShardIteratorAsync(new GetShardIteratorRequest
             {
                 StreamArn = streamArn,
                 ShardId = shard.ShardId,
                 ShardIteratorType = new ShardIteratorType(_config.ShardIteratorType)
             }, stoppingToken);
 
-            _logger.LogInformation("Got iterator for shard {shardId}", shard.ShardId);
-            iterators.Add(iteratorResponse.ShardIterator);
+            _logger.LogInformation("Got new iterator for shard {shardId}", shard.ShardId);
+            iterators[shard.ShardId] = iteratorResponse.ShardIterator;
         }
 
         return iterators;