fix intermittent OperationCancelledException on closing websocket connection (#177)

* fix build.

* fix intermittent OperationCancelledException on closing websocket connection.

* always close connection after RunAsync has finished.

* Make sure token is not created multiple times as ObjectDisposedException could throw.

* added ConnectionId to log messages.

* do not forcefully close connection unless exception is caught.

* capture any exception on CloseAsync.

* Some corrections to OWIN/System.WebSockets implementation

Code format

Removed IsConnected property

Code formatting

Use CancellationToken.None

Fixing if

Checking for multiple disposes in AsyncWampConnection, checking for cancellation requests in WebSocketWrapperConnection

Author: bigbearzhu

src/net45/Extensions/WampSharp.Owin/Owin/OwinWebSocketWrapper.cs

@@ -13,6 +13,7 @@ internal class OwinWebSocketWrapper : IWebSocketWrapper
         private readonly Func<ArraySegment<byte>, int, bool, CancellationToken, Task> mSendAsync;
         private readonly Func<ArraySegment<byte>, CancellationToken, Task<Tuple<int, bool, int>>> mReceiveAsync;
         private readonly Func<int, string, CancellationToken, Task> mCloseAsync;
+        private WebSocketState mState = WebSocketState.Open;
 
         private const string WebSocketSendAsync = "websocket.SendAsync";
         private const string WebSocketReceiveAsync = "websocket.ReceiveAsync";
@@ -95,7 +96,29 @@ public async Task<WebSocketReceiveResult> ReceiveAsync
                 await mReceiveAsync(arraySegment, callCancelled)
                     .ConfigureAwait(false);
 
-            return new WebSocketReceiveResult(count: result.Item3, messageType: GetMessageType(result.Item1), endOfMessage: result.Item2);
+            WebSocketMessageType webSocketMessageType = GetMessageType(result.Item1);
+
+            if (webSocketMessageType == WebSocketMessageType.Close)
+            {
+                ChangeState(actionDone: WebSocketState.CloseReceived,
+                            dualAction: WebSocketState.CloseSent);
+
+                return new WebSocketReceiveResult(count: result.Item3,
+                                                  messageType: webSocketMessageType,
+                                                  endOfMessage: result.Item2,
+                                                  closeStatus: this.ClientCloseStatus,
+                                                  closeStatusDescription: WebSocketClientCloseDescription);
+            }
+
+            return new WebSocketReceiveResult(count: result.Item3, messageType: webSocketMessageType, endOfMessage: result.Item2);
+        }
+
+        public WebSocketState State
+        {
+            get
+            {
+                return mState;
+            }
         }
 
         public Task SendAsync(ArraySegment<byte> data, WebSocketMessageType messageType, bool endOfMessage, CancellationToken cancel)
