From 39f693b9ed9b6869c57af291b4ecaaf15b889e3b Mon Sep 17 00:00:00 2001
From: David Fowler <davidfowl@gmail.com>
Date: Thu, 5 Apr 2018 12:47:39 -0700
Subject: [PATCH] Found these while looking at a flaky test (#1871)

- Fix crashing test
- Dispose the ClientWebsocket in StopAsync
---
 .../Internal/WebSocketsTransport.cs           |  4 ++++
 .../HubConnectionTests.cs                     | 21 +++++++++++++------
 2 files changed, 19 insertions(+), 6 deletions(-)

diff --git a/src/Microsoft.AspNetCore.Http.Connections.Client/Internal/WebSocketsTransport.cs b/src/Microsoft.AspNetCore.Http.Connections.Client/Internal/WebSocketsTransport.cs
index eeafd86822..db06b7dae9 100644
--- a/src/Microsoft.AspNetCore.Http.Connections.Client/Internal/WebSocketsTransport.cs
+++ b/src/Microsoft.AspNetCore.Http.Connections.Client/Internal/WebSocketsTransport.cs
@@ -356,6 +356,10 @@ namespace Microsoft.AspNetCore.Http.Connections.Client.Internal
                 // exceptions have been handled in the Running task continuation by closing the channel with the exception
                 return;
             }
+            finally
+            {
+                _webSocket.Dispose();
+            }
 
             Log.TransportStopped(_logger, null);
         }
diff --git a/test/Microsoft.AspNetCore.SignalR.Client.FunctionalTests/HubConnectionTests.cs b/test/Microsoft.AspNetCore.SignalR.Client.FunctionalTests/HubConnectionTests.cs
index dd31e19714..5278eb42ec 100644
--- a/test/Microsoft.AspNetCore.SignalR.Client.FunctionalTests/HubConnectionTests.cs
+++ b/test/Microsoft.AspNetCore.SignalR.Client.FunctionalTests/HubConnectionTests.cs
@@ -169,13 +169,22 @@ namespace Microsoft.AspNetCore.SignalR.Client.FunctionalTests
                 var restartTcs = new TaskCompletionSource<object>();
                 connection.Closed += async e =>
                 {
-                    logger.LogInformation("Closed event triggered");
-                    if (!restartTcs.Task.IsCompleted)
+                    try
                     {
-                        logger.LogInformation("Restarting connection");
-                        await connection.StartAsync().OrTimeout();
-                        logger.LogInformation("Restarted connection");
-                        restartTcs.SetResult(null);
+                        logger.LogInformation("Closed event triggered");
+                        if (!restartTcs.Task.IsCompleted)
+                        {
+                            logger.LogInformation("Restarting connection");
+                            await connection.StartAsync().OrTimeout();
+                            logger.LogInformation("Restarted connection");
+                            restartTcs.SetResult(null);
+                        }
+                    }
+                    catch (Exception ex)
+                    {
+                        // It's important to try catch here since this happens
+                        // on a thread pool thread
+                        restartTcs.TrySetException(ex);
                     }
                 };