From e40b21803941be1deb1e98474399c48b1152d04e Mon Sep 17 00:00:00 2001
From: Brennan <brecon@microsoft.com>
Date: Thu, 24 Oct 2019 20:54:33 -0700
Subject: [PATCH] Fix race with CTS disposing (#11757)

---
 .../src/Internal/HttpConnectionContext.cs     | 41 ++++++++++++++++++-
 .../src/Internal/HttpConnectionDispatcher.cs  | 27 ++++--------
 .../test/HttpConnectionDispatcherTests.cs     | 16 +++++++-
 3 files changed, 61 insertions(+), 23 deletions(-)

diff --git a/src/SignalR/common/Http.Connections/src/Internal/HttpConnectionContext.cs b/src/SignalR/common/Http.Connections/src/Internal/HttpConnectionContext.cs
index 6d3fe467e9..dac620efa8 100644
--- a/src/SignalR/common/Http.Connections/src/Internal/HttpConnectionContext.cs
+++ b/src/SignalR/common/Http.Connections/src/Internal/HttpConnectionContext.cs
@@ -253,7 +253,7 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
 
                     if (TransportType == HttpTransportType.WebSockets)
                     {
-                        // The websocket transport will close the application output automatically when reading is cancelled
+                        // The websocket transport will close the application output automatically when reading is canceled
                         Cancellation?.Cancel();
                     }
                     else
@@ -443,6 +443,45 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
             }
         }
 
+        internal async Task<bool> CancelPreviousPoll(HttpContext context)
+        {
+            CancellationTokenSource cts;
+            lock (_stateLock)
+            {
+                // Need to sync cts access with DisposeAsync as that will dispose the cts
+                if (Status == HttpConnectionStatus.Disposed)
+                {
+                    cts = null;
+                }
+                else
+                {
+                    cts = Cancellation;
+                    Cancellation = null;
+                }
+            }
+
+            using (cts)
+            {
+                // Cancel the previous request
+                cts?.Cancel();
+
+                try
+                {
+                    // Wait for the previous request to drain
+                    await PreviousPollTask;
+                }
+                catch (OperationCanceledException)
+                {
+                    // Previous poll canceled due to connection closing, close this poll too
+                    context.Response.ContentType = "text/plain";
+                    context.Response.StatusCode = StatusCodes.Status204NoContent;
+                    return false;
+                }
+
+                return true;
+            }
+        }
+
         public void MarkInactive()
         {
             lock (_stateLock)
diff --git a/src/SignalR/common/Http.Connections/src/Internal/HttpConnectionDispatcher.cs b/src/SignalR/common/Http.Connections/src/Internal/HttpConnectionDispatcher.cs
index 6ed9932839..9da1ea0c18 100644
--- a/src/SignalR/common/Http.Connections/src/Internal/HttpConnectionDispatcher.cs
+++ b/src/SignalR/common/Http.Connections/src/Internal/HttpConnectionDispatcher.cs
@@ -164,7 +164,7 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
 
                 Log.EstablishedConnection(_logger);
 
-                // Allow the reads to be cancelled
+                // Allow the reads to be canceled
                 connection.Cancellation = new CancellationTokenSource();
 
                 var ws = new WebSocketsServerTransport(options.WebSockets, connection.Application, connection, _loggerFactory);
@@ -189,28 +189,15 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
                     return;
                 }
 
+                if (!await connection.CancelPreviousPoll(context))
+                {
+                    // Connection closed. It's already set the response status code.
+                    return;
+                }
+
                 // Create a new Tcs every poll to keep track of the poll finishing, so we can properly wait on previous polls
                 var currentRequestTcs = new TaskCompletionSource<object>(TaskCreationOptions.RunContinuationsAsynchronously);
 
-                using (connection.Cancellation)
-                {
-                    // Cancel the previous request
-                    connection.Cancellation?.Cancel();
-
-                    try
-                    {
-                        // Wait for the previous request to drain
-                        await connection.PreviousPollTask;
-                    }
-                    catch (OperationCanceledException)
-                    {
-                        // Previous poll canceled due to connection closing, close this poll too
-                        context.Response.ContentType = "text/plain";
-                        context.Response.StatusCode = StatusCodes.Status204NoContent;
-                        return;
-                    }
-                }
-
                 if (!connection.TryActivateLongPollingConnection(
                         connectionDelegate, context, options.LongPolling.PollTimeout,
                         currentRequestTcs.Task, _loggerFactory, _logger))
diff --git a/src/SignalR/common/Http.Connections/test/HttpConnectionDispatcherTests.cs b/src/SignalR/common/Http.Connections/test/HttpConnectionDispatcherTests.cs
index 6e28f47cc6..d543cd4f9a 100644
--- a/src/SignalR/common/Http.Connections/test/HttpConnectionDispatcherTests.cs
+++ b/src/SignalR/common/Http.Connections/test/HttpConnectionDispatcherTests.cs
@@ -1148,9 +1148,22 @@ namespace Microsoft.AspNetCore.Http.Connections.Tests
                 Assert.True(request1.IsCompleted);
 
                 request1 = dispatcher.ExecuteAsync(context1, options, app);
+                var count = 0;
+                // Wait until the request has started internally
+                while (connection.TransportTask.IsCompleted && count < 50)
+                {
+                    count++;
+                    await Task.Delay(15);
+                }
+                if (count == 50)
+                {
+                    Assert.True(false, "Poll took too long to start");
+                }
+
                 var request2 = dispatcher.ExecuteAsync(context2, options, app);
 
-                await request1;
+                // Wait for poll to be canceled
+                await request1.OrTimeout();
 
                 Assert.Equal(StatusCodes.Status204NoContent, context1.Response.StatusCode);
                 Assert.Equal(HttpConnectionStatus.Active, connection.Status);
@@ -1164,7 +1177,6 @@ namespace Microsoft.AspNetCore.Http.Connections.Tests
         }
 
         [Fact]
-        [Flaky("https://github.com/aspnet/AspNetCore-Internal/issues/2040", "All")]
         public async Task MultipleRequestsToActiveConnectionId409ForLongPolling()
         {
             using (StartVerifiableLog())