Fix race with CTS disposing

Author: BrennanConroy

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

@@ -249,7 +249,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
@@ -439,6 +439,37 @@ 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
+                cts = Status == HttpConnectionStatus.Disposed ? null : Cancellation;
+            }
+
+            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

@@ -155,7 +155,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 WebSocketsTransport(options.WebSockets, connection.Application, connection, _loggerFactory);
@@ -180,28 +180,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))