DSM overhead optimizations

kr-igor · kr-igor · commit 3850cdb84958 · 2026-04-13T15:31:55.000-05:00
diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Kafka/KafkaHelper.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Kafka/KafkaHelper.cs
@@ -142,6 +142,19 @@ private static long GetMessageSize<T>(T message)
             return size;
         }
 
+        // NOTE: tags must be sorted alphabetically — called only on edge-tag cache miss
+        private static string[] BuildProduceEdgeTags(string clusterId, string topic)
+        {
+            if (!StringUtil.IsNullOrEmpty(clusterId))
+            {
+                return StringUtil.IsNullOrEmpty(topic)
+                           ? ["direction:out", $"kafka_cluster_id:{clusterId}", "type:kafka"]
+                           : ["direction:out", $"kafka_cluster_id:{clusterId}", $"topic:{topic}", "type:kafka"];
+            }
+
+            return ["direction:out", $"topic:{topic}", "type:kafka"];
+        }
+
         internal static Scope? CreateConsumerScope(
             Tracer tracer,
             DataStreamsManager dataStreamsManager,
@@ -398,17 +411,14 @@ internal static void TryInjectHeaders<TTopicPartitionMarker, TMessage>(
                     ProducerCache.TryGetProducer(producer, out _, out var producerClusterId);
 
                     string[] edgeTags;
-                    if (!StringUtil.IsNullOrEmpty(producerClusterId))
+                    if (StringUtil.IsNullOrEmpty(topic) && StringUtil.IsNullOrEmpty(producerClusterId))
                     {
-                        edgeTags = StringUtil.IsNullOrEmpty(topic)
-                                       ? ["direction:out", $"kafka_cluster_id:{producerClusterId}", "type:kafka"]
-                                       : ["direction:out", $"kafka_cluster_id:{producerClusterId}", $"topic:{topic}", "type:kafka"];
+                        edgeTags = DefaultProduceEdgeTags;
                     }
                     else
                     {
-                        edgeTags = StringUtil.IsNullOrEmpty(topic)
-                                       ? DefaultProduceEdgeTags
-                                       : ["direction:out", $"topic:{topic}", "type:kafka"];
+                        var cacheKey = new ProduceEdgeTagCacheKey(producerClusterId ?? string.Empty, topic ?? string.Empty);
+                        edgeTags = dataStreamsManager.GetOrCreateEdgeTags(cacheKey, static k => BuildProduceEdgeTags(k.ClusterId, k.Topic));
                     }
 
                     var msgSize = dataStreamsManager.IsInDefaultState ? 0 : GetMessageSize(message);
diff --git a/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Kafka/ProduceEdgeTagCacheKey.cs b/tracer/src/Datadog.Trace/ClrProfiler/AutoInstrumentation/Kafka/ProduceEdgeTagCacheKey.cs
@@ -0,0 +1,43 @@
+// <copyright file="ProduceEdgeTagCacheKey.cs" company="Datadog">
+// 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;
+
+namespace Datadog.Trace.ClrProfiler.AutoInstrumentation.Kafka;
+
+/// <summary>
+/// Value-type cache key for produce edge tags. Using a named struct avoids boxing and
+/// is compatible with all supported target frameworks.
+/// </summary>
+internal readonly struct ProduceEdgeTagCacheKey : IEquatable<ProduceEdgeTagCacheKey>
+{
+    public readonly string ClusterId;
+    public readonly string Topic;
+
+    public ProduceEdgeTagCacheKey(string clusterId, string topic)
+    {
+        ClusterId = clusterId;
+        Topic = topic;
+    }
+
+    public bool Equals(ProduceEdgeTagCacheKey other)
+        => ClusterId == other.ClusterId && Topic == other.Topic;
+
+    public override bool Equals(object? obj)
+        => obj is ProduceEdgeTagCacheKey other && Equals(other);
+
+    public override int GetHashCode()
+    {
+        unchecked
+        {
+            int hash = 17;
+            hash = (hash * 31) + (ClusterId?.GetHashCode() ?? 0);
+            hash = (hash * 31) + (Topic?.GetHashCode() ?? 0);
+            return hash;
+        }
+    }
+}
diff --git a/tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsContextPropagator.cs b/tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsContextPropagator.cs
@@ -34,6 +34,29 @@ internal void Inject<TCarrier>(PathwayContext context, TCarrier headers, bool is
     {
         if (headers is null) { ThrowHelper.ThrowArgumentNullException(nameof(headers)); }
 
+#if NETCOREAPP3_1_OR_GREATER
+        // Encode directly into stack buffers to avoid heap allocations for the intermediate byte arrays.
+        // The only unavoidable allocation is the final ToArray() for headers.Add, since Kafka takes ownership.
+        Span<byte> encodedBytes = stackalloc byte[PathwayContextEncoder.MaxEncodedSize];
+        var encodedLen = PathwayContextEncoder.EncodeInto(context, encodedBytes);
+        var encodedSlice = encodedBytes.Slice(0, encodedLen);
+
+        Span<byte> base64Bytes = stackalloc byte[PathwayContextEncoder.MaxBase64EncodedSize];
+        var status = Base64.EncodeToUtf8(encodedSlice, base64Bytes, out _, out int bytesWritten);
+
+        if (status != OperationStatus.Done)
+        {
+            Log.Error("Failed to encode Data Streams context to Base64. OperationStatus: {Status}", status);
+            return;
+        }
+
+        headers.Add(DataStreamsPropagationHeaders.PropagationKeyBase64, base64Bytes.Slice(0, bytesWritten).ToArray());
+
+        if (isDataStreamsLegacyHeadersEnabled)
+        {
+            headers.Add(DataStreamsPropagationHeaders.PropagationKey, encodedSlice.ToArray());
+        }
+#else
         var encodedBytes = PathwayContextEncoder.Encode(context);
 
         // Calculate the maximum length of the base64 encoded data
@@ -62,6 +85,7 @@ internal void Inject<TCarrier>(PathwayContext context, TCarrier headers, bool is
         {
             headers.Add(DataStreamsPropagationHeaders.PropagationKey, encodedBytes);
         }
+#endif
     }
 
     /// <summary>
diff --git a/tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsManager.cs b/tracer/src/Datadog.Trace/DataStreamsMonitoring/DataStreamsManager.cs
@@ -7,6 +7,7 @@
 using System;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
+using System.Runtime.CompilerServices;
 using System.Threading;
 using System.Threading.Tasks;
 using Datadog.Trace.Agent.DiscoveryService;
@@ -27,6 +28,13 @@ namespace Datadog.Trace.DataStreamsMonitoring;
 /// </summary>
 internal sealed class DataStreamsManager
 {
+    /// <summary>
+    /// Maximum number of distinct keys stored in a single per-type edge-tag cache.
+    /// When the limit is reached, new keys are computed on the fly without caching to
+    /// prevent unbounded memory growth caused by high-cardinality identifiers.
+    /// </summary>
+    internal const int MaxEdgeTagCacheSize = 1000;
+
     private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor<DataStreamsManager>();
     private static readonly AsyncLocal<PathwayContext?> LastConsumePathway = new(); // saves the context on consume checkpointing only
     private readonly object _nodeHashUpdateLock = new();
@@ -36,6 +44,7 @@ internal sealed class DataStreamsManager
     private readonly IDisposable _updateSubscription;
     private readonly bool _isLegacyDsmHeadersEnabled;
     private readonly bool _isInDefaultState;
+    private readonly ConditionalWeakTable<string[], NodeHashCacheEntry> _nodeHashCache = new();
     private long _nodeHashBase; // note that this actually represents a `ulong` that we have done an unsafe cast for
     private MutableSettings _previousMutableSettings;
     private string? _previousContainerTagsHash;
@@ -303,7 +312,24 @@ 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 nodeHash = HashHelper.CalculateNodeHash(nodeHashBase, edgeTags);
+            var cacheEntry = _nodeHashCache.GetOrCreateValue(edgeTags);
+            NodeHash nodeHash;
+
+            // Fast lock-free path: snapshot is an immutable object published via a volatile field.
+            // If the base still matches we avoid taking any lock on the hot path.
+            if (!cacheEntry.TryGetNodeHash(nodeHashBase, out nodeHash))
+            {
+                lock (cacheEntry)
+                {
+                    // Double-check under lock in case another thread raced to update
+                    if (!cacheEntry.TryGetNodeHash(nodeHashBase, out nodeHash))
+                    {
+                        nodeHash = HashHelper.CalculateNodeHash(nodeHashBase, edgeTags);
+                        cacheEntry.Store(nodeHashBase, nodeHash);
+                    }
+                }
+            }
+
             var parentHash = previousContext?.Hash ?? default;
             var pathwayHash = HashHelper.CalculatePathwayHash(nodeHash, parentHash);
 
@@ -351,6 +377,34 @@ public void InjectPathwayContextAsBase64String<TCarrier>(PathwayContext? context
         }
     }
 
+    /// <summary>
+    /// Returns a cached edge-tag array for the given key, creating and caching it on first use.
+    /// On cache hits, zero heap allocations occur. The factory is only invoked on the first call
+    /// per unique key, making this safe to use on high-throughput hot paths.
+    /// Once the cache reaches <see cref="MaxEdgeTagCacheSize"/> entries the result is computed
+    /// fresh each time (no caching) to bound memory usage for high-cardinality key spaces.
+    /// </summary>
+    /// <typeparam name="TKey">A value type (struct) used as the cache key — no boxing.</typeparam>
+    /// <param name="key">The cache key derived from the caller's natural identifiers.</param>
+    /// <param name="factory">A static factory that builds the edge-tag array from the key on cache miss.</param>
+    public string[] GetOrCreateEdgeTags<TKey>(TKey key, Func<TKey, string[]> factory)
+        where TKey : notnull, IEquatable<TKey>
+    {
+        var cache = EdgeTagCache<TKey>.Cache;
+        if (cache.TryGetValue(key, out var existing))
+        {
+            return existing;
+        }
+
+        if (cache.Count >= MaxEdgeTagCacheSize)
+        {
+            // High-cardinality key space — bypass cache to prevent unbounded memory growth
+            return factory(key);
+        }
+
+        return cache.GetOrAdd(key, factory);
+    }
+
     /// <summary>
     /// Make sure we only extract the schema (a costly operation) on select occasions
     /// </summary>
@@ -371,4 +425,58 @@ public bool ShouldExtractSchema(Span span, string operation, out int weight)
         weight = 0;
         return false;
     }
+
+    /// <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
+    /// snapshot without a lock, and only acquire the lock when the base has changed or is missing.
+    /// </summary>
+    private sealed class NodeHashCacheEntry
+    {
+        // Immutable snapshot published via volatile write; null until first computation.
+        private volatile NodeHashSnapshot? _snapshot;
+
+        /// <summary>
+        /// Tries to return the cached <see cref="NodeHash"/> for <paramref name="nodeHashBase"/>
+        /// without acquiring any lock (lock-free read via volatile field).
+        /// </summary>
+        public bool TryGetNodeHash(NodeHashBase nodeHashBase, out NodeHash nodeHash)
+        {
+            var snap = _snapshot; // volatile read — acts as a load-acquire barrier
+            if (snap is not null && snap.Base == nodeHashBase.Value)
+            {
+                nodeHash = snap.Hash;
+                return true;
+            }
+
+            nodeHash = default;
+            return false;
+        }
+
+        /// <summary>
+        /// Stores a newly-computed <see cref="NodeHash"/>. Must be called under a lock held by the caller.
+        /// The volatile write ensures the snapshot is visible to all threads before the lock is released.
+        /// </summary>
+        public void Store(NodeHashBase nodeHashBase, NodeHash nodeHash)
+        {
+            _snapshot = new NodeHashSnapshot(nodeHashBase.Value, nodeHash); // volatile write
+        }
+
+        /// <summary>Immutable payload published atomically via the volatile <see cref="_snapshot"/> field.</summary>
+        private sealed class NodeHashSnapshot
+        {
+            private readonly ulong _base;
+            private readonly NodeHash _hash;
+
+            internal NodeHashSnapshot(ulong @base, NodeHash hash)
+            {
+                _base = @base;
+                _hash = hash;
+            }
+
+            internal ulong Base => _base;
+
+            internal NodeHash Hash => _hash;
+        }
+    }
 }
diff --git a/tracer/src/Datadog.Trace/DataStreamsMonitoring/EdgeTagCache.cs b/tracer/src/Datadog.Trace/DataStreamsMonitoring/EdgeTagCache.cs
@@ -0,0 +1,20 @@
+// <copyright file="EdgeTagCache.cs" company="Datadog">
+// 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>
+
+using System;
+using System.Collections.Concurrent;
+
+namespace Datadog.Trace.DataStreamsMonitoring;
+
+/// <summary>
+/// Process-wide cache of edge tag arrays, keyed by a caller-supplied value type.
+/// One dictionary instance exists per distinct TKey type (static generic class pattern).
+/// This avoids boxing and lets each integration use its own natural key shape.
+/// </summary>
+internal static class EdgeTagCache<TKey>
+    where TKey : notnull, IEquatable<TKey>
+{
+    internal static readonly ConcurrentDictionary<TKey, string[]> Cache = new();
+}
diff --git a/tracer/src/Datadog.Trace/DataStreamsMonitoring/PathwayContextEncoder.cs b/tracer/src/Datadog.Trace/DataStreamsMonitoring/PathwayContextEncoder.cs
@@ -15,6 +15,12 @@ internal static class PathwayContextEncoder
 {
     private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor(typeof(PathwayContextEncoder));
 
+    /// <summary>Maximum byte length produced by EncodeInto: 8 (hash) + 9 (pathway ms) + 9 (edge ms).</summary>
+    internal const int MaxEncodedSize = 26;
+
+    /// <summary>Maximum byte length of the Base64 encoding of <see cref="MaxEncodedSize"/> bytes: ceil(26/3)*4 = 36.</summary>
+    internal const int MaxBase64EncodedSize = 36;
+
     /// <summary>
     /// Encodes a <see cref="PathwayContext"/> as a series of bytes
     /// NOTE: the encoding is lossy, in that we convert <see cref="PathwayContext.PathwayStart"/>
@@ -38,6 +44,25 @@ public static byte[] Encode(PathwayContext pathway)
         return bytes;
     }
 
+#if NETCOREAPP3_1_OR_GREATER
+    /// <summary>
+    /// Zero-allocation alternative to <see cref="Encode"/>: writes the encoded pathway directly into
+    /// a caller-supplied <paramref name="buffer"/> (must be at least <see cref="MaxEncodedSize"/> bytes).
+    /// </summary>
+    /// <returns>Number of bytes written into <paramref name="buffer"/>.</returns>
+    public static int EncodeInto(PathwayContext pathway, Span<byte> buffer)
+    {
+        var pathwayStartMs = ToMilliseconds(pathway.PathwayStart);
+        var edgeStartMs = ToMilliseconds(pathway.EdgeStart);
+
+        BinaryPrimitivesHelper.WriteUInt64LittleEndian(buffer, pathway.Hash.Value);
+        var pathwayBytes = VarEncodingHelper.WriteVarLongZigZag(buffer.Slice(8), pathwayStartMs);
+        var edgeBytes = VarEncodingHelper.WriteVarLongZigZag(buffer.Slice(8 + pathwayBytes), edgeStartMs);
+
+        return 8 + pathwayBytes + edgeBytes;
+    }
+#endif
+
     /// <summary>
     /// Tries to decode a <see cref="PathwayContext"/> from a <c>byte[]</c>.
     /// NOTE: the encoding process is lossy, so the decoded <see cref="PathwayContext"/>
diff --git a/tracer/src/Datadog.Trace/DataStreamsMonitoring/Utils/BinaryPrimitivesHelper.cs b/tracer/src/Datadog.Trace/DataStreamsMonitoring/Utils/BinaryPrimitivesHelper.cs
@@ -13,6 +13,11 @@ namespace Datadog.Trace.DataStreamsMonitoring.Utils;
 
 internal static class BinaryPrimitivesHelper
 {
+#if NETCOREAPP3_1_OR_GREATER
+    public static void WriteUInt64LittleEndian(Span<byte> bytes, ulong value)
+        => System.Buffers.Binary.BinaryPrimitives.WriteUInt64LittleEndian(bytes, value);
+#endif
+
     public static void WriteUInt64LittleEndian(byte[] bytes, ulong value)
     {
 #if NETCOREAPP3_1_OR_GREATER
diff --git a/tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/DataStreamsManagerTests.cs b/tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/DataStreamsManagerTests.cs
diff --git a/tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/PathwayContextEncoderTests.cs b/tracer/test/Datadog.Trace.Tests/DataStreamsMonitoring/PathwayContextEncoderTests.cs

Original file line number	Diff line number	Diff line change
`@@ -13,6 +13,11 @@ namespace Datadog.Trace.DataStreamsMonitoring.Utils;`
`13`	`13`
`14`	`14`	`internal static class BinaryPrimitivesHelper`
`15`	`15`	`{`
	`16`	`+#if NETCOREAPP3_1_OR_GREATER`
	`17`	`+ public static void WriteUInt64LittleEndian(Span<byte> bytes, ulong value)`
	`18`	`+ => System.Buffers.Binary.BinaryPrimitives.WriteUInt64LittleEndian(bytes, value);`
	`19`	`+#endif`
	`20`	`+`
`16`	`21`	`public static void WriteUInt64LittleEndian(byte[] bytes, ulong value)`
`17`	`22`	`{`
`18`	`23`	`#if NETCOREAPP3_1_OR_GREATER`