Change the behavior of ReceiveAsync cancellation to throw OperationCa…

…nceledException (#17444) * fix the behavior of ReceiveAsync cancellation throws WebSocketException
dotnet · Mar 25, 2017 · d6b1125 · d6b1125
1 parent 37d1a01
commit d6b1125
Show file tree

Hide file tree

Showing 6 changed files with 47 additions and 27 deletions.
diff --git a/src/Common/src/System/Net/WebSockets/ManagedWebSocket.cs b/src/Common/src/System/Net/WebSockets/ManagedWebSocket.cs
@@ -732,9 +732,11 @@ private async Task<WebSocketReceiveResult> ReceiveAsyncPrivate(ArraySegment<byte
             }
             catch (Exception exc)
             {
-                throw _state == WebSocketState.Aborted ?
-                    new WebSocketException(WebSocketError.InvalidState, SR.Format(SR.net_WebSockets_InvalidState_ClosedOrAborted, "System.Net.WebSockets.InternalClientWebSocket", "Aborted"), exc) :
-                    new WebSocketException(WebSocketError.ConnectionClosedPrematurely, exc);
+                if (_state == WebSocketState.Aborted)
+                {
+                    throw new OperationCanceledException(nameof(WebSocketState.Aborted), exc);
+                }
+                throw new WebSocketException(WebSocketError.ConnectionClosedPrematurely, exc);
             }
             finally
             {

diff --git a/src/System.Net.WebSockets.Client/src/System/Net/WebSockets/WinHttpWebSocket.cs b/src/System.Net.WebSockets.Client/src/System/Net/WebSockets/WinHttpWebSocket.cs
@@ -795,14 +795,7 @@ private void CancelAllOperations()
 
             if (_operation.TcsReceive != null)
             {
-                var exception = new WebSocketException(
-                    WebSocketError.InvalidState,
-                    SR.Format(
-                        SR.net_WebSockets_InvalidState_ClosedOrAborted,
-                        "System.Net.WebSockets.InternalClientWebSocket",
-                        "Aborted"));
-
-                _operation.TcsReceive.TrySetException(exception);
+                _operation.TcsReceive.TrySetCanceled();
             }
 
             if (_operation.TcsSend != null)

diff --git a/src/System.Net.WebSockets.Client/src/System/Net/WebSockets/WinHttpWebSocketCallback.cs b/src/System.Net.WebSockets.Client/src/System/Net/WebSockets/WinHttpWebSocketCallback.cs
@@ -315,24 +315,16 @@ private static void OnWebSocketError(
 
             if (asyncResult.AsyncResult.dwError == Interop.WinHttp.ERROR_WINHTTP_OPERATION_CANCELLED)
             {
-                var exception = new WebSocketException(
-                    WebSocketError.InvalidState,
-                    SR.Format(
-                        SR.net_WebSockets_InvalidState_ClosedOrAborted,
-                        "System.Net.WebSockets.InternalClientWebSocket",
-                        "Aborted"),
-                    innerException);
-
                 state.UpdateState(WebSocketState.Aborted);
 
                 if (state.TcsReceive != null)
                 {
-                    state.TcsReceive.TrySetException(exception);
+                    state.TcsReceive.TrySetCanceled();
                 }
 
                 if (state.TcsSend != null)
                 {
-                    state.TcsSend.TrySetException(exception);
+                    state.TcsSend.TrySetCanceled();
                 }
 
                 return;

diff --git a/src/System.Net.WebSockets.Client/tests/CancelTest.cs b/src/System.Net.WebSockets.Client/tests/CancelTest.cs
@@ -117,23 +117,52 @@ await cws.SendAsync(
                 await cws.CloseOutputAsync(WebSocketCloseStatus.NormalClosure, "CancelShutdown", cts.Token);
             }, server);
         }
+
+        [OuterLoop] // TODO: Issue #11345
+        [ConditionalTheory(nameof(WebSocketsSupported)), MemberData(nameof(EchoServers))]
+        public async Task ReceiveAsync_CancelThenReceive_ThrowsOperationCanceledException(Uri server)
+        {
+            using (ClientWebSocket cws = await WebSocketHelper.GetConnectedWebSocket(server, TimeOutMilliseconds, _output))
+            {
+                var recvBuffer = new byte[100];
+                var segment = new ArraySegment<byte>(recvBuffer);
+                var cts = new CancellationTokenSource();
 
+                cts.Cancel();
+                Task receive = cws.ReceiveAsync(segment, cts.Token);
+                await Assert.ThrowsAnyAsync<OperationCanceledException>(() => receive);
+            }
+        }
+
         [OuterLoop] // TODO: Issue #11345
         [ConditionalTheory(nameof(WebSocketsSupported)), MemberData(nameof(EchoServers))]
-        public async Task ReceiveAsync_CancelAndReceive_ThrowsWebSocketExceptionWithMessage(Uri server)
+        public async Task ReceiveAsync_ReceiveThenCancel_ThrowsOperationCanceledException(Uri server)
         {
             using (ClientWebSocket cws = await WebSocketHelper.GetConnectedWebSocket(server, TimeOutMilliseconds, _output))
             {
                 var recvBuffer = new byte[100];
                 var segment = new ArraySegment<byte>(recvBuffer);
+                var cts = new CancellationTokenSource();
 
+                Task receive = cws.ReceiveAsync(segment, cts.Token);
+                cts.Cancel();
+                await Assert.ThrowsAnyAsync<OperationCanceledException>(() => receive);
+            }
+        }
+
+        [OuterLoop] // TODO: Issue #11345
+        [ConditionalTheory(nameof(WebSocketsSupported)), MemberData(nameof(EchoServers))]
+        public async Task ReceiveAsync_AfterCancellationDoReceiveAsync_ThrowsWebSocketException(Uri server)
+        {
+            using (ClientWebSocket cws = await WebSocketHelper.GetConnectedWebSocket(server, TimeOutMilliseconds, _output))
+            {
+                var recvBuffer = new byte[100];
+                var segment = new ArraySegment<byte>(recvBuffer);
                 var cts = new CancellationTokenSource();
-                // OperationCancelledException is thrown only if the token is canceled before calling ReceiveAsync
-                // Once it returns (with a Task<>), any cancellation that occurs is treated as a WebSocketException
+
                 Task recieve = cws.ReceiveAsync(segment, cts.Token);
                 cts.Cancel();
-                WebSocketException wse = await Assert.ThrowsAnyAsync<WebSocketException>(() => recieve);
-
+
                 WebSocketException ex = await Assert.ThrowsAsync<WebSocketException>(() =>
                     cws.ReceiveAsync(segment, CancellationToken.None));
                 Assert.Equal(

diff --git a/src/System.Net.WebSockets.Client/tests/CloseTest.cs b/src/System.Net.WebSockets.Client/tests/CloseTest.cs
@@ -314,7 +314,7 @@ public async Task CloseAsync_DuringConcurrentReceiveAsync_ExpectedStates(Uri ser
                     await t;
                     Assert.Equal(WebSocketState.Closed, cws.State);
                 }
-                catch (WebSocketException)
+                catch (OperationCanceledException)
                 {
                     Assert.Equal(WebSocketState.Aborted, cws.State);
                 }

diff --git a/src/System.Net.WebSockets.Client/tests/SendReceiveTest.cs b/src/System.Net.WebSockets.Client/tests/SendReceiveTest.cs
@@ -191,6 +191,10 @@ await cws.SendAsync(
                                 (errCode == WebSocketError.InvalidState) || (errCode == WebSocketError.Success),
                                 "WebSocketErrorCode");
                         }
+                        else if (ex is OperationCanceledException)
+                        {
+                            Assert.Equal(WebSocketState.Aborted, cws.State);
+                        }
                         else
                         {
                             Assert.True(false, "Unexpected exception: " + ex.Message);