CSHARP-4026: Preemptively cancel in progress operations when SDAM heartbeats timeout. (#765)

CSHARP-4026: Preemptively cancel in progress operations when SDAM heartbeats timeout.

Author: DmitryLukyanov

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

@@ -161,66 +161,6 @@ public void ThrowIfNotInitialized()
             public override string ToString() => State.ToString();
         }
 
-        internal sealed class MaintenanceHelper : IDisposable
-        {
-            private CancellationTokenSource _cancellationTokenSource = null;
-            private readonly Action<CancellationToken> _maintenanceAction;
-            private Thread _maintenanceThread;
-            private readonly TimeSpan _interval;
-
-            public MaintenanceHelper(Action<CancellationToken> maintenanceAction, TimeSpan interval)
-            {
-                _interval = interval;
-                _maintenanceAction = Ensure.IsNotNull(maintenanceAction, nameof(maintenanceAction));
-            }
-
-            public bool IsRunning => _maintenanceThread != null;
-
-            public void Cancel()
-            {
-                if (_interval == Timeout.InfiniteTimeSpan)
-                {
-                    return;
-                }
-
-                CancelAndDispose();
-                _cancellationTokenSource = null;
-                _maintenanceThread = null;
-                // the previous _maintenanceThread might not be stopped yet, but it will be soon
-            }
-
-            public void Start()
-            {
-                if (_interval == Timeout.InfiniteTimeSpan)
-                {
-                    return;
-                }
-
-                CancelAndDispose();
-                _cancellationTokenSource = new CancellationTokenSource();
-                var cancellationToken = _cancellationTokenSource.Token;
-
-                _maintenanceThread = new Thread(new ParameterizedThreadStart(ThreadStart)) { IsBackground = true };
-                _maintenanceThread.Start(cancellationToken);
-
-                void ThreadStart(object cancellationToken)
-                {
-                    _maintenanceAction((CancellationToken)cancellationToken);
-                }
-            }
-
-            public void Dispose()
-            {
-                CancelAndDispose();
-            }
-
-            private void CancelAndDispose()
-            {
-                _cancellationTokenSource?.Cancel();
-                _cancellationTokenSource?.Dispose();
-            }
-        }
-
         private sealed class AcquireConnectionHelper : IDisposable
         {
             // private fields
@@ -406,7 +346,7 @@ private void HandleException(Exception ex)
             }
         }
 
-        private sealed class PooledConnection : IConnection, ICheckOutReasonTracker
+        internal sealed class PooledConnection : IConnection, ICheckOutReasonTracker
         {
             private CheckOutReason? _checkOutReason;
             private readonly IConnection _connection;
@@ -712,11 +652,13 @@ private void ThrowIfDisposed()
             }
         }
 
