HTTP2: Send remaining data and trailers together

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

@@ -239,7 +239,7 @@ private void FinishWritingHeaders(int streamId, int payloadLength, bool done)
             }
         }
 
-        public ValueTask<FlushResult> WriteDataAsync(int streamId, StreamOutputFlowControl flowControl, in ReadOnlySequence<byte> data, bool endStream)
+        public ValueTask<FlushResult> WriteDataAsync(int streamId, StreamOutputFlowControl flowControl, in ReadOnlySequence<byte> data, bool endStream, bool forceFlush)
         {
             // The Length property of a ReadOnlySequence can be expensive, so we cache the value.
             var dataLength = data.Length;
@@ -261,7 +261,13 @@ public ValueTask<FlushResult> WriteDataAsync(int streamId, StreamOutputFlowContr
                 // 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);
-                return TimeFlushUnsynchronizedAsync();
+
+                if (forceFlush)
+                {
+                    return TimeFlushUnsynchronizedAsync();
+                }
+
+                return default;
             }
         }
 

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

@@ -375,7 +375,9 @@ private async ValueTask<FlushResult> ProcessDataWrites()
                         // Write any remaining content then write trailers
                         if (readResult.Buffer.Length > 0)
                         {
-                            flushResult = await _frameWriter.WriteDataAsync(_streamId, _flowControl, readResult.Buffer, endStream: false);
+                            // 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, forceFlush: false);
                         }
 
                         _stream.ResponseTrailers.SetReadOnly();
@@ -399,7 +401,7 @@ private async ValueTask<FlushResult> ProcessDataWrites()
                         {
                             _stream.DecrementActiveClientStreamCount();
                         }
-                        flushResult = await _frameWriter.WriteDataAsync(_streamId, _flowControl, readResult.Buffer, endStream);
+                        flushResult = await _frameWriter.WriteDataAsync(_streamId, _flowControl, readResult.Buffer, endStream, forceFlush: true);
                     }
 
                     _pipeReader.AdvanceTo(readResult.Buffer.End);

src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2StreamTests.cs

@@ -2044,6 +2044,127 @@ await ExpectAsync(Http2FrameType.DATA,
             Assert.Contains(CoreStrings.HPackErrorNotEnoughBuffer, message.Exception.Message);
         }
 
