CSHARP-3671: Better wait queue timeout errors for load balanced clusters.

Author: DmitryLukyanov

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

@@ -34,6 +34,7 @@ public sealed class ReadPreferenceBinding : IReadBinding
         private readonly ReadPreference _readPreference;
         private readonly IServerSelector _serverSelector;
         private readonly ICoreSessionHandle _session;
+        private readonly TrackedOperationRunContext _trackedOperationRunContext;
 
         // constructors
         /// <summary>
@@ -43,11 +44,17 @@ public sealed class ReadPreferenceBinding : IReadBinding
         /// <param name="readPreference">The read preference.</param>
         /// <param name="session">The session.</param>
         public ReadPreferenceBinding(ICluster cluster, ReadPreference readPreference, ICoreSessionHandle session)
+            : this(cluster, readPreference, session, trackedOperationRunContext: null)
+        {
+        }
+
+        internal ReadPreferenceBinding(ICluster cluster, ReadPreference readPreference, ICoreSessionHandle session, TrackedOperationRunContext trackedOperationRunContext)
         {
             _cluster = Ensure.IsNotNull(cluster, nameof(cluster));
             _readPreference = Ensure.IsNotNull(readPreference, nameof(readPreference));
             _session = Ensure.IsNotNull(session, nameof(session));
             _serverSelector = new ReadPreferenceServerSelector(readPreference);
+            _trackedOperationRunContext = trackedOperationRunContext; // can be null
         }
 
         // properties
