Fix batcher logic to never create empty batches (#2793)

* Reproduce the batcher bug

* Verify there is something to batch

* Fix batcher logic

---------

Co-authored-by: Daniel Marbach <danielmarbach@users.noreply.github.com>

Author: danielmarbach

src/NServiceBus.Transport.SQS.Tests/SqsPreparedMessageBatcherTests.cs

@@ -90,6 +90,44 @@ public void SingleBatchForLessOrEqual10Entries()
             });
         }
 
+        [Test]
+        public void SingleBatchEvenWhenSingleMessageNotFittingIntoBatchDueToMessageSize()
+        {
+            var preparedMessages = new[]
+            {
+                new SqsPreparedMessage {MessageId = Guid.NewGuid().ToString(), Destination = "destination1", QueueUrl = "https://destination1", Body = GenerateBody(256)},
+            };
+            PrecalculateSize(preparedMessages);
+
+            var batches = SqsPreparedMessageBatcher.Batch(preparedMessages);
+
+            Assert.Multiple(() =>
+            {
+                Assert.That(batches, Has.Count.EqualTo(1));
+                Assert.That(batches.ElementAt(0).BatchRequest.Entries, Has.Count.EqualTo(1));
+            });
+        }
+
+        [Test]
+        public void BatchPerMessageNotFittingIntoBatchDueToMessageSize()
+        {
+            var preparedMessages = new[]
+            {
+                new SqsPreparedMessage {MessageId = Guid.NewGuid().ToString(), Destination = "destination1", QueueUrl = "https://destination1", Body = GenerateBody(256)},
+                new SqsPreparedMessage {MessageId = Guid.NewGuid().ToString(), Destination = "destination1", QueueUrl = "https://destination1", Body = GenerateBody(256)},
+            };
+            PrecalculateSize(preparedMessages);
+
+            var batches = SqsPreparedMessageBatcher.Batch(preparedMessages);
+
+            Assert.Multiple(() =>
+            {
+                Assert.That(batches, Has.Count.EqualTo(2));
+                Assert.That(batches.ElementAt(0).BatchRequest.Entries, Has.Count.EqualTo(1));
+                Assert.That(batches.ElementAt(1).BatchRequest.Entries, Has.Count.EqualTo(1));
+            });
+        }
+
         [Test]
         public void MultipleBatchesForGreaterThan10Entries()
         {
@@ -116,7 +154,7 @@ public void MultipleBatchesForGreaterThan10Entries()
 
             Assert.Multiple(() =>
             {
-                Assert.That(batches.Count(), Is.EqualTo(2));
+                Assert.That(batches, Has.Count.EqualTo(2));
                 Assert.That(batches.ElementAt(0).BatchRequest.Entries, Has.Count.EqualTo(10));
                 Assert.That(batches.ElementAt(1).BatchRequest.Entries, Has.Count.EqualTo(3));
             });
@@ -182,7 +220,7 @@ public void MultipleBatchesForMessagesNotFittingIntoBatchDueToMessageSize()
 
             Assert.Multiple(() =>
             {
-                Assert.That(batches.Count(), Is.EqualTo(7));
+                Assert.That(batches, Has.Count.EqualTo(7));
                 Assert.That(batches.ElementAt(0).BatchRequest.Entries, Has.Count.EqualTo(1));
                 Assert.That(batches.ElementAt(1).BatchRequest.Entries, Has.Count.EqualTo(1));
                 Assert.That(batches.ElementAt(2).BatchRequest.Entries, Has.Count.EqualTo(4));
@@ -253,7 +291,7 @@ public void MultipleBatchesForMessagesWithMessageIdNotFittingIntoBatchDueToMessa
 
             Assert.Multiple(() =>
             {
-                Assert.That(batches.Count(), Is.EqualTo(7));
+                Assert.That(batches, Has.Count.EqualTo(7));
                 Assert.That(batches.ElementAt(0).BatchRequest.Entries, Has.Count.EqualTo(1));
                 Assert.That(batches.ElementAt(1).BatchRequest.Entries, Has.Count.EqualTo(1));
                 Assert.That(batches.ElementAt(2).BatchRequest.Entries, Has.Count.EqualTo(4));
@@ -304,7 +342,7 @@ public void BatchPerDestination_MultipleBatchesForGreaterThan10Entries()
 
             Assert.Multiple(() =>
             {
-                Assert.That(batches.Count(), Is.EqualTo(4));
+                Assert.That(batches, Has.Count.EqualTo(4));
                 Assert.That(batches.ElementAt(0).BatchRequest.Entries, Has.Count.EqualTo(10));
                 Assert.That(batches.ElementAt(1).BatchRequest.Entries, Has.Count.EqualTo(3));
                 Assert.That(batches.ElementAt(2).BatchRequest.Entries, Has.Count.EqualTo(10));

src/NServiceBus.Transport.SQS/SqsPreparedMessageBatcher.cs

@@ -11,47 +11,56 @@ static class SqsPreparedMessageBatcher
         public static IReadOnlyList<SqsBatchEntry> Batch(IEnumerable<SqsPreparedMessage> preparedMessages)
         {
             var allBatches = new List<SqsBatchEntry>();
-            var currentDestinationBatches = new Dictionary<string, SqsPreparedMessage>(TransportConstraints.MaximumItemsInBatch);
+            var currentBatch = new Dictionary<string, SqsPreparedMessage>(TransportConstraints.MaximumItemsInBatch);
 
+            // Group messages by destination to ensure batches only contain messages for the same queue
             var groupByDestination = preparedMessages.GroupBy(m => m.QueueUrl, StringComparer.Ordinal);
-            foreach (var group in groupByDestination)
+
+            foreach (var destinationGroup in groupByDestination)
             {
-                SqsPreparedMessage? firstMessage = null;
-                var payloadSize = 0L;
-                foreach (var message in group)
+                SqsPreparedMessage? referenceMessage = null;
+                var currentBatchSize = 0L;
+
+                foreach (var message in destinationGroup)
                 {
-                    firstMessage ??= message;
+                    referenceMessage ??= message;
+                    var messageSize = message.Size; // Size calculation is assumed to be done previously
 
-                    // Assumes the size was already calculated by the dispatcher
-                    var size = message.Size;
-                    payloadSize += size;
+                    // Check if this message would push the batch over the size limit
+                    if (currentBatchSize + messageSize > TransportConstraints.MaximumMessageSize)
+                    {
+                        // Finalize current batch if it has any messages
+                        if (currentBatch.Count > 0)
+                        {
+                            allBatches.Add(referenceMessage.ToBatchRequest(currentBatch));
+                            currentBatch.Clear();
+                        }
 
-                    if (payloadSize > TransportConstraints.MaximumMessageSize)
+                        currentBatchSize = messageSize;
+                    }
+                    else
                     {
-                        allBatches.Add(message.ToBatchRequest(currentDestinationBatches));
-                        currentDestinationBatches.Clear();
-                        payloadSize = size;
+                        // Message will fit within the current batch
+                        currentBatchSize += messageSize;
                     }
 
-                    // we don't have to recheck payload size here because the support layer checks that a request can always fit 256 KB size limit
-                    // we can't take MessageId because batch request ID can only contain alphanumeric characters, hyphen and underscores, message id could be overloaded
-                    currentDestinationBatches.Add(Guid.NewGuid().ToString(), message);
+                    // Add message to the current batch with a unique batch ID
+                    currentBatch.Add(Guid.NewGuid().ToString(), message);
 
-                    var currentCount = currentDestinationBatches.Count;
-                    if (currentCount != TransportConstraints.MaximumItemsInBatch)
+                    // Check if we've reached the maximum items per batch
+                    if (currentBatch.Count == TransportConstraints.MaximumItemsInBatch)
                     {
-                        continue;
+                        allBatches.Add(message.ToBatchRequest(currentBatch));
+                        currentBatch.Clear();
+                        currentBatchSize = 0;
                     }
-
-                    allBatches.Add(message.ToBatchRequest(currentDestinationBatches));
-                    currentDestinationBatches.Clear();
-                    payloadSize = 0;
                 }
 
-                if (currentDestinationBatches.Count > 0)
+                // Finalize any remaining messages in the batch
+                if (currentBatch.Count > 0)
                 {
-                    allBatches.Add(firstMessage!.ToBatchRequest(currentDestinationBatches));
-                    currentDestinationBatches.Clear();
+                    allBatches.Add(referenceMessage!.ToBatchRequest(currentBatch));
+                    currentBatch.Clear();
                 }
             }