Fix batching Azure ServiceBus messages to partitioned endpoints

Fixes #2000

Author: PerfectlyNormal

src/Transports/Azure/Wolverine.AzureServiceBus.Tests/end_to_end.cs

@@ -4,6 +4,7 @@
 using Wolverine.AzureServiceBus.Internal;
 using Wolverine.Configuration;
 using Wolverine.Tracking;
+using Wolverine.Transports;
 using Xunit;
 
 namespace Wolverine.AzureServiceBus.Tests;
@@ -44,6 +45,18 @@ public async Task InitializeAsync()
                     // Require sessions on this subscription
                     .RequireSessions(1)
 
+                    .ProcessInline();
+
+                opts.PublishMessage<AsbMessage4>().ToAzureServiceBusTopic("asb4").BufferedInMemory();
+                opts.ListenToAzureServiceBusSubscription("asb4")
+                    .FromTopic("asb4", cfg =>
+                    {
+                        cfg.EnablePartitioning = true;
+                    })
+
+                    // Require sessions on this subscription
+                    .RequireSessions(1)
+
                     .ProcessInline();
             }).StartAsync();
 
@@ -155,11 +168,39 @@ public async Task send_and_receive_multiple_messages_to_subscription_with_sessio
         session.Received.MessagesOf<AsbMessage3>().Select(x => x.Name)
             .ShouldBe(new string[]{"Red", "Green", "Refactor"});
     }
+
+    [Fact]
+    public async Task split_messages_with_different_sessionids_into_separate_batches()
+    {
+        Func<IMessageContext, Task> sendMany = async bus =>
+        {
+            await bus.SendAsync(new AsbMessage4("Dummy 1.1"), new DeliveryOptions { GroupId = "1" });
+            await bus.SendAsync(new AsbMessage4("Dummy 1.2"), new DeliveryOptions { GroupId = "1" });
+            await bus.SendAsync(new AsbMessage4("Dummy 2.1"), new DeliveryOptions { GroupId = "2" });
+            await bus.SendAsync(new AsbMessage4("Dummy 3.1"), new DeliveryOptions { GroupId = "3" });
+            await bus.SendAsync(new AsbMessage4("Dummy 4.1"), new DeliveryOptions { GroupId = "4" });
+            await bus.SendAsync(new AsbMessage4("Dummy 4.2"), new DeliveryOptions { GroupId = "4" });
+        };
+
+        var session = await _host.TrackActivity()
+            .IncludeExternalTransports()
+            .Timeout(30.Seconds())
+            .ExecuteAndWaitAsync(sendMany);
+
+        // Verify that all messages were received and processed in order inside the session
+        var names = session.Received.MessagesOf<AsbMessage4>().Select(x => x.Name).ToList();
+        names.Count.ShouldBe(6);
+        names.Where(x => x.StartsWith("Dummy 1")).ShouldBe(["Dummy 1.1", "Dummy 1.2"]);
+        names.Where(x => x.StartsWith("Dummy 2")).ShouldBe(["Dummy 2.1"]);
+        names.Where(x => x.StartsWith("Dummy 3")).ShouldBe(["Dummy 3.1"]);
+        names.Where(x => x.StartsWith("Dummy 4")).ShouldBe(["Dummy 4.1", "Dummy 4.2"]);
+    }
 }
 
 public record AsbMessage1(string Name);
 public record AsbMessage2(string Name);
 public record AsbMessage3(string Name);
+public record AsbMessage4(string Name);
 
 public static class AsbMessageHandler
 {
@@ -177,4 +218,9 @@ public static void Handle(AsbMessage3 message)
     {
         // nothing
     }
+
+    public static void Handle(AsbMessage4 message)
+    {
+        // nothing
+    }
 }

src/Transports/Azure/Wolverine.AzureServiceBus/Internal/AzureServiceBusEndpoint.cs

@@ -50,6 +50,7 @@ public AzureServiceBusEndpoint(AzureServiceBusTransport parent, Uri uri, Endpoin
     public abstract ValueTask<bool> CheckAsync();
     public abstract ValueTask TeardownAsync(ILogger logger);
     public abstract ValueTask SetupAsync(ILogger logger);
+    public abstract bool IsPartitioned { get; }
 
     protected override bool supportsMode(EndpointMode mode)
     {

src/Transports/Azure/Wolverine.AzureServiceBus/Internal/AzureServiceBusQueue.cs

@@ -102,6 +102,8 @@ await Parent.WithServiceBusClientAsync(async client =>
         });
     }
 
