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.

Author: ajcvickers

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

@@ -34,7 +34,7 @@ namespace MongoDB.Driver.Core.Connections
     internal class CommandEventHelper
     {
         private readonly EventLogger<LogCategories.Command> _eventLogger;
-        private readonly ConcurrentDictionary<int, CommandState> _state;
+        private ConcurrentDictionary<int, CommandState> _state;
 
         private readonly bool _shouldProcessRequestMessages;
         private readonly bool _shouldTrackStarted;
@@ -57,7 +57,7 @@ public CommandEventHelper(EventLogger<LogCategories.Command> eventLogger, Tracin
             _tracingDisabled = tracingOptions?.Disabled == true;
             _queryTextMaxLength = tracingOptions?.QueryTextMaxLength ?? 0;
 
-            _shouldTrackState = _shouldTrackSucceeded || _shouldTrackFailed || !_tracingDisabled;
+            _shouldTrackState = _shouldTrackSucceeded || _shouldTrackFailed;
             _shouldProcessRequestMessages = _shouldTrackStarted || _shouldTrackState;
 
             if (_shouldTrackState)
@@ -68,30 +68,18 @@ public CommandEventHelper(EventLogger<LogCategories.Command> eventLogger, Tracin
             }
         }
 
-        public bool ShouldCallBeforeSending
-        {
-            get { return _shouldProcessRequestMessages; }
-        }
+        public bool ShouldCallBeforeSending => _shouldProcessRequestMessages || ShouldTraceWithActivityListener();
 
-        public bool ShouldCallAfterSending
-        {
-            get { return _shouldTrackState; }
-        }
+        public bool ShouldCallAfterSending => _shouldTrackState || ShouldTraceWithActivityListener();
 
-        public bool ShouldCallErrorSending
-        {
-            get { return _shouldTrackState; }
-        }
+        public bool ShouldCallErrorSending => _shouldTrackState || ShouldTraceWithActivityListener();
 
-        public bool ShouldCallAfterReceiving
-        {
-            get { return _shouldTrackState; }
-        }
+        public bool ShouldCallAfterReceiving => _shouldTrackState || ShouldTraceWithActivityListener();
 
-        public bool ShouldCallErrorReceiving
-        {
-            get { return _shouldTrackState; }
-        }
+        public bool ShouldCallErrorReceiving => _shouldTrackState || ShouldTraceWithActivityListener();
+
+        private bool ShouldTraceWithActivityListener()
+            => !_tracingDisabled && MongoTelemetry.ActivitySource.HasListeners();
 
         public void CompleteFailedCommandActivity(Exception exception)
         {
@@ -129,8 +117,12 @@ public void BeforeSending(
 
         public void AfterSending(RequestMessage message, ConnectionId connectionId, ObjectId? serviceId, bool skipLogging)
         {
-            CommandState state;
-            if (_state.TryGetValue(message.RequestId, out state) &&
+            if (_state == null)
+            {
+                return;
+            }
+
+            if (_state.TryGetValue(message.RequestId, out var state) &&
                 state.ExpectedResponseType == ExpectedResponseType.None)
             {
                 state.Stopwatch.Stop();
@@ -157,12 +149,16 @@ public void AfterSending(RequestMessage message, ConnectionId connectionId, Obje
 
         public void ErrorSending(RequestMessage message, ConnectionId connectionId, ObjectId? serviceId, Exception exception, bool skipLogging)
         {
-            CommandState state;
-            if (_state.TryRemove(message.RequestId, out state))
+            CompleteCommandActivityWithException(exception);
+
+            if (_state == null)
             {
-                state.Stopwatch.Stop();
+                return;
+            }
 
-                CompleteCommandActivityWithException(exception);
+            if (_state.TryRemove(message.RequestId, out var state))
+            {
+                state.Stopwatch.Stop();
 
                 _eventLogger.LogAndPublish(new CommandFailedEvent(
                         state.CommandName,
@@ -179,8 +175,12 @@ public void ErrorSending(RequestMessage message, ConnectionId connectionId, Obje
 
         public void AfterReceiving(ResponseMessage message, IByteBuffer buffer, ConnectionId connectionId, ObjectId? serviceId, MessageEncoderSettings encoderSettings, bool skipLogging)
         {
-            CommandState state;
-            if (!_state.TryRemove(message.ResponseTo, out state))
+            if (_state == null)
+            {
+                return;
+            }
+
+            if (!_state.TryRemove(message.ResponseTo, out var state))
             {
                 // this indicates a bug in the sending portion...
                 return;
@@ -198,17 +198,21 @@ public void AfterReceiving(ResponseMessage message, IByteBuffer buffer, Connecti
 
         public void ErrorReceiving(int responseTo, ConnectionId connectionId, ObjectId? serviceId, Exception exception, bool skipLogging)
         {
-            CommandState state;
-            if (!_state.TryRemove(responseTo, out state))
+            CompleteCommandActivityWithException(exception);
+
+            if (_state == null)
+            {
+                return;
+            }
+
+            if (!_state.TryRemove(responseTo, out var state))
             {
                 // this indicates a bug in the sending portion...
                 return;
             }
 
             state.Stopwatch.Stop();
 
-            CompleteCommandActivityWithException(exception);
-
             _eventLogger.LogAndPublish(new CommandFailedEvent(
                 state.CommandName,
                 state.QueryNamespace.DatabaseNamespace,
@@ -223,13 +227,18 @@ public void ErrorReceiving(int responseTo, ConnectionId connectionId, ObjectId?
 
         public void ConnectionFailed(ConnectionId connectionId, ObjectId? serviceId, Exception exception, bool skipLogging)
         {
-            if (!_shouldTrackFailed && (_tracingDisabled || !MongoTelemetry.ActivitySource.HasListeners()))
+            if (!_shouldTrackFailed && !ShouldTraceWithActivityListener())
             {
                 return;
             }
 
             CompleteCommandActivityWithException(exception);
 
+            if (_state == null)
+            {
+                return;
+            }
+
             var requestIds = _state.Keys;
             foreach (var requestId in requestIds)
             {
@@ -731,11 +740,19 @@ private void TrackCommandState(
             BsonDocument sessionId,
             long? transactionNumber)
         {
-            if (!_shouldTrackState)
+            var shouldTraceCommand = ShouldTraceWithActivityListener() && !shouldRedactCommand && !skipLogging;
+
+            if (!_shouldTrackState && !shouldTraceCommand)
             {
                 return;
             }
 
+            if (!_shouldTrackState && shouldTraceCommand)
+            {
+                // Lazy initialize state dictionary when tracing is needed but event tracking is not
+                _state ??= new ConcurrentDictionary<int, CommandState>();
+            }
+
             var commandState = new CommandState
             {
                 CommandName = commandName,
@@ -746,7 +763,7 @@ private void TrackCommandState(
                 ShouldRedactReply = shouldRedactCommand
             };
 
-            if (!_tracingDisabled && MongoTelemetry.ActivitySource.HasListeners() && !shouldRedactCommand && !skipLogging)
+            if (shouldTraceCommand)
             {
                 _currentCommandActivity = MongoTelemetry.StartCommandActivity(
                     commandName,

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

@@ -119,7 +119,9 @@ public void ShouldTrackState_should_be_correct_with_activity_listener(
                 var tracingOptions = traceCommands ? new TracingOptions() : new TracingOptions { Disabled = true };
                 var commandHelper = new CommandEventHelper(eventLogger, tracingOptions);
 
-                commandHelper._shouldTrackState().Should().Be(logCommands || captureCommandSucceeded || captureCommandFailed || traceCommands);
+                // With the new implementation, _shouldTrackState only reflects event tracking,
+                // not tracing with listeners (which is checked dynamically at execution time)
+                commandHelper._shouldTrackState().Should().Be(logCommands || captureCommandSucceeded || captureCommandFailed);
             }
             finally
             {

tests/SmokeTests/MongoDB.Driver.SmokeTests.Sdk/LibmongocryptTests.cs

@@ -59,7 +59,7 @@ public void Explicit_encryption_with_libmongocrypt_package_works()
                 kmsProviders.Add("local", localKey);
 
                 var keyVaultNamespace = CollectionNamespace.FromFullName("encryption.__keyVault");
-                var keyVaultMongoClient = new MongoClient(InfrastructureUtilities.MongoUri);
+                var keyVaultMongoClient = new MongoClient();
                 var clientEncryptionSettings = new ClientEncryptionOptions(
                     keyVaultMongoClient,
                     keyVaultNamespace,