Skip to content

Commit e11bcc8

Browse files
ajcvickersajcvickers
authored
CSHARP-5935: Command activities may be skipped when using pooled connection (#1918)
* CSHARP-5935: Command activities may be skipped when using pooled connection I was looking at flaky tests, and found MongoClient_should_create_activities_when_tracing_enabled was failing occasionally. Turns out it was an actual bug. Junie (Opus 4.6) says: ### Root Cause `_shouldTrace` in `CommandEventHelper` was set **once** at connection construction time via `MongoTelemetry.ActivitySource.HasListeners()`. Since connections are pooled and reused, if a connection was created before an `ActivityListener` was registered (or by a test with tracing disabled), `_shouldTrace` remained `false` permanently for that connection — command activities were never created even when a listener was later active. * Updates from copilot feedback. * Summary I've successfully refactored the code to fix the bug where activity listeners were being checked when the connection was created (at which point none may be set), but not checked again when the connection was pulled from the pool. Key Changes: Constructor (CommandEventHelper.cs:60): Removed the listener check from _shouldTrackState initialization. Now _shouldTrackState only depends on whether event tracking is needed (_shouldTrackSucceeded || _shouldTrackFailed), not on tracing with listeners. Property Getters (CommandEventHelper.cs:71-94): Updated all property getters (ShouldCallBeforeSending, ShouldCallAfterSending, etc.) to dynamically check MongoTelemetry.ActivitySource.HasListeners() at execution time, ensuring that activity tracing is detected when connections are obtained from the pool. Lazy State Initialization (CommandEventHelper.cs:37, 741-745): Changed _state from readonly to allow lazy initialization. The state dictionary is now only created when needed - either for event tracking or when tracing with listeners is detected at runtime. Null Safety (CommandEventHelper.cs:132, 165, 192, 216, 233): Added null checks before accessing _state in all methods to handle cases where state tracking hasn't been initialized. TrackCommandState (CommandEventHelper.cs:734-770): Updated to check for listeners at execution time and lazily initialize the state dictionary only when tracing is actually needed. Test Update (CommandEventHelperTests.cs:122-124): Updated the test to reflect the new behavior where _shouldTrackState only reflects event tracking, not tracing with listeners (which is now checked dynamically). Benefits: Bug Fixed: Activity listeners are now checked when connections are obtained from the pool, not just when they're created No Unnecessary State Tracking: State tracking is only enabled when actually needed (either for events or when listeners are registered), avoiding the performance overhead of always tracking state Backward Compatible: The behavior is the same for all scenarios, just more efficient The changes ensure that the driver properly detects activity listeners whenever a connection is used, whether it's newly created or retrieved from the pool, while avoiding unnecessary state tracking when no listeners are registered. * Updates based on review feedback. * Minor fixes. * Updates after comments from Boris. * Added an explicit API for ShouldCallConnectionFailed to make it look like the other methods, even though it is not really needed.
1 parent ccca05c commit e11bcc8

File tree

4 files changed

+87
-77
lines changed

4 files changed

+87
-77
lines changed

‎src/MongoDB.Driver/Core/Connections/BinaryConnection.cs‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -155,7 +155,11 @@ private void ConnectionFailed(Exception exception)
155155
{
156156
_failedEventHasBeenRaised = true;
157157
_eventLogger.LogAndPublish(new ConnectionFailedEvent(_connectionId, exception));
158-
_commandEventHelper.ConnectionFailed(_connectionId, _description?.ServiceId, exception, IsInitializing);
158+
159+
if (_commandEventHelper.ShouldCallConnectionFailed)
160+
{
161+
_commandEventHelper.ConnectionFailed(_connectionId, _description?.ServiceId, exception, IsInitializing);
162+
}
159163
}
160164
}
161165

‎src/MongoDB.Driver/Core/Connections/CommandEventHelper.cs‎

