Moved MessageGroupId header add from FairQueuesFeature to InputQueuePump and removed additional references to DoNotAutomaticallyPropagateMessageGroupId from DelayedMessagesPump.

Author: ngallegos

src/NServiceBus.Transport.SQS.Tests/Receiving/DelayedMessagesPumpTests.cs

@@ -84,9 +84,9 @@ public void Initialize_redrive_policy_used_throws()
         Assert.That(exception.Message, Is.EqualTo("Delayed delivery queue 'queue-delay.fifo' should not have Redrive Policy enabled."));
     }
 
-    async Task SetupInitializedPump(bool doNotAutomaticallyPropagateMessageGroupId = false)
+    async Task SetupInitializedPump()
     {
-        pump = new DelayedMessagesPump(FakeInputQueueQueueUrl, mockSqsClient, new QueueCache(mockSqsClient, q => QueueCache.GetSqsQueueName(q, "")), 15 * 60, doNotAutomaticallyPropagateMessageGroupId: doNotAutomaticallyPropagateMessageGroupId);
+        pump = new DelayedMessagesPump(FakeInputQueueQueueUrl, mockSqsClient, new QueueCache(mockSqsClient, q => QueueCache.GetSqsQueueName(q, "")), 15 * 60);
 
         mockSqsClient.GetAttributeNamesRequestsResponse = (queue, attributes) => new GetQueueAttributesResponse
         {

src/NServiceBus.Transport.SQS/Extensions/MessageExtensions.cs

@@ -19,7 +19,7 @@ public static DateTimeOffset GetAdjustedDateTimeFromServerSetAttributes(this Mes
 
     public static string? ExtractMessageGroupId(this Message messsage)
     {
-        var messageGroupId = messsage.Attributes.GetValueOrDefault("MessageGroupId");
+        var messageGroupId = messsage.Attributes?.GetValueOrDefault("MessageGroupId");
 
         return !string.IsNullOrEmpty(messageGroupId) ? messageGroupId : null;
     }

src/NServiceBus.Transport.SQS/Features/FairQueuesFeature.cs

@@ -16,25 +16,10 @@ protected override void Setup(FeatureConfigurationContext context)
             return;
         }
 
-        context.Pipeline.Register(new PersistIncomingMessageGroupIdToHeadersBehavior(), "SQS persist incoming MessageGroupId to headers behavior");
         context.Pipeline.Register(new ApplyMessageGroupIdFromHeadersToOutgoingMessageBehavior(), "SQS apply MessageGroupId from headers to outgoing message behavior");
     }
 }
 
