chore: Ensure consistent DT header insertion/replacement logic across all wrappers (#3512)

* chore: Ensure consistent DT header insertion/replacement logic across all wrappers.

* Fix flaky container test

* test: Fix Kafka container test race condition and byte count range

Move existing-headers produce+consume before cancellation-token
consumers to prevent race where cancellation consumer could grab the
larger existing-headers message. Widen byte count assertion range to
accommodate non-deterministic span ordering from FirstOrDefault.

* Rework dt propagation tests for SQS

Author: tippmar-nr

src/Agent/NewRelic/Agent/Extensions/NewRelic.Agent.Extensions/AwsSdk/SqsHelper.cs

@@ -78,10 +78,6 @@ public static void InsertDistributedTraceHeaders(ITransaction transaction, objec
             var getMessageAttributes = _getMessageAttributes.GetOrAdd(smr.GetType(), t => VisibilityBypasser.Instance.GeneratePropertyAccessor<IDictionary>(t, "MessageAttributes"));
             var messageAttributes = getMessageAttributes(smr);
 
-            // if we can't add all DT headers, don't add any
-            if ((messageAttributes.Count + dtHeaderCount - headersInserted) > MaxSQSMessageAttributes)
-                return;
-
             // create a new MessageAttributeValue instance
             var messageAttributeValueTypeFactory = _messageAttributeValueTypeFactory ??= VisibilityBypasser.Instance.GenerateTypeFactory(smr.GetType().Assembly.FullName, "Amazon.SQS.Model.MessageAttributeValue");
             object newMessageAttributeValue = messageAttributeValueTypeFactory.Invoke();
@@ -92,7 +88,19 @@ public static void InsertDistributedTraceHeaders(ITransaction transaction, objec
             var stringValuePropertySetter = VisibilityBypasser.Instance.GeneratePropertySetter<string>(newMessageAttributeValue, "StringValue");
             stringValuePropertySetter(value);
 
-            messageAttributes.Add(key, newMessageAttributeValue);
+            if (messageAttributes.Contains(key))
+            {
+                // Replace existing header in place — no count change, no limit check needed
+                messageAttributes[key] = newMessageAttributeValue;
+            }
+            else
+            {
+                // New header — check that adding all remaining DT headers won't exceed the SQS limit
+                if ((messageAttributes.Count + dtHeaderCount - headersInserted) > MaxSQSMessageAttributes)
+                    return;
+
+                messageAttributes.Add(key, newMessageAttributeValue);
+            }
 
             ++headersInserted;
         });

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/RestSharp/AppendHeaders.cs

@@ -26,14 +26,17 @@ public AfterWrappedMethodDelegate BeforeWrappedMethod(InstrumentedMethodCall ins
     {
         var httpWebRequest = (HttpWebRequest)instrumentedMethodCall.MethodCall.MethodArguments[0];
 
-        var setHeaders = new Action<HttpWebRequest, string, string>((carrier, key, value) =>
+        // Insert DT headers after AppendHeaders completes, so they overwrite any
+        // pre-existing DT headers that RestSharp copied from the RestRequest.
+        return Delegates.GetDelegateFor(onComplete: () =>
         {
-            // 'Set' will replace an existing value
-            httpWebRequest.Headers?.Set(key, value);
-        });
-
-        agent.CurrentTransaction.InsertDistributedTraceHeaders(httpWebRequest, setHeaders);
+            var setHeaders = new Action<HttpWebRequest, string, string>((carrier, key, value) =>
+            {
+                // 'Set' will replace an existing value
+                carrier.Headers?.Set(key, value);
+            });
 
-        return Delegates.NoOp;
+            agent.CurrentTransaction.InsertDistributedTraceHeaders(httpWebRequest, setHeaders);
+        });
     }
 }

src/Agent/NewRelic/Agent/Extensions/Providers/Wrapper/Wcf3/ChannelFactoryWrapper.cs

