Skip to content

Commit bfc5ae5

Browse files
rkargMsftReubenBond
rkargMsft
authored and
ReubenBond
committed
Refactor tracing: add granular ActivitySources and IAsyncEnumerable tracing
Introduce distinct ActivitySources for application, runtime, lifecycle, and storage tracing in Orleans. Update grain call filter and runtime logic to use the correct source based on operation type. Enhance async enumerable tracing with session activities and proper context propagation. Add new grain and tests to verify async enumerable activity spans. Improve test coverage and assertions for new tracing structure, increasing observability and diagnostic precision across subsystems.
1 parent ba90e7c commit bfc5ae5

File tree

13 files changed

+281
-109
lines changed

13 files changed

+281
-109
lines changed

‎src/Orleans.Core.Abstractions/Diagnostics/ActivitySources.cs‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
using System.Diagnostics;
2+
3+
namespace Orleans.Diagnostics;
4+
5+
public static class ActivitySources
6+
{
7+
/// <summary>
8+
/// Spans triggered from application level code
9+
/// </summary>
10+
public static string ApplicationGrainActivitySourceName = "Microsoft.Orleans.Application";
11+
/// <summary>
12+
/// Spans triggered from Orleans runtime code
13+
/// </summary>
14+
public static string RuntimeActivitySourceName = "Microsoft.Orleans.Runtime";
15+
/// <summary>
16+
/// Spans tied to lifecycle operations such as activation, migration, and deactivation.
17+
/// </summary>
18+
public static string LifecycleActivitySourceName = "Microsoft.Orleans.Lifecycle";
19+
/// <summary>
20+
/// Spans tied to persistent storage operations.
21+
/// </summary>
22+
public static string StorageActivitySourceName = "Microsoft.Orleans.Storage";
23+
/// <summary>
24+
/// A wildcard name to match all Orleans activity sources.
25+
/// </summary>
26+
public static string AllActivitySourceName = "Microsoft.Orleans.*";
27+
28+
internal static readonly ActivitySource ApplicationGrainSource = new(ApplicationGrainActivitySourceName, "1.0.0");
29+
internal static readonly ActivitySource RuntimeGrainSource = new(RuntimeActivitySourceName, "1.0.0");
30+
internal static readonly ActivitySource LifecycleGrainSource = new(LifecycleActivitySourceName, "1.0.0");
31+
internal static readonly ActivitySource StorageGrainSource = new(StorageActivitySourceName, "1.0.0");
32+
}

‎src/Orleans.Core.Abstractions/Properties/AssemblyInfo.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
[assembly: InternalsVisibleTo("ServiceBus.Tests")]
1313
[assembly: InternalsVisibleTo("Tester.AzureUtils")]
1414
[assembly: InternalsVisibleTo("AWSUtils.Tests")]
15+
[assembly: InternalsVisibleTo("Tester")]
1516
[assembly: InternalsVisibleTo("TesterInternal")]
1617
[assembly: InternalsVisibleTo("TestInternalGrainInterfaces")]
1718
[assembly: InternalsVisibleTo("TestInternalGrains")]

‎src/Orleans.Core.Abstractions/Runtime/AsyncEnumerableRequest.cs‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@
99
using Microsoft.Extensions.DependencyInjection;
1010
using Microsoft.Extensions.Logging;
1111
using Orleans.Concurrency;
12+
using Orleans.Diagnostics;
1213
using Orleans.Invocation;
1314
using Orleans.Serialization.Invocation;
1415

@@ -199,6 +200,7 @@ internal sealed class AsyncEnumeratorProxy<T> : IAsyncEnumerator<T>
199200
private int _batchIndex;
200201
private bool _disposed;
201202
private bool _isInitialized;
203+
private Activity? _sessionActivity;
202204

203205
private bool IsBatch => (_current.State & EnumerationResult.Batch) != 0;
204206
private bool IsElement => (_current.State & EnumerationResult.Element) != 0;
@@ -270,6 +272,9 @@ public async ValueTask DisposeAsync()
270272