-class PersistIncomingMessageGroupIdToHeadersBehavior : IBehavior<IIncomingPhysicalMessageContext, IIncomingPhysicalMessageContext>
-{
-    public Task Invoke(IIncomingPhysicalMessageContext context, Func<IIncomingPhysicalMessageContext, Task> next)
-    {
-        if (context.Extensions.TryGet<Amazon.SQS.Model.Message>(out var nativeMessage)
-            && nativeMessage.Attributes.TryGetValue("MessageGroupId", out var messageGroupId)
-            && !string.IsNullOrWhiteSpace(messageGroupId))
-        {
-            context.Message.Headers[TransportHeaders.MessageGroupId] = messageGroupId;
-        }
-        return next(context);
-    }
-}
-
 class ApplyMessageGroupIdFromHeadersToOutgoingMessageBehavior : IBehavior<IOutgoingPhysicalMessageContext, IOutgoingPhysicalMessageContext>
 {
     public Task Invoke(IOutgoingPhysicalMessageContext context, Func<IOutgoingPhysicalMessageContext, Task> next)

src/NServiceBus.Transport.SQS/Receiving/DelayedMessagesPump.cs

@@ -12,7 +12,7 @@ namespace NServiceBus.Transport.SQS;
 using Extensions;
 using Logging;
 
-class DelayedMessagesPump(string receiveAddress, IAmazonSQS sqsClient, QueueCache queueCache, int queueDelayTimeSeconds, bool doNotAutomaticallyPropagateMessageGroupId = false)
+class DelayedMessagesPump(string receiveAddress, IAmazonSQS sqsClient, QueueCache queueCache, int queueDelayTimeSeconds)
 {
     public async Task Initialize(CancellationToken cancellationToken = default)
     {
@@ -225,14 +225,10 @@ IReadOnlyCollection<SqsReceivedDelayedMessage> PrepareMessages(ReceiveMessageRes
                 // Copy over all the message attributes so we don't lose part of the message when moving to the delayed delivery queue
                 preparedMessage.CopyMessageAttributes(receivedMessage.MessageAttributes);
 
-                // Preserve fair queue MessageGroupId if it was set
-                if (!doNotAutomaticallyPropagateMessageGroupId)
+                var existingMessageGroupId = receivedMessage.ExtractMessageGroupId();
+                if (!string.IsNullOrEmpty(existingMessageGroupId))
                 {
-                    var existingMessageGroupId = receivedMessage.ExtractMessageGroupId();
-                    if (existingMessageGroupId != null)
-                    {
-                        preparedMessage.MessageGroupId = existingMessageGroupId;
-                    }
+                    preparedMessage.MessageGroupId = existingMessageGroupId;
                 }
 
                 preparedMessage.MessageAttributes.Remove(TransportHeaders.DelaySeconds);

src/NServiceBus.Transport.SQS/Receiving/InputQueuePump.cs

@@ -29,7 +29,8 @@ class InputQueuePump(
         Action<string, Exception, CancellationToken> criticalErrorAction,
         int? configuredVisibilityTimeoutInSeconds,
         TimeSpan maxAutoMessageVisibilityRenewalDuration,
-        bool setupInfrastructure = true)
+        bool setupInfrastructure = true,
+        bool doNotAutomaticallyPropagateMessageGroupId = false)
         : IMessageReceiver
     {
         public async Task Initialize(PushRuntimeSettings limitations, OnMessage onMessage, OnError onError, CancellationToken cancellationToken = default)
@@ -471,6 +472,15 @@ async Task<bool> InnerProcessMessage(Dictionary<string, string> headers, string
             transportTransaction.Set(nativeMessage);
             transportTransaction.Set("IncomingMessageId", headers[Headers.MessageId]);
 
+            if (doNotAutomaticallyPropagateMessageGroupId is false)
+            {
+                var messageGroupId = nativeMessage.ExtractMessageGroupId();
+                if (!string.IsNullOrEmpty(messageGroupId))
+                {
+                    headers[TransportHeaders.MessageGroupId] = messageGroupId;
+                }
+            }
+
             try
             {
                 var messageContext = new MessageContext(

src/NServiceBus.Transport.SQS/Receiving/MessagePump.cs

@@ -28,11 +28,11 @@ public MessagePump(string receiverId,
         bool doNotAutomaticallyPropagateMessageGroupId)
     {
         this.disableDelayedDelivery = disableDelayedDelivery;
-        inputQueuePump = new InputQueuePump(receiverId, receiveAddress, errorQueueAddress, purgeOnStartup, sqsClient, queueCache, s3Settings, subscriptionManager, criticalErrorAction, visibilityTimeoutInSeconds, maxAutoMessageVisibilityRenewalDuration, setupInfrastructure);
+        inputQueuePump = new InputQueuePump(receiverId, receiveAddress, errorQueueAddress, purgeOnStartup, sqsClient, queueCache, s3Settings, subscriptionManager, criticalErrorAction, visibilityTimeoutInSeconds, maxAutoMessageVisibilityRenewalDuration, setupInfrastructure, doNotAutomaticallyPropagateMessageGroupId);
         if (!disableDelayedDelivery)
         {
             delayedMessagesPump =
-                new DelayedMessagesPump(receiveAddress, sqsClient, queueCache, queueDelayTimeSeconds, doNotAutomaticallyPropagateMessageGroupId);
+                new DelayedMessagesPump(receiveAddress, sqsClient, queueCache, queueDelayTimeSeconds);
         }
     }