Fix the code that adds OTel and publish sequence number headers.

Author: lukebakken

projects/RabbitMQ.Client/Impl/ChannelBase.cs

@@ -1885,93 +1885,70 @@ await _confirmSemaphore.WaitAsync(cancellationToken)
                 return null;
             }
 
-            BasicProperties? rv = null;
-            if (sendActivity is not null)
-            {
-                // This activity is marked as recorded, so let's propagate the trace and span ids.
-                if (sendActivity.IsAllDataRequested)
-                {
-                    if (!string.IsNullOrEmpty(basicProperties.CorrelationId))
-                    {
-                        sendActivity.SetTag(RabbitMQActivitySource.MessageConversationId, basicProperties.CorrelationId);
-                    }
-                }
+            var newHeaders = new Dictionary<string, object?>();
+            MaybeAddActivityToHeaders(newHeaders, basicProperties.CorrelationId, sendActivity);
+            MaybeAddPublishSequenceNumberToHeaders(newHeaders);
 
-                IDictionary<string, object?>? headers = basicProperties.Headers;
-                if (headers is null)
-                {
-                    rv = AddActivityHeaders(basicProperties, sendActivity);
-                }
-                else
-                {
-                    // Inject the ActivityContext into the message headers to propagate trace context to the receiving service.
-                    RabbitMQActivitySource.ContextInjector(sendActivity, headers);
-                    rv = null;
-                }
+            switch (basicProperties)
+            {
+                case BasicProperties writableProperties:
+                    MergeHeaders(newHeaders, writableProperties);
+                    return null;
+                case EmptyBasicProperty:
+                    return new BasicProperties { Headers = newHeaders };
+                default:
+                    return new BasicProperties(basicProperties) { Headers = newHeaders };
             }
 
-            if (_publisherConfirmationTrackingEnabled)
+            static void MergeHeaders(IDictionary<string, object?> newHeaders, BasicProperties props)
             {
-                byte[] publishSequenceNumberBytes;
-                if (BitConverter.IsLittleEndian)
-                {
-                    publishSequenceNumberBytes = BitConverter.GetBytes(BinaryPrimitives.ReverseEndianness(publishSequenceNumber));
-                }
-                else
+                if (props.Headers is null)
                 {
-                    publishSequenceNumberBytes = BitConverter.GetBytes(publishSequenceNumber);
-                }
-
-                IDictionary<string, object?>? headers = basicProperties.Headers;
-                if (headers is null)
-                {
-                    rv = AddPublishSequenceNumberHeader(basicProperties, publishSequenceNumberBytes);
+                    props.Headers = newHeaders;
                 }
                 else
                 {
-                    headers[Constants.PublishSequenceNumberHeader] = publishSequenceNumberBytes;
-                    rv = null;
+                    foreach (KeyValuePair<string, object?> val in newHeaders)
+                    {
+                        props.Headers[val.Key] = val.Value;
+                    }
                 }
             }
 
-            return rv;
-
-            static BasicProperties? AddActivityHeaders(TProperties basicProperties, Activity sendActivity)
+            void MaybeAddActivityToHeaders(IDictionary<string, object?> headers,
+                string? correlationId, Activity? sendActivity)
             {
-                var headers = new Dictionary<string, object?>();
-
-                // Inject the ActivityContext into the message headers to propagate trace context to the receiving service.
-                RabbitMQActivitySource.ContextInjector(sendActivity, headers);
-
-                switch (basicProperties)
+                if (sendActivity is not null)
                 {
-                    case BasicProperties writableProperties:
-                        writableProperties.Headers = headers;
-                        return null;
-                    case EmptyBasicProperty:
-                        return new BasicProperties { Headers = headers };
-                    default:
-                        return new BasicProperties(basicProperties) { Headers = headers };
+                    // This activity is marked as recorded, so let's propagate the trace and span ids.
+                    if (sendActivity.IsAllDataRequested)
+                    {
+                        if (!string.IsNullOrEmpty(correlationId))
+                        {
+                            sendActivity.SetTag(RabbitMQActivitySource.MessageConversationId, correlationId);
+                        }
+                    }
+
+                    // Inject the ActivityContext into the message headers to propagate trace context to the receiving service.
+                    RabbitMQActivitySource.ContextInjector(sendActivity, headers);
                 }
             }
 
-            static BasicProperties? AddPublishSequenceNumberHeader(TProperties basicProperties, byte[] publishSequenceNumberBytes)
+            void MaybeAddPublishSequenceNumberToHeaders(IDictionary<string, object?> headers)
             {
-                var headers = new Dictionary<string, object?>()
+                if (_publisherConfirmationTrackingEnabled)
                 {
-                    [Constants.PublishSequenceNumberHeader] = publishSequenceNumberBytes
-                };
-
-                switch (basicProperties)
-                {
-                    case BasicProperties writableProperties:
-                        writableProperties.Headers = headers;
-                        return null;
-                    case EmptyBasicProperty:
-                        return new BasicProperties { Headers = headers };
-                    default:
-                        return new BasicProperties(basicProperties) { Headers = headers };
+                    byte[] publishSequenceNumberBytes;
+                    if (BitConverter.IsLittleEndian)
+                    {
+                        publishSequenceNumberBytes = BitConverter.GetBytes(BinaryPrimitives.ReverseEndianness(publishSequenceNumber));
+                    }
+                    else
+                    {
+                        publishSequenceNumberBytes = BitConverter.GetBytes(publishSequenceNumber);
+                    }
 
+                    headers[Constants.PublishSequenceNumberHeader] = publishSequenceNumberBytes;
                 }
             }
         }

projects/Test/SequentialIntegration/TestActivitySource.cs

@@ -99,7 +99,7 @@ public async Task TestPublisherAndConsumerActivityTags(bool useRoutingKeyAsOpera
             };
 
             string consumerTag = await _channel.BasicConsumeAsync(queueName, autoAck: true, consumer: consumer);
-            await _channel.BasicPublishAsync("", q.QueueName, true, sendBody);
+            Assert.True(await _channel.BasicPublishAsync("", q.QueueName, true, sendBody));
 
             await consumerReceivedTcs.Task.WaitAsync(TimeSpan.FromSeconds(5));
             Assert.True(await consumerReceivedTcs.Task);