Fix race when connection is canceled and new poll comes in (#2697)

This commit is contained in:
BrennanConroy 2018-07-31 16:36:54 -07:00 committed by GitHub
parent 651b39bc90
commit c8eedb3540
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 57 additions and 55 deletions

View File

@ -88,6 +88,8 @@ namespace Microsoft.AspNetCore.Http.Connections.Internal
public Task TransportTask { get; set; } public Task TransportTask { get; set; }
public Task PreviousPollTask { get; set; } = Task.CompletedTask;
public Task ApplicationTask { get; set; } public Task ApplicationTask { get; set; }
public DateTime LastSeenUtc { get; set; } public DateTime LastSeenUtc { get; set; }

View File

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