THRIFT-2081 Specified timeout should be used in TSocket.Open()

Patch: Jens Geyer

diff --git a/lib/csharp/src/Transport/TSocket.cs b/lib/csharp/src/Transport/TSocket.cs
index c05b6c2..1d50968 100644 (file)
--- a/lib/csharp/src/Transport/TSocket.cs
+++ b/lib/csharp/src/Transport/TSocket.cs
@@ -131,11 +131,78 @@ namespace Thrift.Transport
                                InitSocket();
                        }
 
-                       client.Connect(host, port);
+                       if( timeout == 0)                       // no timeout -> infinite
+                       {
+                               client.Connect(host, port);
+                       }
+                       else                        // we have a timeout -> use it
+                       {
+                ConnectHelper hlp = new ConnectHelper(client);
+                IAsyncResult asyncres = client.BeginConnect(host, port, new AsyncCallback(ConnectCallback), hlp);
+                bool bConnected = asyncres.AsyncWaitHandle.WaitOne(timeout) && client.Connected;
+                if (!bConnected)
+                {
+                    lock (hlp.Mutex)
+                    {
+                        if( hlp.CallbackDone)
+                        {
+                            asyncres.AsyncWaitHandle.Close();
+                            client.Close();
+                        }
+                        else
+                        {
+                            hlp.DoCleanup = true;
+                            client = null;
+                        }
+                    }
+                    throw new TTransportException(TTransportException.ExceptionType.TimedOut, "Connect timed out");
+                }
+                       }
+
                        inputStream = client.GetStream();
                        outputStream = client.GetStream();
                }
 
+
+        static void ConnectCallback(IAsyncResult asyncres)
+        {
+            ConnectHelper hlp = asyncres.AsyncState as ConnectHelper;
+            lock (hlp.Mutex)
+            {
+                hlp.CallbackDone = true;
+                    
+                try
+                {
+                    if( hlp.Client.Client != null)
+                        hlp.Client.EndConnect(asyncres);
+                }
+                catch (SocketException)
+                {
+                    // catch that away
+                }
+
+                if (hlp.DoCleanup) 
+                {
+                    asyncres.AsyncWaitHandle.Close();
+                    if (hlp.Client is IDisposable)
+                        ((IDisposable)hlp.Client).Dispose();
+                    hlp.Client = null;
+                }
+            }
+        }
+
+        private class ConnectHelper
+        {
+            public object Mutex = new object();
+            public bool DoCleanup = false;
+            public bool CallbackDone = false;
+            public TcpClient Client;
+            public ConnectHelper(TcpClient client)
+            {
+                Client = client;
+            }
+        }
+
                public override void Close()
                {
                        base.Close();