@@ -119,7 +119,13 @@ private void SetHeaders (HttpRequestMessageProperty carrier, string key, string
 
     private void SetHeaders(Message carrier, string key, string value)
     {
-        // 'Set' will replace an existing value
+        // Remove existing header if present, then add the new one
+        var headerIndex = carrier.Headers.FindHeader(key, string.Empty);
+        if (headerIndex >= 0)
+        {
+            carrier.Headers.RemoveAt(headerIndex);
+        }
+
         carrier.Headers.Add(MessageHeader.CreateHeader(key, string.Empty, value));
     }
 

tests/Agent/IntegrationTests/Applications/BasicMvcApplication/Controllers/DefaultController.cs

@@ -160,6 +160,21 @@ public string ChainedWebRequest(string chainedServerName, string chainedPortNumb
         return "Worked";
     }
 
+    public string ChainedWebRequestWithExistingHeaders(string chainedServerName, string chainedPortNumber, string chainedAction)
+    {
+        var address = $"http://{chainedServerName}:{chainedPortNumber}/Default/{chainedAction}";
+#pragma warning disable SYSLIB0014 // obsolete usage is ok here
+        var httpWebRequest = (HttpWebRequest)WebRequest.Create(address);
+        // Pre-populate with stale DT headers — agent should replace them
+        httpWebRequest.Headers.Set("traceparent", "00-stale0000000000000000000000000-stale000000000-01");
+        httpWebRequest.Headers.Set("tracestate", "stale=value");
+        httpWebRequest.Headers.Set("newrelic", "stale-newrelic-payload");
+        httpWebRequest.GetResponse();
+#pragma warning restore SYSLIB0014
+
+        return "Worked";
+    }
+
     public async Task<string> ChainedHttpClient(string chainedServerName, string chainedPortNumber, string chainedAction)
     {
         var address = $"http://{chainedServerName}:{chainedPortNumber}/Default/{chainedAction}";

tests/Agent/IntegrationTests/Applications/BasicMvcApplication/Controllers/DistributedTracingController.cs

@@ -62,4 +62,23 @@ public async Task<string> MakeExternalCallUsingRestClient(string externalCallUrl
         }
         return "Worked";
     }
+
+    [Route("DistributedTracing/MakeExternalCallUsingRestClientWithExistingHeaders")]
+    public async Task<string> MakeExternalCallUsingRestClientWithExistingHeaders(string externalCallUrl)
+    {
+        var uri = new Uri(externalCallUrl);
+        var client = new RestClient($"http://{uri.Host}:{uri.Port}");
+        var restRequest = new RestRequest(uri.PathAndQuery);
+        // Pre-populate with stale DT headers — agent should replace them
+        restRequest.AddHeader("traceparent", "00-stale0000000000000000000000000-stale000000000-01");
+        restRequest.AddHeader("tracestate", "stale=value");
+        restRequest.AddHeader("newrelic", "stale-newrelic-payload");
+        var response = await client.ExecuteTaskAsync<IEnumerable<Bird>>(restRequest);
+
+        if ((response.StatusCode != HttpStatusCode.OK) && (response.StatusCode != HttpStatusCode.NoContent))
+        {
+            return $"Unexpected HTTP status code {response.StatusCode}";
+        }
+        return "Worked";
+    }
 }

tests/Agent/IntegrationTests/Applications/Owin4WebApi/Controllers/DistributedTracingSenderController.cs

@@ -30,4 +30,29 @@ public async Task<string> CallNext(string nextUrl)
             return result;
         }
     }
+
+    [HttpGet]
+    [Route("api/CallNextWithExistingHeaders")]
+    public async Task<string> CallNextWithExistingHeaders(string nextUrl)
+    {
+        try
+        {
+            using (var client = new HttpClient())
+            {
+                var request = new HttpRequestMessage(HttpMethod.Get, nextUrl);
+                // Pre-populate with stale DT headers — agent should replace them
+                request.Headers.Add("traceparent", "00-stale0000000000000000000000000-stale000000000-01");
+                request.Headers.Add("tracestate", "stale=value");
+                request.Headers.Add("newrelic", "stale-newrelic-payload");
+                var response = await client.SendAsync(request);
+                var result = await response.Content.ReadAsStringAsync();
+                return result;
+            }
+        }
+        catch (Exception ex)
+        {
+            var result = $"Exception occurred in {nameof(DistributedTracingSenderController)} calling [{nextUrl}]: {ex}";
+            return result;
+        }
+    }
 }

tests/Agent/IntegrationTests/ContainerApplications/AwsSdkTestApp/AwsSdkExercisers/AwsSdkSQSExerciser.cs

@@ -142,6 +142,29 @@ await _amazonSqsClient.DeleteMessageAsync(new DeleteMessageRequest
         }
     }
 
