Fix several WebSocket tests blocking xunit's SynchronizationContext (#…

…32591) * Fix several WebSocket tests that were blocking xunit's SynchronizationContext * Address PR feedback
dotnet · Oct 3, 2018 · 7685f1c · 7685f1c
1 parent 9cb4ea0
commit 7685f1c
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 35 deletions.
diff --git a/src/Common/tests/System/Threading/Tasks/TaskTimeoutExtensions.cs b/src/Common/tests/System/Threading/Tasks/TaskTimeoutExtensions.cs
@@ -12,6 +12,19 @@ namespace System.Threading.Tasks
 {
     public static class TaskTimeoutExtensions
     {
+        public static async Task WithCancellation(this Task task, CancellationToken cancellationToken)
+        {
+            var tcs = new TaskCompletionSource<bool>();
+            using (cancellationToken.Register(s => ((TaskCompletionSource<bool>)s).TrySetResult(true), tcs))
+            {
+                if (task != await Task.WhenAny(task, tcs.Task).ConfigureAwait(false))
+                {
+                    throw new OperationCanceledException(cancellationToken);
+                }
+                await task; // already completed; propagate any exception
+            }
+        }
+
         public static async Task TimeoutAfter(this Task task, int millisecondsTimeout)
         {
             var cts = new CancellationTokenSource();

diff --git a/src/System.Net.WebSockets.Client/tests/SendReceiveTest.cs b/src/System.Net.WebSockets.Client/tests/SendReceiveTest.cs
@@ -161,7 +161,7 @@ public async Task SendAsync_MultipleOutstandingSendOperations_Throws(Uri server)
                             cts.Token);
                     }
 
-                    Task.WaitAll(tasks);
+                    await Task.WhenAll(tasks);
 
                     Assert.Equal(WebSocketState.Open, cws.State);
                 }
@@ -225,41 +225,38 @@ await SendAsync(
                         tasks[i] = ReceiveAsync(cws, recvSegment, cts.Token);
                     }
 
-                    Task.WaitAll(tasks);
+                    await Task.WhenAll(tasks);
                     Assert.Equal(WebSocketState.Open, cws.State);
                 }
-                catch (AggregateException ag)
+                catch (Exception ex)
                 {
-                    foreach (var ex in ag.InnerExceptions)
+                    if (ex is InvalidOperationException)
                     {
-                        if (ex is InvalidOperationException)
-                        {
-                            Assert.Equal(
-                                ResourceHelper.GetExceptionMessage(
-                                    "net_Websockets_AlreadyOneOutstandingOperation",
-                                    "ReceiveAsync"),
-                                ex.Message);
+                        Assert.Equal(
+                            ResourceHelper.GetExceptionMessage(
+                                "net_Websockets_AlreadyOneOutstandingOperation",
+                                "ReceiveAsync"),
+                            ex.Message);
 
-                            Assert.Equal(WebSocketState.Aborted, cws.State);
-                        }
-                        else if (ex is WebSocketException)
-                        {
-                            // Multiple cases.
-                            Assert.Equal(WebSocketState.Aborted, cws.State);
+                        Assert.Equal(WebSocketState.Aborted, cws.State);
+                    }
+                    else if (ex is WebSocketException)
+                    {
+                        // Multiple cases.
+                        Assert.Equal(WebSocketState.Aborted, cws.State);
 
-                            WebSocketError errCode = (ex as WebSocketException).WebSocketErrorCode;
-                            Assert.True(
-                                (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);
-                        }
+                        WebSocketError errCode = (ex as WebSocketException).WebSocketErrorCode;
+                        Assert.True(
+                            (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);
                     }
                 }
             }
@@ -381,7 +378,7 @@ public async Task SendReceive_ConnectionClosedPrematurely_ReceiveAsyncFailsAndWe
 
             Func<ClientWebSocket, LoopbackServer, Uri, Task> connectToServerThatAbortsConnection = async (clientSocket, server, url) =>
             {
-                AutoResetEvent pendingReceiveAsyncPosted =  new AutoResetEvent(false);
+                var pendingReceiveAsyncPosted = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
 
                 // Start listening for incoming connections on the server side.
                 Task acceptTask = server.AcceptConnectionAsync(async connection =>
@@ -390,7 +387,7 @@ public async Task SendReceive_ConnectionClosedPrematurely_ReceiveAsyncFailsAndWe
                     Assert.True(await LoopbackHelper.WebSocketHandshakeAsync(connection));
 
                     // Wait for client-side ConnectAsync to complete and for a pending ReceiveAsync to be posted.
-                    pendingReceiveAsyncPosted.WaitOne(TimeOutMilliseconds);
+                    await pendingReceiveAsyncPosted.Task.TimeoutAfter(TimeOutMilliseconds);
 
                     // Close the underlying connection prematurely (without sending a WebSocket Close frame).
                     connection.Socket.Shutdown(SocketShutdown.Both);
@@ -405,10 +402,10 @@ public async Task SendReceive_ConnectionClosedPrematurely_ReceiveAsyncFailsAndWe
                 var recvBuffer = new byte[100];
                 var recvSegment = new ArraySegment<byte>(recvBuffer);
                 Task pendingReceiveAsync = ReceiveAsync(clientSocket, recvSegment, cts.Token);
-                pendingReceiveAsyncPosted.Set();
+                pendingReceiveAsyncPosted.SetResult(true);
 
                 // Wait for the server to close the underlying connection.
-                acceptTask.Wait(cts.Token);
+                await acceptTask.WithCancellation(cts.Token);
 
                 WebSocketException pendingReceiveException = await Assert.ThrowsAsync<WebSocketException>(() => pendingReceiveAsync);
 

diff --git a/src/System.Net.WebSockets.Client/tests/SendReceiveTest.netcoreapp.cs b/src/System.Net.WebSockets.Client/tests/SendReceiveTest.netcoreapp.cs
@@ -16,7 +16,7 @@ protected override async Task<WebSocketReceiveResult> ReceiveAsync(WebSocket ws,
         {
             ValueWebSocketReceiveResult r = await ws.ReceiveAsync(
                 (Memory<byte>)arraySegment,
-                cancellationToken);
+                cancellationToken).ConfigureAwait(false);
             return new WebSocketReceiveResult(r.Count, r.MessageType, r.EndOfMessage, ws.CloseStatus, ws.CloseStatusDescription);
         }