Skip to content

Commit e5ae882

Browse files
committed
ActivityTraceListenerManager now uses an OpenTelemetry sampler instead of ActivityListener.
1 parent d465930 commit e5ae882

10 files changed

+160
-24
lines changed

ClassLibraries/Macross.OpenTelemetry.Extensions/Code/ActivityEnrichmentScope/ActivityEnrichmentScopeProcessor.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
namespace Macross.OpenTelemetry.Extensions
77
{
8-
internal class ActivityEnrichmentScopeProcessor : BaseProcessor<Activity>
8+
internal sealed class ActivityEnrichmentScopeProcessor : BaseProcessor<Activity>
99
{
1010
/// <inheritdoc/>
1111
public override void OnEnd(Activity activity)

ClassLibraries/Macross.OpenTelemetry.Extensions/Code/ActivityTraceListenerManager.cs

Lines changed: 25 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,15 @@
11
using System.Collections.Concurrent;
22
using System.Collections.Generic;
3+
using System.Reflection;
34
using System.Runtime.CompilerServices;
45
using System.Threading;
56

67
using Microsoft.Extensions.Options;
78

9+
using OpenTelemetry.Trace;
10+
11+
using Macross.OpenTelemetry.Extensions;
12+
813
namespace System.Diagnostics
914
{
1015
/// <summary>
@@ -16,23 +21,35 @@ public class ActivityTraceListenerManager : IDisposable
1621

1722
private readonly ConcurrentDictionary<ActivityTraceId, TraceListener> _TraceListeners = new ConcurrentDictionary<ActivityTraceId, TraceListener>();
1823
private readonly EventWaitHandle _StopHandle = new EventWaitHandle(false, EventResetMode.ManualReset);
24+
private readonly ActivityTraceListenerSampler _ActivityTraceListenerSampler;
1925
private readonly IDisposable _OptionsChangeToken;
2026
private ActivityTraceListenerManagerOptions? _Options;
2127
private Thread? _CleanupThread;
22-
private ActivityListener? _ActivityListener;
2328
private bool _HasInitialized;
2429
private long _LastRequestedListenerDateTimeBinary;
2530

2631
/// <summary>
2732
/// Initializes a new instance of the <see cref="ActivityTraceListenerManager"/> class.
2833
/// </summary>
34+
/// <param name="tracerProvider"><see cref="TracerProvider"/>.</param>
2935
/// <param name="options"><see cref="ActivityTraceListenerManagerOptions"/>.</param>
3036
public ActivityTraceListenerManager(
37+
TracerProvider tracerProvider,
3138
IOptionsMonitor<ActivityTraceListenerManagerOptions> options)
3239
{
40+
if (tracerProvider == null)
41+
throw new ArgumentNullException(nameof(tracerProvider));
3342
if (options == null)
3443
throw new ArgumentNullException(nameof(options));
3544

45+
FieldInfo? samplerFieldInfo = tracerProvider.GetType().GetField("sampler", BindingFlags.Instance | BindingFlags.NonPublic);
46+
if (samplerFieldInfo == null)
47+
throw new NotSupportedException($"sampler field could not be read reflectively on tracerProvider of type {tracerProvider.GetType()}.");
48+
49+
_ActivityTraceListenerSampler = (samplerFieldInfo.GetValue(tracerProvider) as ActivityTraceListenerSampler)!;
50+
if (_ActivityTraceListenerSampler == null)
51+
throw new NotSupportedException("ActivityTraceListenerManager cannot be used without the ActivityTraceListenerSampler. Call SetActivityTraceListenerSampler on TracerProviderBuilder during startup.");
52+
3653
ApplyOptions(options.CurrentValue);
3754
_OptionsChangeToken = options.OnChange(ApplyOptions);
3855
}
@@ -82,6 +99,10 @@ public IActivityTraceListener RegisterTraceListener(ActivityTraceId traceId)
8299
: listener;
83100
}
84101

102+
[MethodImpl(MethodImplOptions.AggressiveInlining)]
103+
internal bool IsTraceIdRegistered(ActivityTraceId activityTraceId)
104+
=> _TraceListeners.ContainsKey(activityTraceId);
105+
85106
/// <summary>
86107
/// Releases the unmanaged resources used by this class and optionally releases the managed resources.
87108
/// </summary>
@@ -99,16 +120,10 @@ protected virtual void Dispose(bool isDisposing)
99120
{
100121
_OptionsChangeToken.Dispose();
101122
_StopHandle.Dispose();
102-
_ActivityListener?.Dispose();
123+
_ActivityTraceListenerSampler.ActivityTraceListenerManager = null;
103124
}
104125
}
105126

106-
private void OnActivityStopped(Activity activity)
107-
{
108-
if (_TraceListeners.TryGetValue(activity.TraceId, out TraceListener? listener))
109-
listener.Add(activity);
110-
}
111-
112127
[MethodImpl(MethodImplOptions.AggressiveInlining)]
113128
private void EnsureInitialized()
114129
{
@@ -126,14 +141,7 @@ private void EnsureInitialized()
126141
};
127142
_CleanupThread.Start();
128143

129-
// Watch out doing this in prod, it's expensive.
130-
_ActivityListener = new ActivityListener
131-
{
132-
ShouldListenTo = source => true, // Listens to all sources.
133-
Sample = (ref ActivityCreationOptions<ActivityContext> options) => ActivitySamplingResult.AllData, // Causes all Activity objects to be created and populated.
134-
ActivityStopped = OnActivityStopped
135-
};
136-
ActivitySource.AddActivityListener(_ActivityListener);
144+
_ActivityTraceListenerSampler.ActivityTraceListenerManager = this;
137145

138146
_HasInitialized = true;
139147
}
@@ -155,8 +163,7 @@ private void CleanupThreadBody()
155163
lock (_TraceListeners)
156164
{
157165
_HasInitialized = false;
158-
_ActivityListener?.Dispose();
159-
_ActivityListener = null;
166+
_ActivityTraceListenerSampler.ActivityTraceListenerManager = null;
160167
}
161168

