Switch .NET tracer to injecting both base64 & binary headers (#6448)

## Summary of changes
- Switch .NET tracer to injecting both base64 & binary headers
- Binary headers can be disabled via a config variable:
DD_DATA_STREAMS_LEGACY_HEADERS=false (defaults to true)
- .NET tracer, when extracting the context from headers looks at both
base64 & binary headers

## Reason for change
Data Streams previously used binary encoding in Kafka headers. This was
causing issues in cross language communication because of the difference
in negative byte handling. That’s why we switched to base64 encoding.
Today, .NET is the only remaining tracer using binary encoding for
Kafka, SQS & RabbitMQ.

This is causing 3 issues:
- Communication between .NET & Java breaks, since Java doesn’t support
negative bytes in Kafka headers
- Manual instrumentation breaks between .NET and any other language,
since manual instrumentation uses base64 encoding
- .NET can’t consume payloads from any other tracers, since the other
tracers encode in base64

Also, byte headers are causing a crash of the [.NET
application](https://datadog.zendesk.com/agent/tickets/1824351) (not
reproduced yet).

## Implementation details
- Injects base64-encoded headers by default to ensure cross-language
compatibility.
- Binary headers are injected only if `DD_DATA_STREAMS_LEGACY_HEADERS`
(default: true) is enabled for backward compatibility.
- Prefers base64 headers for extraction. If base64 is unavailable or
malformed, it falls back to binary headers when legacy headers are
enabled.
- A new configuration setting DD_DATA_STREAMS_LEGACY_HEADERS controls
whether binary headers are included.

## Test coverage
I've added tests for the following cases:
- Inject_WhenLegacyHeadersDisabled_DoesNotIncludeBinaryHeader 
- Ensures that when legacy headers are disabled, the binary header is
not included in the injected headers.

- Extract_WhenBothHeadersPresent_PrefersBase64Header
- Confirms that when both Base64 and binary headers are present, the
Base64 header is preferred for extraction.

- InjectedHeaders_HaveCorrectFormat
- Validates that injected headers are in the correct format, especially
checking if the Base64-encoded header is valid and correctly decoded.

- InjectHeaders_WhenLegacyHeadersDisabled_DoesNotIncludeLegacyHeader 
- Ensures that legacy headers are excluded when legacy header support is
disabled.

- Inject_WhenLegacyHeadersEnabled_IncludesBothHeaders 
- Confirms that both Base64 and binary headers are injected when legacy
header support is enabled.

- Extract_WhenBase64HeaderIsMalformed_ReturnsFallbackToBinary 
- Verifies that if the Base64 header is malformed, the system falls back
to using the binary header for extraction.

## Other details
<!-- Fixes #{issue} -->

<!-- ⚠️ Note: where possible, please obtain 2 approvals prior to
merging. Unless CODEOWNERS specifies otherwise, for external teams it is
typically best to have one review from a team member, and one review
from apm-dotnet. Trivial changes do not require 2 reviews. -->

---------

Co-authored-by: Andrew Lock <andrew.lock@datadoghq.com>
Co-authored-by: Steven Bouwkamp <steven.bouwkamp@datadoghq.com>

Author: 3 people

tracer/src/Datadog.Trace/Configuration/ConfigurationKeys.cs

@@ -819,6 +819,13 @@ internal static class DataStreamsMonitoring
             /// </summary>
             /// <see cref="TracerSettings.IsDataStreamsMonitoringEnabled"/>
             public const string Enabled = "DD_DATA_STREAMS_ENABLED";
+
+            /// <summary>
+            /// Configuration key for enabling legacy binary headers in Data Streams Monitoring.
+            /// Default is true.
+            /// </summary>
+             /// <see cref="TracerSettings.IsDataStreamsLegacyHeadersEnabled"/>
+            public const string LegacyHeadersEnabled = "DD_DATA_STREAMS_LEGACY_HEADERS";
         }
     }
 }

tracer/src/Datadog.Trace/Configuration/TracerSettings.cs

@@ -520,6 +520,10 @@ _ when x.ToBoolean() is { } boolean => boolean,
                                             .WithKeys(ConfigurationKeys.DataStreamsMonitoring.Enabled)
                                             .AsBool(false);
 
