Improve ConnectWithTimeout in PooledSocket

cnblogs-dudu · cnblogs-dudu · commit 44bfae8a7d1b · 2019-04-30T18:40:28.000+08:00
diff --git a/Enyim.Caching/Memcached/PooledSocket.cs b/Enyim.Caching/Memcached/PooledSocket.cs
@@ -44,16 +44,20 @@ public PooledSocket(DnsEndPoint endpoint, TimeSpan connectionTimeout, TimeSpan r
             socket.ReceiveTimeout = rcv;
             socket.SendTimeout = rcv;
 
-            ConnectWithTimeout(socket, endpoint, timeout);
+            if (!ConnectWithTimeout(socket, endpoint, timeout))
+            {
+                throw new TimeoutException($"Could not connect to {endpoint.Host}:{endpoint.Port}.");
+            }
 
             this.socket = socket;
             this.endpoint = endpoint;
 
             this.inputStream = new NetworkStream(socket);
         }
 
-        private void ConnectWithTimeout(Socket socket, DnsEndPoint endpoint, int timeout)
+        private bool ConnectWithTimeout(Socket socket, DnsEndPoint endpoint, int timeout)
         {
+            bool connected = false;
             socket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.KeepAlive, true);
             var args = new SocketAsyncEventArgs();
 
@@ -64,25 +68,31 @@ private void ConnectWithTimeout(Socket socket, DnsEndPoint endpoint, int timeout
                     .FirstOrDefault(ip => ip.AddressFamily == AddressFamily.InterNetwork);
                 if (address == null)
                     throw new ArgumentException(String.Format("Could not resolve host '{0}'.", endpoint.Host));
-                args.RemoteEndPoint = new IPEndPoint(address, endpoint.Port);
-            }
-            else
-            {
-                //DnsEndPoint is not working on linux
-                args.RemoteEndPoint = new IPEndPoint(address, endpoint.Port);
             }
 
-            using (var mres = new ManualResetEventSlim())
+            //Learn from https://github.com/dotnet/corefx/blob/release/2.2/src/System.Data.SqlClient/src/System/Data/SqlClient/SNI/SNITcpHandle.cs#L180
+            var cts = new CancellationTokenSource();
+            cts.CancelAfter(timeout);
+            void Cancel()
             {
-                args.Completed += (s, e) => mres.Set();
-                if (socket.ConnectAsync(args))
+                if (!socket.Connected)
                 {
-                    if (!mres.Wait(timeout))
-                    {
-                        throw new TimeoutException("Could not connect to " + endpoint);
-                    }
+                    socket.Dispose();
                 }
             }
+            cts.Token.Register(Cancel);
+
+            socket.Connect(address, endpoint.Port);
+            if (socket.Connected)
+            {
+                connected = true;
+            }
+            else
+            {
+                socket.Dispose();
+            }
+
+            return connected;
         }
 
         public Action<PooledSocket> CleanupCallback { get; set; }