+        [Fact]
+        public async Task ResponseTrailers_WithLargeUnflushedData_DataExceedsFlowControlAvailableAndNotSentWithTrailers()
+        {
+            const int windowSize = (int)Http2PeerSettings.DefaultMaxFrameSize;
+            _clientSettings.InitialWindowSize = windowSize;
+
+            var headers = new[]
+            {
+                new KeyValuePair<string, string>(HeaderNames.Method, "GET"),
+                new KeyValuePair<string, string>(HeaderNames.Path, "/"),
+                new KeyValuePair<string, string>(HeaderNames.Scheme, "http"),
+            };
+            await InitializeConnectionAsync(async context =>
+            {
+                await context.Response.StartAsync();
+
+                // Body exceeds flow control available and requires the client to allow more
+                // data via updating the window
+                context.Response.BodyWriter.GetMemory(windowSize + 1);
+                context.Response.BodyWriter.Advance(windowSize + 1);
+
+                context.Response.AppendTrailer("CustomName", "Custom Value");
+            }).DefaultTimeout();
+
+            await StartStreamAsync(1, headers, endStream: true).DefaultTimeout();
+
+            var headersFrame = await ExpectAsync(Http2FrameType.HEADERS,
+                withLength: 37,
+                withFlags: (byte)Http2HeadersFrameFlags.END_HEADERS,
+                withStreamId: 1).DefaultTimeout();
+
+            await ExpectAsync(Http2FrameType.DATA,
+                withLength: 16384,
+                withFlags: (byte)Http2DataFrameFlags.NONE,
+                withStreamId: 1).DefaultTimeout();
+
+            var dataTask = ExpectAsync(Http2FrameType.DATA,
+                withLength: 1,
+                withFlags: (byte)Http2DataFrameFlags.NONE,
+                withStreamId: 1).DefaultTimeout();
+
+            // Reading final frame of data requires window update
+            // Verify this data task is waiting on window update
+            Assert.False(dataTask.IsCompletedSuccessfully);
+
+            await SendWindowUpdateAsync(1, 1);
+
+            await dataTask;
+
+            var trailersFrame = await ExpectAsync(Http2FrameType.HEADERS,
+                withLength: 25,
+                withFlags: (byte)(Http2HeadersFrameFlags.END_HEADERS | Http2HeadersFrameFlags.END_STREAM),
+                withStreamId: 1).DefaultTimeout();
+
+            await StopConnectionAsync(expectedLastStreamId: 1, ignoreNonGoAwayFrames: false).DefaultTimeout();
+
+            _hpackDecoder.Decode(headersFrame.PayloadSequence, endHeaders: false, handler: this);
+
+            Assert.Equal(2, _decodedHeaders.Count);
+            Assert.Contains("date", _decodedHeaders.Keys, StringComparer.OrdinalIgnoreCase);
+            Assert.Equal("200", _decodedHeaders[HeaderNames.Status]);
+
+            _decodedHeaders.Clear();
+            _hpackDecoder.Decode(trailersFrame.PayloadSequence, endHeaders: true, handler: this);
+
+            Assert.Single(_decodedHeaders);
+            Assert.Equal("Custom Value", _decodedHeaders["CustomName"]);
+        }
+
+        [Fact]
+        public async Task ResponseTrailers_WithUnflushedData_DataSentWithTrailers()
+        {
+            var headers = new[]
+            {
+                new KeyValuePair<string, string>(HeaderNames.Method, "GET"),
+                new KeyValuePair<string, string>(HeaderNames.Path, "/"),
+                new KeyValuePair<string, string>(HeaderNames.Scheme, "http"),
+            };
+            await InitializeConnectionAsync(async context =>
+            {
+                await context.Response.StartAsync();
+
+                var s = context.Response.BodyWriter.GetMemory(1);
+                s.Span[0] = byte.MaxValue;
+                context.Response.BodyWriter.Advance(1);
+
+                context.Response.AppendTrailer("CustomName", "Custom Value");
+            });
+
+            await StartStreamAsync(1, headers, endStream: true);
+
+            var headersFrame = await ExpectAsync(Http2FrameType.HEADERS,
+                withLength: 37,
+                withFlags: (byte)Http2HeadersFrameFlags.END_HEADERS,
+                withStreamId: 1);
+
+            await ExpectAsync(Http2FrameType.DATA,
+                withLength: 1,
+                withFlags: (byte)Http2DataFrameFlags.NONE,
+                withStreamId: 1);
+
+            var trailersFrame = await ExpectAsync(Http2FrameType.HEADERS,
+                withLength: 25,
+                withFlags: (byte)(Http2HeadersFrameFlags.END_HEADERS | Http2HeadersFrameFlags.END_STREAM),
+                withStreamId: 1);
+
+            await StopConnectionAsync(expectedLastStreamId: 1, ignoreNonGoAwayFrames: false);
+
+            _hpackDecoder.Decode(headersFrame.PayloadSequence, endHeaders: false, handler: this);
+
+            Assert.Equal(2, _decodedHeaders.Count);
+            Assert.Contains("date", _decodedHeaders.Keys, StringComparer.OrdinalIgnoreCase);
+            Assert.Equal("200", _decodedHeaders[HeaderNames.Status]);
+
+            _decodedHeaders.Clear();
+            _hpackDecoder.Decode(trailersFrame.PayloadSequence, endHeaders: true, handler: this);
+
+            Assert.Single(_decodedHeaders);
+            Assert.Equal("Custom Value", _decodedHeaders["CustomName"]);
+        }
+
         [Fact]
         public async Task ApplicationException_BeforeFirstWrite_Sends500()
         {