271273
if (_isInitialized)
272274
{
275+
// Restore the session activity as the current activity so that DisposeAsync RPC is parented to it
276+
var previousActivity = Activity.Current;
277+
Activity.Current = _sessionActivity;
273278
try
274279
{
275280
await _target.DisposeAsync(_requestId);
@@ -279,9 +284,18 @@ public async ValueTask DisposeAsync()
279284
var logger = ((GrainReference)_target).Shared.ServiceProvider.GetRequiredService<ILogger<AsyncEnumerableRequest<T>>>();
280285
logger.LogWarning(exception, "Failed to dispose async enumerator.");
281286
}
287+
finally
288+
{
289+
Activity.Current = previousActivity;
290+
}
282291
}
283292

284293
_cancellationTokenSource?.Dispose();
294+
295+
// Stop the session activity after DisposeAsync completes
296+
_sessionActivity?.Stop();
297+
_sessionActivity?.Dispose();
298+
285299
_disposed = true;
286300
}
287301

@@ -302,6 +316,14 @@ public async ValueTask<bool> MoveNextAsync()
302316
}
303317

304318
var isActive = _isInitialized;
319+
320+
// Restore the session activity as the current activity so that RPC calls are parented to it
321+
var previousActivity = Activity.Current;
322+
if (_sessionActivity is not null)
323+
{
324+
Activity.Current = _sessionActivity;
325+
}
326+
305327
try
306328
{
307329
(EnumerationResult Status, object Value) result;
@@ -311,6 +333,11 @@ public async ValueTask<bool> MoveNextAsync()
311333

312334
if (!_isInitialized)
313335
{
336+
// Start the session activity on first enumeration call
337+
// This span wraps the entire enumeration session
338+
_sessionActivity = ActivitySources.ApplicationGrainSource.StartActivity(_request.GetActivityName(), ActivityKind.Client);
339+
_sessionActivity?.SetTag("orleans.async_enumerable.request_id", _requestId.ToString());
340+
314341
// Assume the enumerator is active as soon as the call begins.
315342
isActive = true;
316343
result = await _target.StartEnumeration(_requestId, _request, _cancellationToken);
@@ -324,10 +351,12 @@ public async ValueTask<bool> MoveNextAsync()
324351
isActive = result.Status.IsActive();
325352
if (result.Status is EnumerationResult.Error)
326353
{
354+
_sessionActivity?.SetStatus(ActivityStatusCode.Error);
327355
ExceptionDispatchInfo.Capture((Exception)result.Value).Throw();
328356
}
329357
else if (result.Status is EnumerationResult.Canceled)
330358
{
359+
_sessionActivity?.SetStatus(ActivityStatusCode.Error, "Canceled");
331360
throw new OperationCanceledException();
332361
}
333362

@@ -339,6 +368,7 @@ public async ValueTask<bool> MoveNextAsync()
339368

340369
if (result.Status is EnumerationResult.MissingEnumeratorError)
341370
{
371+
_sessionActivity?.SetStatus(ActivityStatusCode.Error, "MissingEnumerator");
342372
throw new EnumerationAbortedException("Enumeration aborted: the remote target does not have a record of this enumerator."
343373
+ " This likely indicates that the remote grain was deactivated since enumeration begun or that the enumerator was idle for longer than the expiration period.");
344374
}
@@ -356,6 +386,11 @@ await _target.DisposeAsync(_requestId).AsTask()
356386
.ConfigureAwait(ConfigureAwaitOptions.ContinueOnCapturedContext | ConfigureAwaitOptions.SuppressThrowing);
357387
throw;
358388
}
389+
finally
390+
{
391+
// Restore the previous activity after each call
392+
Activity.Current = previousActivity;
393+
}
359394
}
360395
}
361396

‎src/Orleans.Core/Diagnostics/ActivityPropagationGrainCallFilter.cs‎

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
using System.Diagnostics;
2+
using Orleans.Diagnostics;
23

