HTTP2: Write data and trailers under one lock (#24822)

Author: JamesNK

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

@@ -198,7 +198,7 @@ public void WriteResponseHeaders(int streamId, int statusCode, Http2HeadersFrame
             }
         }
 
-        public ValueTask<FlushResult> WriteResponseTrailers(int streamId, HttpResponseTrailers headers)
+        public ValueTask<FlushResult> WriteResponseTrailersAsync(int streamId, HttpResponseTrailers headers)
         {
             lock (_writeLock)
             {
@@ -256,6 +256,9 @@ private void FinishWritingHeaders(int streamId, int payloadLength, bool done)
 
         public ValueTask<FlushResult> WriteDataAsync(int streamId, StreamOutputFlowControl flowControl, in ReadOnlySequence<byte> data, bool endStream, bool firstWrite, bool forceFlush)
         {
+            // Logic in this method is replicated in WriteDataAndTrailersAsync.
+            // Changes here may need to be mirrored in WriteDataAndTrailersAsync.
+
             // The Length property of a ReadOnlySequence can be expensive, so we cache the value.
             var dataLength = data.Length;
 
@@ -286,6 +289,43 @@ public ValueTask<FlushResult> WriteDataAsync(int streamId, StreamOutputFlowContr
             }
         }
 
+        public ValueTask<FlushResult> WriteDataAndTrailersAsync(int streamId, StreamOutputFlowControl flowControl, in ReadOnlySequence<byte> data, bool firstWrite, HttpResponseTrailers headers)
+        {
+            // This method combines WriteDataAsync and WriteResponseTrailers.
+            // Changes here may need to be mirrored in WriteDataAsync.
+
+            // The Length property of a ReadOnlySequence can be expensive, so we cache the value.
+            var dataLength = data.Length;
+
+            lock (_writeLock)
+            {
+                if (_completed || flowControl.IsAborted)
+                {
+                    return default;
+                }
+
+                // Zero-length data frames are allowed to be sent immediately even if there is no space available in the flow control window.
+                // https://httpwg.org/specs/rfc7540.html#rfc.section.6.9.1
+                if (dataLength != 0 && dataLength > flowControl.Available)
+                {
+                    return WriteDataAndTrailersAsyncCore(this, streamId, flowControl, data, dataLength, firstWrite, headers);
+                }
+
+                // This cast is safe since if dataLength would overflow an int, it's guaranteed to be greater than the available flow control window.
+                flowControl.Advance((int)dataLength);
+                WriteDataUnsynchronized(streamId, data, dataLength, endStream: false);
+
+                return WriteResponseTrailersAsync(streamId, headers);
+            }
+
+            static async ValueTask<FlushResult> WriteDataAndTrailersAsyncCore(Http2FrameWriter writer, int streamId, StreamOutputFlowControl flowControl, ReadOnlySequence<byte> data, long dataLength, bool firstWrite, HttpResponseTrailers headers)
+            {
+                await writer.WriteDataAsync(streamId, flowControl, data, dataLength, endStream: false, firstWrite);
+
+                return await writer.WriteResponseTrailersAsync(streamId, headers);
+            }
+        }
+
         /*  Padding is not implemented
             +---------------+
             |Pad Length? (8)|

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

@@ -423,16 +423,19 @@ private async ValueTask ProcessDataWrites()
                         {
                             // Output is ending and there are trailers to write
                             // Write any remaining content then write trailers
-                            if (readResult.Buffer.Length > 0)
-                            {
-                                // Only flush if required (i.e. content length exceeds flow control availability)
-                                // Writing remaining content without flushing allows content and trailers to be sent in the same packet
-                                await _frameWriter.WriteDataAsync(StreamId, _flowControl, readResult.Buffer, endStream: false, firstWrite, forceFlush: false);
-                            }
 
                             _stream.ResponseTrailers.SetReadOnly();
                             _stream.DecrementActiveClientStreamCount();
-                            flushResult = await _frameWriter.WriteResponseTrailers(StreamId, _stream.ResponseTrailers);
+
+                            if (readResult.Buffer.Length > 0)
+                            {
+                                // It is faster to write data and trailers together. Locking once reduces lock contention.
+                                flushResult = await _frameWriter.WriteDataAndTrailersAsync(StreamId, _flowControl, readResult.Buffer, firstWrite, _stream.ResponseTrailers);
+                            }
+                            else
+                            {
+                                flushResult = await _frameWriter.WriteResponseTrailersAsync(StreamId, _stream.ResponseTrailers);
+                            }
                         }
                         else if (readResult.IsCompleted && _streamEnded)
                         {