+    public override bool IsPartitioned { get => Options.EnablePartitioning; }
+
     private async Task purgeWithSessions(ServiceBusClient client)
     {
         var cancellation = new CancellationTokenSource();

src/Transports/Azure/Wolverine.AzureServiceBus/Internal/AzureServiceBusSenderProtocol.cs

@@ -46,6 +46,14 @@ public async Task SendBatchAsync(ISenderCallback callback, OutgoingMessageBatch
             }
         }
 
+        if (_endpoint.IsPartitioned)
+            await sendPartitionedBatches(callback, messages, batch);
+        else
+            await sendBatches(callback, messages, batch);
+    }
+
+    private async Task sendBatches(ISenderCallback callback, List<ServiceBusMessage> messages, OutgoingMessageBatch batch)
+    {
         try
         {
             var serviceBusMessageBatch = await _sender.CreateMessageBatchAsync();
@@ -55,7 +63,7 @@ public async Task SendBatchAsync(ISenderCallback callback, OutgoingMessageBatch
                 if (!serviceBusMessageBatch.TryAddMessage(message))
                 {
                     _logger.LogInformation("Wolverine had to break up outgoing message batches at {Uri}, you may want to reduce the MaximumMessagesToReceive configuration. No messages were lost, this is strictly informative", _endpoint.Uri);
-                    
+
                     // Send the currently full batch
                     await _sender.SendMessagesAsync(serviceBusMessageBatch, _runtime.Cancellation);
                     serviceBusMessageBatch.Dispose();
@@ -77,4 +85,53 @@ public async Task SendBatchAsync(ISenderCallback callback, OutgoingMessageBatch
             await callback.MarkProcessingFailureAsync(batch, e);
         }
     }
+
+    private async Task sendPartitionedBatches(ISenderCallback callback, List<ServiceBusMessage> messages, OutgoingMessageBatch batch)
+    {
+        try
+        {
+            var serviceBusMessageBatch = await _sender.CreateMessageBatchAsync();
+
+            var groupedMessages = messages
+                .GroupBy(x => x.SessionId)
+                .ToDictionary(x => x.Key, x => x.ToList());
+
+            foreach (var (sessionId, messageGroup) in groupedMessages)
+            {
+                _logger.LogDebug("Processing batch with session id '{SessionId}'", sessionId);
+
+                foreach (var message in messageGroup)
+                {
+                    if (!serviceBusMessageBatch.TryAddMessage(message))
+                    {
+                        _logger.LogInformation("Wolverine had to break up outgoing message batches at {Uri}, you may want to reduce the MaximumMessagesToReceive configuration. No messages were lost, this is strictly informative", _endpoint.Uri);
+
+                        // Send the currently full batch
+                        await _sender.SendMessagesAsync(serviceBusMessageBatch, _runtime.Cancellation);
+                        serviceBusMessageBatch.Dispose();
+
+                        // Create a new batch and add the message to it
+                        serviceBusMessageBatch = await _sender.CreateMessageBatchAsync();
+                        serviceBusMessageBatch.TryAddMessage(message);
+                    }
+                }
+
+                await _sender.SendMessagesAsync(serviceBusMessageBatch, _runtime.Cancellation);
+                serviceBusMessageBatch = await _sender.CreateMessageBatchAsync();
+            }
+
+            // Send the final batch
+            if (serviceBusMessageBatch.Count > 0)
+            {
+                await _sender.SendMessagesAsync(serviceBusMessageBatch, _runtime.Cancellation);
+            }
+            serviceBusMessageBatch.Dispose();
+
+            await callback.MarkSuccessfulAsync(batch);
+        }
+        catch (Exception e)
+        {
+            await callback.MarkProcessingFailureAsync(batch, e);
+        }
+    }
 }

src/Transports/Azure/Wolverine.AzureServiceBus/Internal/AzureServiceBusSubscription.cs

@@ -175,6 +175,8 @@ public override async ValueTask InitializeAsync(ILogger logger)
         _hasInitialized = true;
     }
 
+    public override bool IsPartitioned { get => Topic.IsPartitioned; }
+
     internal async ValueTask InitializeAsync(ServiceBusAdministrationClient client, ILogger logger)
     {
         if (Parent.AutoProvision)

src/Transports/Azure/Wolverine.AzureServiceBus/Internal/AzureServiceBusTopic.cs

@@ -116,4 +116,6 @@ public AzureServiceBusSubscription FindOrCreateSubscription(string subscriptionN
 
         return subscription;
     }
+
+    public override bool IsPartitioned { get => Options.EnablePartitioning; }
 }