+    [MethodImpl(MethodImplOptions.NoOptimization | MethodImplOptions.NoInlining)]
+    public async Task SQS_SendMessageWithExistingDTHeadersAsync(string message)
+    {
+        if (_sqsQueueUrl == null)
+        {
+            throw new InvalidOperationException("Queue URL is not set. Call SQS_Initialize or SQS_SetQueueUrl first.");
+        }
+
+        var request = new SendMessageRequest
+        {
+            QueueUrl = _sqsQueueUrl,
+            MessageBody = message,
+            MessageAttributes = new Dictionary<string, MessageAttributeValue>
+            {
+                { "traceparent", new MessageAttributeValue { DataType = "String", StringValue = "00-stale0000000000000000000000000-stale000000000-01" } },
+                { "tracestate", new MessageAttributeValue { DataType = "String", StringValue = "stale=value" } },
+                { "newrelic", new MessageAttributeValue { DataType = "String", StringValue = "stale-newrelic-payload" } },
+            }
+        };
+
+        await _amazonSqsClient.SendMessageAsync(request);
+    }
+
     // send message batch
     [MethodImpl(MethodImplOptions.NoOptimization | MethodImplOptions.NoInlining)]
     public async Task SQS_SendMessageBatchAsync(string[] messages)

tests/Agent/IntegrationTests/ContainerApplications/AwsSdkTestApp/Controllers/AwsSdkSQSController.cs

@@ -82,6 +82,18 @@ public async Task SQS_SendMessageToQueueAsync([Required]string message, [Require
         _logger.LogInformation("Message {Message} sent to {Queue}", message, messageQueueUrl);
     }
 
+    // GET: /AwsSdk/SQS_SendMessageWithExistingDTHeadersToQueue?message=Hello&messageQueueUrl=MyQueue
+    [HttpGet("SQS_SendMessageWithExistingDTHeadersToQueue")]
+    public async Task SQS_SendMessageWithExistingDTHeadersToQueueAsync([Required]string message, [Required]string messageQueueUrl)
+    {
+        _logger.LogInformation("Sending message with existing DT headers {Message} to {Queue}", message, messageQueueUrl);
+        using var awsSdkSQSExerciser = new AwsSdkSQSExerciser();
+        awsSdkSQSExerciser.SQS_SetQueueUrl(messageQueueUrl);
+
+        await awsSdkSQSExerciser.SQS_SendMessageWithExistingDTHeadersAsync(message);
+        _logger.LogInformation("Message with existing DT headers {Message} sent to {Queue}", message, messageQueueUrl);
+    }
+
     // GET: /AwsSdk/SQS_SendMessageBatchToQueue?messageQueueUrl=MyQueue
     [HttpGet("SQS_ReceiveMessageFromQueue")]
     public async Task<IEnumerable<Message>> SQS_ReceiveMessageFromQueueAsync([Required]string messageQueueUrl)

tests/Agent/IntegrationTests/ContainerApplications/KafkaTestApp/Controllers/KafkaController.cs

@@ -38,6 +38,13 @@ public async Task<string> ProduceAsync()
         return "Complete";
     }
 
+    [HttpGet("produceasyncwithexistingheaders")]
+    public async Task<string> ProduceAsyncWithExistingHeaders()
+    {
+        await _producer.ProduceAsyncWithExistingHeaders();
+        return "Complete";
+    }
+
     [HttpGet("bootstrap_server")]
     public string GetBootstrapServer() => Program.GetBootstrapServer();
 

tests/Agent/IntegrationTests/ContainerApplications/KafkaTestApp/Producer.cs

@@ -2,6 +2,7 @@
 // SPDX-License-Identifier: Apache-2.0
 
 using System.Runtime.CompilerServices;
+using System.Text;
 using System.Threading.Tasks;
 using Confluent.Kafka;
 using Confluent.Kafka.Admin;
@@ -66,4 +67,25 @@ public async Task ProduceAsync()
 
         await _producer.ProduceAsync(_topic, new Message<string, string> { Key = user, Value = item });
     }
+
+    [MethodImpl(MethodImplOptions.NoInlining)]
+    public async Task ProduceAsyncWithExistingHeaders()
+    {
+        var user = "asyncExistingHeadersUser";
+        var item = "asyncExistingHeadersItem";
+
+        var message = new Message<string, string>
+        {
+            Key = user,
+            Value = item,
+            Headers = new Headers
+            {
+                { "traceparent", Encoding.ASCII.GetBytes("00-stale0000000000000000000000000-stale000000000-01") },
+                { "tracestate", Encoding.ASCII.GetBytes("stale=value") },
+                { "newrelic", Encoding.ASCII.GetBytes("stale-newrelic-payload") }
+            }
+        };
+
+        await _producer.ProduceAsync(_topic, message);
+    }
 }