Skip to content

Commit b549899

Browse files
adelinowonaadelinowona
committed
Implement operation-level tracing
1 parent fc9f61f commit b549899

File tree

2 files changed

+194
-8
lines changed

2 files changed

+194
-8
lines changed

‎src/MongoDB.Driver/OperationContext.cs‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,12 @@ internal OperationContext(IClock clock, long initialTimestamp, TimeSpan? timeout
5151

5252
public OperationContext RootContext { get; private init; }
5353

54+
// OpenTelemetry operation metadata
55+
internal string OperationName { get; init; }
56+
internal string DatabaseName { get; init; }
57+
internal string CollectionName { get; init; }
58+
internal bool IsTracingEnabled { get; init; }
59+
5460
public TimeSpan RemainingTimeout
5561
{
5662
get
@@ -213,5 +219,17 @@ public OperationContext WithTimeout(TimeSpan timeout)
213219
RootContext = RootContext
214220
};
215221
}
222+
223+
internal OperationContext WithOperationMetadata(string operationName, string databaseName, string collectionName, bool isTracingEnabled)
224+
{
225+
return new OperationContext(Clock, InitialTimestamp, Timeout, CancellationToken)
226+
{
227+
RootContext = RootContext,
228+
OperationName = operationName,
229+
DatabaseName = databaseName,
230+
CollectionName = collectionName,
231+
IsTracingEnabled = isTracingEnabled
232+
};
233+
}
216234
}
217235
}

‎src/MongoDB.Driver/OperationExecutor.cs‎

Lines changed: 176 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
*/
1515

1616
using System;
17+
using System.Diagnostics;
1718
using System.Threading.Tasks;
1819
using MongoDB.Driver.Core;
1920
using MongoDB.Driver.Core.Bindings;
@@ -50,8 +51,41 @@ public TResult ExecuteReadOperation<TResult>(
5051
Ensure.IsNotNull(readPreference, nameof(readPreference));
5152
ThrowIfDisposed();
5253

53-
using var binding = CreateReadBinding(session, readPreference, allowChannelPinning);
54-
return operation.Execute(operationContext, binding);
54+
using var activityScope = TransactionActivityScope.CreateIfNeeded(session);
55+
using var activity = operationContext.IsTracingEnabled
56+
? MongoTelemetry.StartOperationActivity(
57+
operationContext.OperationName,
58+
operationContext.DatabaseName,
59+
operationContext.CollectionName)
60+
: null;
61+
62+
try
63+
{
64+
using var binding = CreateReadBinding(session, readPreference, allowChannelPinning);
65+
var result = operation.Execute(operationContext, binding);
66+
activity?.SetStatus(ActivityStatusCode.Ok);
67+
return result;
68+
}
69+
catch (Exception ex)
70+
{
71+
// Only record exceptions that originate at the operation level
72+
// Command-level exceptions (MongoCommandException, MongoWriteException, etc.)
73+
// are already recorded by CommandEventHelper on the command span
74+
if (activity != null)
75+
{
76+
if (!IsCommandLevelException(ex))
77+
{
78+
MongoTelemetry.RecordException(activity, ex);
79+
}
80+
else
81+
{
82+
// For command-level exceptions, only set error status without recording details
83+
activity.SetStatus(ActivityStatusCode.Error);
84+
}
85+
}
86+
87+
throw;
88+
}
5589
}
5690

