CSHARP-3840: Unresponsive/deadlocked cluster.Dispose() (#1557)

Author: sanych-sun

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

@@ -122,103 +122,62 @@ private void ConfigureConnectedSocket(Socket socket)
 
         private void Connect(Socket socket, EndPoint endPoint, CancellationToken cancellationToken)
         {
-            var state = 1; // 1 == connecting, 2 == connected, 3 == timedout, 4 == cancelled
+            IAsyncResult connectOperation;
 
-            using (new Timer(_ => ChangeState(3), null, _settings.ConnectTimeout, Timeout.InfiniteTimeSpan))
-            using (cancellationToken.Register(() => ChangeState(4)))
+            if (endPoint is DnsEndPoint dnsEndPoint)
             {
-                try
-                {
-                    var dnsEndPoint = endPoint as DnsEndPoint;
-                    if (dnsEndPoint != null)
-                    {
-                        // mono doesn't support DnsEndPoint in its BeginConnect method.
-                        socket.Connect(dnsEndPoint.Host, dnsEndPoint.Port);
-                    }
-                    else
-                    {
-                        socket.Connect(endPoint);
-                    }
-                    ChangeState(2); // note: might not actually go to state 2 if already in state 3 or 4
-                }
-                catch when (state == 1)
-                {
-                    try { socket.Dispose(); } catch { }
-                    throw;
-                }
-                catch when (state >= 3)
-                {
-                    // a timeout or operation cancelled exception will be thrown instead
-                }
+                // mono doesn't support DnsEndPoint in its BeginConnect method.
+                connectOperation = socket.BeginConnect(dnsEndPoint.Host, dnsEndPoint.Port, null, null);
+            }
+            else
+            {
+                connectOperation = socket.BeginConnect(endPoint, null, null);
+            }
 
-                if (state == 3)
-                {
-                    var message = string.Format("Timed out connecting to {0}. Timeout was {1}.", endPoint, _settings.ConnectTimeout);
-                    throw new TimeoutException(message);
-                }
-                if (state == 4) { throw new OperationCanceledException(); }
+            WaitHandle.WaitAny([connectOperation.AsyncWaitHandle, cancellationToken.WaitHandle], _settings.ConnectTimeout);
+
+            if (!connectOperation.IsCompleted)
+            {
+                try { socket.Dispose(); } catch { }
+
+                cancellationToken.ThrowIfCancellationRequested();
+                throw new TimeoutException($"Timed out connecting to {endPoint}. Timeout was {_settings.ConnectTimeout}.");
             }
 
-            void ChangeState(int to)
+            try
             {
-                var from = Interlocked.CompareExchange(ref state, to, 1);
-                if (from == 1 && to >= 3)
-                {
-                    try { socket.Dispose(); } catch { } // disposing the socket aborts the connection attempt
-                }
+                socket.EndConnect(connectOperation);
+            }
+            catch
+            {
+                try { socket.Dispose(); } catch { }
+                throw;
             }
         }
 
         private async Task ConnectAsync(Socket socket, EndPoint endPoint, CancellationToken cancellationToken)
         {
-            var state = 1; // 1 == connecting, 2 == connected, 3 == timedout, 4 == cancelled
+            var timeoutTask = Task.Delay(_settings.ConnectTimeout, cancellationToken);
+            var connectTask = socket.ConnectAsync(endPoint);
+
+            await Task.WhenAny(connectTask, timeoutTask).ConfigureAwait(false);
 
-            using (new Timer(_ => ChangeState(3), null, _settings.ConnectTimeout, Timeout.InfiniteTimeSpan))
-            using (cancellationToken.Register(() => ChangeState(4)))
+            if (!connectTask.IsCompleted)
             {
-                try
-                {
-                    var dnsEndPoint = endPoint as DnsEndPoint;
-#if !NET472
-                    await socket.ConnectAsync(endPoint).ConfigureAwait(false);
-#else
-                    if (dnsEndPoint != null)
-                    {
-                        // mono doesn't support DnsEndPoint in its BeginConnect method.
-                        await Task.Factory.FromAsync(socket.BeginConnect(dnsEndPoint.Host, dnsEndPoint.Port, null, null), socket.EndConnect).ConfigureAwait(false);
-                    }
-                    else
-                    {
-                        await Task.Factory.FromAsync(socket.BeginConnect(endPoint, null, null), socket.EndConnect).ConfigureAwait(false);
-                    }
-#endif
-                    ChangeState(2); // note: might not actually go to state 2 if already in state 3 or 4
-                }
-                catch when (state == 1)
-                {
-                    try { socket.Dispose(); } catch { }
-                    throw;
-                }
-                catch when (state >= 3)
-                {
-                    // a timeout or operation cancelled exception will be thrown instead
-                }
+                try { socket.Dispose(); } catch { }
 
-                if (state == 3)
-                {
-                    var message = string.Format("Timed out connecting to {0}. Timeout was {1}.", endPoint, _settings.ConnectTimeout);
-                    throw new TimeoutException(message);
-                }
-                if (state == 4) { throw new OperationCanceledException(); }
+                cancellationToken.ThrowIfCancellationRequested();
+                throw new TimeoutException($"Timed out connecting to {endPoint}. Timeout was {_settings.ConnectTimeout}.");
             }
 
-            void ChangeState(int to)
+            try
             {
-                var from = Interlocked.CompareExchange(ref state, to, 1);
-                if (from == 1 && to >= 3)
-                {
-                    try { socket.Dispose(); } catch { } // disposing the socket aborts the connection attempt
-                }
+                await connectTask.ConfigureAwait(false);
+            }
+            catch
+            {
+                try { socket.Dispose(); } catch { }
+                throw;
             }
         }