Work around potential race in PipeWriter (IIS Edition) (#11165)

This commit is contained in:
Stephen Halter 2019-06-20 17:38:19 -07:00 committed by GitHub
parent e95c3a08eb
commit 0d92fee9d0
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 22 additions and 44 deletions

View File

@ -13,25 +13,22 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
{ {
internal class OutputProducer internal class OutputProducer
{ {
// This locks access to to all of the below fields // This locks access to _completed.
private readonly object _contextLock = new object(); private readonly object _contextLock = new object();
private ValueTask<FlushResult> _flushTask;
private bool _completed = false; private bool _completed = false;
private readonly Pipe _pipe; private readonly Pipe _pipe;
// https://github.com/dotnet/corefxlab/issues/1334 // https://github.com/dotnet/corefxlab/issues/1334
// Pipelines don't support multiple awaiters on flush // https://github.com/aspnet/AspNetCore/issues/8843
// this is temporary until it does // Pipelines don't support multiple awaiters on flush. This is temporary until it does.
private TaskCompletionSource<object> _flushTcs; // _lastFlushTask field should only be get or set under the _flushLock.
private readonly object _flushLock = new object(); private readonly object _flushLock = new object();
private Action _flushCompleted; private Task _lastFlushTask = Task.CompletedTask;
public OutputProducer(Pipe pipe) public OutputProducer(Pipe pipe)
{ {
_pipe = pipe; _pipe = pipe;
_flushCompleted = OnFlushCompleted;
} }
public PipeReader Reader => _pipe.Reader; public PipeReader Reader => _pipe.Reader;
@ -90,35 +87,27 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
private Task FlushAsync(PipeWriter pipeWriter, CancellationToken cancellationToken) private Task FlushAsync(PipeWriter pipeWriter, CancellationToken cancellationToken)
{ {
var awaitable = pipeWriter.FlushAsync(cancellationToken);
if (awaitable.IsCompleted)
{
// The flush task can't fail today
return Task.CompletedTask;
}
return FlushAsyncAwaited(awaitable, cancellationToken);
}
private async Task FlushAsyncAwaited(ValueTask<FlushResult> awaitable, CancellationToken cancellationToken)
{
// https://github.com/dotnet/corefxlab/issues/1334
// Since the flush awaitable doesn't currently support multiple awaiters
// we need to use a task to track the callbacks.
// All awaiters get the same task
lock (_flushLock) lock (_flushLock)
{ {
_flushTask = awaitable; _lastFlushTask = _lastFlushTask.IsCompleted ?
if (_flushTcs == null || _flushTcs.Task.IsCompleted) FlushNowAsync(pipeWriter, cancellationToken) :
{ AwaitLastFlushAndThenFlushAsync(_lastFlushTask, pipeWriter, cancellationToken);
_flushTcs = new TaskCompletionSource<object>();
_flushTask.GetAwaiter().OnCompleted(_flushCompleted); return _lastFlushTask;
}
} }
}
private Task FlushNowAsync(PipeWriter pipeWriter, CancellationToken cancellationToken)
{
var awaitable = pipeWriter.FlushAsync(cancellationToken);
return awaitable.IsCompleted ? Task.CompletedTask : FlushNowAsyncAwaited(awaitable, cancellationToken);
}
private async Task FlushNowAsyncAwaited(ValueTask<FlushResult> awaitable, CancellationToken cancellationToken)
{
try try
{ {
await _flushTcs.Task; await awaitable;
cancellationToken.ThrowIfCancellationRequested(); cancellationToken.ThrowIfCancellationRequested();
} }
catch (OperationCanceledException ex) catch (OperationCanceledException ex)
@ -132,21 +121,10 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
} }
} }
private void OnFlushCompleted() private async Task AwaitLastFlushAndThenFlushAsync(Task lastFlushTask, PipeWriter pipeWriter, CancellationToken cancellationToken)
{ {
try await lastFlushTask;
{ await FlushNowAsync(pipeWriter, cancellationToken);
_flushTask.GetAwaiter().GetResult();
_flushTcs.TrySetResult(null);
}
catch (Exception exception)
{
_flushTcs.TrySetResult(exception);
}
finally
{
_flushTask = default;
}
} }
} }
} }