Fixing #155

Author: Elad Zelingher

src/net45/WampSharp/WAMP2/V2/Api/CalleeProxy/Callbacks/Async/AsyncOperationCallback.cs

@@ -85,7 +85,7 @@ public void Error<TMessage>(IWampFormatter<TMessage> formatter, TMessage details
 
         public void SetException(Exception exception)
         {
-            mTask.SetException(exception);
+            mTask.TrySetException(exception);
         }
     }
 }

src/net45/WampSharp/WAMP2/V2/Api/CalleeProxy/ClientInvocationHandler.cs

@@ -21,7 +21,7 @@ internal class ClientInvocationHandler : WampCalleeProxyInvocationHandler
         private readonly IWampRpcOperationCatalogProxy mCatalogProxy;
         private readonly IWampClientConnectionMonitor mMonitor;
 
-        private TaskCompletionSource<object> mDisconnectionTaskCompletionSource;
+        private TaskCompletionSource<Exception> mDisconnectionTaskCompletionSource;
 
         private ManualResetEvent mDisconnectionWaitHandle;
 
@@ -37,10 +37,10 @@ public ClientInvocationHandler(IWampRpcOperationCatalogProxy catalogProxy,
         {
             mCatalogProxy = catalogProxy;
             mMonitor = monitor;
-            mDisconnectionTaskCompletionSource = new TaskCompletionSource<object>();
+
+            mDisconnectionTaskCompletionSource = new TaskCompletionSource<Exception>();
             mDisconnectionWaitHandle = new ManualResetEvent(false);
 
-            mMonitor.ConnectionEstablished += OnConnectionEstablished;
             mMonitor.ConnectionError += OnConnectionError;
             mMonitor.ConnectionBroken += OnConnectionBroken;
         }
@@ -54,16 +54,14 @@ public ClientInvocationHandler(IWampRealmProxy realmProxy) :
 
         #region Private Methods
 
-        private void OnConnectionEstablished(object sender, WampSessionCreatedEventArgs e)
-        {
-            mDisconnectionTaskCompletionSource = new TaskCompletionSource<object>();
-            mDisconnectionWaitHandle = new ManualResetEvent(false);
-        }
 
         private void OnConnectionBroken(object sender, WampSessionCloseEventArgs e)
         {
             Exception exception = new WampConnectionBrokenException(e);
             SetException(exception);
+
+            mDisconnectionTaskCompletionSource = new TaskCompletionSource<Exception>();
+            mDisconnectionWaitHandle = new ManualResetEvent(false);
         }
 
         private void OnConnectionError(object sender, WampConnectionErrorEventArgs e)
@@ -75,7 +73,7 @@ private void OnConnectionError(object sender, WampConnectionErrorEventArgs e)
         private void SetException(Exception exception)
         {
             mDisconnectionException = exception;
-            mDisconnectionTaskCompletionSource.TrySetException(exception);
+            mDisconnectionTaskCompletionSource.TrySetResult(exception);
             mDisconnectionWaitHandle.Set();
         }
 
@@ -92,17 +90,29 @@ protected override Task<T> AwaitForResult<T>(AsyncOperationCallback<T> asyncOper
 #if ASYNC
             Task<T> operationTask = asyncOperationCallback.Task;
 
+            Task<Exception> disconnectionTask = mDisconnectionTaskCompletionSource.Task;
+
             Task task = await Task.WhenAny(operationTask,
-                                           mDisconnectionTaskCompletionSource.Task)
+                                           disconnectionTask)
                                   .ConfigureAwait(false);
 
+            if (!operationTask.IsCompleted)
+            {
+                Exception exception = await disconnectionTask.ConfigureAwait(false);
+
+                asyncOperationCallback.SetException(exception);
+            }
+
             T result = await operationTask.ConfigureAwait(false);
 
             return result;
+
 #else
             IObservable<T> merged =
-                Observable.Amb(asyncOperationCallback.Task.ToObservable(),
-                                    mDisconnectionTaskCompletionSource.Task.ToObservable().Cast<T>());
+                Observable.Amb
+                (asyncOperationCallback.Task.ToObservable(),
+                 mDisconnectionTaskCompletionSource.Task.ToObservable()
+                                                   .SelectMany(x => Observable.Throw<T>(x)));
 
             Task<T> task = merged.ToTask();
 
@@ -114,7 +124,7 @@ protected override Task<T> AwaitForResult<T>(AsyncOperationCallback<T> asyncOper
         protected override void WaitForResult<T>(SyncCallback<T> callback)
         {
             int signaledIndex =
-                WaitHandle.WaitAny(new[] {mDisconnectionWaitHandle, callback.WaitHandle},
+                WaitHandle.WaitAny(new[] { mDisconnectionWaitHandle, callback.WaitHandle },
                                    Timeout.Infinite);
 
 
@@ -137,6 +147,6 @@ protected override void Invoke(ICalleeProxyInterceptor interceptor, IWampRawRpcO
                                  arguments);
         }
 
-        #endregion
+#endregion
     }
 }

src/net45/WampSharp/WAMP2/V2/Client/Session/WampSessionClient.cs

@@ -194,6 +194,7 @@ public void OnConnectionClosed()
             }
 
             mConnectionBrokenRaised = false;
+            mOpenTask = new TaskCompletionSource<bool>();
         }
 
         public void OnConnectionError(Exception exception)
@@ -207,13 +208,7 @@ public void OnConnectionError(Exception exception)
 
         private void SetOpenTaskErrorIfNeeded(Exception exception)
         {
-            TaskCompletionSource<bool> openTask;
-
-            lock (mLock)
-            {
-                openTask = mOpenTask;
-                mOpenTask = new TaskCompletionSource<bool>();
-            }
+            TaskCompletionSource<bool> openTask = mOpenTask;
 
             if (openTask != null)
             {