Fix race with CTS disposing

src/SignalR/common/Http.Connections/src/Internal/HttpConnectionContext.cs

@@ -253,7 +253,7 @@ private async Task WaitOnTasks(Task applicationTask, Task transportTask, bool cl
 
                     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 @@ private void FailActivationUnsynchronized(HttpContext nonClonedContext, ILogger
             }
         }
 
+        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)

src/SignalR/common/Http.Connections/src/Internal/HttpConnectionDispatcher.cs

@@ -164,7 +164,7 @@ private async Task ExecuteAsync(HttpContext context, ConnectionDelegate connecti
 
                 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 @@ private async Task ExecuteAsync(HttpContext context, ConnectionDelegate connecti
                     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)
+                if (!await connection.CancelPreviousPoll(context))
                 {
-                    // 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;
-                    }
+                    // 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);
+
                 if (!connection.TryActivateLongPollingConnection(
                         connectionDelegate, context, options.LongPolling.PollTimeout,
                         currentRequestTcs.Task, _loggerFactory, _logger))

src/SignalR/common/Http.Connections/test/HttpConnectionDispatcherTests.cs

@@ -1148,9 +1148,22 @@ public async Task RequestToActiveConnectionIdKillsPreviousConnectionLongPolling(
                 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 @@ public async Task RequestToActiveConnectionIdKillsPreviousConnectionLongPolling(
         }
 
         [Fact]
-        [Flaky("https://github.com/aspnet/AspNetCore-Internal/issues/2040", "All")]
         public async Task MultipleRequestsToActiveConnectionId409ForLongPolling()
         {
             using (StartVerifiableLog())