Implement Connect and ConnectAsync in PooledSocket

Author: cnblogs-dudu

Enyim.Caching/Memcached/MemcachedNode.cs

@@ -702,7 +702,7 @@ protected internal virtual PooledSocket CreateSocket()
             try
             {
                 var ps = new PooledSocket(this.endPoint, this.config.ConnectionTimeout, this.config.ReceiveTimeout, _logger);
-                //ps.Connect();
+                ps.Connect();
                 return ps;
             }
             catch (Exception ex)
@@ -717,7 +717,7 @@ protected internal virtual async Task<PooledSocket> CreateSocketAsync()
             try
             {
                 var ps = new PooledSocket(this.endPoint, this.config.ConnectionTimeout, this.config.ReceiveTimeout, _logger);
-                //await ps.ConnectAsync();
+                await ps.ConnectAsync();
                 return ps;
             }
             catch (Exception ex)

Enyim.Caching/Memcached/PooledSocket.cs

@@ -19,21 +19,22 @@ public partial class PooledSocket : IDisposable
 
         private bool _isAlive;
         private Socket _socket;
-        private EndPoint _endpoint;
+        private readonly EndPoint _endpoint;
+        private readonly int _connectionTimeout;
 
         private Stream _inputStream;
         private AsyncSocketHelper _helper;
 
         public PooledSocket(EndPoint endpoint, TimeSpan connectionTimeout, TimeSpan receiveTimeout, ILogger logger)
         {
             _logger = logger;
-
             _isAlive = true;
 
             var socket = new Socket(AddressFamily.InterNetwork, SocketType.Stream, ProtocolType.Tcp);
+            socket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.KeepAlive, true);
             socket.NoDelay = true;
 
-            var timeout = connectionTimeout == TimeSpan.MaxValue
+            _connectionTimeout = connectionTimeout == TimeSpan.MaxValue
                 ? Timeout.Infinite
                 : (int)connectionTimeout.TotalMilliseconds;
 
@@ -44,45 +45,73 @@ public PooledSocket(EndPoint endpoint, TimeSpan connectionTimeout, TimeSpan rece
             socket.ReceiveTimeout = rcv;
             socket.SendTimeout = rcv;
 
-            if (!ConnectWithTimeout(socket, endpoint, timeout))
-            {
-                throw new TimeoutException($"Could not connect to {endpoint}.");
-            }
-
             _socket = socket;
             _endpoint = endpoint;
-
-            _inputStream = new NetworkStream(socket);
         }
 
-        private bool ConnectWithTimeout(Socket socket, EndPoint endpoint, int timeout)
+        public void Connect()
         {
-            bool connected = false;
-            socket.SetSocketOption(SocketOptionLevel.Socket, SocketOptionName.KeepAlive, true);
+            bool success = false;
 
             //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);
+            cts.CancelAfter(_connectionTimeout);
             void Cancel()
             {
-                if (!socket.Connected)
+                if (!_socket.Connected)
                 {
-                    socket.Dispose();
+                    _socket.Dispose();
                 }
             }
             cts.Token.Register(Cancel);
 
-            socket.Connect(endpoint);
-            if (socket.Connected)
+            _socket.Connect(_endpoint);
+            if (_socket.Connected)
+            {
+                success = true;
+            }
+            else
+            {
+                _socket.Dispose();
+            }
+
+            if (success)
+            {
+                _inputStream = new NetworkStream(_socket);
+            }
+            else
+            {
+                throw new TimeoutException($"Could not connect to {_endpoint}.");
+            }
+        }
+
+        public async Task ConnectAsync()
+        {
+            bool success = false;
+
+            var connTask = _socket.ConnectAsync(_endpoint);
+            if (await Task.WhenAny(connTask, Task.Delay(_connectionTimeout)) == connTask)
             {
-                connected = true;
+                await connTask;
+            }
+
+            if (_socket.Connected)
+            {
+                success = true;
             }
             else
             {
-                socket.Dispose();
+                _socket.Dispose();
             }
 
-            return connected;
+            if (success)
+            {
+                _inputStream = new NetworkStream(_socket);
+            }
+            else
+            {
+                throw new TimeoutException($"Could not connect to {_endpoint}.");
+            }
         }
 
         public Action<PooledSocket> CleanupCallback { get; set; }
@@ -159,16 +188,14 @@ protected void Dispose(bool disposing)
                 try
                 {
                     if (_socket != null)
-                        try
-                        {
-                            _socket.Dispose();
-                        }
-                        catch
-                        {
-                        }
+                    {
+                        try { _socket.Dispose(); } catch { }
+                    }
 
                     if (_inputStream != null)
+                    {
                         _inputStream.Dispose();
+                    }
 
                     _inputStream = null;
                     _socket = null;
@@ -325,7 +352,7 @@ public void Write(IList<ArraySegment<byte>> buffers)
             if (_socket.Send(buffers, SocketFlags.None, out status) != total)
                 System.Diagnostics.Debugger.Break();
 #else
-            this.socket.Send(buffers, SocketFlags.None, out status);
+            _socket.Send(buffers, SocketFlags.None, out status);
 #endif
 
             if (status != SocketError.Success)