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

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

@@ -88,6 +88,8 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
 
         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 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
         {
             Task disposeTask;
 
+            Cancellation?.Dispose();
+
             await StateLock.WaitAsync();
             try
             {

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

@@ -188,6 +188,9 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
                     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,6 +208,7 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
                     {
                         var existing = connection.GetHttpContext();
                         Log.ConnectionAlreadyActive(_logger, connection.ConnectionId, existing.TraceIdentifier);
+                    }
 
                     using (connection.Cancellation)
                     {
@@ -212,10 +216,9 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
                         connection.Cancellation?.Cancel();
 
                         // Wait for the previous request to drain
-                            await connection.TransportTask;
+                        await connection.PreviousPollTask;
 
-                            Log.PollCanceled(_logger, connection.ConnectionId, existing.TraceIdentifier);
-                        }
+                        connection.PreviousPollTask = currentRequestTcs.Task;
                     }
 
                     // Mark the connection as active
@@ -267,6 +270,8 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
 
                 var resultTask = await Task.WhenAny(connection.ApplicationTask, connection.TransportTask);
 
+                try
+                {
                     var pollAgain = true;
 
                     // If the application ended before the transport task then we potentially need to end the connection
@@ -296,28 +301,18 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
                     }
 
                     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)
                     {
                         // Mark the connection as inactive
                         connection.LastSeenUtc = DateTime.UtcNow;
 
                         connection.Status = HttpConnectionStatus.Inactive;
-
-                            // Dispose the cancellation token
-                            connection.Cancellation?.Dispose();
-
-                            connection.Cancellation = null;
                     }
                 }
                 finally
                 {
-                        connection.StateLock.Release();
-                    }
+                    // Artificial task queue
+                    // This will cause incoming polls to wait until the previous poll has finished updating internal state info
+                    currentRequestTcs.TrySetResult(null);
                 }
             }
         }