-        private sealed class ListConnectionHolder
+        internal sealed class ListConnectionHolder
         {
             private readonly SemaphoreSlimSignalable _semaphoreSlimSignalable;
             private readonly object _lock = new object();
+            private readonly object _lockInUse = new object();
             private readonly List<PooledConnection> _connections;
+            private readonly List<PooledConnection> _connectionsInUse;
 
             private readonly Action<ConnectionPoolRemovingConnectionEvent> _removingConnectionEventHandler;
             private readonly Action<ConnectionPoolRemovedConnectionEvent> _removedConnectionEventHandler;
@@ -725,6 +667,7 @@ public ListConnectionHolder(IEventSubscriber eventSubscriber, SemaphoreSlimSigna
             {
                 _semaphoreSlimSignalable = semaphoreSlimSignalable;
                 _connections = new List<PooledConnection>();
+                _connectionsInUse = new List<PooledConnection>();
 
                 eventSubscriber.TryGetEventHandler(out _removingConnectionEventHandler);
                 eventSubscriber.TryGetEventHandler(out _removedConnectionEventHandler);
@@ -745,6 +688,7 @@ public void Clear()
             {
                 lock (_lock)
                 {
+                    // In use Connections MUST be closed when they are checked in to the closed pool.
                     foreach (var connection in _connections)
                     {
                         RemoveConnection(connection);
@@ -755,30 +699,44 @@ public void Clear()
                 }
             }
 
-            public void Prune(CancellationToken cancellationToken)
+            public void Prune(int? maxExpiredGenerationInUse, CancellationToken cancellationToken)
             {
-                PooledConnection[] expiredConnections;
-                lock (_lock)
+                RemoveExpiredConnections(_connections, generation: null, _lock, signal: true);
+
+                if (maxExpiredGenerationInUse.HasValue)
                 {
-                    expiredConnections = _connections.Where(c => c.IsExpired).ToArray();
+                    RemoveExpiredConnections(_connectionsInUse, generation: maxExpiredGenerationInUse.Value, _lockInUse, signal: false);
                 }
 
-                foreach (var connection in expiredConnections)
+                void RemoveExpiredConnections(List<PooledConnection> connections, int? generation, object @lock, bool signal)
                 {
-                    cancellationToken.ThrowIfCancellationRequested();
+                    PooledConnection[] expiredConnections;
+                    lock (@lock)
+                    {
+                        expiredConnections = connections.Where(c => c.IsExpired && (generation == null || c.Generation <= generation)).ToArray();
+                    }
 
-                    lock (_lock)
+                    foreach (var connection in expiredConnections)
                     {
-                        // At this point connection is always expired and might be disposed
-                        // If connection is already disposed the removal logic was already executed
-                        if (connection.IsDisposed)
+                        cancellationToken.ThrowIfCancellationRequested();
+
+                        lock (@lock)
                         {
-                            continue;
-                        }
+                            // At this point connection is always expired and might be disposed
+                            // If connection is already disposed the removal logic was already executed
+                            if (connection.IsDisposed)
+                            {
+                                continue;
+                            }
 
-                        RemoveConnection(connection);
-                        _connections.Remove(connection);
-                        SignalOrReset();
+                            RemoveConnection(connection);
+                            connections.Remove(connection);
+
+                            if (signal)
+                            {
+                                SignalOrReset();
+                            }
+                        }
                     }
                 }
             }
@@ -806,11 +764,26 @@ public PooledConnection Acquire()
 
                     SignalOrReset();
                 }
+
+                if (result != null)
+                {
+                    TrackInUseConnection(result);
+
+                    // This connection can be expired and not disposed by Prune. Dispose if needed
+                    if (result.IsExpired)
+                    {
+                        RemoveConnection(result);
+                        result = null;
+                    }
+                }
+
                 return result;
             }
 
             public void Return(PooledConnection connection)
             {
+                UntrackInUseConnection(connection);
+
                 lock (_lock)
                 {
                     _connections.Add(connection);
@@ -826,6 +799,7 @@ public void RemoveConnection(PooledConnection connection)
                 }
 
                 var stopwatch = Stopwatch.StartNew();
+                UntrackInUseConnection(connection); // no op if connection is not in use
                 connection.Dispose();
                 stopwatch.Stop();
 
@@ -849,9 +823,25 @@ private void SignalOrReset()
                     _semaphoreSlimSignalable.Signal();
                 }
             }
+
+            public void TrackInUseConnection(PooledConnection connection)
+            {
+                lock (_lockInUse)
+                {
+                    _connectionsInUse.Add(connection);
+                }
+            }
+
+            public void UntrackInUseConnection(PooledConnection connection)
+            {
+                lock (_lockInUse)
+                {
+                    _connectionsInUse.Remove(connection);
+                }
+            }
         }
 
-        private sealed class ConnectionCreator : IDisposable
+        internal sealed class ConnectionCreator : IDisposable
         {
             private readonly ExclusiveConnectionPool _pool;
             private readonly TimeSpan _connectingTimeout;
@@ -888,8 +878,7 @@ public PooledConnection CreateOpened(CancellationToken cancellationToken)
                         _pool.CreateTimeoutException(stopwatch, $"Timed out waiting for in connecting queue after {stopwatch.ElapsedMilliseconds}ms.");
                     }
 
-                    var connection = CreateOpenedInternal(cancellationToken);
-                    return connection;
+                    return CreateOpenedInternal(cancellationToken);
                 }
                 catch (Exception ex)
                 {
@@ -992,9 +981,10 @@ public void Dispose()
                     _pool._maxConnectingQueue.Release();
                 }
 
-                if (_disposeConnection)
+                if (_disposeConnection && _connection != null)
                 {
-                    _connection?.Dispose();
+                    _pool.ConnectionHolder.UntrackInUseConnection(_connection);
+                    _connection.Dispose();
                 }
             }
 
@@ -1029,7 +1019,10 @@ private void StartCreating(CancellationToken cancellationToken)
                 cancellationToken.ThrowIfCancellationRequested();
 
                 _stopwatch = Stopwatch.StartNew();
-                _connection = _pool.CreateNewConnection();
+
+                var connection = _pool.CreateNewConnection();
+                _pool.ConnectionHolder.TrackInUseConnection(connection);
+                _connection = connection;
             }
 
             private void FinishCreating(ConnectionDescription description)

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

