Fix Abort Lock (#55738)

* Fix Abort lock

* better

* some updates

* comment + rename

---------

Co-authored-by: Brennan <[email protected]>

Author: wtgodbe

src/Servers/Kestrel/Core/src/Internal/Http2/Http2FrameWriter.cs

@@ -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()
@@ -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)
@@ -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.
@@ -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.
@@ -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;
@@ -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)

src/Servers/Kestrel/Core/src/Internal/Http2/Http2OutputProducer.cs

@@ -590,6 +590,7 @@ public void Reset()
 
     internal void OnRequestProcessingEnded()
     {
+        var shouldCompleteStream = false;
         lock (_dataWriterLock)
         {
             if (_requestProcessingComplete)
@@ -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)
@@ -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)