Skip to content

Commit

Permalink
Fix Abort Lock (dotnet#55738)
Browse files Browse the repository at this point in the history
* Fix Abort lock

* better

* some updates

* comment + rename

---------

Co-authored-by: Brennan <brecon@microsoft.com>
  • Loading branch information
wtgodbe and BrennanConroy authored May 15, 2024
1 parent 8198eeb commit 6c95ad9
Show file tree
Hide file tree
Showing 2 changed files with 62 additions and 20 deletions.
47 changes: 38 additions & 9 deletions src/Servers/Kestrel/Core/src/Internal/Http2/Http2FrameWriter.cs
Original file line number Diff line number Diff line change
Expand Up @@ -379,19 +379,36 @@ public void UpdateMaxFrameSize(uint maxFrameSize)
}
}

/// <summary>
/// Call while in the <see cref="_writeLock"/>.
/// </summary>
/// <returns><c>true</c> if already completed.</returns>
private bool CompleteUnsynchronized()
{
if (_completed)
{
return true;
}

_completed = true;
_outputWriter.Abort();

return false;
}

public void Complete()
{
lock (_writeLock)
{
if (_completed)
if (CompleteUnsynchronized())
{
return;
}

_completed = true;
AbortConnectionFlowControl();
_outputWriter.Abort();
}

// Call outside of _writeLock as this can call Http2OutputProducer.Stop which can acquire Http2OutputProducer._dataWriterLock
// which is not the desired lock order
AbortConnectionFlowControl();
}

public Task ShutdownAsync()
Expand All @@ -413,8 +430,15 @@ public void Abort(ConnectionAbortedException error)
_aborted = true;
_connectionContext.Abort(error);

Complete();
if (CompleteUnsynchronized())
{
return;
}
}

// Call outside of _writeLock as this can call Http2OutputProducer.Stop which can acquire Http2OutputProducer._dataWriterLock
// which is not the desired lock order
AbortConnectionFlowControl();
}

private ValueTask<FlushResult> FlushEndOfStreamHeadersAsync(Http2Stream stream)
Expand Down Expand Up @@ -487,7 +511,7 @@ private void WriteResponseHeadersUnsynchronized(int streamId, int statusCode, Ht
_outgoingFrame.PrepareHeaders(headerFrameFlags, streamId);
var buffer = _headerEncodingBuffer.AsSpan();
var done = HPackHeaderWriter.BeginEncodeHeaders(statusCode, _hpackEncoder, _headersEnumerator, buffer, out var payloadLength);
FinishWritingHeaders(streamId, payloadLength, done);
FinishWritingHeadersUnsynchronized(streamId, payloadLength, done);
}
// Any exception from the HPack encoder can leave the dynamic table in a corrupt state.
// Since we allow custom header encoders we don't know what type of exceptions to expect.
Expand Down Expand Up @@ -528,7 +552,7 @@ private ValueTask<FlushResult> WriteDataAndTrailersAsync(Http2Stream stream, in
_outgoingFrame.PrepareHeaders(Http2HeadersFrameFlags.END_STREAM, streamId);
var buffer = _headerEncodingBuffer.AsSpan();
var done = HPackHeaderWriter.BeginEncodeHeaders(_hpackEncoder, _headersEnumerator, buffer, out var payloadLength);
FinishWritingHeaders(streamId, payloadLength, done);
FinishWritingHeadersUnsynchronized(streamId, payloadLength, done);
}
// Any exception from the HPack encoder can leave the dynamic table in a corrupt state.
// Since we allow custom header encoders we don't know what type of exceptions to expect.
Expand All @@ -542,7 +566,7 @@ private ValueTask<FlushResult> WriteDataAndTrailersAsync(Http2Stream stream, in
}
}

private void FinishWritingHeaders(int streamId, int payloadLength, bool done)
private void FinishWritingHeadersUnsynchronized(int streamId, int payloadLength, bool done)
{
var buffer = _headerEncodingBuffer.AsSpan();
_outgoingFrame.PayloadLength = payloadLength;
Expand Down Expand Up @@ -934,6 +958,11 @@ private void ConsumeConnectionWindow(long bytes)
}
}

/// <summary>
/// Do not call this method under the _writeLock.
/// This method can call Http2OutputProducer.Stop which can acquire Http2OutputProducer._dataWriterLock
/// which is not the desired lock order
/// </summary>
private void AbortConnectionFlowControl()
{
lock (_windowUpdateLock)
Expand Down
35 changes: 24 additions & 11 deletions src/Servers/Kestrel/Core/src/Internal/Http2/Http2OutputProducer.cs
Original file line number Diff line number Diff line change
Expand Up @@ -590,6 +590,7 @@ public void Reset()

internal void OnRequestProcessingEnded()
{
var shouldCompleteStream = false;
lock (_dataWriterLock)
{
if (_requestProcessingComplete)
Expand All @@ -600,15 +601,24 @@ internal void OnRequestProcessingEnded()

_requestProcessingComplete = true;

if (_completedResponse)
{
Stream.CompleteStream(errored: false);
}
shouldCompleteStream = _completedResponse;
}

// Complete outside of lock, anything this method does that needs a lock will acquire a lock itself.
// Additionally, this method should only be called once per Reset so calling outside of the lock is fine from the perspective
// of multiple threads calling OnRequestProcessingEnded.
if (shouldCompleteStream)
{
Stream.CompleteStream(errored: false);
}

}

internal ValueTask<FlushResult> CompleteResponseAsync()
{
var shouldCompleteStream = false;
ValueTask<FlushResult> task = default;

lock (_dataWriterLock)
{
if (_completedResponse)
Expand All @@ -619,22 +629,25 @@ internal ValueTask<FlushResult> CompleteResponseAsync()

_completedResponse = true;

ValueTask<FlushResult> task = default;

if (_resetErrorCode is { } error)
{
// If we have an error code to write, write it now that we're done with the response.
// Always send the reset even if the response body is completed. The request body may not have completed yet.
task = _frameWriter.WriteRstStreamAsync(StreamId, error);
}

if (_requestProcessingComplete)
{
Stream.CompleteStream(errored: false);
}
shouldCompleteStream = _requestProcessingComplete;
}

return task;
// Complete outside of lock, anything this method does that needs a lock will acquire a lock itself.
// CompleteResponseAsync also should never be called in parallel so calling this outside of the lock doesn't
// cause any weirdness with parallel threads calling this method and no longer waiting on the stream completion call.
if (shouldCompleteStream)
{
Stream.CompleteStream(errored: false);
}

return task;
}

internal Memory<byte> GetFakeMemory(int minSize)
Expand Down

0 comments on commit 6c95ad9

Please sign in to comment.