@@ -105,16 +128,21 @@ public Task SendAsync(ArraySegment<byte> data, WebSocketMessageType messageType,
 
         public Task CloseAsync(WebSocketCloseStatus closeStatus, string closeDescription, CancellationToken cancel)
         {
+            ChangeState(actionDone: WebSocketState.CloseSent,
+                        dualAction: WebSocketState.CloseReceived);
+
             return mCloseAsync((int) closeStatus, closeDescription, cancel);
         }
 
-        public bool IsConnected
+        private void ChangeState(WebSocketState actionDone, WebSocketState dualAction)
         {
-            get
+            if (State == WebSocketState.Open)
             {
-                WebSocketCloseStatus? closeStatus = ClientCloseStatus;
-
-                return ((closeStatus == null) || (closeStatus == 0));
+                mState = actionDone;
+            }
+            else if (mState == dualAction)
+            {
+                mState = WebSocketState.Closed;
             }
         }
 

src/net45/Extensions/WampSharp.Owin/WampSharp.Owin.project.json

@@ -1,8 +1,11 @@
-{
+{
   "frameworks": {
     "net45": {}
   },
   "runtimes": {
     "win": {}
+  },
+  "dependencies": {
+    "Microsoft.Owin": "3.0.1"
   }
 }

src/net45/Extensions/WampSharp.WebSockets/WebSockets/WebSocketWrapperConnection.cs

@@ -5,6 +5,7 @@
 using System.Threading.Tasks;
 using WampSharp.Core.Listener;
 using WampSharp.Core.Message;
+using WampSharp.Logging;
 using WampSharp.V2.Authentication;
 using WampSharp.V2.Binding.Parsers;
 
@@ -16,15 +17,17 @@ public abstract class WebSocketWrapperConnection<TMessage> : AsyncWebSocketWampC
     {
         private readonly IWampStreamingMessageParser<TMessage> mParser;
         private readonly IWebSocketWrapper mWebSocket;
-        private readonly CancellationTokenSource mCancellationTokenSource;
+        private CancellationTokenSource mCancellationTokenSource;
         private readonly Uri mAddressUri;
+        private CancellationToken mCancellationToken;
 
         public WebSocketWrapperConnection(IWebSocketWrapper webSocket, IWampStreamingMessageParser<TMessage> parser, ICookieProvider cookieProvider, ICookieAuthenticatorFactory cookieAuthenticatorFactory) :
             base(cookieProvider, cookieAuthenticatorFactory)
         {
             mWebSocket = webSocket;
             mParser = parser;
             mCancellationTokenSource = new CancellationTokenSource();
+            mCancellationToken = mCancellationTokenSource.Token;
         }
 
         protected WebSocketWrapperConnection(IClientWebSocketWrapper clientWebSocket, Uri addressUri, string protocolName, IWampStreamingMessageParser<TMessage> parser) :
@@ -36,8 +39,9 @@ protected WebSocketWrapperConnection(IClientWebSocketWrapper clientWebSocket, Ur
 
         protected override Task SendAsync(WampMessage<object> message)
         {
+            mLogger.Debug("Attempting to send a message");
             ArraySegment<byte> messageToSend = GetMessageInBytes(message);
-            return mWebSocket.SendAsync(messageToSend, WebSocketMessageType, true, mCancellationTokenSource.Token);
+            return mWebSocket.SendAsync(messageToSend, WebSocketMessageType, true, mCancellationToken);
         }
 
         protected abstract ArraySegment<byte> GetMessageInBytes(WampMessage<object> message);
@@ -48,12 +52,12 @@ protected async void Connect()
         {
             try
             {
-                await this.ClientWebSocket.ConnectAsync(mAddressUri, mCancellationTokenSource.Token)
+                await this.ClientWebSocket.ConnectAsync(mAddressUri, mCancellationToken)
                           .ConfigureAwait(false);
 
                 RaiseConnectionOpen();
 
-                Task task = Task.Run((Func<Task>) this.RunAsync, mCancellationTokenSource.Token);
+                Task task = Task.Run((Func<Task>) this.RunAsync, mCancellationToken);
             }
             catch (Exception ex)
             {
@@ -87,36 +91,13 @@ data is very small.
                 MemoryStream memoryStream = new MemoryStream();
 
                 // Checks WebSocket state.
-                while (mWebSocket.IsConnected)
+                while (IsConnected && !mCancellationToken.IsCancellationRequested)
                 {
                     // Reads data.
-                    WebSocketReceiveResult webSocketReceiveResult;
+                    WebSocketReceiveResult webSocketReceiveResult =
+                        await ReadMessage(receivedDataBuffer, memoryStream);
 
-                    long length = 0;
-                    do
-                    {
-                        webSocketReceiveResult =
-                            await mWebSocket.ReceiveAsync(receivedDataBuffer, mCancellationTokenSource.Token)
-                            .ConfigureAwait(false);
-
-                        length += webSocketReceiveResult.Count;
-
-                        await memoryStream.WriteAsync(receivedDataBuffer.Array, receivedDataBuffer.Offset,
-                                                      webSocketReceiveResult.Count, mCancellationTokenSource.Token)
-                                          .ConfigureAwait(false);
-
-                    } while (!webSocketReceiveResult.EndOfMessage);
-
-                    // If input frame is cancelation frame, send close command.
-                    if (webSocketReceiveResult.MessageType == WebSocketMessageType.Close)
-                    {
-                        this.RaiseConnectionClosed();
-
-                        await mWebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure,
-                                                    String.Empty, mCancellationTokenSource.Token)
-                                        .ConfigureAwait(false);
-                    }
-                    else
+                    if (webSocketReceiveResult.MessageType != WebSocketMessageType.Close)
                     {
                         memoryStream.Position = 0;
                         OnNewMessage(memoryStream);
@@ -128,9 +109,62 @@ await mWebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure,
             }
             catch (Exception ex)
             {
-                RaiseConnectionError(ex);
-                RaiseConnectionClosed();
+                // Cancellation token could be cancelled in Dispose if a 
+                // Goodbye message has been received.
+                if (!(ex is OperationCanceledException) ||
+                    !mCancellationToken.IsCancellationRequested)
+                {
+                    RaiseConnectionError(ex);
+                }
             }
+
+            if (mWebSocket.State != WebSocketState.CloseReceived &&
+                mWebSocket.State != WebSocketState.Closed)
+            {
+                await CloseWebSocket().ConfigureAwait(false);
+            }
+
+            RaiseConnectionClosed();
+        }
+
+        private async Task CloseWebSocket()
+        {
+            try
+            {
+                await mWebSocket.CloseAsync(WebSocketCloseStatus.NormalClosure,
+                                            String.Empty,
+                                            CancellationToken.None)
+                                .ConfigureAwait(false);
+            }
+            catch (Exception ex)
+            {
+                mLogger.WarnException("Failed sending a close message to client", ex);
+            }
+        }
+
+        private async Task<WebSocketReceiveResult> ReadMessage(ArraySegment<byte> receivedDataBuffer, MemoryStream memoryStream)
+        {
+            WebSocketReceiveResult webSocketReceiveResult;
+
+            long length = 0;
+
+            do
+            {
+                webSocketReceiveResult =
+                    await mWebSocket.ReceiveAsync(receivedDataBuffer, mCancellationToken)
+                                    .ConfigureAwait(false);
+
+                length += webSocketReceiveResult.Count;
+
+                await memoryStream.WriteAsync(receivedDataBuffer.Array,
+                                              receivedDataBuffer.Offset,
+                                              webSocketReceiveResult.Count,
+                                              mCancellationToken)
+                                  .ConfigureAwait(false);
+            }
+            while (!webSocketReceiveResult.EndOfMessage);
+
+            return webSocketReceiveResult;
         }
 
         private void OnNewMessage(MemoryStream payloadData)
@@ -142,14 +176,16 @@ private void OnNewMessage(MemoryStream payloadData)
         protected override void Dispose()
         {
             mCancellationTokenSource.Cancel();
+            mCancellationTokenSource.Dispose();
+            mCancellationTokenSource = null;
         }
 
         protected override bool IsConnected
         {
             get
             {
-                return mWebSocket.IsConnected;
+                return mWebSocket.State == WebSocketState.Open;
             }
         }
     }
-}
+}

src/net45/Extensions/WampSharp.WebSockets/WebSockets/Wrapper/IWebSocketWrapper.cs

@@ -7,7 +7,7 @@ namespace WampSharp.WebSockets
 {
     public interface IWebSocketWrapper
     {
-        bool IsConnected { get; }
+        WebSocketState State { get; }
         Task<WebSocketReceiveResult> ReceiveAsync(ArraySegment<byte> arraySegment, CancellationToken callCancelled);
         Task SendAsync(ArraySegment<byte> data, WebSocketMessageType messageType, bool endOfMessage, CancellationToken cancel);
         Task CloseAsync(WebSocketCloseStatus closeStatus, string closeDescription, CancellationToken cancel);

src/net45/Extensions/WampSharp.WebSockets/WebSockets/Wrapper/WebSocketWrapper.cs

@@ -14,14 +14,6 @@ public WebSocketWrapper(WebSocket webSocket)
             mWebSocket = webSocket;
         }
 
-        public bool IsConnected
-        {
-            get
-            {
-                return mWebSocket.State == WebSocketState.Open;
-            }
-        }
-
         public Task<WebSocketReceiveResult> ReceiveAsync(ArraySegment<byte> arraySegment, CancellationToken callCancelled)
         {
             return mWebSocket.ReceiveAsync(arraySegment, callCancelled);
@@ -36,5 +28,10 @@ public Task CloseAsync(WebSocketCloseStatus closeStatus, string closeDescription
         {
             return mWebSocket.CloseAsync(closeStatus, closeDescription, cancel);
         }
+
+        public WebSocketState State
+        {
+            get { return mWebSocket.State; }
+        }
     }
 }

src/net45/WampSharp/Core/Listener/Connections/AsyncConnection/AsyncWampConnection.cs

@@ -12,10 +12,11 @@ public abstract class AsyncWampConnection<TMessage> : IWampConnection<TMessage>,
     {
         private readonly ActionBlock<WampMessage<object>> mSendBlock;
         protected readonly ILog mLogger;
+        private int mDisposeCalled = 0;
 
         protected AsyncWampConnection()
         {
-            mLogger = LogProvider.GetLogger(this.GetType());
+            mLogger = new LoggerWithConnectionId(LogProvider.GetLogger(this.GetType()));
             mSendBlock = new ActionBlock<WampMessage<object>>(x => InnerSend(x));
         }
 
@@ -112,6 +113,7 @@ protected virtual void RaiseMessageArrived(WampMessage<TMessage> message)
 
         protected virtual void RaiseConnectionClosed()
         {
+            mLogger.Debug("Connection has been closed");
             var handler = ConnectionClosed;
             if (handler != null) handler(this, EventArgs.Empty);
         }
@@ -125,9 +127,12 @@ protected virtual void RaiseConnectionError(Exception ex)
 
         void IDisposable.Dispose()
         {
-            mSendBlock.Complete();
-            mSendBlock.Completion.Wait();
-            this.Dispose();
+            if (Interlocked.CompareExchange(ref mDisposeCalled, 1, 0) == 0)
+            {
+                mSendBlock.Complete();
+                mSendBlock.Completion.Wait();
+                this.Dispose();
+            }
         }
 
         protected abstract void Dispose();
@@ -136,9 +141,12 @@ void IDisposable.Dispose()
 
         async Task IAsyncDisposable.DisposeAsync()
         {
-            mSendBlock.Complete();
-            await mSendBlock.Completion;
-            this.Dispose();
+            if (Interlocked.CompareExchange(ref mDisposeCalled, 1, 0) == 0)
+            {
+                mSendBlock.Complete();
+                await mSendBlock.Completion;
+                this.Dispose();
+            }
         }
 
 #else
@@ -151,5 +159,26 @@ Task IAsyncDisposable.DisposeAsync()
 
 #endif
 
+        // TODO: move this to another file (after making it more generic)
+        // TODO: or get rid of this.
+        private class LoggerWithConnectionId : ILog
+        {
+            private readonly ILog mLogger;
+            private readonly string mConnectionId;
+
+            public LoggerWithConnectionId(ILog logger)
+            {
+                mConnectionId = Guid.NewGuid().ToString();
+                mLogger = logger;
+            }
+
+            public bool Log(LogLevel logLevel, Func<string> messageFunc, Exception exception = null, params object[] formatParameters)
+            {
+                using (LogProvider.OpenMappedContext("ConncetionId", mConnectionId))
+                {
+                    return mLogger.Log(logLevel, messageFunc, exception, formatParameters);
+                }
+            }
+        }
     }
 }