[api] Optimise TraceContextPropagator.Extract (#5749)

Co-authored-by: Vishwesh Bankwar <[email protected]>
Co-authored-by: Mikel Blanchard <[email protected]>

Author: 3 people

src/OpenTelemetry.Api/CHANGELOG.md

@@ -11,6 +11,9 @@ Notes](../../RELEASENOTES.md).
   returned an empty set.
   ([#5745](https://github.com/open-telemetry/opentelemetry-dotnet/pull/5745))
 
+* Optimize performance of `TraceContextPropagator.Extract`.
+  ([#5749](https://github.com/open-telemetry/opentelemetry-dotnet/pull/5749))
+
 ## 1.9.0
 
 Released 2024-Jun-14

src/OpenTelemetry.Api/Context/Propagation/TraceContextPropagator.cs

@@ -1,9 +1,9 @@
 // Copyright The OpenTelemetry Authors
 // SPDX-License-Identifier: Apache-2.0
 
+using System.Buffers;
 using System.Diagnostics;
 using System.Runtime.CompilerServices;
-using System.Text;
 using OpenTelemetry.Internal;
 
 namespace OpenTelemetry.Context.Propagation;
@@ -76,7 +76,7 @@ public override PropagationContext Extract<T>(PropagationContext context, T carr
             var tracestateCollection = getter(carrier, TraceState);
             if (tracestateCollection?.Any() ?? false)
             {
-                TryExtractTracestate(tracestateCollection.ToArray(), out tracestate);
+                TryExtractTracestate(tracestateCollection, out tracestate);
             }
 
             return new PropagationContext(
@@ -220,31 +220,37 @@ internal static bool TryExtractTraceparent(string traceparent, out ActivityTrace
         return true;
     }
 
-    internal static bool TryExtractTracestate(string[] tracestateCollection, out string tracestateResult)
+    internal static bool TryExtractTracestate(IEnumerable<string> tracestateCollection, out string tracestateResult)
     {
         tracestateResult = string.Empty;
 
-        if (tracestateCollection != null)
+        char[] rentedArray = null;
+        Span<char> traceStateBuffer = stackalloc char[128]; // 256B
+        Span<char> keyLookupBuffer = stackalloc char[96]; // 192B (3x32 keys)
+        int keys = 0;
+        int charsWritten = 0;
+
+        try
         {
-            var keySet = new HashSet<string>();
-            var result = new StringBuilder();
-            for (int i = 0; i < tracestateCollection.Length; ++i)
+            foreach (var tracestateItem in tracestateCollection)
             {
-                var tracestate = tracestateCollection[i].AsSpan();
-                int begin = 0;
-                while (begin < tracestate.Length)
+                var tracestate = tracestateItem.AsSpan();
+                int position = 0;
+
+                while (position < tracestate.Length)
                 {
-                    int length = tracestate.Slice(begin).IndexOf(',');
+                    int length = tracestate.Slice(position).IndexOf(',');
                     ReadOnlySpan<char> listMember;
+
                     if (length != -1)
                     {
-                        listMember = tracestate.Slice(begin, length).Trim();
-                        begin += length + 1;
+                        listMember = tracestate.Slice(position, length).Trim();
+                        position += length + 1;
                     }
                     else
                     {
-                        listMember = tracestate.Slice(begin).Trim();
-                        begin = tracestate.Length;
+                        listMember = tracestate.Slice(position).Trim();
+                        position = tracestate.Length;
                     }
 
                     // https://github.com/w3c/trace-context/blob/master/spec/20-http_request_header_format.md#tracestate-header-field-values
@@ -255,7 +261,7 @@ internal static bool TryExtractTracestate(string[] tracestateCollection, out str
                         continue;
                     }
 
-                    if (keySet.Count >= 32)
+                    if (keys >= 32)
                     {
                         // https://github.com/w3c/trace-context/blob/master/spec/20-http_request_header_format.md#list
                         // test_tracestate_member_count_limit
@@ -286,25 +292,107 @@ internal static bool TryExtractTracestate(string[] tracestateCollection, out str
                     }
 
                     // ValidateKey() call above has ensured the key does not contain upper case letters.
-                    if (!keySet.Add(key.ToString()))
+
+                    var duplicationCheckLength = Math.Min(key.Length, 3);
+
+                    if (keys > 0)
                     {
-                        // test_tracestate_duplicated_keys
-                        return false;
+                        // Fast path check of first three chars for potential duplicated keys
+                        var potentialMatchingKeyPosition = 1;
+                        var found = false;
+                        for (int i = 0; i < keys * 3; i += 3)
+                        {
+                            if (keyLookupBuffer.Slice(i, duplicationCheckLength).SequenceEqual(key.Slice(0, duplicationCheckLength)))
+                            {
+                                found = true;
+                                break;
+                            }
+
+                            potentialMatchingKeyPosition++;
+                        }
+
+                        // If the fast check has found a possible duplicate, we need to do a full check
+                        if (found)
+                        {
+                            var bufferToCompare = traceStateBuffer.Slice(0, charsWritten);
+
+                            // We know which key is the first possible duplicate, so skip to that key
+                            // by slicing to the position after the appropriate comma.
+                            for (int i = 1; i < potentialMatchingKeyPosition; i++)
+                            {
+                                var commaIndex = bufferToCompare.IndexOf(',');
+
+                                if (commaIndex > -1)
+                                {
+                                    bufferToCompare.Slice(commaIndex);
+                                }
+                            }
+
+                            int existingIndex = -1;
+                            while ((existingIndex = bufferToCompare.IndexOf(key)) > -1)
+                            {
+                                if ((existingIndex > 0 && bufferToCompare[existingIndex - 1] != ',') || bufferToCompare[existingIndex + key.Length] != '=')
+                                {
+                                    continue; // this is not a key
+                                }
+
+                                return false; // test_tracestate_duplicated_keys
+                            }
+                        }
                     }
 
-                    if (result.Length > 0)
+                    // Store up to the first three characters of the key for use in the duplicate lookup fast path
+                    var startKeyLookupIndex = keys > 0 ? keys * 3 : 0;
+                    key.Slice(0, duplicationCheckLength).CopyTo(keyLookupBuffer.Slice(startKeyLookupIndex));
+
+                    // Check we have capacity to write the key and value
+                    var requiredCapacity = charsWritten > 0 ? listMember.Length + 1 : listMember.Length;
+
+                    while (charsWritten + requiredCapacity > traceStateBuffer.Length)
                     {
-                        result.Append(',');
+                        GrowBuffer(ref rentedArray, ref traceStateBuffer);
                     }
 
-                    result.Append(listMember.ToString());
+                    if (charsWritten > 0)
+                    {
+                        traceStateBuffer[charsWritten++] = ',';
+                    }
+
+                    listMember.CopyTo(traceStateBuffer.Slice(charsWritten));
+                    charsWritten += listMember.Length;
+
+                    keys++;
                 }
             }
 
-            tracestateResult = result.ToString();
+            tracestateResult = traceStateBuffer.Slice(0, charsWritten).ToString();
+
+            return true;
+        }
+        finally
+        {
+            if (rentedArray is not null)
+            {
+                ArrayPool<char>.Shared.Return(rentedArray);
+                rentedArray = null;
+            }
         }
 
-        return true;
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        static void GrowBuffer(ref char[] array, ref Span<char> buffer)
+        {
+            var newBuffer = ArrayPool<char>.Shared.Rent(buffer.Length * 2);
+
+            buffer.CopyTo(newBuffer.AsSpan());
+
+            if (array is not null)
+            {
+                ArrayPool<char>.Shared.Return(array);
+            }
+
+            array = newBuffer;
+            buffer = array.AsSpan();
+        }
     }
 
     private static byte HexCharToByte(char c)

test/Benchmarks/Context/Propagation/TraceContextPropagatorBenchmarks.cs

@@ -0,0 +1,72 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+using BenchmarkDotNet.Attributes;
+using OpenTelemetry.Context.Propagation;
+
+namespace Benchmarks.Context.Propagation;
+
+public class TraceContextPropagatorBenchmarks
+{
+    private const string TraceParent = "traceparent";
+    private const string TraceState = "tracestate";
+    private const string TraceId = "0af7651916cd43dd8448eb211c80319c";
+    private const string SpanId = "b9c7c989f97918e1";
+
+    private static readonly Random Random = new(455946);
+    private static readonly TraceContextPropagator TraceContextPropagator = new();
+
+    private static readonly Func<IReadOnlyDictionary<string, string>, string, IEnumerable<string>> Getter = (headers, name) =>
+    {
+        if (headers.TryGetValue(name, out var value))
+        {
+            return [value];
+        }
+
+        return [];
+    };
+
+    private Dictionary<string, string> headers;
+
+    [Params(true, false)]
+    public bool LongListMember { get; set; }
+
+    [Params(0, 4, 32)]
+    public int MembersCount { get; set; }
+
+    public Dictionary<string, string> Headers => this.headers;
+
+    [GlobalSetup]
+    public void Setup()
+    {
+        var length = this.LongListMember ? 256 : 20;
+
+        var value = new string('a', length);
+
+        Span<char> keyBuffer = stackalloc char[length - 2];
+
+        string traceState = string.Empty;
+        for (var i = 0; i < this.MembersCount; i++)
+        {
+            // We want a unique key for each member
+            for (var j = 0; j < length - 2; j++)
+            {
+                keyBuffer[j] = (char)('a' + Random.Next(0, 26));
+            }
+
+            var key = keyBuffer.ToString();
+
+            var listMember = $"{key}{i:00}={value}";
+            traceState += i < this.MembersCount - 1 ? $"{listMember}," : listMember;
+        }
+
+        this.headers = new Dictionary<string, string>
+        {
+            { TraceParent, $"00-{TraceId}-{SpanId}-01" },
+            { TraceState, traceState },
+        };
+    }
+
+    [Benchmark(Baseline = true)]
+    public void Extract() => _ = TraceContextPropagator!.Extract(default, this.headers, Getter);
+}

test/OpenTelemetry.Tests/Trace/Propagation/TraceContextPropagatorTest.cs

@@ -13,12 +13,12 @@ public class TraceContextPropagatorTest
     private const string TraceId = "0af7651916cd43dd8448eb211c80319c";
     private const string SpanId = "b9c7c989f97918e1";
 
-    private static readonly string[] Empty = Array.Empty<string>();
+    private static readonly string[] Empty = [];
     private static readonly Func<IDictionary<string, string>, string, IEnumerable<string>> Getter = (headers, name) =>
     {
         if (headers.TryGetValue(name, out var value))
         {
-            return new[] { value };
+            return [value];
         }
 
         return Empty;
@@ -31,7 +31,7 @@ public class TraceContextPropagatorTest
             return value;
         }
 
-        return Array.Empty<string>();
+        return [];
     };
 
     private static readonly Action<IDictionary<string, string>, string, string> Setter = (carrier, name, value) =>
@@ -183,8 +183,18 @@ public void DuplicateKeys()
         // test_tracestate_duplicated_keys
         Assert.Empty(CallTraceContextPropagator("foo=1,foo=1"));
         Assert.Empty(CallTraceContextPropagator("foo=1,foo=2"));
-        Assert.Empty(CallTraceContextPropagator(new[] { "foo=1", "foo=1" }));
-        Assert.Empty(CallTraceContextPropagator(new[] { "foo=1", "foo=2" }));
+        Assert.Empty(CallTraceContextPropagator(["foo=1", "foo=1"]));
+        Assert.Empty(CallTraceContextPropagator(["foo=1", "foo=2"]));
+        Assert.Empty(CallTraceContextPropagator("foo=1,bar=2,baz=3,foo=4"));
+    }
+
+    [Fact]
+    public void NoDuplicateKeys()
+    {
+        Assert.Equal("foo=1,bar=foo,baz=2", CallTraceContextPropagator("foo=1,bar=foo,baz=2"));
+        Assert.Equal("foo=1,bar=2,baz=foo", CallTraceContextPropagator("foo=1,bar=2,baz=foo"));
+        Assert.Equal("foo=1,foo@tenant=2", CallTraceContextPropagator("foo=1,foo@tenant=2"));
+        Assert.Equal("foo=1,tenant@foo=2", CallTraceContextPropagator("foo=1,tenant@foo=2"));
     }
 
     [Fact]