Fix test flakiness with OnTimeout (#29359)

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

@@ -423,8 +423,11 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
                         break;
                     }
 
+                    // Observe HTTP/2 backpressure
                     var actual = flowControl.AdvanceUpToAndWait(dataLength, out availabilityTask);
 
+                    var shouldFlush = false;
+
                     if (actual > 0)
                     {
                         if (actual < dataLength)
@@ -439,25 +442,32 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2
                             dataLength = 0;
                         }
 
-                        // Don't call TimeFlushUnsynchronizedAsync() since we time this write while also accounting for
+                        // Don't call FlushAsync() with the min data rate, since we time this write while also accounting for
                         // flow control induced backpressure below.
-                        writeTask = _flusher.FlushAsync();
+                        shouldFlush = true;
                     }
                     else if (firstWrite)
                     {
                         // If we're facing flow control induced backpressure on the first write for a given stream's response body,
                         // we make sure to flush the response headers immediately.
-                        writeTask = _flusher.FlushAsync();
+                        shouldFlush = true;
                     }
 
-                    firstWrite = false;
-
-                    if (_minResponseDataRate != null)
+                    if (shouldFlush)
                     {
-                        _timeoutControl.BytesWrittenToBuffer(_minResponseDataRate, _unflushedBytes);
+                        if (_minResponseDataRate != null)
+                        {
+                            // Call BytesWrittenToBuffer before FlushAsync() to make testing easier, otherwise the Flush can cause test code to run before the timeout
+                            // control updates and if the test checks for a timeout it can fail
+                            _timeoutControl.BytesWrittenToBuffer(_minResponseDataRate, _unflushedBytes);
+                        }
+
+                        _unflushedBytes = 0;
+
+                        writeTask = _flusher.FlushAsync();
                     }
 
-                    _unflushedBytes = 0;
+                    firstWrite = false;
                 }
 
                 // Avoid timing writes that are already complete. This is likely to happen during the last iteration.

src/Servers/Kestrel/Core/src/Internal/Infrastructure/PipeWriterHelpers/TimingPipeFlusher.cs

@@ -50,13 +50,15 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure.PipeW
 
         public ValueTask<FlushResult> FlushAsync(MinDataRate? minRate, long count, IHttpOutputAborter? outputAborter, CancellationToken cancellationToken)
         {
-            var pipeFlushTask = _writer.FlushAsync(cancellationToken);
-
             if (minRate is object)
             {
+                // Call BytesWrittenToBuffer before FlushAsync() to make testing easier, otherwise the Flush can cause test code to run before the timeout
+                // control updates and if the test checks for a timeout it can fail
                 _timeoutControl!.BytesWrittenToBuffer(minRate, count);
             }
 
+            var pipeFlushTask = _writer.FlushAsync(cancellationToken);
+
             if (pipeFlushTask.IsCompletedSuccessfully)
             {
                 var flushResult = pipeFlushTask.Result;

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

@@ -550,7 +550,8 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             _mockConnectionContext.VerifyNoOtherCalls();
         }
 
-        [Fact(Skip = "https://github.com/dotnet/aspnetcore-internal/issues/2197")]
+        [Fact]
+        [QuarantinedTest("https://github.com/dotnet/aspnetcore-internal/issues/2197")]
         public async Task DATA_Sent_TooSlowlyDueToOutputFlowControlOnMultipleStreams_AbortsConnectionAfterAdditiveRateTimeout()
         {
             var mockSystemClock = _serviceContext.MockSystemClock;
@@ -582,7 +583,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             await SendDataAsync(3, _maxData, endStream: true);
 
             await ExpectAsync(Http2FrameType.HEADERS,
-                withLength: 32,
+                withLength: 2,
                 withFlags: (byte)Http2HeadersFrameFlags.END_HEADERS,
                 withStreamId: 3);
             await ExpectAsync(Http2FrameType.DATA,