LongPolling: Setting connection to inactive while there is still an active poll (#2769)

Author: BrennanConroy

src/Microsoft.AspNetCore.Http.Connections/Internal/HttpConnectionContext.cs

@@ -88,6 +88,8 @@ public HttpConnectionContext(string id, IDuplexPipe transport, IDuplexPipe appli
 
         public Task TransportTask { get; set; }
 
+        public Task PreviousPollTask { get; set; } = Task.CompletedTask;
+
         public Task ApplicationTask { get; set; }
 
         public DateTime LastSeenUtc { get; set; }
@@ -180,6 +182,8 @@ public async Task DisposeAsync(bool closeGracefully = false)
         {
             Task disposeTask;
 
+            Cancellation?.Dispose();
+
             await StateLock.WaitAsync();
             try
             {

src/Microsoft.AspNetCore.Http.Connections/Internal/HttpConnectionDispatcher.cs

@@ -188,6 +188,9 @@ 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);
+
                 await connection.StateLock.WaitAsync();
                 try
                 {
@@ -205,17 +208,17 @@ private async Task ExecuteAsync(HttpContext context, ConnectionDelegate connecti
                     {
                         var existing = connection.GetHttpContext();
                         Log.ConnectionAlreadyActive(_logger, connection.ConnectionId, existing.TraceIdentifier);
+                    }
 
-                        using (connection.Cancellation)
-                        {
-                            // Cancel the previous request
-                            connection.Cancellation?.Cancel();
+                    using (connection.Cancellation)
+                    {
+                        // Cancel the previous request
+                        connection.Cancellation?.Cancel();
 
-                            // Wait for the previous request to drain
-                            await connection.TransportTask;
+                        // Wait for the previous request to drain
+                        await connection.PreviousPollTask;
 
-                            Log.PollCanceled(_logger, connection.ConnectionId, existing.TraceIdentifier);
-                        }
+                        connection.PreviousPollTask = currentRequestTcs.Task;
                     }
 
                     // Mark the connection as active
@@ -267,57 +270,49 @@ private async Task ExecuteAsync(HttpContext context, ConnectionDelegate connecti
 
                 var resultTask = await Task.WhenAny(connection.ApplicationTask, connection.TransportTask);
 
-                var pollAgain = true;
-
-                // If the application ended before the transport task then we potentially need to end the connection
-                if (resultTask == connection.ApplicationTask)
+                try
                 {
-                    // Complete the transport (notifying it of the application error if there is one)
-                    connection.Transport.Output.Complete(connection.ApplicationTask.Exception);
+                    var pollAgain = true;
 
-                    // Wait for the transport to run
-                    await connection.TransportTask;
-
-                    // If the status code is a 204 it means the connection is done
-                    if (context.Response.StatusCode == StatusCodes.Status204NoContent)
+                    // If the application ended before the transport task then we potentially need to end the connection
+                    if (resultTask == connection.ApplicationTask)
                     {
-                        // We should be able to safely dispose because there's no more data being written
-                        // We don't need to wait for close here since we've already waited for both sides
-                        await _manager.DisposeAndRemoveAsync(connection, closeGracefully: false);
+                        // Complete the transport (notifying it of the application error if there is one)
+                        connection.Transport.Output.Complete(connection.ApplicationTask.Exception);
 
-                        // Don't poll again if we've removed the connection completely
-                        pollAgain = false;
-                    }
-                }
-                else if (context.Response.StatusCode == StatusCodes.Status204NoContent)
-                {
-                    // Don't poll if the transport task was canceled
-                    pollAgain = false;
-                }
+                        // Wait for the transport to run
+                        await connection.TransportTask;
 
-                if (pollAgain)
-                {
-                    // Otherwise, we update the state to inactive again and wait for the next poll
-                    await connection.StateLock.WaitAsync();
-                    try
-                    {
-                        if (connection.Status == HttpConnectionStatus.Active)
+                        // If the status code is a 204 it means the connection is done
+                        if (context.Response.StatusCode == StatusCodes.Status204NoContent)
                         {
-                            // Mark the connection as inactive
-                            connection.LastSeenUtc = DateTime.UtcNow;
-
-                            connection.Status = HttpConnectionStatus.Inactive;
+                            // We should be able to safely dispose because there's no more data being written
+                            // We don't need to wait for close here since we've already waited for both sides
+                            await _manager.DisposeAndRemoveAsync(connection, closeGracefully: false);
 
-                            // Dispose the cancellation token
-                            connection.Cancellation?.Dispose();
-
-                            connection.Cancellation = null;
+                            // Don't poll again if we've removed the connection completely
+                            pollAgain = false;
                         }
                     }
-                    finally
+                    else if (context.Response.StatusCode == StatusCodes.Status204NoContent)
                     {
-                        connection.StateLock.Release();
+                        // Don't poll if the transport task was canceled
+                        pollAgain = false;
                     }
+
+                    if (pollAgain)
+                    {
+                        // Mark the connection as inactive
+                        connection.LastSeenUtc = DateTime.UtcNow;
+
+                        connection.Status = HttpConnectionStatus.Inactive;
+                    }
+                }
+                finally
+                {
+                    // Artificial task queue
+                    // This will cause incoming polls to wait until the previous poll has finished updating internal state info
+                    currentRequestTcs.TrySetResult(null);
                 }
             }
         }