Merge branch 'normj/lambda-ddbstreams' of https://github.com/aws/aws-lambda-dotnet into normj/lambda-ddbstreams

Author: normj

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

@@ -100,36 +100,50 @@ protected override async Task ExecuteAsync(CancellationToken stoppingToken)
 
     private async Task PollStream(string streamArn, CancellationToken stoppingToken)
     {
-        // Track iterators by shard ID to preserve position across re-discoveries
-        var shardIterators = await GetShardIterators(streamArn, null, stoppingToken);
-        var emptyPollCount = 0;
+        // Shard polling strategy:
+        //
+        // Goal: Only deliver records to Lambda that were written AFTER the test tool started.
+        //
+        // 1. At startup, discover all shards. Open shards get a LATEST iterator (future records only).
+        //    Closed shards are recorded in a "closed at startup" set and never polled — they contain
+        //    only historical data from before the tool started.
+        //
+        // 2. Every 30 seconds (or immediately when a shard is exhausted), re-discover shards:
+        //    - Shards already being polled: leave their iterator alone (preserves position).
+        //    - Shards in the "closed at startup" set: skip (pre-existing historical data).
+        //    - Any other shard (new since startup): poll with TRIM_HORIZON to read all its records,
+        //      since the shard was created after the tool started and all its data is relevant.
+
+        var closedAtStartup = new HashSet<string>();
+        var shardIterators = await DiscoverInitialShards(streamArn, closedAtStartup, stoppingToken);
+
+        _logger.LogInformation("Initial discovery: {openCount} open shard(s), {closedCount} closed shard(s) at startup",
+            shardIterators.Count, closedAtStartup.Count);
+
+        var lastDiscoveryTime = DateTime.UtcNow;
+        const int ShardRediscoveryIntervalSeconds = 30;
 
         while (!stoppingToken.IsCancellationRequested)
         {
-            // Poll all shards concurrently
-            var shardIds = shardIterators.Keys.ToList();
+            // Poll all active shards concurrently
             var tasks = new List<Task<(string ShardId, GetRecordsResponse? Response)>>();
-            foreach (var shardId in shardIds)
+            foreach (var (shardId, iterator) in shardIterators)
             {
-                var iterator = shardIterators[shardId];
                 if (iterator == null)
                     continue;
-
                 tasks.Add(PollShard(shardId, iterator, stoppingToken));
             }
 
             var activeCount = tasks.Count;
-            _logger.LogInformation("Polling {activeShardCount} active shard(s) out of {totalCount} total", activeCount, shardIterators.Count);
+            _logger.LogInformation("Polling {activeShardCount} active shard(s)", activeCount);
 
             if (activeCount == 0)
             {
-                // All iterators exhausted — re-discover shards but preserve existing positions
-                _logger.LogInformation("All shard iterators exhausted, re-discovering shards");
-                shardIterators = await GetShardIterators(streamArn, shardIterators, stoppingToken);
-                emptyPollCount = 0;
+                // No active shards — re-discover
+                shardIterators = await DiscoverNewShards(streamArn, shardIterators, closedAtStartup, stoppingToken);
+                lastDiscoveryTime = DateTime.UtcNow;
                 if (shardIterators.Count == 0)
                 {
-                    _logger.LogInformation("No shards found, sleeping 1000ms before retry");
                     await Task.Delay(1000, stoppingToken);
                 }
                 continue;
@@ -138,16 +152,18 @@ private async Task PollStream(string streamArn, CancellationToken stoppingToken)
             var results = await Task.WhenAll(tasks);
 
             var hasRecords = false;
+            var shardExhausted = false;
             foreach (var (shardId, response) in results)
             {
                 if (response == null)
                     continue;
 
                 if (response.NextShardIterator == null)
                 {
-                    _logger.LogInformation("Shard {shardId} has been closed (NextShardIterator is null), records in final batch: {count}",
+                    _logger.LogInformation("Shard {shardId} exhausted (closed), records in final batch: {count}",
                         shardId, response.Records.Count);
                     shardIterators.Remove(shardId);
+                    shardExhausted = true;
                 }
                 else
                 {
@@ -161,11 +177,7 @@ private async Task PollStream(string streamArn, CancellationToken stoppingToken)
                 _logger.LogInformation("Retrieved {recordCount} record(s) from shard {shardId}", response.Records.Count, shardId);
                 var lambdaRecords = ConvertToLambdaRecords(response.Records, streamArn);
 
-                var lambdaPayload = new DynamoDBEvent
-                {
-                    Records = lambdaRecords
-                };
-
+                var lambdaPayload = new DynamoDBEvent { Records = lambdaRecords };
                 var invokeRequest = new InvokeRequest
                 {
                     InvocationType = InvocationType.RequestResponse,
@@ -185,31 +197,19 @@ private async Task PollStream(string streamArn, CancellationToken stoppingToken)
                 }
             }
 
-            // If any shards were removed (closed), re-discover to pick up child shards
-            if (results.Any(r => r.Response?.NextShardIterator == null))
+            // Re-discover if a shard was exhausted or 30 seconds have elapsed
+            var timeSinceDiscovery = (DateTime.UtcNow - lastDiscoveryTime).TotalSeconds;
+            if (shardExhausted || timeSinceDiscovery >= ShardRediscoveryIntervalSeconds)
             {
-                _logger.LogInformation("Closed shard(s) detected, re-discovering to find child shards");
-                shardIterators = await GetShardIterators(streamArn, shardIterators, stoppingToken);
-                emptyPollCount = 0;
+                _logger.LogInformation("Re-discovering shards (exhausted={shardExhausted}, elapsed={elapsed}s)",
+                    shardExhausted, (int)timeSinceDiscovery);
+                shardIterators = await DiscoverNewShards(streamArn, shardIterators, closedAtStartup, stoppingToken);
+                lastDiscoveryTime = DateTime.UtcNow;
                 continue;
             }
 
-            if (hasRecords)
+            if (!hasRecords)
             {
-                emptyPollCount = 0;
-            }
-            else
-            {
-                emptyPollCount++;
-                if (emptyPollCount >= 30)
-                {
-                    _logger.LogInformation("No records after {count} consecutive polls, re-discovering shards", emptyPollCount);
-                    shardIterators = await GetShardIterators(streamArn, shardIterators, stoppingToken);
-                    emptyPollCount = 0;
-                    continue;
-                }
-
-                _logger.LogInformation("No records found (empty poll #{count}), sleeping {pollingInterval}ms", emptyPollCount, _config.PollingIntervalMs);
                 await Task.Delay(_config.PollingIntervalMs, stoppingToken);
             }
         }
@@ -227,57 +227,93 @@ private async Task PollStream(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.
+    /// Initial shard discovery at startup. Uses LATEST for open shards and records closed shard IDs.
     /// </summary>
-    private async Task<Dictionary<string, string?>> GetShardIterators(string streamArn, Dictionary<string, string?>? existingIterators, CancellationToken stoppingToken)
+    private async Task<Dictionary<string, string?>> DiscoverInitialShards(string streamArn, HashSet<string> closedAtStartup, CancellationToken stoppingToken)
     {
-        _logger.LogInformation("Discovering shards for stream {streamArn}", streamArn);
-        var shards = new List<Shard>();
-        string? lastEvaluatedShardId = null;
+        var shards = await GetAllShards(streamArn, stoppingToken);
+        var iterators = new Dictionary<string, string?>();
 
-        do
+        foreach (var shard in shards)
         {
-            var describeResponse = await _streamsClient.DescribeStreamAsync(new DescribeStreamRequest
+            var isClosed = shard.SequenceNumberRange?.EndingSequenceNumber != null;
+            if (isClosed)
+            {
+                closedAtStartup.Add(shard.ShardId);
+                continue;
+            }
+
+            // Open shard — use LATEST to only get records created after startup
+            var iteratorResponse = await _streamsClient.GetShardIteratorAsync(new GetShardIteratorRequest
             {
                 StreamArn = streamArn,
-                ExclusiveStartShardId = lastEvaluatedShardId
+                ShardId = shard.ShardId,
+                ShardIteratorType = ShardIteratorType.LATEST
             }, stoppingToken);
 
-            shards.AddRange(describeResponse.StreamDescription.Shards);
-            lastEvaluatedShardId = describeResponse.StreamDescription.LastEvaluatedShardId;
-        } while (lastEvaluatedShardId != null);
-
-        _logger.LogInformation("Discovered {shardCount} total shard(s) for stream", shards.Count);
+            _logger.LogInformation("Got LATEST iterator for startup shard {shardId}", shard.ShardId);
+            iterators[shard.ShardId] = iteratorResponse.ShardIterator;
+        }
 
-        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);
+        return iterators;
+    }
 
-        var iterators = new Dictionary<string, string?>();
+    /// <summary>
+    /// Ongoing shard discovery. Preserves existing iterators, skips shards closed at startup,
+    /// and starts TRIM_HORIZON pollers for any new shards (even if closed).
+    /// </summary>
+    private async Task<Dictionary<string, string?>> DiscoverNewShards(string streamArn, Dictionary<string, string?> existingIterators, HashSet<string> closedAtStartup, CancellationToken stoppingToken)
+    {
+        var shards = await GetAllShards(streamArn, stoppingToken);
+        var iterators = new Dictionary<string, string?>(existingIterators);
 
-        foreach (var shard in openShards)
+        foreach (var shard in shards)
         {
-            // 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;
+            // Already being polled — leave iterator alone
+            if (iterators.ContainsKey(shard.ShardId))
                 continue;
-            }
 
+            // Was closed at startup — skip
+            if (closedAtStartup.Contains(shard.ShardId))
+                continue;
+
+            // New shard discovered after startup — use TRIM_HORIZON to read all its records
             var iteratorResponse = await _streamsClient.GetShardIteratorAsync(new GetShardIteratorRequest
             {
                 StreamArn = streamArn,
                 ShardId = shard.ShardId,
-                ShardIteratorType = new ShardIteratorType(_config.ShardIteratorType)
+                ShardIteratorType = ShardIteratorType.TRIM_HORIZON
             }, stoppingToken);
 
-            _logger.LogInformation("Got new iterator for shard {shardId}", shard.ShardId);
+            _logger.LogInformation("Got TRIM_HORIZON iterator for new shard {shardId}", shard.ShardId);
             iterators[shard.ShardId] = iteratorResponse.ShardIterator;
         }
 
         return iterators;
     }
 
+    private async Task<List<Shard>> GetAllShards(string streamArn, CancellationToken stoppingToken)
+    {
+        _logger.LogInformation("Discovering shards for stream {streamArn}", streamArn);
+        var shards = new List<Shard>();
+        string? lastEvaluatedShardId = null;
+
+        do
+        {
+            var describeResponse = await _streamsClient.DescribeStreamAsync(new DescribeStreamRequest
+            {
+                StreamArn = streamArn,
+                ExclusiveStartShardId = lastEvaluatedShardId
+            }, stoppingToken);
+
+            shards.AddRange(describeResponse.StreamDescription.Shards);
+            lastEvaluatedShardId = describeResponse.StreamDescription.LastEvaluatedShardId;
+        } while (lastEvaluatedShardId != null);
+
+        _logger.LogInformation("Discovered {shardCount} shard(s)", shards.Count);
+        return shards;
+    }
+
     /// <summary>
     /// Convert from the SDK's DynamoDB Streams records to the Lambda event's DynamoDB record type.
     /// </summary>

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

@@ -28,11 +28,6 @@ public class DynamoDBStreamsEventSourceBackgroundServiceConfig
     /// </summary>
     public required string TableName { get; init; }
 
-    /// <summary>
-    /// The shard iterator type to use when reading from the stream.
-    /// </summary>
-    public required string ShardIteratorType { get; init; } = "LATEST";
-
     /// <summary>
     /// The polling interval in milliseconds between stream reads when no records are found.
     /// </summary>

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

@@ -40,12 +40,6 @@ internal class DynamoDBStreamsEventSourceConfig
     /// </summary>
     public string? TableName { get; set; }
 
-    /// <summary>
-    /// The shard iterator type to use when reading from the stream.
-    /// Valid values: LATEST, TRIM_HORIZON. Default is TRIM_HORIZON.
-    /// </summary>
-    public string? ShardIteratorType { get; set; }
-
     /// <summary>
     /// The polling interval in milliseconds between stream reads when no records are found.
     /// Default is 1000.

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

@@ -92,7 +92,6 @@ public static DynamoDBStreamsEventSourceProcess Startup(RunCommandSettings setti
                 FunctionName = config.FunctionName ?? LambdaRuntimeApi.DefaultFunctionName,
                 LambdaRuntimeApi = lambdaRuntimeApi,
                 TableName = tableName,
-                ShardIteratorType = config.ShardIteratorType ?? "TRIM_HORIZON",
                 PollingIntervalMs = config.PollingIntervalMs ?? DefaultPollingIntervalMs
             };
 
@@ -201,9 +200,6 @@ internal static List<DynamoDBStreamsEventSourceConfig> LoadDynamoDBStreamsEventS
                     case "tablename":
                         config.TableName = keyValuePair[1].Trim();
                         break;
-                    case "sharditeratortype":
-                        config.ShardIteratorType = keyValuePair[1].Trim();
-                        break;
                     case "pollingintervalms":
                         if (!int.TryParse(keyValuePair[1].Trim(), out var pollingInterval))
                         {