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
{
// This locks access to to all of the below fields
// This locks access to _completed.
private readonly object _contextLock = new object();
private ValueTask<FlushResult> _flushTask;
private bool _completed = false;
private readonly Pipe _pipe;
// https://github.com/dotnet/corefxlab/issues/1334
// Pipelines don't support multiple awaiters on flush
// this is temporary until it does
private TaskCompletionSource<object> _flushTcs;
// https://github.com/aspnet/AspNetCore/issues/8843
// Pipelines don't support multiple awaiters on flush. This is temporary until it does.
// _lastFlushTask field should only be get or set under the _flushLock.
private readonly object _flushLock = new object();
private Action _flushCompleted;
private Task _lastFlushTask = Task.CompletedTask;
public OutputProducer(Pipe pipe)
{
_pipe = pipe;
_flushCompleted = OnFlushCompleted;
}
public PipeReader Reader => _pipe.Reader;
@ -90,35 +87,27 @@ namespace Microsoft.AspNetCore.Server.IIS.Core
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)
{
_flushTask = awaitable;
if (_flushTcs == null || _flushTcs.Task.IsCompleted)
{
_flushTcs = new TaskCompletionSource<object>();
_lastFlushTask = _lastFlushTask.IsCompleted ?
FlushNowAsync(pipeWriter, cancellationToken) :
AwaitLastFlushAndThenFlushAsync(_lastFlushTask, pipeWriter, cancellationToken);
_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
{
await _flushTcs.Task;
await awaitable;
cancellationToken.ThrowIfCancellationRequested();
}
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
{
_flushTask.GetAwaiter().GetResult();
_flushTcs.TrySetResult(null);
}
catch (Exception exception)
{
_flushTcs.TrySetResult(exception);
}
finally
{
_flushTask = default;
}
await lastFlushTask;
await FlushNowAsync(pipeWriter, cancellationToken);
}
}
}