@@ -76,7 +76,7 @@ public ExclusiveConnectionPool(
             _connectionExceptionHandler = Ensure.IsNotNull(connectionExceptionHandler, nameof(connectionExceptionHandler));
             Ensure.IsNotNull(eventSubscriber, nameof(eventSubscriber));
 
-            _maintenanceHelper = new MaintenanceHelper(MaintainSize, _settings.MaintenanceInterval);
+            _maintenanceHelper = new MaintenanceHelper(this, _settings.MaintenanceInterval);
             _poolState = new PoolState(EndPointHelper.ToString(_endPoint));
             _checkOutReasonCounter = new CheckOutReasonCounter();
 
@@ -155,6 +155,8 @@ public ServerId ServerId
             get { return _serverId; }
         }
 
+        public ConnectionPoolSettings Settings => _settings;
+
         public int UsedCount
         {
             get
@@ -164,6 +166,9 @@ public int UsedCount
             }
         }
 
+        // internal properties
+        internal ListConnectionHolder ConnectionHolder => _connectionHolder;
+
         // public methods
         public IConnectionHandle AcquireConnection(CancellationToken cancellationToken)
         {
@@ -177,7 +182,7 @@ public async Task<IConnectionHandle> AcquireConnectionAsync(CancellationToken ca
             return await helper.AcquireConnectionAsync(cancellationToken).ConfigureAwait(false);
         }
 
-        public void Clear()
+        public void Clear(bool closeInUseConnections = false)
         {
             lock (_poolState)
             {
@@ -187,8 +192,9 @@ public void Clear()
                 {
                     _clearingEventHandler?.Invoke(new ConnectionPoolClearingEvent(_serverId, _settings));
 
-                    _maintenanceHelper.Cancel();
+                    int? maxGenerationToReap = closeInUseConnections ? _generation : null;
                     _generation++;
+                    _maintenanceHelper.Stop(maxGenerationToReap);
 
                     _maxConnectionsQueue.Signal();
                     _maxConnectingQueue.Signal();
@@ -267,8 +273,8 @@ public void Dispose()
                     _closingEventHandler(new ConnectionPoolClosingEvent(_serverId));
                 }
 
-                _connectionHolder.Clear();
                 _maintenanceHelper.Dispose();
+                _connectionHolder.Clear();
                 _maxConnectionsQueue.Dispose();
                 _maxConnectingQueue.Dispose();
                 if (_closedEventHandler != null)
@@ -290,56 +296,13 @@ public int GetGeneration(ObjectId? serviceId)
             return Generation;
         }
 
-        // private methods
-        private void MaintainSize(CancellationToken cancellationToken)
+        // internal methods
+        internal SemaphoreSlimSignalable.SemaphoreSlimSignalableAwaiter CreateMaxConnectionsAwaiter()
         {
-            try
-            {
-                while (!cancellationToken.IsCancellationRequested)
-                {
-                    try
-                    {
-                        _connectionHolder.Prune(cancellationToken);
-                        EnsureMinSize(cancellationToken);
-                    }
-                    catch
-                    {
-                        // ignore exceptions
-                    }
-                    ThreadHelper.Sleep(_settings.MaintenanceInterval, cancellationToken);
-                }
-            }
-            catch
-            {
-                // ignore exceptions
-            }
-        }
-
-        private void EnsureMinSize(CancellationToken cancellationToken)
-        {
-            var minTimeout = TimeSpan.FromMilliseconds(20);
-
-            while (CreatedCount < _settings.MinConnections && !cancellationToken.IsCancellationRequested)
-            {
-                using (var poolAwaiter = _maxConnectionsQueue.CreateAwaiter())
-                {
-                    var entered = poolAwaiter.WaitSignaled(minTimeout, cancellationToken);
-                    if (!entered)
-                    {
-                        return;
-                    }
-
-                    using (var connectionCreator = new ConnectionCreator(this, minTimeout))
-                    {
-                        var connection = connectionCreator.CreateOpened(cancellationToken);
-                        _connectionHolder.Return(connection);
-                    }
-                }
-
-                cancellationToken.ThrowIfCancellationRequested();
-            }
+            return _maxConnectionsQueue.CreateAwaiter();
         }
 
+        // private methods
         private void ReleaseConnection(PooledConnection connection)
         {
             if (_checkingInConnectionEventHandler != null)

src/MongoDB.Driver.Core/Core/ConnectionPools/IConnectionPool.cs

@@ -63,7 +63,8 @@ public interface IConnectionPool : IDisposable
         /// <summary>
         /// Clears the connection pool and marks it as paused.
         /// </summary>
-        void Clear();
+        /// <param name="closeInUseConnections">Whether in use connections should be closed.</param>
+        void Clear(bool closeInUseConnections = false);
 
         /// <summary>
         /// Clears the connection pool for the specified service.