5791
public async Task<TResult> ExecuteReadOperationAsync<TResult>(
@@ -67,8 +101,41 @@ public async Task<TResult> ExecuteReadOperationAsync<TResult>(
67101
Ensure.IsNotNull(readPreference, nameof(readPreference));
68102
ThrowIfDisposed();
69103

70-
using var binding = CreateReadBinding(session, readPreference, allowChannelPinning);
71-
return await operation.ExecuteAsync(operationContext, binding).ConfigureAwait(false);
104+
using var activityScope = TransactionActivityScope.CreateIfNeeded(session);
105+
using var activity = operationContext.IsTracingEnabled
106+
? MongoTelemetry.StartOperationActivity(
107+
operationContext.OperationName,
108+
operationContext.DatabaseName,
109+
operationContext.CollectionName)
110+
: null;
111+
112+
try
113+
{
114+
using var binding = CreateReadBinding(session, readPreference, allowChannelPinning);
115+
var result = await operation.ExecuteAsync(operationContext, binding).ConfigureAwait(false);
116+
activity?.SetStatus(ActivityStatusCode.Ok);
117+
return result;
118+
}
119+
catch (Exception ex)
120+
{
121+
// Only record exceptions that originate at the operation level
122+
// Command-level exceptions (MongoCommandException, MongoWriteException, etc.)
123+
// are already recorded by CommandEventHelper on the command span
124+
if (activity != null)
125+
{
126+
if (!IsCommandLevelException(ex))
127+
{
128+
MongoTelemetry.RecordException(activity, ex);
129+
}
130+
else
131+
{
132+
// For command-level exceptions, only set error status without recording details
133+
activity.SetStatus(ActivityStatusCode.Error);
134+
}
135+
}
136+
137+
throw;
138+
}
72139
}
73140

74141
public TResult ExecuteWriteOperation<TResult>(
@@ -82,8 +149,41 @@ public TResult ExecuteWriteOperation<TResult>(
82149
Ensure.IsNotNull(operation, nameof(operation));
83150
ThrowIfDisposed();
84151

85-
using var binding = CreateReadWriteBinding(session, allowChannelPinning);
86-
return operation.Execute(operationContext, binding);
152+
using var activityScope = TransactionActivityScope.CreateIfNeeded(session);
153+
using var activity = operationContext.IsTracingEnabled
154+
? MongoTelemetry.StartOperationActivity(
155+
operationContext.OperationName,
156+
operationContext.DatabaseName,
157+
operationContext.CollectionName)
158+
: null;
159+
160+
try
161+
{
162+
using var binding = CreateReadWriteBinding(session, allowChannelPinning);
163+
var result = operation.Execute(operationContext, binding);
164+
activity?.SetStatus(ActivityStatusCode.Ok);
165+
return result;
166+
}
167+
catch (Exception ex)
168+
{
169+
// Only record exceptions that originate at the operation level
170+
// Command-level exceptions (MongoCommandException, MongoWriteException, etc.)
171+
// are already recorded by CommandEventHelper on the command span
172+
if (activity != null)
173+
{
174+
if (!IsCommandLevelException(ex))
175+
{
176+
MongoTelemetry.RecordException(activity, ex);
177+
}
178+
else
179+
{
180+
// For command-level exceptions, only set error status without recording details
181+
activity.SetStatus(ActivityStatusCode.Error);
182+
}
183+
}
184+
185+
throw;
186+
}
87187
}
88188

89189
public async Task<TResult> ExecuteWriteOperationAsync<TResult>(
@@ -97,8 +197,41 @@ public async Task<TResult> ExecuteWriteOperationAsync<TResult>(
97197
Ensure.IsNotNull(operation, nameof(operation));
98198
ThrowIfDisposed();
99199

100-
using var binding = CreateReadWriteBinding(session, allowChannelPinning);
101-
return await operation.ExecuteAsync(operationContext, binding).ConfigureAwait(false);
200+
using var activityScope = TransactionActivityScope.CreateIfNeeded(session);
201+
using var activity = operationContext.IsTracingEnabled
202+
? MongoTelemetry.StartOperationActivity(
203+
operationContext.OperationName,
204+
operationContext.DatabaseName,
205+
operationContext.CollectionName)
206+
: null;
207+
208+
try
209+
{
210+
using var binding = CreateReadWriteBinding(session, allowChannelPinning);
211+
var result = await operation.ExecuteAsync(operationContext, binding).ConfigureAwait(false);
212+
activity?.SetStatus(ActivityStatusCode.Ok);
213+
return result;
214+
}
215+
catch (Exception ex)
216+
{
217+
// Only record exceptions that originate at the operation level
218+
// Command-level exceptions (MongoCommandException, MongoWriteException, etc.)
219+
// are already recorded by CommandEventHelper on the command span
220+
if (activity != null)
221+
{
222+
if (!IsCommandLevelException(ex))
223+
{
224+
MongoTelemetry.RecordException(activity, ex);
225+
}
226+
else
227+
{
228+
// For command-level exceptions, only set error status without recording details
229+
activity.SetStatus(ActivityStatusCode.Error);
230+
}
231+
}
232+
233+
throw;
234+
}
102235
}
103236

104237
public IClientSessionHandle StartImplicitSession()
@@ -143,5 +276,40 @@ private void ThrowIfDisposed()
143276
throw new ObjectDisposedException(nameof(OperationExecutor));
144277
}
145278
}
279+
280+
private static bool IsCommandLevelException(Exception ex)
281+
{
282+
// Command-level exceptions are those that originate from MongoDB server
283+
return ex is MongoServerException;
284+
}
285+
286+
/// <summary>
287+
/// Manages Activity.Current scope when executing operations within a transaction.
288+
/// Temporarily sets Activity.Current to the transaction activity so operation activities nest correctly,
289+
/// then restores to the original value to prevent AsyncLocal flow issues.
290+
/// </summary>
291+
private sealed class TransactionActivityScope : IDisposable
292+
{
293+
private readonly Activity _originalActivity;
294+
295+
private TransactionActivityScope(Activity transactionActivity)
296+
{
297+
_originalActivity = Activity.Current;
298+
Activity.Current = transactionActivity;
299+
}
300+
301+
public static TransactionActivityScope CreateIfNeeded(IClientSessionHandle session)
302+
{
303+
var transactionActivity = session?.WrappedCoreSession?.CurrentTransaction?.TransactionActivity;
304+
return transactionActivity != null
305+
? new TransactionActivityScope(transactionActivity)
306+
: null;
307+
}
308+
309+
public void Dispose()
310+
{
311+
Activity.Current = _originalActivity;
312+
}
313+
}
146314
}
147315
}

0 commit comments

Comments
 (0)