[DSM] Reduce per-message overhead in Kafka produce and consume hot paths

Author: kr-igor

tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Kafka/KafkaHelper.cs

@@ -299,7 +299,15 @@ private static string[] BuildProduceEdgeTags(string clusterId, string topic)
                         tags.MessageQueueTimeMs == null ? 0 : (long)tags.MessageQueueTimeMs,
                         pathwayContext);
 
-                    message?.Headers?.Remove(DataStreamsPropagationHeaders.TemporaryBase64PathwayContext); // remove eventual junk
+                    // TemporaryBase64PathwayContext is only written by our consumer code when
+                    // KafkaCreateConsumerScopeEnabled=false. When it is true (the default), the
+                    // header is never present so the unconditional Remove performs a wasted O(n)
+                    // linear header scan on every message. Skip it when we know it can't be there.
+                    if (!tracer.CurrentTraceSettings.Settings.KafkaCreateConsumerScopeEnabled)
+                    {
+                        message?.Headers?.Remove(DataStreamsPropagationHeaders.TemporaryBase64PathwayContext);
+                    }
+
                     if (!tracer.CurrentTraceSettings.Settings.KafkaCreateConsumerScopeEnabled && message?.Headers is not null && span.Context.PathwayContext != null)
                     {
                         // write the _new_ pathway (the "consume" checkpoint that we just set above) to the headers as a way to pass its value to an eventual

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

@@ -110,6 +110,19 @@ internal void Inject<TCarrier>(PathwayContext context, TCarrier headers, bool is
         {
             try
             {
+#if NETCOREAPP3_1_OR_GREATER
+                // Decode directly into a stack buffer to avoid a heap allocation per consume.
+                Span<byte> decodedBytes = stackalloc byte[PathwayContextEncoder.MaxEncodedSize];
+                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;
+                }
+
+                return PathwayContextEncoder.Decode(decodedBytes.Slice(0, bytesWritten));
+#else
                 // 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
@@ -134,6 +147,7 @@ internal void Inject<TCarrier>(PathwayContext context, TCarrier headers, bool is
                         return PathwayContextEncoder.Decode(decodedBytes.AsSpan(0, bytesWritten).ToArray());
                     }
                 }
+#endif
             }
             catch (Exception ex)
             {

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

@@ -44,7 +44,11 @@ internal sealed class DataStreamsManager
     private readonly IDisposable _updateSubscription;
     private readonly bool _isLegacyDsmHeadersEnabled;
     private readonly bool _isInDefaultState;
-    private readonly ConditionalWeakTable<string[], NodeHashCacheEntry> _nodeHashCache = new();
+    // Keyed by string[] identity (reference equality) — safe because EdgeTagCache holds strong
+    // references to the cached arrays (bounded by MaxEdgeTagCacheSize).
+    private readonly ConcurrentDictionary<string[], NodeHashCacheEntry> _nodeHashCache =
+        new(NodeHashCacheKeyComparer.Instance);
+
     private long _nodeHashBase; // note that this actually represents a `ulong` that we have done an unsafe cast for
     private MutableSettings _previousMutableSettings;
     private string? _previousContainerTagsHash;
@@ -312,7 +316,7 @@ public void InjectPathwayContextAsBase64String<TCarrier>(PathwayContext? context
 
             // Don't blame me, blame the fact we can't do Volatile.Read with a ulong in .NET FX...
             var nodeHashBase = new NodeHashBase(unchecked((ulong)Volatile.Read(ref _nodeHashBase)));
-            var cacheEntry = _nodeHashCache.GetOrCreateValue(edgeTags);
+            var cacheEntry = _nodeHashCache.GetOrAdd(edgeTags, static _ => new NodeHashCacheEntry());
             NodeHash nodeHash;
 
             // Fast lock-free path: snapshot is an immutable object published via a volatile field.
@@ -426,6 +430,20 @@ public bool ShouldExtractSchema(Span span, string operation, out int weight)
         return false;
     }
 
+    /// <summary>
+    /// Reference-equality comparer for string[] keys in <see cref="_nodeHashCache"/>.
+    /// Two string[] objects are considered equal only when they are the same instance,
+    /// which is always true for the cached arrays held by <see cref="EdgeTagCache{TKey}"/>.
+    /// </summary>
+    private sealed class NodeHashCacheKeyComparer : IEqualityComparer<string[]>
+    {
+        internal static readonly NodeHashCacheKeyComparer Instance = new();
+
+        public bool Equals(string[]? x, string[]? y) => ReferenceEquals(x, y);
+
+        public int GetHashCode(string[] obj) => System.Runtime.CompilerServices.RuntimeHelpers.GetHashCode(obj);
+    }
+
     /// <summary>
     /// Memoized NodeHash associated with a specific edge-tag array instance and nodeHashBase value.
     /// The volatile <see cref="_snapshot"/> field enables a lock-free fast path: callers read the

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

@@ -112,6 +112,50 @@ public static int EncodeInto(PathwayContext pathway, Span<byte> buffer)
         return new PathwayContext(new PathwayHash(hash), pathwayStartNs, edgeStartNs);
     }
 
+#if NETCOREAPP3_1_OR_GREATER
+    /// <summary>
+    /// Zero-allocation alternative to <see cref="Decode(byte[])"/>: decodes directly from a
+    /// caller-supplied <paramref name="bytes"/> span (e.g. a stackalloc buffer).
+    /// </summary>
+    public static PathwayContext? Decode(Span<byte> bytes)
+    {
+        if (bytes.Length < 10)
+        {
+            Log.Warning<int>("Error decoding Data Stream PathwayContext from bytes: insufficient bytes ({ByteCount})", bytes.Length);
+            return null;
+        }
+
+        var hash = System.Buffers.Binary.BinaryPrimitives.ReadUInt64LittleEndian(bytes);
+
+        var pathwayStartMs = VarEncodingHelper.ReadVarLongZigZag(bytes.Slice(8), out var bytesRead);
+        if (pathwayStartMs is null)
+        {
+            Log.Warning("Error decoding Data Stream PathwayContext from bytes: invalid pathway start");
+            return null;
+        }
+
+        var edgeStartMs = VarEncodingHelper.ReadVarLongZigZag(bytes.Slice(8 + bytesRead), out _);
+        if (edgeStartMs is null)
+        {
+            Log.Warning("Error decoding Data Stream PathwayContext from bytes: invalid edge start");
+            return null;
+        }
+
+        var pathwayStartNs = ToNanoseconds(pathwayStartMs.Value);
+        var edgeStartNs = ToNanoseconds(edgeStartMs.Value);
+        if (pathwayStartMs > pathwayStartNs || edgeStartMs.Value > edgeStartNs)
+        {
+            Log.Warning(
+                "Overflow detected in Data Stream PathwayContext from bytes: invalid pathway {PathwayMs}ms or edge {EdgeMs}ms",
+                pathwayStartMs,
+                edgeStartMs);
+            return null;
+        }
+
+        return new PathwayContext(new PathwayHash(hash), pathwayStartNs, edgeStartNs);
+    }
+#endif
+
     private static long ToNanoseconds(long milliseconds)
         => milliseconds * 1_000_000;