Implement session and transaction tracing

Author: adelinowona

src/MongoDB.Driver/ClientSessionHandle.cs

@@ -158,7 +158,12 @@ public IClientSessionHandle Fork()
         public void StartTransaction(TransactionOptions transactionOptions = null)
         {
             var effectiveTransactionOptions = GetEffectiveTransactionOptions(transactionOptions);
-            _coreSession.StartTransaction(effectiveTransactionOptions);
+
+            // Check if tracing is enabled for this client
+            var tracingOptions = _client.Settings.TracingOptions;
+            var isTracingEnabled = tracingOptions == null || !tracingOptions.Disabled;
+
+            ((ICoreSessionInternal)_coreSession).StartTransaction(effectiveTransactionOptions, isTracingEnabled);
         }
 
         /// <inheritdoc />

src/MongoDB.Driver/Core/Bindings/CoreSession.cs

@@ -14,10 +14,12 @@
 */
 
 using System;
+using System.Diagnostics;
 using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
 using MongoDB.Bson;
+using MongoDB.Driver;
 using MongoDB.Driver.Core.Clusters;
 using MongoDB.Driver.Core.Misc;
 using MongoDB.Driver.Core.Operations;
@@ -190,6 +192,20 @@ void ICoreSessionInternal.AbortTransaction(AbortTransactionOptions options, Canc
                 _currentTransaction.SetState(CoreTransactionState.Aborted);
                 // The transaction is aborted.The session MUST be unpinned regardless
                 // of whether the abortTransaction command succeeds or fails
+                if (_currentTransaction.TransactionActivity != null)
+                {
+                    var transactionActivity = _currentTransaction.TransactionActivity;
+                    _currentTransaction.TransactionActivity = null;
+
+                    // Set status to Ok for successfully aborted transaction
+                    transactionActivity.SetStatus(ActivityStatusCode.Ok);
+
+                    // Dispose the transaction activity. Note: Activity.Current was already restored to the
+                    // parent in StartTransaction() to prevent AsyncLocal flow issues, so the transaction
+                    // activity was never persisted in Activity.Current. Set it explicitly to be defensive.
+                    Activity.Current = _currentTransaction.ParentActivity;
+                    transactionActivity.Dispose();
+                }
                 _currentTransaction.UnpinAll();
             }
         }
@@ -244,6 +260,20 @@ async Task ICoreSessionInternal.AbortTransactionAsync(AbortTransactionOptions op
                 _currentTransaction.SetState(CoreTransactionState.Aborted);
                 // The transaction is aborted.The session MUST be unpinned regardless
                 // of whether the abortTransaction command succeeds or fails
+                if (_currentTransaction.TransactionActivity != null)
+                {
+                    var transactionActivity = _currentTransaction.TransactionActivity;
+                    _currentTransaction.TransactionActivity = null;
+
+                    // Set status to Ok for successfully aborted transaction
+                    transactionActivity.SetStatus(ActivityStatusCode.Ok);
+
+                    // Dispose the transaction activity. Note: Activity.Current was already restored to the
+                    // parent in StartTransaction() to prevent AsyncLocal flow issues, so the transaction
+                    // activity was never persisted in Activity.Current. Set it explicitly to be defensive.
+                    Activity.Current = _currentTransaction.ParentActivity;
+                    transactionActivity.Dispose();
+                }
                 _currentTransaction.UnpinAll();
             }
         }
@@ -334,6 +364,21 @@ void ICoreSessionInternal.CommitTransaction(CommitTransactionOptions options, Ca
             {
                 _isCommitTransactionInProgress = false;
                 _currentTransaction.SetState(CoreTransactionState.Committed);
+                // Stop the transaction span immediately so it's captured for testing
+                if (_currentTransaction.TransactionActivity != null)
+                {
+                    var transactionActivity = _currentTransaction.TransactionActivity;
+                    _currentTransaction.TransactionActivity = null;
+
+                    // Set status to Ok for successfully committed transaction
+                    transactionActivity.SetStatus(ActivityStatusCode.Ok);
+
+                    // Dispose the transaction activity. Note: Activity.Current was already restored to the
+                    // parent in StartTransaction() to prevent AsyncLocal flow issues, so the transaction
+                    // activity was never persisted in Activity.Current. Set it explicitly to be defensive.
+                    Activity.Current = _currentTransaction.ParentActivity;
+                    transactionActivity.Dispose();
+                }
             }
         }
 
@@ -374,6 +419,21 @@ async Task ICoreSessionInternal.CommitTransactionAsync(CommitTransactionOptions
             {
                 _isCommitTransactionInProgress = false;
                 _currentTransaction.SetState(CoreTransactionState.Committed);
+                // Stop the transaction span immediately so it's captured for testing
+                if (_currentTransaction.TransactionActivity != null)
+                {
+                    var transactionActivity = _currentTransaction.TransactionActivity;
+                    _currentTransaction.TransactionActivity = null;
+
+                    // Set status to Ok for successfully committed transaction
+                    transactionActivity.SetStatus(ActivityStatusCode.Ok);
+
+                    // Dispose the transaction activity. Note: Activity.Current was already restored to the
+                    // parent in StartTransaction() to prevent AsyncLocal flow issues, so the transaction
+                    // activity was never persisted in Activity.Current. Set it explicitly to be defensive.
+                    Activity.Current = _currentTransaction.ParentActivity;
+                    transactionActivity.Dispose();
+                }
             }
         }
 
