Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ namespace Datadog.Trace.Activity.Handlers
internal sealed class ActivityHandlerCommon
{
private static readonly IDatadogLogger Log = DatadogLogging.GetLoggerFor(typeof(ActivityHandlerCommon));
internal static readonly ConcurrentDictionary<string, ActivityMapping> ActivityMappingById = new();
internal static readonly ConcurrentDictionary<ActivityKey, ActivityMapping> ActivityMappingById = new();
private static readonly IntegrationId IntegrationId = IntegrationId.OpenTelemetry;

/// <summary>
Expand All @@ -48,10 +48,10 @@ public static void ActivityStarted<T>(string sourceName, T activity, OpenTelemet
string? rawTraceId = null;
string? rawSpanId = null;

// for non-IW3CActivity interfaces we'll use Activity.Id as the key as they don't have a guaranteed TraceId+SpanId
// for non-IW3CActivity interfaces we'll use Activity.Id and string.Empty as the key as they don't have a guaranteed TraceId+SpanId
// for IW3CActivity interfaces we'll use the Activity.TraceId + Activity.SpanId as the key
// have to also validate that the TraceId and SpanId actually exist and aren't null - as they can be in some cases
string? activityKey = null;
ActivityKey? activityKey = null;

if (activity is IW3CActivity w3cActivity)
{
Expand All @@ -67,7 +67,7 @@ public static void ActivityStarted<T>(string sourceName, T activity, OpenTelemet
// Doing a lookup on just the TraceId+ParentSpanId seems to be more resilient.
if (activityTraceId != null!)
{
if (ActivityMappingById.TryGetValue(activityTraceId + parentSpanId, out ActivityMapping mapping))
if (ActivityMappingById.TryGetValue(new ActivityKey(activityTraceId, parentSpanId), out ActivityMapping mapping))
{
parent = mapping.Scope.Span.Context;
}
Expand All @@ -88,7 +88,7 @@ public static void ActivityStarted<T>(string sourceName, T activity, OpenTelemet
else
{
// we have a ParentSpanId/ParentId, but no TraceId/SpanId, so default to use the ParentId for lookup
if (ActivityMappingById.TryGetValue(parentId, out ActivityMapping mapping))
if (ActivityMappingById.TryGetValue(new ActivityKey(parentId), out ActivityMapping mapping))
{
parent = mapping.Scope.Span.Context;
}
Expand Down Expand Up @@ -136,7 +136,7 @@ public static void ActivityStarted<T>(string sourceName, T activity, OpenTelemet

if (activityTraceId != null! && activitySpanId != null!)
{
activityKey = activityTraceId + activitySpanId;
activityKey = new(traceId: activityTraceId, spanId: activitySpanId);
}
}

Expand Down Expand Up @@ -164,9 +164,9 @@ public static void ActivityStarted<T>(string sourceName, T activity, OpenTelemet
return;
}

activityKey ??= activity.Id;
activityKey ??= new(activity.Id);

if (activityKey is null)
if (!activityKey.Value.IsValid())
{
// identified by Error Tracking
// unsure how exactly this occurs after reading through the Activity source code
Expand All @@ -176,7 +176,7 @@ public static void ActivityStarted<T>(string sourceName, T activity, OpenTelemet
return;
}

activityMapping = ActivityMappingById.GetOrAdd(activityKey, _ => new(activity.Instance!, CreateScopeFromActivity(activity, tags, parent, traceId, spanId, rawTraceId, rawSpanId)));
activityMapping = ActivityMappingById.GetOrAdd(activityKey.Value, _ => new(activity.Instance!, CreateScopeFromActivity(activity, tags, parent, traceId, spanId, rawTraceId, rawSpanId)));
}
catch (Exception ex)
{
Expand Down Expand Up @@ -213,17 +213,20 @@ public static void ActivityStopped<T>(string sourceName, T activity)
return;
}

string key;
if (activity is IW3CActivity w3cActivity)
// Non-w3c activities will have null trace/span IDs, even though they implement IW3CActivity
ActivityKey key;
if (activity is IW3CActivity w3cActivity
&& w3cActivity.TraceId != null!
&& w3cActivity.SpanId != null!)
Comment on lines +216 to +220
Copy link
Member Author

Choose a reason for hiding this comment

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

I think this is the bug we were hitting with hierarchical IDs:

  • We were always landing in this branch (IW3CActivity is about the Version of Activity, not the type of ID in use)
  • The key was an empty string
  • We'd do the lookup on the dictionary, and it would fail
  • We'd then enumerate the dictionary, duck typing repeatedly, trying to find the activity that closed
  • massive allocation 💥

Copy link
Collaborator

Choose a reason for hiding this comment

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

Ohhh wow nice catch

{
key = w3cActivity.TraceId + w3cActivity.SpanId;
key = new(w3cActivity.TraceId, w3cActivity.SpanId);
}
else
{
key = activity.Id;
key = new(activity.Id);
}

if (key is null)
if (!key.IsValid())
{
// Adding this as a protective measure as Error Tracking identified
// instances where StartActivity had an Activity with null Id, SpanId, TraceId
Expand Down Expand Up @@ -259,9 +262,11 @@ public static void ActivityStopped<T>(string sourceName, T activity)
Log.Information($"DefaultActivityHandler.ActivityStopped: [Missing Activity]");
}

List<string>? toDelete = null;
foreach (var (activityId, item) in ActivityMappingById)
List<ActivityKey>? toDelete = null;
foreach (var kvp in ActivityMappingById)
{
var activityId = kvp.Key;
var item = kvp.Value;
var activityObject = item.Activity;
var scope = item.Scope;
var hasClosed = false;
Expand Down Expand Up @@ -293,7 +298,7 @@ public static void ActivityStopped<T>(string sourceName, T activity)

if (hasClosed)
{
toDelete ??= new List<string>();
toDelete ??= new List<ActivityKey>();
toDelete.Add(activityId);
}
}
Expand Down
50 changes: 50 additions & 0 deletions tracer/src/Datadog.Trace/Activity/Handlers/ActivityKey.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,50 @@
// <copyright file="ActivityKey.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;
using System.Collections.Generic;

namespace Datadog.Trace.Activity.Handlers;

internal readonly struct ActivityKey : IEquatable<ActivityKey>
{
private readonly string _traceId;
private readonly string _spanId;

public ActivityKey(string traceId, string spanId)
{
_traceId = traceId;
_spanId = spanId;
}

public ActivityKey(string id)
{
// Arbitrary which way around these are
_spanId = id;
_traceId = string.Empty;
}

public bool IsValid() => _traceId is not null && _spanId is not null;

public bool Equals(ActivityKey other) =>
string.Equals(_traceId, other._traceId, StringComparison.Ordinal) &&
string.Equals(_spanId, other._spanId, StringComparison.Ordinal);

public override bool Equals(object? obj) =>
obj is ActivityKey other && Equals(other);

public override int GetHashCode()
{
// TraceId and SpanId must be non-null, so can just just the values directly.
// If we allow this to be called with any null values, then we must add null checks here
unchecked
{
return (StringComparer.Ordinal.GetHashCode(_traceId) * 397)
^ StringComparer.Ordinal.GetHashCode(_spanId);
}
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -46,17 +46,17 @@ public void ActivityStopped<T>(string sourceName, T activity)
// then we can retrieve the active message object using our mapping. With access to the message
// object, we can accurately calculate the payload size for the DataStreamsCheckpoint

string key;
ActivityKey key;
if (activity is IW3CActivity w3cActivity)
{
key = w3cActivity.TraceId + w3cActivity.SpanId;
key = new(w3cActivity.TraceId, w3cActivity.SpanId);
}
else
{
key = activity.Id;
key = new ActivityKey(activity.Id);
}

if (ActivityHandlerCommon.ActivityMappingById.TryRemove(key, out ActivityMapping activityMapping)
if (key.IsValid() && ActivityHandlerCommon.ActivityMappingById.TryRemove(key, out ActivityMapping activityMapping)
&& activityMapping.Scope?.Span is Span span)
{
// Copy over the data to our Span object so we can do an efficient tags lookup
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -45,13 +45,13 @@ public void ActivityStopped<T>(string sourceName, T activity)
where T : IActivity
{
// Find the span and update it before the common handler processes it
string key = activity switch
ActivityKey key = activity switch
{
IW3CActivity w3cActivity => w3cActivity.TraceId + w3cActivity.SpanId,
_ => activity.Id
IW3CActivity w3cActivity => new(w3cActivity.TraceId, w3cActivity.SpanId),
_ => new(activity.Id)
};

if (key != null && ActivityHandlerCommon.ActivityMappingById.TryRemove(key, out var activityMapping) && activityMapping.Scope.Span is Span span)
if (key.IsValid() && ActivityHandlerCommon.ActivityMappingById.TryRemove(key, out var activityMapping) && activityMapping.Scope.Span is Span span)
{
Log.Debug("ActivityStopped: Processing span for activity '{ActivityId}'", activity.Id);

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -31,17 +31,17 @@ public static void OnStart(object processor, object activityData)
return;
}

string key;
ActivityKey key;
if (activityData.TryDuckCast<IW3CActivity>(out var w3cActivity) && w3cActivity.TraceId is { } traceId && w3cActivity.SpanId is { } spanId)
{
key = traceId + spanId;
key = new(traceId, spanId);
}
else
{
key = activity.Id;
key = new(activity.Id);
}

if (ActivityHandlerCommon.ActivityMappingById.TryGetValue(key, out var activityMapping))
if (key.IsValid() && ActivityHandlerCommon.ActivityMappingById.TryGetValue(key, out var activityMapping))
{
if (baseProcessor.ParentProvider is not null)
{
Expand Down
73 changes: 73 additions & 0 deletions tracer/test/Datadog.Trace.Tests/Activity/ActivityKeyTests.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
// <copyright file="ActivityKeyTests.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.Collections.Generic;
using Datadog.Trace.Activity.Handlers;
using FluentAssertions;
using Xunit;

namespace Datadog.Trace.Tests.Activity;

public class ActivityKeyTests
{
[Theory]
[InlineData("spanId", "traceId")]
[InlineData("id", "")]
public void Equality(string spanId, string traceId)
{
var activity1 = new ActivityKey(traceId: traceId, spanId: spanId);
var activity2 = new ActivityKey(traceId: traceId, spanId: spanId);
activity1.Equals(activity2).Should().BeTrue();
activity1.GetHashCode().Should().Be(activity2.GetHashCode());

var other = new ActivityKey(traceId: traceId, spanId: "newId");
activity1.Equals(other).Should().BeFalse();
activity1.GetHashCode().Should().NotBe(other.GetHashCode());
activity2.Equals(other).Should().BeFalse();
activity2.GetHashCode().Should().NotBe(other.GetHashCode());
}

[Theory]
[InlineData("spanId", "traceId")]
[InlineData("id", "")]
public void Equality_DefaultComparer(string spanId, string traceId)
{
var activity1 = new ActivityKey(traceId: traceId, spanId: spanId);
var activity2 = new ActivityKey(traceId: traceId, spanId: spanId);

EqualityComparer<ActivityKey>.Default.Equals(activity1, activity2).Should().BeTrue();
var hashCode1 = EqualityComparer<ActivityKey>.Default.GetHashCode(activity1).GetHashCode();
var hashCode2 = EqualityComparer<ActivityKey>.Default.GetHashCode(activity2);
hashCode1.Should().Be(hashCode2);

var other = new ActivityKey(traceId: traceId, spanId: "newId");
var hashCodeOther = EqualityComparer<ActivityKey>.Default.GetHashCode(other).GetHashCode();
EqualityComparer<ActivityKey>.Default.Equals(activity1, other).Should().BeFalse();
hashCode1.Should().NotBe(hashCodeOther);
EqualityComparer<ActivityKey>.Default.Equals(activity2, other).Should().BeFalse();
hashCode2.Should().NotBe(hashCodeOther);
}

[Theory]
[InlineData("spanId", "traceId", true)]
[InlineData("spanId", "", true)]
[InlineData("", "traceId", true)]
[InlineData("", "", true)]
[InlineData("spanId", null, false)]
[InlineData(null, "traceId", false)]
public void IsValid_SpanId_TraceId(string traceId, string spanId, bool isValid)
{
new ActivityKey(traceId: traceId, spanId: spanId).IsValid().Should().Be(isValid);
}

[Theory]
[InlineData("some_id", true)]
[InlineData("", true)]
[InlineData(null, false)]
public void IsValid_SpanId_Id(string id, bool isValid)
{
new ActivityKey(id).IsValid().Should().Be(isValid);
}
}
4 changes: 2 additions & 2 deletions tracer/test/Datadog.Trace.Tests/ActivityTests.cs
Original file line number Diff line number Diff line change
Expand Up @@ -207,12 +207,12 @@ public class ActivityFixture : IDisposable

public ActivityFixture()
{
Activity.ActivityListener.Initialize();
Datadog.Trace.Activity.ActivityListener.Initialize();
}

public void Dispose()
{
Activity.ActivityListener.StopListeners();
Datadog.Trace.Activity.ActivityListener.StopListeners();
}

public void StartActivity(SD.Activity activity)
Expand Down
Loading