[api] Optimise TraceContextPropagator.Extract (#5749)

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

src/OpenTelemetry.Api/CHANGELOG.md

@@ -11,6 +11,9 @@ Notes](../../RELEASENOTES.md).
  returned an empty set.
  ([#5745](https:/open-telemetry/opentelemetry-dotnet/pull/5745))
 
+* Optimize performance of `TraceContextPropagator.Extract`.
+ ([#5749](https:/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:/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:/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]