+            IsDataStreamsLegacyHeadersEnabled = config
+                                               .WithKeys(ConfigurationKeys.DataStreamsMonitoring.LegacyHeadersEnabled)
+                                               .AsBool(true);
+
             IsRareSamplerEnabled = config
                                   .WithKeys(ConfigurationKeys.RareSamplerEnabled)
                                   .AsBool(false);
@@ -985,6 +989,11 @@ public bool DiagnosticSourceEnabled
         /// </summary>
         internal bool IsDataStreamsMonitoringEnabled => DynamicSettings.DataStreamsMonitoringEnabled ?? _isDataStreamsMonitoringEnabled;
 
+        /// <summary>
+        /// Gets a value indicating whether to inject legacy binary headers for Data Streams.
+        /// </summary>
+        internal bool IsDataStreamsLegacyHeadersEnabled { get; }
+
         /// <summary>
         /// Gets a value indicating whether the rare sampler is enabled or not.
         /// </summary>

tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsContextPropagator.cs

@@ -2,13 +2,14 @@
 // Unless explicitly stated otherwise all files in this repository are licensed under the Apache 2 License.
 // This product includes software developed at Datadog (https://www.datadoghq.com/). Copyright 2017 Datadog, Inc.
 // </copyright>
-
 #nullable enable
-
 using System;
 using System.Text;
 using Datadog.Trace.Headers;
+using Datadog.Trace.Logging;
 using Datadog.Trace.Util;
+using Datadog.Trace.VendoredMicrosoftCode.System.Buffers;
+using Datadog.Trace.VendoredMicrosoftCode.System.Buffers.Text;
 
 namespace Datadog.Trace.DataStreamsMonitoring;
 