@@ -414,6 +474,9 @@ public void MarkDirty()
 
         /// <inheritdoc />
         public void StartTransaction(TransactionOptions transactionOptions = null)
+            => ((ICoreSessionInternal)this).StartTransaction(transactionOptions, isTracingEnabled: false);
+
+        void ICoreSessionInternal.StartTransaction(TransactionOptions transactionOptions, bool isTracingEnabled)
         {
             EnsureStartTransactionCanBeCalled();
 
@@ -425,7 +488,19 @@ public void StartTransaction(TransactionOptions transactionOptions = null)
             }
 
             _currentTransaction?.UnpinAll(); // unpin data if any when a new transaction is started
-            _currentTransaction = new CoreTransaction(transactionNumber, effectiveTransactionOptions);
+            _currentTransaction = new CoreTransaction(transactionNumber, effectiveTransactionOptions, isTracingEnabled);
+
+            // Start transaction span for OpenTelemetry tracing (if enabled)
+            if (isTracingEnabled)
+            {
+                // Store the parent activity to restore later
+                _currentTransaction.ParentActivity = Activity.Current;
+                _currentTransaction.TransactionActivity = MongoTelemetry.StartTransactionActivity();
+
+                // Immediately restore Activity.Current to the parent to prevent AsyncLocal flow issues.
+                // The transaction activity will be explicitly set as parent for operations within the transaction.
+                Activity.Current = _currentTransaction.ParentActivity;
+            }
         }
 
         /// <inheritdoc />
@@ -559,19 +634,87 @@ private void EnsureTransactionsAreSupported()
 
         private TResult ExecuteEndTransactionOnPrimary<TResult>(OperationContext operationContext, IReadOperation<TResult> operation)
         {
-            using (var sessionHandle = new NonDisposingCoreSessionHandle(this))
-            using (var binding = ChannelPinningHelper.CreateReadWriteBinding(_cluster, sessionHandle))
+            // Determine operation name and create operation-level span if tracing is enabled
+            string operationName = operation switch
+            {
+                CommitTransactionOperation => "commitTransaction",
+                AbortTransactionOperation => "abortTransaction",
+                _ => null
+            };
+
+            // Temporarily set Activity.Current to transaction activity so the operation nests under it
+            var transactionActivity = _currentTransaction?.TransactionActivity;
+            var previousActivity = Activity.Current;
+            if (transactionActivity != null)
             {
-                return operation.Execute(operationContext, binding);
+                Activity.Current = transactionActivity;
+            }
+
+            using var activity = _currentTransaction?.IsTracingEnabled == true && operationName != null
+                ? MongoTelemetry.StartOperationActivity(operationName, "admin", collectionName: null)
+                : null;
+
+            // Don't restore Activity.Current yet - let it stay as the operation activity
+            // so command activities nest under it. We'll restore after the operation completes.
+
+            try
+            {
+                using (var sessionHandle = new NonDisposingCoreSessionHandle(this))
+                using (var binding = ChannelPinningHelper.CreateReadWriteBinding(_cluster, sessionHandle))
+                {
+                    return operation.Execute(operationContext, binding);
+                }
+            }
+            finally
+            {
+                // Restore Activity.Current after operation completes
+                if (transactionActivity != null)
+                {
+                    Activity.Current = previousActivity;
+                }
             }
         }
 
         private async Task<TResult> ExecuteEndTransactionOnPrimaryAsync<TResult>(OperationContext operationContext, IReadOperation<TResult> operation)
         {
-            using (var sessionHandle = new NonDisposingCoreSessionHandle(this))
-            using (var binding = ChannelPinningHelper.CreateReadWriteBinding(_cluster, sessionHandle))
+            // Determine operation name and create operation-level span if tracing is enabled
+            string operationName = operation switch
+            {
+                CommitTransactionOperation => "commitTransaction",
+                AbortTransactionOperation => "abortTransaction",
+                _ => null
+            };
+
+            // Temporarily set Activity.Current to transaction activity so the operation nests under it
+            var transactionActivity = _currentTransaction?.TransactionActivity;
+            var previousActivity = Activity.Current;
+            if (transactionActivity != null)
+            {
+                Activity.Current = transactionActivity;
+            }
+
+            using var activity = _currentTransaction?.IsTracingEnabled == true && operationName != null
+                ? MongoTelemetry.StartOperationActivity(operationName, "admin", collectionName: null)
+                : null;
+
+            // Don't restore Activity.Current yet - let it stay as the operation activity
+            // so command activities nest under it. We'll restore after the operation completes.
+
+            try
+            {
+                using (var sessionHandle = new NonDisposingCoreSessionHandle(this))
+                using (var binding = ChannelPinningHelper.CreateReadWriteBinding(_cluster, sessionHandle))
+                {
+                    return await operation.ExecuteAsync(operationContext, binding).ConfigureAwait(false);
+                }
+            }
+            finally
             {
-                return await operation.ExecuteAsync(operationContext, binding).ConfigureAwait(false);
+                // Restore Activity.Current after operation completes
+                if (transactionActivity != null)
+                {
+                    Activity.Current = previousActivity;
+                }
             }
         }
 