34
namespace Orleans.Runtime
45
{
@@ -10,10 +11,33 @@ internal abstract class ActivityPropagationGrainCallFilter
1011
internal const string RpcSystem = "orleans";
1112
internal const string OrleansNamespacePrefix = "Orleans";
1213

13-
protected static ActivitySource GetActivitySource(IGrainCallContext context) =>
14-
context.Request.GetInterfaceType().Namespace?.StartsWith(OrleansNamespacePrefix) == true
15-
? ActivitySources.RuntimeGrainSource
16-
: ActivitySources.ApplicationGrainSource;
14+
protected static ActivitySource GetActivitySource(IGrainCallContext context)
15+
{
16+
var interfaceType = context.Request.GetInterfaceType();
17+
var interfaceTypeName = interfaceType.Name;
18+
19+
switch (interfaceTypeName)
20+
{
21+
// Memory storage uses grains for its implementation
22+
case "IMemoryStorageGrain":
23+
return ActivitySources.StorageGrainSource;
24+
25+
// This extension is for explicit migrate/deactivate calls
26+
case "IGrainManagementExtension":
27+
// This target is for accepting migration batches
28+
case "IActivationMigrationManagerSystemTarget":
29+
return ActivitySources.LifecycleGrainSource;
30+
31+
// These extensions are for async stream subscriptions
32+
case "IAsyncEnumerableGrainExtension":
33+
return ActivitySources.ApplicationGrainSource;
34+
35+
default:
36+
return interfaceType.Namespace?.StartsWith(OrleansNamespacePrefix) == true
37+
? ActivitySources.RuntimeGrainSource
38+
: ActivitySources.ApplicationGrainSource;
39+
}
40+
}
1741

1842
protected static void GetRequestContextValue(object carrier, string fieldName, out string fieldValue, out IEnumerable<string> fieldValues)
1943
{

‎src/Orleans.Core/Diagnostics/ActivitySources.cs‎

Lines changed: 0 additions & 12 deletions
This file was deleted.

‎src/Orleans.Runtime/Catalog/ActivationData.cs‎

Lines changed: 9 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,17 +1,14 @@
11
#nullable enable
2-
using System;
32
using System.Buffers;
4-
using System.Collections.Generic;
53
using System.Diagnostics;
64
using System.Diagnostics.CodeAnalysis;
75
using System.Runtime.CompilerServices;
86
using System.Runtime.InteropServices;
9-
using System.Threading;
10-
using System.Threading.Tasks;
117
using Microsoft.Extensions.DependencyInjection;
128
using Microsoft.Extensions.Logging;
139
using Orleans.Configuration;
1410
using Orleans.Core.Internal;
11+
using Orleans.Diagnostics;
1512
using Orleans.GrainDirectory;
1613
using Orleans.Internal;
1714
using Orleans.Runtime.Placement;
@@ -1212,9 +1209,9 @@ private void RehydrateInternal(IRehydrationContext context)
12121209
{
12131210
// Start a span for rehydration
12141211
rehydrateSpan = _activationActivity is not null
1215-
? ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.ActivationRehydrate,
1212+
? ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.ActivationRehydrate,
12161213
ActivityKind.Internal, _activationActivity.Context)
1217-
: ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.ActivationRehydrate,
1214+
: ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.ActivationRehydrate,
12181215
ActivityKind.Internal);
12191216
rehydrateSpan?.SetTag("orleans.grain.id", GrainId.ToString());
12201217
rehydrateSpan?.SetTag("orleans.grain.type", _shared.GrainTypeName);
@@ -1290,9 +1287,9 @@ private void OnDehydrate(IDehydrationContext context)
12901287
{
12911288
// Start a span for dehydration, parented to the migration request that triggered it
12921289
dehydrateSpan = parentContext.HasValue
1293-
? ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.ActivationDehydrate,
1290+
? ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.ActivationDehydrate,
12941291
ActivityKind.Internal, parentContext.Value)
1295-
: ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.ActivationDehydrate,
1292+
: ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.ActivationDehydrate,
12961293
ActivityKind.Internal);
12971294
dehydrateSpan?.SetTag("orleans.grain.id", GrainId.ToString());
12981295
dehydrateSpan?.SetTag("orleans.grain.type", _shared.GrainTypeName);
@@ -1556,8 +1553,8 @@ private async Task ActivateAsync(Dictionary<string, object>? requestContextData,
15561553
{
15571554
// Start directory registration activity as a child of the activation activity
15581555
Activity? registerSpan = _activationActivity is not null
1559-
? ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.RegisterDirectoryEntry, ActivityKind.Internal, _activationActivity.Context)
1560-
: ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.RegisterDirectoryEntry, ActivityKind.Internal);
1556+
? ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.RegisterDirectoryEntry, ActivityKind.Internal, _activationActivity.Context)
1557+
: ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.RegisterDirectoryEntry, ActivityKind.Internal);
15611558
registerSpan?.SetTag("orleans.grain.id", GrainId.ToString());
15621559
registerSpan?.SetTag("orleans.silo.id", _shared.Runtime.SiloAddress.ToString());
15631560
registerSpan?.SetTag("orleans.activation.id", ActivationId.ToString());
@@ -1676,8 +1673,8 @@ private async Task ActivateAsync(Dictionary<string, object>? requestContextData,
16761673
{
16771674
// Start a span for OnActivateAsync execution
16781675
using var onActivateSpan = _activationActivity is not null
1679-
? ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.OnActivate, ActivityKind.Internal, _activationActivity.Context)
1680-
: ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.OnActivate, ActivityKind.Internal);
1676+
? ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.OnActivate, ActivityKind.Internal, _activationActivity.Context)
1677+
: ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.OnActivate, ActivityKind.Internal);
16811678
onActivateSpan?.SetTag("orleans.grain.id", GrainId.ToString());
16821679
onActivateSpan?.SetTag("orleans.grain.type", _shared.GrainTypeName ?? GrainInstance.GetType().FullName);
16831680
onActivateSpan?.SetTag("orleans.silo.id", _shared.Runtime.SiloAddress.ToString());