@@ -17,6 +18,8 @@ namespace Datadog.Trace.DataStreamsMonitoring;
 /// </summary>
 internal class DataStreamsContextPropagator
 {
+    private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor<DataStreamsContextPropagator>();
+
     public static DataStreamsContextPropagator Instance { get; } = new();
 
     /// <summary>
@@ -28,10 +31,42 @@ internal class DataStreamsContextPropagator
     /// <typeparam name="TCarrier">Type of header collection</typeparam>
     public void Inject<TCarrier>(PathwayContext context, TCarrier headers)
         where TCarrier : IBinaryHeadersCollection
+        => Inject(context, headers, Tracer.Instance.Settings.IsDataStreamsLegacyHeadersEnabled);
+
+    // Internal for testing
+    internal void Inject<TCarrier>(PathwayContext context, TCarrier headers, bool isDataStreamsLegacyHeadersEnabled)
+        where TCarrier : IBinaryHeadersCollection
     {
         if (headers is null) { ThrowHelper.ThrowArgumentNullException(nameof(headers)); }
 
-        headers.Add(DataStreamsPropagationHeaders.PropagationKey, PathwayContextEncoder.Encode(context));
+        var encodedBytes = PathwayContextEncoder.Encode(context);
+
+        // Calculate the maximum length of the base64 encoded data
+        // Base64 encoding encodes 3 bytes of data into 4 bytes of encoded data
+        // So the maximum length is ceil(encodedBytes.Length / 3) * 4 and using integer arithmetic it's ((encodedBytes.Length + 2) / 3) * 4
+        int base64Length = ((encodedBytes.Length + 2) / 3) * 4;
+        byte[] base64EncodedContextBytes = new byte[base64Length];
+        var status = Base64.EncodeToUtf8(encodedBytes, base64EncodedContextBytes, out _, out int bytesWritten);
+
+        if (status != OperationStatus.Done)
+        {
+            Log.Error("Failed to encode Data Streams context to Base64. OperationStatus: {Status}", status);
+            return;
+        }
+
+        if (bytesWritten == base64EncodedContextBytes.Length)
+        {
+            headers.Add(DataStreamsPropagationHeaders.PropagationKeyBase64, base64EncodedContextBytes);
+        }
+        else
+        {
+            headers.Add(DataStreamsPropagationHeaders.PropagationKeyBase64, base64EncodedContextBytes.AsSpan(0, bytesWritten).ToArray());
+        }
+
+        if (isDataStreamsLegacyHeadersEnabled)
+        {
+            headers.Add(DataStreamsPropagationHeaders.PropagationKey, encodedBytes);
+        }
     }
 
     /// <summary>
@@ -42,12 +77,68 @@ public void Inject<TCarrier>(PathwayContext context, TCarrier headers)
     /// <returns>A new <see cref="PathwayContext"/> that contains the values obtained from <paramref name="headers"/>.</returns>
     public PathwayContext? Extract<TCarrier>(TCarrier headers)
         where TCarrier : IBinaryHeadersCollection
+        => Extract(headers, Tracer.Instance.Settings.IsDataStreamsLegacyHeadersEnabled);
+
+    // internal for testing
+    internal PathwayContext? Extract<TCarrier>(TCarrier headers, bool isDataStreamsLegacyHeadersEnabled)
+        where TCarrier : IBinaryHeadersCollection
     {
         if (headers is null) { ThrowHelper.ThrowArgumentNullException(nameof(headers)); }
 
-        var bytes = headers.TryGetLastBytes(DataStreamsPropagationHeaders.PropagationKey);
+        // Try to extract from the base64 header first
+        var base64Bytes = headers.TryGetLastBytes(DataStreamsPropagationHeaders.PropagationKeyBase64);
+        if (base64Bytes is { Length: > 0 })
+        {
+            try
+            {
+                // Calculate the maximum decoded length
+                // Base64 encoding encodes 3 bytes of data into 4 bytes of encoded data
+                // So the maximum decoded length is (base64Bytes.Length * 3) / 4
+                int decodedLength = (base64Bytes.Length * 3) / 4;
+                byte[] decodedBytes = new byte[decodedLength];
+
+                var status = Base64.DecodeFromUtf8(base64Bytes, decodedBytes, out _, out int bytesWritten);
+
+                if (status != OperationStatus.Done)
+                {
+                    Log.Error("Failed to decode Base64 data streams context. OperationStatus: {Status}", status);
+                    return null;
+                }
+                else
+                {
+                    if (bytesWritten == decodedBytes.Length)
+                    {
+                        return PathwayContextEncoder.Decode(decodedBytes);
+                    }
+                    else
+                    {
+                        return PathwayContextEncoder.Decode(decodedBytes.AsSpan(0, bytesWritten).ToArray());
+                    }
+                }
+            }
+            catch (Exception ex)
+            {
+                Log.Error(ex, "Failed to decode base64 Data Streams context.");
+            }
+        }
 
-        return bytes is { } ? PathwayContextEncoder.Decode(bytes) : null;
+        if (isDataStreamsLegacyHeadersEnabled)
+        {
+            var binaryBytes = headers.TryGetLastBytes(DataStreamsPropagationHeaders.PropagationKey);
+            if (binaryBytes is { Length: > 0 })
+            {
+                try
+                {
+                    return PathwayContextEncoder.Decode(binaryBytes);
+                }
+                catch (Exception ex)
+                {
+                    Log.Error(ex, "Failed to decode binary Data Streams context.");
+                }
+            }
+        }
+
+        return null;
     }
 
     /// <summary>

tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsManager.cs

@@ -83,21 +83,12 @@ public async Task DisposeAsync()
     public void InjectPathwayContext<TCarrier>(PathwayContext? context, TCarrier headers)
         where TCarrier : IBinaryHeadersCollection
     {
-        if (!IsEnabled)
-        {
-            return;
-        }
-
-        if (context is not null)
+        if (!IsEnabled || context is null)
         {
-            DataStreamsContextPropagator.Instance.Inject(context.Value, headers);
             return;
         }
 
-        // This shouldn't happen normally, as you should call SetCheckpoint before calling InjectPathwayContext
-        // But if data streams was disabled, you call SetCheckpoint, and then data streams is enabled
-        // you will hit this code path
-        Log.Debug("Attempted to inject null pathway context");
+        DataStreamsContextPropagator.Instance.Inject(context.Value, headers);
     }
 
     public void TrackBacklog(string tags, long value)

tracer/src/Datadog.Trace/TracerManager.cs

@@ -528,6 +528,9 @@ void WriteDictionary(IReadOnlyDictionary<string, string> dictionary)
                     writer.WritePropertyName("data_streams_enabled");
                     writer.WriteValue(instanceSettings.IsDataStreamsMonitoringEnabled);
 
+                    writer.WritePropertyName("data_streams_legacy_headers_enabled");
+                    writer.WriteValue(instanceSettings.IsDataStreamsLegacyHeadersEnabled);
+
                     writer.WritePropertyName("span_sampling_rules");
                     writer.WriteValue(instanceSettings.SpanSamplingRules);
 

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/DataStreamsMonitoringKafkaTests.cs

@@ -4,6 +4,7 @@
 // </copyright>
 
 using System;
+using System.Collections.Generic;
 using System.Threading.Tasks;
 using Datadog.Trace.Configuration;
 using Datadog.Trace.TestHelpers;
@@ -27,6 +28,14 @@ public DataStreamsMonitoringKafkaTests(ITestOutputHelper output)
         SetServiceVersion("1.0.0");
     }
 