Lines changed: 53 additions & 58 deletions
Original file line numberDiff line numberDiff line change
@@ -34,14 +34,14 @@ namespace MongoDB.Driver.Core.Connections
3434
internal class CommandEventHelper
3535
{
3636
private readonly EventLogger<LogCategories.Command> _eventLogger;
37-
private readonly ConcurrentDictionary<int, CommandState> _state;
37+
private ConcurrentDictionary<int, CommandState> _state;
3838

39-
private readonly bool _shouldProcessRequestMessages;
39+
private readonly bool _eventsNeedBeforeSending;
4040
private readonly bool _shouldTrackStarted;
41-
private readonly bool _shouldTrackState;
41+
private readonly bool _eventsNeedState;
4242
private readonly bool _shouldTrackFailed;
4343
private readonly bool _shouldTrackSucceeded;
44-
private readonly bool _shouldTrace;
44+
private readonly bool _tracingDisabled;
4545
private readonly int _queryTextMaxLength;
4646

4747
private Activity _currentCommandActivity;
@@ -54,44 +54,34 @@ public CommandEventHelper(EventLogger<LogCategories.Command> eventLogger, Tracin
5454
_shouldTrackFailed = _eventLogger.IsEventTracked<CommandFailedEvent>();
5555
_shouldTrackStarted = _eventLogger.IsEventTracked<CommandStartedEvent>();
5656

57-
_shouldTrace = tracingOptions?.Disabled != true && MongoTelemetry.ActivitySource.HasListeners();
57+
_tracingDisabled = tracingOptions?.Disabled == true;
5858
_queryTextMaxLength = tracingOptions?.QueryTextMaxLength ?? 0;
5959

60-
_shouldTrackState = _shouldTrackSucceeded || _shouldTrackFailed || _shouldTrace;
61-
_shouldProcessRequestMessages = _shouldTrackStarted || _shouldTrackState;
60+
_eventsNeedState = _shouldTrackSucceeded || _shouldTrackFailed;
61+
_eventsNeedBeforeSending = _shouldTrackStarted || _eventsNeedState;
6262

63-
if (_shouldTrackState)
63+
if (_eventsNeedState)
6464
{
6565
// we only need to track state if we have to raise
6666
// a succeeded or failed event or for tracing
6767
_state = new ConcurrentDictionary<int, CommandState>();
6868
}
6969
}
7070

71-
public bool ShouldCallBeforeSending
72-
{
73-
get { return _shouldProcessRequestMessages; }
74-
}
71+
public bool ShouldCallBeforeSending => _eventsNeedBeforeSending || ShouldTraceWithActivityListener();
7572

76-
public bool ShouldCallAfterSending
77-
{
78-
get { return _shouldTrackState; }
79-
}
73+
public bool ShouldCallAfterSending => _eventsNeedState || ShouldTraceWithActivityListener();
8074

81-
public bool ShouldCallErrorSending
82-
{
83-
get { return _shouldTrackState; }
84-
}
75+
public bool ShouldCallErrorSending => _eventsNeedState || ShouldTraceWithActivityListener();
8576

86-
public bool ShouldCallAfterReceiving
87-
{
88-
get { return _shouldTrackState; }
89-
}
77+
public bool ShouldCallAfterReceiving => _eventsNeedState || ShouldTraceWithActivityListener();
9078

91-
public bool ShouldCallErrorReceiving
92-
{
93-
get { return _shouldTrackState; }
94-
}
79+
public bool ShouldCallErrorReceiving => _eventsNeedState || ShouldTraceWithActivityListener();
80+
81+
public bool ShouldCallConnectionFailed => (_shouldTrackFailed || ShouldTraceWithActivityListener()) && _state != null;
82+
83+
private bool ShouldTraceWithActivityListener()
84+
=> !_tracingDisabled && MongoTelemetry.ActivitySource.HasListeners();
9585

9686
public void CompleteFailedCommandActivity(Exception exception)
9787
{
@@ -129,8 +119,12 @@ public void BeforeSending(
129119

130120
public void AfterSending(RequestMessage message, ConnectionId connectionId, ObjectId? serviceId, bool skipLogging)
131121
{
132-
CommandState state;
133-
if (_state.TryGetValue(message.RequestId, out state) &&
122+
if (_state == null)
123+
{
124+
return;
125+
}
126+
127+
if (_state.TryGetValue(message.RequestId, out var state) &&
134128
state.ExpectedResponseType == ExpectedResponseType.None)
135129
{
136130
state.Stopwatch.Stop();
@@ -157,12 +151,16 @@ public void AfterSending(RequestMessage message, ConnectionId connectionId, Obje
157151

158152
public void ErrorSending(RequestMessage message, ConnectionId connectionId, ObjectId? serviceId, Exception exception, bool skipLogging)
159153
{
160-
CommandState state;
161-
if (_state.TryRemove(message.RequestId, out state))
154+
if (_state == null)
162155
{
163-
state.Stopwatch.Stop();
156+
return;
157+
}
158+
159+
CompleteFailedCommandActivity(exception);
164160

165-
CompleteCommandActivityWithException(exception);
161+
if (_state.TryRemove(message.RequestId, out var state))
162+
{
163+
state.Stopwatch.Stop();
166164

167165
_eventLogger.LogAndPublish(new CommandFailedEvent(
168166
state.CommandName,
@@ -179,8 +177,12 @@ public void ErrorSending(RequestMessage message, ConnectionId connectionId, Obje
179177

180178
public void AfterReceiving(ResponseMessage message, IByteBuffer buffer, ConnectionId connectionId, ObjectId? serviceId, MessageEncoderSettings encoderSettings, bool skipLogging)
181179
{
182-
CommandState state;
183-
if (!_state.TryRemove(message.ResponseTo, out state))
180+
if (_state == null)
181+
{
182+
return;
183+
}
184+
185+
if (!_state.TryRemove(message.ResponseTo, out var state))
184186
{
185187
// this indicates a bug in the sending portion...
186188
return;
@@ -198,17 +200,21 @@ public void AfterReceiving(ResponseMessage message, IByteBuffer buffer, Connecti
198200

199201
public void ErrorReceiving(int responseTo, ConnectionId connectionId, ObjectId? serviceId, Exception exception, bool skipLogging)
200202
{
201-
CommandState state;
202-
if (!_state.TryRemove(responseTo, out state))
203+
if (_state == null)
204+
{
205+
return;
206+
}
207+
208+
CompleteFailedCommandActivity(exception);
209+
210+
if (!_state.TryRemove(responseTo, out var state))
203211
{
204212
// this indicates a bug in the sending portion...
205213
return;
206214
}
207215

208216
state.Stopwatch.Stop();
209217

210-
CompleteCommandActivityWithException(exception);
211-
212218
_eventLogger.LogAndPublish(new CommandFailedEvent(
213219
state.CommandName,
214220
state.QueryNamespace.DatabaseNamespace,
@@ -223,12 +229,7 @@ public void ErrorReceiving(int responseTo, ConnectionId connectionId, ObjectId?
223229

224230
public void ConnectionFailed(ConnectionId connectionId, ObjectId? serviceId, Exception exception, bool skipLogging)
225231
{
226-
if (!_shouldTrackFailed && !_shouldTrace)
227-
{
228-
return;
229-
}
230-
231-
CompleteCommandActivityWithException(exception);
232+
CompleteFailedCommandActivity(exception);
232233

233234
var requestIds = _state.Keys;
234235
foreach (var requestId in requestIds)
@@ -731,11 +732,15 @@ private void TrackCommandState(
731732
BsonDocument sessionId,
732733
long? transactionNumber)
733734
{
734-
if (!_shouldTrackState)
735+
var shouldTraceCommand = ShouldTraceWithActivityListener() && !shouldRedactCommand && !skipLogging;
736+
737+
if (!_eventsNeedState && !shouldTraceCommand)
735738
{
736739
return;
737740
}
738741

742+
_state ??= new ConcurrentDictionary<int, CommandState>();
743+
739744
var commandState = new CommandState
740745
{
741746
CommandName = commandName,
@@ -746,7 +751,7 @@ private void TrackCommandState(
746751
ShouldRedactReply = shouldRedactCommand
747752
};
748753

749-
if (_shouldTrace && !shouldRedactCommand && !skipLogging)
754+
if (shouldTraceCommand)
750755
{
751756
_currentCommandActivity = MongoTelemetry.StartCommandActivity(
752757
commandName,
@@ -770,16 +775,6 @@ private void CompleteCommandActivityWithSuccess(BsonDocument reply = null)
770775
}
771776
}
772777

773-
private void CompleteCommandActivityWithException(Exception exception)
774-
{
775-
if (_currentCommandActivity is not null)
776-
{
777-
MongoTelemetry.RecordException(_currentCommandActivity, exception);
778-
_currentCommandActivity.Dispose();
779-
_currentCommandActivity = null;
780-
}
781-
}
782-
783778
private void HandleCommandFailure(
784779
CommandState state,
785780
BsonDocument reply,

‎tests/MongoDB.Driver.Tests/Core/Connections/CommandEventHelperTests.cs‎

Lines changed: 28 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,7 @@ public void ShouldRedactCommand_should_return_expected_result(string commandJson
5858
public void ShouldTrackState_should_be_correct(
5959
[Values(false, true)] bool logCommands,
6060
[Values(false, true)] bool captureCommandSucceeded,
61-
[Values(false, true)] bool captureCommandFailed,
62-
[Values(false, true)] bool traceCommands)
61+
[Values(false, true)] bool captureCommandFailed)
6362
{
6463
var mockLogger = new Mock<ILogger<LogCategories.Command>>();
6564
mockLogger.Setup(m => m.IsEnabled(LogLevel.Debug)).Returns(logCommands);
@@ -76,16 +75,16 @@ public void ShouldTrackState_should_be_correct(
7675
}
7776

7877
var eventLogger = new EventLogger<LogCategories.Command>(eventCapturer, mockLogger.Object);
79-
var tracingOptions = traceCommands ? new TracingOptions() : new TracingOptions { Disabled = true };
78+
var tracingOptions = new TracingOptions { Disabled = true };
8079
var commandHelper = new CommandEventHelper(eventLogger, tracingOptions);
8180

82-
// No ActivityListener, so tracing doesn't contribute to _shouldTrackState
83-
commandHelper._shouldTrackState().Should().Be(logCommands || captureCommandSucceeded || captureCommandFailed);
81+
// No ActivityListener, so tracing doesn't contribute to _eventsNeedState
82+
commandHelper._eventsNeedState().Should().Be(logCommands || captureCommandSucceeded || captureCommandFailed);
8483
}
8584

8685
[Theory]
8786
[ParameterAttributeData]
88-
public void ShouldTrackState_should_be_correct_with_activity_listener(
87+
public void Callbacks_turn_on_when_listener_is_added_even_if_no_events(
8988
[Values(false, true)] bool logCommands,
9089
[Values(false, true)] bool captureCommandSucceeded,
9190
[Values(false, true)] bool captureCommandFailed,
@@ -94,13 +93,6 @@ public void ShouldTrackState_should_be_correct_with_activity_listener(
9493
ActivityListener listener = null;
9594
try
9695
{
97-
listener = new ActivityListener
98-
{
99-
ShouldListenTo = source => source.Name == "MongoDB.Driver",
100-
Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.AllData
101-
};
102-
ActivitySource.AddActivityListener(listener);
103-
10496
var mockLogger = new Mock<ILogger<LogCategories.Command>>();
10597
mockLogger.Setup(m => m.IsEnabled(LogLevel.Debug)).Returns(logCommands);
10698

@@ -119,7 +111,27 @@ public void ShouldTrackState_should_be_correct_with_activity_listener(
119111
var tracingOptions = traceCommands ? new TracingOptions() : new TracingOptions { Disabled = true };
120112
var commandHelper = new CommandEventHelper(eventLogger, tracingOptions);
121113

122-
commandHelper._shouldTrackState().Should().Be(logCommands || captureCommandSucceeded || captureCommandFailed || traceCommands);
114+
// When there are no listeners, these only return true if logging is enabled or an event is registered,
115+
// regardless of whether tracing is enabled.
116+
commandHelper.ShouldCallBeforeSending.Should().Be(captureCommandSucceeded || captureCommandFailed || logCommands);
117+
commandHelper.ShouldCallAfterSending.Should().Be(captureCommandSucceeded || captureCommandFailed || logCommands);
118+
commandHelper.ShouldCallErrorSending.Should().Be(captureCommandSucceeded || captureCommandFailed || logCommands);
119+
commandHelper.ShouldCallAfterReceiving.Should().Be(captureCommandSucceeded || captureCommandFailed || logCommands);
120+
commandHelper.ShouldCallErrorReceiving.Should().Be(captureCommandSucceeded || captureCommandFailed || logCommands);
121+
122+
listener = new ActivityListener
123+
{
124+
ShouldListenTo = source => source.Name == "MongoDB.Driver",
125+
Sample = (ref ActivityCreationOptions<ActivityContext> _) => ActivitySamplingResult.AllData
126+
};
127+
ActivitySource.AddActivityListener(listener);
128+
129+
// With listeners registered, these always return true unless everything is disabled.
130+
commandHelper.ShouldCallBeforeSending.Should().Be(captureCommandSucceeded || captureCommandFailed || logCommands || traceCommands);
131+
commandHelper.ShouldCallAfterSending.Should().Be(captureCommandSucceeded || captureCommandFailed || logCommands || traceCommands);
132+
commandHelper.ShouldCallErrorSending.Should().Be(captureCommandSucceeded || captureCommandFailed || logCommands || traceCommands);
133+
commandHelper.ShouldCallAfterReceiving.Should().Be(captureCommandSucceeded || captureCommandFailed || logCommands || traceCommands);
134+
commandHelper.ShouldCallErrorReceiving.Should().Be(captureCommandSucceeded || captureCommandFailed || logCommands || traceCommands);
123135
}
124136
finally
125137
{
@@ -130,8 +142,8 @@ public void ShouldTrackState_should_be_correct_with_activity_listener(
130142

131143
internal static class CommandEventHelperReflector
132144
{
133-
public static bool _shouldTrackState(this CommandEventHelper commandEventHelper) =>
134-
(bool)Reflector.GetFieldValue(commandEventHelper, nameof(_shouldTrackState));
145+
public static bool _eventsNeedState(this CommandEventHelper commandEventHelper) =>
146+
(bool)Reflector.GetFieldValue(commandEventHelper, nameof(_eventsNeedState));
135147

136148

137149
public static bool ShouldRedactCommand(BsonDocument command) =>

‎tests/SmokeTests/MongoDB.Driver.SmokeTests.Sdk/OpenTelemetryTests.cs‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,6 @@ namespace MongoDB.Driver.SmokeTests.Sdk
2626
[Trait("Category", "Integration")]
2727
public sealed class OpenTelemetryTests
2828
{
29-
3029
[Fact]
3130
public void MongoClient_should_create_activities_when_tracing_enabled()
3231
{
@@ -84,7 +83,7 @@ public void MongoClient_should_not_create_activities_when_tracing_disabled()
8483
capturedActivities.Should().BeEmpty();
8584
}
8685

87-
private static ActivityListener CreateActivityListener(out List<Activity> capturedActivities)
86+
private static ActivityListener CreateActivityListener(out IReadOnlyCollection<Activity> capturedActivities)
8887
{
8988
var activities = new List<Activity>();
9089
var listener = new ActivityListener

0 commit comments

Comments
 (0)