‎src/Orleans.Runtime/Catalog/Catalog.cs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
using Orleans.GrainDirectory;
55
using Orleans.Runtime.GrainDirectory;
66
using System.Diagnostics;
7+
using Orleans.Diagnostics;
78

89
namespace Orleans.Runtime
910
{
@@ -171,8 +172,8 @@ public IGrainContext GetOrCreateActivation(
171172
// Start activation span with parent context from request if available
172173
var parentContext = TryGetActivityContext(requestContextData);
173174
Activity activationActivity = parentContext.HasValue
174-
? ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.ActivateGrain, ActivityKind.Internal, parentContext.Value)
175-
: ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.ActivateGrain, ActivityKind.Internal);
175+
? ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.ActivateGrain, ActivityKind.Internal, parentContext.Value)
176+
: ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.ActivateGrain, ActivityKind.Internal);
176177
if (activationActivity is not null)
177178
{
178179
activationActivity.SetTag("orleans.grain.id", grainId.ToString());

‎src/Orleans.Runtime/Placement/PlacementService.cs‎

Lines changed: 3 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,14 +1,10 @@
1-
using System;
2-
using System.Collections.Generic;
31
using System.Diagnostics;
4-
using System.Linq;
52
using System.Runtime.CompilerServices;
63
using System.Runtime.InteropServices;
7-
using System.Threading.Tasks;
84
using Microsoft.Extensions.Logging;
95
using Microsoft.Extensions.Options;
106
using Orleans.Configuration;
11-
using Orleans.Placement;
7+
using Orleans.Diagnostics;
128
using Orleans.Runtime.GrainDirectory;
139
using Orleans.Runtime.Internal;
1410
using Orleans.Runtime.Placement.Filtering;
@@ -128,7 +124,7 @@ public SiloAddress[] GetCompatibleSilos(PlacementTarget target)
128124
var director = _placementFilterDirectoryResolver.GetFilterDirector(placementFilter);
129125

130126
// Create a span for each filter invocation
131-
using var filterSpan = ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.FilterPlacementCandidates);
127+
using var filterSpan = ActivitySources.LifecycleGrainSource.StartActivity(ActivityNames.FilterPlacementCandidates);
132128
filterSpan?.SetTag("orleans.placement.filter.type", placementFilter.GetType().Name);
133129
filterSpan?.SetTag("orleans.grain.type", grainType.ToString());
134130