+    public static IEnumerable<object[]> GetKafkaTestData()
+    {
+        yield return new object[] { true, true };
+        yield return new object[] { true, false };
+        yield return new object[] { false, true };
+        yield return new object[] { false, false };
+    }
+
     /// <returns>A <see cref="Task"/> representing the asynchronous operation.</returns>
     /// <summary>
     /// This sample does a series of produces and consumes to create two pipelines:
@@ -57,15 +66,16 @@ public DataStreamsMonitoringKafkaTests(ITestOutputHelper output)
     ///    C2->>+T3: Produce
     /// </summary>
     /// <param name="enableConsumerScopeCreation">Is the scope created manually or using built-in support</param>
+    /// <param name="enableLegacyHeaders">Should legacy headers be enabled?</param>
     [SkippableTheory]
-    [InlineData(true)]
-    [InlineData(false)]
+    [MemberData(nameof(GetKafkaTestData))]
     [Trait("Category", "EndToEnd")]
     [Trait("Category", "ArmUnsupported")]
-    public async Task SubmitsDataStreams(bool enableConsumerScopeCreation)
+    public async Task SubmitsDataStreams(bool enableConsumerScopeCreation, bool enableLegacyHeaders)
     {
         SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.Enabled, "1");
         SetEnvironmentVariable(ConfigurationKeys.KafkaCreateConsumerScopeEnabled, enableConsumerScopeCreation ? "1" : "0");
+        SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.LegacyHeadersEnabled, enableLegacyHeaders ? "1" : "0");
 
         using var agent = EnvironmentHelper.GetMockAgent(useTelemetry: true);
 
@@ -100,6 +110,7 @@ public async Task HandlesBatchProcessing()
         SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.Enabled, "1");
         // set variable to create short spans on receive instead of spans that last until the next consume
         SetEnvironmentVariable(ConfigurationKeys.KafkaCreateConsumerScopeEnabled, "0");
+        SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.LegacyHeadersEnabled, "1");
 
         using var agent = EnvironmentHelper.GetMockAgent(useTelemetry: true);
 

tracer/test/Datadog.Trace.ClrProfiler.IntegrationTests/DataStreamsMonitoringRabbitMQTests.cs

@@ -35,6 +35,7 @@ public DataStreamsMonitoringRabbitMQTests(ITestOutputHelper output)
     public async Task HandleProduceAndConsume(string packageVersion)
     {
         SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.Enabled, "1");
+        SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.LegacyHeadersEnabled, "1");
 
         using var assertionScope = new AssertionScope();
         using var agent = EnvironmentHelper.GetMockAgent();
@@ -58,6 +59,7 @@ await Verifier.Verify(PayloadsToPoints(agent.DataStreams), settings)
     public async Task ValidateSpanTags(string packageVersion)
     {
         SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.Enabled, "1");
+        SetEnvironmentVariable(ConfigurationKeys.DataStreamsMonitoring.LegacyHeadersEnabled, "1");
 
         using var assertionScope = new AssertionScope();
         using var agent = EnvironmentHelper.GetMockAgent();

tracer/test/Datadog.Trace.Tests/Configuration/TracerSettingsTests.cs

@@ -974,6 +974,16 @@ public void IsDataStreamsMonitoringEnabled(string value, bool expected)
             settings.IsDataStreamsMonitoringEnabled.Should().Be(expected);
         }
 
+        [Theory]
+        [MemberData(nameof(BooleanTestCases), true)]
+        public void IsDataStreamsLegacyHeadersEnabled(string value, bool expected)
+        {
+            var source = CreateConfigurationSource((ConfigurationKeys.DataStreamsMonitoring.LegacyHeadersEnabled, value));
+            var settings = new TracerSettings(source);
+
+            settings.IsDataStreamsLegacyHeadersEnabled.Should().Be(expected);
+        }
+
         [Theory]
         [MemberData(nameof(BooleanTestCases), false)]
         public void IsRareSamplerEnabled(string value, bool expected)