162169
_Options.ClosedAction?.Invoke();
Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
using System;
2+
using System.Diagnostics;
3+
4+
using OpenTelemetry.Trace;
5+
6+
namespace Macross.OpenTelemetry.Extensions
7+
{
8+
internal sealed class ActivityTraceListenerSampler : Sampler
9+
{
10+
private readonly Sampler _InnerSampler;
11+
12+
internal ActivityTraceListenerManager? ActivityTraceListenerManager { get; set; }
13+
14+
public ActivityTraceListenerSampler(Sampler innerSampler)
15+
{
16+
_InnerSampler = innerSampler ?? throw new ArgumentNullException(nameof(innerSampler));
17+
}
18+
19+
public override SamplingResult ShouldSample(in SamplingParameters samplingParameters)
20+
{
21+
ActivityTraceListenerManager? activityTraceListenerManager = ActivityTraceListenerManager;
22+
return activityTraceListenerManager != null
23+
&& activityTraceListenerManager.IsTraceIdRegistered(samplingParameters.TraceId)
24+
? new SamplingResult(SamplingDecision.RecordAndSample)
25+
: _InnerSampler.ShouldSample(samplingParameters);
26+
}
27+
}
28+
}

ClassLibraries/Macross.OpenTelemetry.Extensions/Code/Extensions/TracerProviderBuilderExtensions.cs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,5 +18,14 @@ public static TracerProviderBuilder AddActivityEnrichmentScopeProcessor(this Tra
1818
#pragma warning disable CA2000 // Dispose objects before losing scope
1919
=> tracerProviderBuilder.AddProcessor(new ActivityEnrichmentScopeProcessor());
2020
#pragma warning restore CA2000 // Dispose objects before losing scope
21+
22+
/// <summary>
23+
/// Adds a <see cref="Sampler"/> which will automatically sample any span under a traceId registered through <see cref="ActivityTraceListenerManager"/>.
24+
/// </summary>
25+
/// <param name="tracerProviderBuilder"><see cref="TracerProviderBuilder"/>.</param>
26+
/// <param name="innerSampler"><see cref="Sampler"/>.</param>
27+
/// <returns>Returns the supplied <see cref="TracerProviderBuilder"/> for chaining.</returns>
28+
public static TracerProviderBuilder SetActivityTraceListenerSampler(this TracerProviderBuilder tracerProviderBuilder, Sampler innerSampler)
29+
=> tracerProviderBuilder.SetSampler(new ActivityTraceListenerSampler(innerSampler));
2130
}
2231
}