@@ -442,7 +438,7 @@ private static Activity TryRestoreActivityContext(Dictionary<string, object> req
442438
if (!string.IsNullOrEmpty(traceParent) && ActivityContext.TryParse(traceParent, traceState, isRemote: true, out var parentContext))
443439
{
444440
// Start the activity from the Catalog's ActivitySource to properly associate it with activation tracing
445-
return ActivitySources.RuntimeGrainSource.StartActivity(operationName, ActivityKind.Internal, parentContext);
441+
return ActivitySources.LifecycleGrainSource.StartActivity(operationName, ActivityKind.Internal, parentContext);
446442
}
447443

448444
return null;

‎src/Orleans.Runtime/Storage/StateStorageBridge.cs‎

Lines changed: 7 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,11 @@
11
#nullable enable
2-
using System;
32
using System.Collections.Concurrent;
43
using System.Diagnostics;
54
using System.Diagnostics.CodeAnalysis;
65
using System.Runtime.ExceptionServices;
7-
using System.Threading.Tasks;
86
using Microsoft.Extensions.DependencyInjection;
97
using Microsoft.Extensions.Logging;
10-
using Orleans.Runtime;
8+
using Orleans.Diagnostics;
119
using Orleans.Serialization.Activators;
1210
using Orleans.Serialization.Serializers;
1311
using Orleans.Storage;
@@ -93,8 +91,8 @@ public async Task ReadStateAsync()
9391
}
9492

9593
using var activity = parentContext.HasValue
96-
? ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.StorageRead, ActivityKind.Client, parentContext.Value)
97-
: ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.StorageRead, ActivityKind.Client);
94+
? ActivitySources.StorageGrainSource.StartActivity(ActivityNames.StorageRead, ActivityKind.Client, parentContext.Value)
95+
: ActivitySources.StorageGrainSource.StartActivity(ActivityNames.StorageRead, ActivityKind.Client);
9896
activity?.SetTag("orleans.grain.id", _grainContext.GrainId.ToString());
9997
activity?.SetTag("orleans.storage.provider", _shared.ProviderTypeName);
10098
activity?.SetTag("orleans.storage.state.name", _shared.Name);
@@ -127,8 +125,8 @@ public async Task WriteStateAsync()
127125
}
128126

129127
using var activity = parentContext.HasValue
130-
? ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.StorageWrite, ActivityKind.Client, parentContext.Value)
131-
: ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.StorageWrite, ActivityKind.Client);
128+
? ActivitySources.StorageGrainSource.StartActivity(ActivityNames.StorageWrite, ActivityKind.Client, parentContext.Value)
129+
: ActivitySources.StorageGrainSource.StartActivity(ActivityNames.StorageWrite, ActivityKind.Client);
132130
activity?.SetTag("orleans.grain.id", _grainContext.GrainId.ToString());
133131
activity?.SetTag("orleans.storage.provider", _shared.ProviderTypeName);
134132
activity?.SetTag("orleans.storage.state.name", _shared.Name);
@@ -160,8 +158,8 @@ public async Task ClearStateAsync()
160158
}
161159

162160
using var activity = parentContext.HasValue
163-
? ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.StorageClear, ActivityKind.Client, parentContext.Value)
164-
: ActivitySources.RuntimeGrainSource.StartActivity(ActivityNames.StorageClear, ActivityKind.Client);
161+
? ActivitySources.StorageGrainSource.StartActivity(ActivityNames.StorageClear, ActivityKind.Client, parentContext.Value)
162+
: ActivitySources.StorageGrainSource.StartActivity(ActivityNames.StorageClear, ActivityKind.Client);
165163
activity?.SetTag("orleans.grain.id", _grainContext.GrainId.ToString());
166164
activity?.SetTag("orleans.storage.provider", _shared.ProviderTypeName);
167165
activity?.SetTag("orleans.storage.state.name", _shared.Name);

0 commit comments

Comments
 (0)