src/MongoDB.Driver/Core/Bindings/CoreTransaction.cs

@@ -13,6 +13,7 @@
 * limitations under the License.
 */
 
+using System.Diagnostics;
 using MongoDB.Bson;
 using MongoDB.Driver.Core.Servers;
 
@@ -32,6 +33,9 @@ public class CoreTransaction
         private readonly long _transactionNumber;
         private readonly TransactionOptions _transactionOptions;
         private readonly object _lock = new object();
+        private Activity _transactionActivity;
+        private Activity _parentActivity;
+        private readonly bool _isTracingEnabled;
 
         // public constructors
         /// <summary>
@@ -40,9 +44,21 @@ public class CoreTransaction
         /// <param name="transactionNumber">The transaction number.</param>
         /// <param name="transactionOptions">The transaction options.</param>
         public CoreTransaction(long transactionNumber, TransactionOptions transactionOptions)
+            : this(transactionNumber, transactionOptions, isTracingEnabled: false)
+        {
+        }
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="CoreTransaction" /> class.
+        /// </summary>
+        /// <param name="transactionNumber">The transaction number.</param>
+        /// <param name="transactionOptions">The transaction options.</param>
+        /// <param name="isTracingEnabled">Whether OpenTelemetry tracing is enabled for this transaction.</param>
+        internal CoreTransaction(long transactionNumber, TransactionOptions transactionOptions, bool isTracingEnabled)
         {
             _transactionNumber = transactionNumber;
             _transactionOptions = transactionOptions;
+            _isTracingEnabled = isTracingEnabled;
             _state = CoreTransactionState.Starting;
             _isEmpty = true;
         }
@@ -58,6 +74,29 @@ public CoreTransaction(long transactionNumber, TransactionOptions transactionOpt
 
         internal OperationContext OperationContext { get; set; }
 
+        /// <summary>
+        /// Gets or sets the transaction activity (for OpenTelemetry tracing).
+        /// </summary>
+        internal Activity TransactionActivity
+        {
+            get => _transactionActivity;
+            set => _transactionActivity = value;
+        }
+
+        /// <summary>
+        /// Gets or sets the parent activity to restore after the transaction completes.
+        /// </summary>
+        internal Activity ParentActivity
+        {
+            get => _parentActivity;
+            set => _parentActivity = value;
+        }
+
+        /// <summary>
+        /// Gets whether OpenTelemetry tracing is enabled for this transaction.
+        /// </summary>
+        internal bool IsTracingEnabled => _isTracingEnabled;
+
         /// <summary>
         /// Gets the transaction state.
         /// </summary>
@@ -138,6 +177,8 @@ internal void UnpinAll()
                 _pinnedChannel?.Dispose();
                 _pinnedChannel = null;
                 _pinnedServer = null;
+                _transactionActivity?.Dispose();
+                _transactionActivity = null;
             }
         }
     }

src/MongoDB.Driver/Core/Bindings/ICoreSessionInternal.cs

@@ -25,4 +25,5 @@ internal interface ICoreSessionInternal
     Task AbortTransactionAsync(AbortTransactionOptions options, CancellationToken cancellationToken = default);
     void CommitTransaction(CommitTransactionOptions options, CancellationToken cancellationToken = default);
     Task CommitTransactionAsync(CommitTransactionOptions options, CancellationToken cancellationToken = default);
+    void StartTransaction(TransactionOptions transactionOptions, bool isTracingEnabled);
 }

src/MongoDB.Driver/Core/Bindings/NoCoreSession.cs

@@ -173,6 +173,11 @@ public void StartTransaction(TransactionOptions transactionOptions = null)
             throw new NotSupportedException("NoCoreSession does not support StartTransaction.");
         }
 
+        void ICoreSessionInternal.StartTransaction(TransactionOptions transactionOptions, bool isTracingEnabled)
+        {
+            throw new NotSupportedException("NoCoreSession does not support StartTransaction.");
+        }
+
         /// <inheritdoc />
         public void SetSnapshotTimeIfNeeded(BsonTimestamp snapshotTime)
         {

src/MongoDB.Driver/Core/Bindings/WrappingCoreSession.cs

@@ -286,6 +286,12 @@ public virtual void StartTransaction(TransactionOptions transactionOptions = nul
             _wrapped.StartTransaction(transactionOptions);
         }
 
+        void ICoreSessionInternal.StartTransaction(TransactionOptions transactionOptions, bool isTracingEnabled)
+        {
+            ThrowIfDisposed();
+            ((ICoreSessionInternal)_wrapped).StartTransaction(transactionOptions, isTracingEnabled);
+        }
+
         /// <inheritdoc />
         public void SetSnapshotTimeIfNeeded(BsonTimestamp snapshotTime)
         {