@@ -82,7 +89,7 @@ public async Task<IChannelSourceHandle> GetReadChannelSourceAsync(CancellationTo
 
         private IChannelSourceHandle GetChannelSourceHelper(IServer server)
         {
-            return new ChannelSourceHandle(new ServerChannelSource(server, _session.Fork()));
+            return new ChannelSourceHandle(new ServerChannelSource(server, _session.Fork(), _trackedOperationRunContext));
         }
 
         /// <inheritdoc/>

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

@@ -16,7 +16,6 @@
 using System;
 using System.Threading;
 using System.Threading.Tasks;
-using MongoDB.Driver.Core.Connections;
 using MongoDB.Driver.Core.Misc;
 using MongoDB.Driver.Core.Servers;
 
@@ -31,6 +30,7 @@ public sealed class ServerChannelSource : IChannelSource
         private bool _disposed;
         private readonly IServer _server;
         private readonly ICoreSessionHandle _session;
+        private readonly TrackedOperationRunContext _trackedOperationRunContext;
 
         // constructors
         /// <summary>
@@ -39,9 +39,15 @@ public sealed class ServerChannelSource : IChannelSource
         /// <param name="server">The server.</param>
         /// <param name="session">The session.</param>
         public ServerChannelSource(IServer server, ICoreSessionHandle session)
+            : this(server, session, trackedOperationRunContext: null)
+        {
+        }
+
+        internal ServerChannelSource(IServer server, ICoreSessionHandle session, TrackedOperationRunContext trackedOperationRunContext)
         {
             _server = Ensure.IsNotNull(server, nameof(server));
             _session = Ensure.IsNotNull(session, nameof(session));
+            _trackedOperationRunContext = trackedOperationRunContext; // can be null
         }
 
         // properties
@@ -78,14 +84,28 @@ public void Dispose()
         public IChannelHandle GetChannel(CancellationToken cancellationToken)
         {
             ThrowIfDisposed();
-            return _server.GetChannel(cancellationToken);
+            if (_trackedOperationRunContext != null && _server is IServerWithTrackedGetChannel trackedServer)
+            {
+                return trackedServer.GetChannel(_trackedOperationRunContext, cancellationToken);
+            }
+            else
+            {
+                return _server.GetChannel(cancellationToken);
+            }
         }
 
         /// <inheritdoc/>
         public Task<IChannelHandle> GetChannelAsync(CancellationToken cancellationToken)
         {
             ThrowIfDisposed();
-            return _server.GetChannelAsync(cancellationToken);
+            if (_trackedOperationRunContext != null && _server is IServerWithTrackedGetChannel trackedServer)
+            {
+                return trackedServer.GetChannelAsync(_trackedOperationRunContext, cancellationToken);
+            }
+            else
+            {
+                return _server.GetChannelAsync(cancellationToken);
+            }
         }
 
         private void ThrowIfDisposed()

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

@@ -0,0 +1,54 @@
+/* Copyright 2021-present MongoDB Inc.
+*
+* Licensed under the Apache License, Version 2.0 (the "License");
+* you may not use this file except in compliance with the License.
+* You may obtain a copy of the License at
+*
+* http://www.apache.org/licenses/LICENSE-2.0
+*
+* Unless required by applicable law or agreed to in writing, software
+* distributed under the License is distributed on an "AS IS" BASIS,
+* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+* See the License for the specific language governing permissions and
+* limitations under the License.
+*/
+
+namespace MongoDB.Driver.Core.Bindings
+{
+    /// <summary>
+    /// Shows running context for tracking operation.
+    /// </summary>
+    public class TrackedOperationRunContext
+    {
+        #region static
+        /// <summary>
+        /// Create not tracked running context.
+        /// </summary>
+        public static TrackedOperationRunContext CreateEmpty() => new TrackedOperationRunContext(false, false);
+        #endregion
+
+        private readonly bool _isInTransaction;
+        private readonly bool _withCursorResult;
+
+        /// <summary>
+        /// Initializes a new instance of the <see cref="TrackedOperationRunContext"/> class.
+        /// </summary>
+        /// <param name="isInTransaction">Determine whether the operation is under transaction.</param>
+        /// <param name="withCursorResult">Determine whether the operation initiates a cursor result.</param>
+        public TrackedOperationRunContext(bool isInTransaction, bool withCursorResult = false)
+        {
+            _isInTransaction = isInTransaction;
+            _withCursorResult = withCursorResult;
+        }
+
+        /// <summary>
+        /// Determine whether the operation is under transaction.
+        /// </summary>
+        public bool IsInTransaction => _isInTransaction;
+
+        /// <summary>
+        /// Determine whether the operation initiates a cursor result.
+        /// </summary>
+        public bool WithCursorResult => _withCursorResult;
+    }
+}

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

@@ -32,6 +32,7 @@ public sealed class WritableServerBinding : IReadWriteBinding
         private readonly ICluster _cluster;
         private bool _disposed;
         private readonly ICoreSessionHandle _session;
+        private readonly TrackedOperationRunContext _trackedOperationRunContext;
 
         // constructors
         /// <summary>
@@ -40,9 +41,15 @@ public sealed class WritableServerBinding : IReadWriteBinding
         /// <param name="cluster">The cluster.</param>
         /// <param name="session">The session.</param>
         public WritableServerBinding(ICluster cluster, ICoreSessionHandle session)
+            : this(cluster, session, trackedOperationRunContext: null)
+        {
+        }
+
+        internal WritableServerBinding(ICluster cluster, ICoreSessionHandle session, TrackedOperationRunContext trackedOperationRunContext)
         {
             _cluster = Ensure.IsNotNull(cluster, nameof(cluster));
             _session = Ensure.IsNotNull(session, nameof(session));
+            _trackedOperationRunContext = trackedOperationRunContext; // can be null
         }
 
         // properties
@@ -94,7 +101,7 @@ public async Task<IChannelSourceHandle> GetWriteChannelSourceAsync(CancellationT
 
         private IChannelSourceHandle GetChannelSourceHelper(IServer server)
         {
-            return new ChannelSourceHandle(new ServerChannelSource(server, _session.Fork()));
+            return new ChannelSourceHandle(new ServerChannelSource(server, _session.Fork(), _trackedOperationRunContext));
         }
 
         /// <inheritdoc/>

src/MongoDB.Driver.Core/Core/ChannelPinningHelper.cs

@@ -32,8 +32,9 @@ public static class ChannelPinningHelper
         /// <param name="cluster">The cluster,</param>
         /// <param name="session">The session.</param>
         /// <param name="readPreference">The read preference.</param>
+        /// <param name="trackedOperationRunContext">The tracked operation run context.</param>
         /// <returns>An effective read binging.</returns>
-        public static IReadBindingHandle CreateReadBinding(ICluster cluster, ICoreSessionHandle session, ReadPreference readPreference)
+        public static IReadBindingHandle CreateReadBinding(ICluster cluster, ICoreSessionHandle session, ReadPreference readPreference, TrackedOperationRunContext trackedOperationRunContext = null)
         {
             IReadBinding readBinding;
             if (session.IsInTransaction &&
@@ -52,7 +53,9 @@ public static IReadBindingHandle CreateReadBinding(ICluster cluster, ICoreSessio
                     // unpin if the next operation is not under transaction
                     session.CurrentTransaction.UnpinAll();
                 }
-                readBinding = new ReadPreferenceBinding(cluster, readPreference, session);
+
+                trackedOperationRunContext = trackedOperationRunContext ?? new TrackedOperationRunContext(session.IsInTransaction, withCursorResult: false); // the default
+                readBinding = new ReadPreferenceBinding(cluster, readPreference, session, trackedOperationRunContext);
             }
 
             return new ReadBindingHandle(readBinding);
@@ -63,8 +66,9 @@ public static IReadBindingHandle CreateReadBinding(ICluster cluster, ICoreSessio
         /// </summary>
         /// <param name="cluster">The cluster.</param>
         /// <param name="session">The session.</param>
+        /// <param name="trackedOperationRunContext">The tracked operation run context.</param>
         /// <returns>An effective read write binging.</returns>
-        public static IReadWriteBindingHandle CreateReadWriteBinding(ICluster cluster, ICoreSessionHandle session)
+        public static IReadWriteBindingHandle CreateReadWriteBinding(ICluster cluster, ICoreSessionHandle session, TrackedOperationRunContext trackedOperationRunContext = null)
         {
             IReadWriteBinding readWriteBinding;
             if (session.IsInTransaction &&
@@ -83,7 +87,9 @@ public static IReadWriteBindingHandle CreateReadWriteBinding(ICluster cluster, I
                     // unpin if the next operation is not under transaction
                     session.CurrentTransaction.UnpinAll();
                 }
-                readWriteBinding = new WritableServerBinding(cluster, session);
+
+                trackedOperationRunContext = trackedOperationRunContext ?? new TrackedOperationRunContext(session.IsInTransaction, withCursorResult: false); // the default
+                readWriteBinding = new WritableServerBinding(cluster, session, trackedOperationRunContext);
             }
 
             return new ReadWriteBindingHandle(readWriteBinding);

src/MongoDB.Driver.Core/Core/ConnectionPools/ExclusiveConnectionPool.Helpers.cs

@@ -29,7 +29,7 @@
 
 namespace MongoDB.Driver.Core.ConnectionPools
 {
-    internal sealed partial class ExclusiveConnectionPool : IConnectionPool
+    internal sealed partial class ExclusiveConnectionPool
     {
         // nested classes
         private static class State
@@ -42,14 +42,16 @@ private static class State
         private sealed class AcquireConnectionHelper
         {
             // private fields
+            private readonly CheckedOutReason _checkedOutReason;
             private readonly ExclusiveConnectionPool _pool;
             private bool _enteredPool;
             private bool _enteredWaitQueue;
             private Stopwatch _stopwatch;
 
             // constructors
-            public AcquireConnectionHelper(ExclusiveConnectionPool pool)
+            public AcquireConnectionHelper(ExclusiveConnectionPool pool, CheckedOutReason checkedOutReason)
             {
+                _checkedOutReason = checkedOutReason;
                 _pool = pool;
             }
 
@@ -82,7 +84,7 @@ public IConnectionHandle EnteredPool(bool enteredPool, CancellationToken cancell
                 if (enteredPool)
                 {
                     var timeSpentInWaitQueue = _stopwatch.Elapsed;
-                    using (var connectionCreator = new ConnectionCreator(_pool, _pool._settings.WaitQueueTimeout - timeSpentInWaitQueue))
+                    using (var connectionCreator = new ConnectionCreator(_pool, _pool._settings.WaitQueueTimeout - timeSpentInWaitQueue, _checkedOutReason))
                     {
                         connection = connectionCreator.CreateOpenedOrReuse(cancellationToken);
                     }
@@ -99,7 +101,7 @@ public async Task<IConnectionHandle> EnteredPoolAsync(bool enteredPool, Cancella
                 if (enteredPool)
                 {
                     var timeSpentInWaitQueue = _stopwatch.Elapsed;
-                    using (var connectionCreator = new ConnectionCreator(_pool, _pool._settings.WaitQueueTimeout - timeSpentInWaitQueue))
+                    using (var connectionCreator = new ConnectionCreator(_pool, _pool._settings.WaitQueueTimeout - timeSpentInWaitQueue, _checkedOutReason))
                     {
                         connection = await connectionCreator.CreateOpenedOrReuseAsync(cancellationToken).ConfigureAwait(false);
                     }
@@ -118,13 +120,20 @@ private AcquiredConnection FinalizePoolEnterance(PooledConnection pooledConnecti
                     var checkedOutConnectionEvent = new ConnectionPoolCheckedOutConnectionEvent(connectionHandle.ConnectionId, _stopwatch.Elapsed, EventContext.OperationId);
                     _pool._checkedOutConnectionEventHandler?.Invoke(checkedOutConnectionEvent);
 
+                    _pool._checkedOutTracker.CheckOut(_checkedOutReason);
+
                     return connectionHandle;
                 }
                 else
                 {
                     _stopwatch.Stop();
 
-                    var message = $"Timed out waiting for a connection after {_stopwatch.ElapsedMilliseconds}ms.";
+                    var message =
+                        "Timeout waiting for connection from the connection pool. " +
+                        $"maxPoolSize: {_pool._settings.MaxConnections}, " +
+                        $"connections in use by cursors: {_pool._checkedOutTracker.GetCheckedOutNumber(CheckedOutReason.Cursor)}, " +
+                        $"connections in use by transactions: {_pool._checkedOutTracker.GetCheckedOutNumber(CheckedOutReason.Transaction)}, " +
+                        $"connections in use by other operations: {_pool._checkedOutTracker.GetCheckedOutNumber(CheckedOutReason.NotSet)}";
                     throw new TimeoutException(message);
                 }
             }
@@ -175,18 +184,25 @@ public void HandleException(Exception ex)
 
         private sealed class PooledConnection : IConnection
         {
+            private readonly CheckedOutReason _checkedOutReason;
             private readonly IConnection _connection;
             private readonly ExclusiveConnectionPool _connectionPool;
             private int _generation;
             private bool _disposed;
 
-            public PooledConnection(ExclusiveConnectionPool connectionPool, IConnection connection)
+            public PooledConnection(ExclusiveConnectionPool connectionPool, IConnection connection, CheckedOutReason checkedOutReason)
             {
+                _checkedOutReason = checkedOutReason;
                 _connectionPool = connectionPool;
                 _connection = connection;
                 _generation = connectionPool._generation;
             }
 
+            public CheckedOutReason CheckedOutReason
+            {
+                get { return _checkedOutReason; }
+            }
+
             public ConnectionId ConnectionId
             {
                 get { return _connection.ConnectionId; }
@@ -621,6 +637,7 @@ private void SignalOrReset()
 
         private sealed class ConnectionCreator : IDisposable
         {
+            private readonly CheckedOutReason _checkedOutReason;
             private readonly ExclusiveConnectionPool _pool;
             private readonly TimeSpan _connectingTimeout;
 
@@ -631,8 +648,9 @@ private sealed class ConnectionCreator : IDisposable
 
             private Stopwatch _stopwatch;
 
-            public ConnectionCreator(ExclusiveConnectionPool pool, TimeSpan connectingTimeout)
+            public ConnectionCreator(ExclusiveConnectionPool pool, TimeSpan connectingTimeout, CheckedOutReason checkedOutReason)
             {
+                _checkedOutReason = checkedOutReason;
                 _pool = pool;
                 _connectingTimeout = connectingTimeout;
                 _connectingWaitStatus = SemaphoreSlimSignalable.SemaphoreWaitResult.None;
@@ -768,7 +786,7 @@ private void StartCreating(CancellationToken cancellationToken)
                 cancellationToken.ThrowIfCancellationRequested();
 
                 _stopwatch = Stopwatch.StartNew();
-                _connection = _pool.CreateNewConnection();
+                _connection = _pool.CreateNewConnection(_checkedOutReason);
             }
 
             private void FinishCreating(ConnectionDescription description)