Ensure properties from the first message needed for a batch are included in the safe size calculation.

Author: SeanFeldman

src/Microsoft.Azure.ServiceBus/Core/MessageBatch.cs

@@ -20,7 +20,6 @@ public class MessageBatch : IDisposable
         private AmqpMessage firstMessage;
         private readonly List<Data> datas;
         private AmqpMessage result;
-        private (string messageId, string sessionId, string partitionKey, string viaPartitionKey) originalMessageData;
 
         /// <summary>
         /// Construct a new batch with a maximum batch size and outgoing plugins callback.
@@ -56,8 +55,27 @@ public async Task<bool> TryAdd(Message message)
 
             if (firstMessage == null)
             {
-                originalMessageData = (processedMessage.MessageId, processedMessage.SessionId, processedMessage.PartitionKey, processedMessage.ViaPartitionKey);
                 firstMessage = amqpMessage;
+
+                if (processedMessage.MessageId != null)
+                {
+                    result.Properties.MessageId = processedMessage.MessageId;
+                }
+
+                if (processedMessage.SessionId != null)
+                {
+                    result.Properties.GroupId = processedMessage.SessionId;
+                }
+
+                if (processedMessage.PartitionKey != null)
+                {
+                    result.MessageAnnotations.Map[AmqpMessageConverter.PartitionKeyName] = processedMessage.PartitionKey;
+                }
+
+                if (processedMessage.ViaPartitionKey != null)
+                {
+                    result.MessageAnnotations.Map[AmqpMessageConverter.ViaPartitionKeyName] = processedMessage.ViaPartitionKey;
+                }
             }
 
             var data = AmqpMessageConverter.ToData(amqpMessage);
@@ -95,26 +113,6 @@ internal AmqpMessage ToAmqpMessage()
                 return firstMessage;
             }
 
-            if (originalMessageData.messageId != null)
-            {
-                result.Properties.MessageId = originalMessageData.messageId;
-            }
-
-            if (originalMessageData.sessionId != null)
-            {
-                result.Properties.GroupId = originalMessageData.sessionId;
-            }
-
-            if (originalMessageData.partitionKey != null)
-            {
-                result.MessageAnnotations.Map[AmqpMessageConverter.PartitionKeyName] = originalMessageData.partitionKey;
-            }
-
-            if (originalMessageData.viaPartitionKey != null)
-            {
-                result.MessageAnnotations.Map[AmqpMessageConverter.ViaPartitionKeyName] = originalMessageData.viaPartitionKey;
-            }
-
             result.MessageFormat = AmqpConstants.AmqpBatchedMessageFormat;
             result.Batchable = true;
             return result;