ClassLibraries/Macross.OpenTelemetry.Extensions/Code/OpenTelemetryEventListener.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
namespace Macross.OpenTelemetry.Extensions
1010
{
11-
internal class OpenTelemetryEventListener : EventListener
11+
internal sealed class OpenTelemetryEventListener : EventListener
1212
{
1313
private readonly List<EventSource> _UnitializedEventSources = new List<EventSource>();
1414
private readonly List<EventSource> _EventSources = new List<EventSource>();

ClassLibraries/Macross.OpenTelemetry.Extensions/Code/OpenTelemetryEventLoggingHostedService.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010
namespace Macross.OpenTelemetry.Extensions
1111
{
1212
#pragma warning disable CA1812 // Remove class never instantiated
13-
internal class OpenTelemetryEventLoggingHostedService : IHostedService, IDisposable
13+
internal sealed class OpenTelemetryEventLoggingHostedService : IHostedService, IDisposable
1414
#pragma warning restore CA1812 // Remove class never instantiated
1515
{
1616
private readonly OpenTelemetryEventListener _OpenTelemetryEventListener;

ClassLibraries/Macross.OpenTelemetry.Extensions/Code/OpenTelemetryExtensionsEventSource.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@ namespace Macross.OpenTelemetry.Extensions
99
/// EventSource implementation for OpenTelemetry extensions.
1010
/// </summary>
1111
[EventSource(Name = "OpenTelemetry-Extensions")]
12-
internal class OpenTelemetryExtensionsEventSource : EventSource
12+
internal sealed class OpenTelemetryExtensionsEventSource : EventSource
1313
{
1414
public static OpenTelemetryExtensionsEventSource Log { get; } = new OpenTelemetryExtensionsEventSource();
1515

ClassLibraries/Macross.OpenTelemetry.Extensions/README.md

Lines changed: 15 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -214,13 +214,26 @@ is provided.
214214
services.AddActivityTraceListener();
215215
}
216216
}
217-
218217
```
219218

220219
This will register the `ActivityTraceListenerManager` with the
221220
`IServiceProvider`.
222221

223-
2) Use the `ActivityTraceListenerManager` to register a listener:
222+
2) Call the `SetActivityTraceListenerSampler` extension where you configure
223+
OpenTelemetry:
224+
225+
```csharp
226+
services.AddOpenTelemetryTracing((serviceProvider, builder) =>
227+
{
228+
builder.SetActivityTraceListenerSampler(new ParentBasedSampler(new AlwaysOnSampler()));
229+
};
230+
```
231+
232+
The `innerSampler` parameter is the sampler which will be used when a trace
233+
listener is NOT registered. The default behavior is shown
234+
(`ParentBasedSampler` w/ `AlwaysOnSampler`).
235+
236+
3) Use the `ActivityTraceListenerManager` to register a listener:
224237

225238
```csharp
226239
using System.Diagnostics;
Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
using System;
2+
using System.Diagnostics;
3+
4+
using Microsoft.VisualStudio.TestTools.UnitTesting;
5+
6+
using Microsoft.Extensions.DependencyInjection;
7+
8+
using OpenTelemetry;
9+
using OpenTelemetry.Trace;
10+
11+
namespace Macross.OpenTelemetry.Extensions.Tests
12+
{
13+
[TestClass]
14+
public sealed class ActivityTraceListenerManagerTests : IDisposable
15+
{
16+
private readonly ActivityContext _ActivityContext = new ActivityContext(
17+
ActivityTraceId.CreateRandom(),
18+
ActivitySpanId.CreateRandom(),
19+
ActivityTraceFlags.Recorded,
20+
isRemote: true);
21+
22+
private readonly ActivitySource _ActivitySource = new ActivitySource(nameof(ActivityEnrichmentScopeProcessorTests));
23+
private readonly TestActivityProcessor _TestActivityProcessor = new TestActivityProcessor(collectEndedSpans: true);
24+
private readonly TracerProvider _TracerProvider;
25+
private readonly ServiceProvider _ServiceProvider;
26+
private readonly ActivityTraceListenerManager _ActivityTraceListenerManager;
27+
28+
public ActivityTraceListenerManagerTests()
29+
{
30+
_TracerProvider = Sdk.CreateTracerProviderBuilder()
31+
.AddSource(_ActivitySource.Name)
32+
.SetActivityTraceListenerSampler(new AlwaysOffSampler())
33+
.AddProcessor(_TestActivityProcessor)
34+
.Build();
35+
36+
IServiceCollection services = new ServiceCollection();
37+
services.AddSingleton(typeof(TracerProvider), _TracerProvider);
38+
services.AddOptions();
39+
services.AddActivityTraceListener();
40+
41+
_ServiceProvider = services.BuildServiceProvider();
42+
43+
_ActivityTraceListenerManager = _ServiceProvider.GetRequiredService<ActivityTraceListenerManager>();
44+
}
45+
46+
public void Dispose()
47+
{
48+
_TracerProvider.Dispose();
49+
_TestActivityProcessor.Dispose();
50+
_ActivitySource.Dispose();
51+
_ServiceProvider.Dispose();
52+
_ActivityTraceListenerManager.Dispose();
53+
}
54+
55+
[TestMethod]
56+
public void ActivityNotSampledWithoutRegistration()
57+
{
58+
Activity? activity = _ActivitySource.StartActivity("Test", ActivityKind.Server, _ActivityContext);
59+
60+
Assert.IsNull(activity);
61+
}
62+
63+
[TestMethod]
64+
public void ActivitySampledWithRegistration()
65+
{
66+
using (IDisposable registration = _ActivityTraceListenerManager.RegisterTraceListener(_ActivityContext.TraceId))
67+
{
68+
Activity? activity = _ActivitySource.StartActivity("Test", ActivityKind.Server, _ActivityContext);
69+
70+
Assert.IsNotNull(activity);
71+
Assert.IsTrue(activity.IsAllDataRequested);
72+
Assert.IsTrue(activity.Recorded);
73+
}
74+
75+
ActivityNotSampledWithoutRegistration();
76+
}
77+
}
78+
}

ClassLibraries/Macross.OpenTelemetry.Extensions/Test/Macross.OpenTelemetry.Extensions.Tests.csproj

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@
1818
</PropertyGroup>
1919

2020
<ItemGroup>
21+
<PackageReference Include="Microsoft.Extensions.DependencyInjection" Version="3.1.1" />
2122
<PackageReference Include="Microsoft.NET.Test.Sdk" Version="16.9.0-preview-20201123-03" />
2223
<PackageReference Include="MSTest.TestAdapter" Version="2.2.0-preview-20210115-03" />
2324
<PackageReference Include="MSTest.TestFramework" Version="2.2.0-preview-20210115-03" />

0 commit comments

Comments
 (0)