Skip to content

[api] Optimise TraceContextPropagator.Extract #5749

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 5 commits into from
Aug 12, 2024
Merged

[api] Optimise TraceContextPropagator.Extract #5749

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
d01fa96
Optimise TraceContextPropagator.Extract
stevejgordon Jul 15, 2024
538592e
Update changelog
stevejgordon Jul 15, 2024
a24b5a9
Reduce initial stackalloc buffer size and prefer try/finally
stevejgordon Jul 16, 2024
35af52d
Fix typo
stevejgordon Aug 8, 2024
d43e26f
Merge from main.
CodeBlanch Aug 12, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions src/OpenTelemetry.Api/CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,9 @@

## Unreleased

* Optimize performance of `TraceContextPropagator.Extract`.
([#5749](https://github.com/open-telemetry/opentelemetry-dotnet/pull/5749))

## 1.9.0

Released 2024-Jun-14
Expand Down
217 changes: 153 additions & 64 deletions src/OpenTelemetry.Api/Context/Propagation/TraceContextPropagator.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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;
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -220,91 +220,180 @@ 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[] array = null;
Span<char> buffer = stackalloc char[256];
Span<char> keyLookupBuffer = stackalloc char[96]; // 3 x 32 keys
Copy link
Member

@CodeBlanch CodeBlanch Jul 15, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Love what this PR is doing but I have some concern here with allocating 704 bytes on the stack ((256 + 96) * 2).

System.Text.Json for example will at max stackalloc 256 bytes (128 chars): https://github.com/dotnet/runtime/blob/a86987ccab917433d065fe5dc8870fc261f79d14/src/libraries/System.Text.Json/Common/JsonConstants.cs#L12-L13

Not sure why that number, but I'm sure a lot of thought went into it 😄 /cc @stephentoub

I think this pattern for "stackalloc with fallback to rented array" is solid but the version here seems a bit off from what I have seen. IMO it more commonly looks like this:

char[]? rentedBuffer = null;
Span<char> destination = length <= Constants.StackallocCharThreshold ?
   stackalloc char[Constants.StackallocCharThreshold] :
   (rentedBuffer = ArrayPool<char>.Shared.Rent(length));

try
{
   ...
}
finally
{
   if (rentedBuffer != null)
      ArrayPool<char>.Shared.Return(rentedBuffer );
}

Would something like that work/help simplify things here?

Copy link

@stephentoub stephentoub Jul 15, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not sure why that number, but I'm sure a lot of thought went into it 😄

It's a bit squishy. We almost never go above 1K. We typically use 256 or 512 bytes, but it varies case-to-case based on knowledge of that particular location and how likely longer buffers are expected to be needed.

Copy link
Contributor Author

@stevejgordon stevejgordon Jul 16, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@CodeBlanch I think I'd seen somewhere that 1K was a "reasonable" max value. I went for 256 chars as, realistically, that should cover most trace state scenarios. We could drop to 128 (256B) instead, though, and still cover most shorter trace state strings. Combined with the 96 chars (192B) for the duplicate lookup, that might be more reasonable.

I can switch to the try/finally here. The Return method initially had some extra logic, but I refactored that before the PR. I left it in that form because I preferred avoiding one extra level of indentation introduced with the extra blocks.

Copy link

@gfoidl gfoidl Jul 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In the pattern the try-finally can be ommited. In case of exceptions the buffer (if rented) will just be dropped instead of returned to the pool, but that isn't a problem to the pool.

In first incarnations of this pattern try-finally got used, but later on that pattern evolved.

Copy link
Member

@CodeBlanch CodeBlanch Jul 17, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't care so much about the try-finally we could drop that. Just want to make sure the stackalloc is good.

We can't really control where Extract runs so it is probably a good idea to be conservative. For incoming request (think AspNetCore instrumentation) it will probably be ~early. But something like processing messages from a queue, who knows! What I would really like to see is us avoid the stackalloc if we know there is a lot of tracestate but I guess hard to do because there could be multiple headers needing to be processed?

Copy link
Contributor Author

@stevejgordon stevejgordon Jul 23, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@CodeBlanch We could just be safe and go with ArrayPool. I preferred stackalloc since we might reasonably expect the state to be small, and we can avoid a small amount of overhead that we incur by renting.

Copy link
Contributor Author

@stevejgordon stevejgordon Jul 23, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@gfoidl I think we need the finally here, though, as there are multiple branches where this code could return from this method on the non-exception path, and we need to ensure that the array is returned in those cases, too.

int keys = 0;
int charsWritten = 0;

foreach (var tracestateItem in tracestateCollection)
{
var keySet = new HashSet<string>();
var result = new StringBuilder();
for (int i = 0; i < tracestateCollection.Length; ++i)
var tracestate = tracestateItem.AsSpan();
int position = 0;

while (position < tracestate.Length)
{
var tracestate = tracestateCollection[i].AsSpan();
int begin = 0;
while (begin < tracestate.Length)
int length = tracestate.Slice(position).IndexOf(',');
ReadOnlySpan<char> listMember;

if (length != -1)
{
int length = tracestate.Slice(begin).IndexOf(',');
ReadOnlySpan<char> listMember;
if (length != -1)
{
listMember = tracestate.Slice(begin, length).Trim();
begin += length + 1;
}
else
{
listMember = tracestate.Slice(begin).Trim();
begin = tracestate.Length;
}
listMember = tracestate.Slice(position, length).Trim();
position += length + 1;
}
else
{
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
if (listMember.IsEmpty)
{
// Empty and whitespace - only list members are allowed.
// Vendors MUST accept empty tracestate headers but SHOULD avoid sending them.
continue;
}
// https://github.com/w3c/trace-context/blob/master/spec/20-http_request_header_format.md#tracestate-header-field-values
if (listMember.IsEmpty)
{
// Empty and whitespace - only list members are allowed.
// Vendors MUST accept empty tracestate headers but SHOULD avoid sending them.
continue;
}

if (keySet.Count >= 32)
{
// https://github.com/w3c/trace-context/blob/master/spec/20-http_request_header_format.md#list
// test_tracestate_member_count_limit
return false;
}
if (keys >= 32)
{
// https://github.com/w3c/trace-context/blob/master/spec/20-http_request_header_format.md#list
// test_tracestate_member_count_limit
return Return(false, ref array);
}

int keyLength = listMember.IndexOf('=');
if (keyLength == listMember.Length || keyLength == -1)
{
// Missing key or value in tracestate
return false;
}
int keyLength = listMember.IndexOf('=');
if (keyLength == listMember.Length || keyLength == -1)
{
// Missing key or value in tracestate
return Return(false, ref array);
}

var key = listMember.Slice(0, keyLength);
if (!ValidateKey(key))
{
// test_tracestate_key_illegal_characters in https://github.com/w3c/trace-context/blob/master/test/test.py
// test_tracestate_key_length_limit
// test_tracestate_key_illegal_vendor_format
return false;
}
var key = listMember.Slice(0, keyLength);
if (!ValidateKey(key))
{
// test_tracestate_key_illegal_characters in https://github.com/w3c/trace-context/blob/master/test/test.py
// test_tracestate_key_length_limit
// test_tracestate_key_illegal_vendor_format
return Return(false, ref array);
}

var value = listMember.Slice(keyLength + 1);
if (!ValidateValue(value))
{
// test_tracestate_value_illegal_characters
return false;
}
var value = listMember.Slice(keyLength + 1);
if (!ValidateValue(value))
{
// test_tracestate_value_illegal_characters
return Return(false, ref array);
}

// ValidateKey() call above has ensured the key does not contain upper case letters.

// 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)
{
// 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)
{
// test_tracestate_duplicated_keys
return false;
if (keyLookupBuffer.Slice(i, duplicationCheckLength).SequenceEqual(key.Slice(0, duplicationCheckLength)))
{
found = true;
break;
}

potentialMatchingKeyPosition++;
}

if (result.Length > 0)
// If the fast check has found a possible duplicate, we need to do a full check
if (found)
{
result.Append(',');
var bufferToCompare = buffer.Slice(0, charsWritten);

// We know which key is the first posible 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 Return(false, ref array); // test_tracestate_duplicated_keys
}
}
}

// 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;

result.Append(listMember.ToString());
while (charsWritten + requiredCapacity > buffer.Length)
{
GrowBuffer(ref array, ref buffer);
}

if (charsWritten > 0)
{
buffer[charsWritten++] = ',';
}

listMember.CopyTo(buffer.Slice(charsWritten));
charsWritten += listMember.Length;

keys++;
}
}

tracestateResult = result.ToString();
tracestateResult = buffer.Slice(0, charsWritten).ToString();

return Return(true, ref array);

[MethodImpl(MethodImplOptions.AggressiveInlining)]
static bool Return(bool returnValue, ref char[] array)
{
if (array is not null)
{
ArrayPool<char>.Shared.Return(array);
array = null;
}

return returnValue;
}

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)
Expand Down
72 changes: 72 additions & 0 deletions test/Benchmarks/Context/Propagation/TraceContextPropagatorBenchmarks.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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);
}
14 changes: 12 additions & 2 deletions test/OpenTelemetry.Tests/Trace/Propagation/TraceContextPropagatorTest.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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]
Expand Down