HTTP/2: Fix flakey response rate timeout tests (#32139)

Author: JamesNK

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

@@ -6,6 +6,7 @@
 using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
+using Microsoft.AspNetCore.Http;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Features;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
@@ -332,6 +333,29 @@ await WaitForConnectionErrorAsyncDoNotCloseTransport<Http2ConnectionErrorExcepti
             _pair.Application.Output.Complete();
         }
 
+        private class EchoAppWithNotification
+        {
+            private readonly TaskCompletionSource _writeStartedTcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
+
+            public Task WriteStartedTask => _writeStartedTcs.Task;
+
+            public async Task RunApp(HttpContext context)
+            {
+                await context.Response.Body.FlushAsync();
+
+                var buffer = new byte[Http2PeerSettings.MinAllowedMaxFrameSize];
+                int received;
+
+                while ((received = await context.Request.Body.ReadAsync(buffer, 0, buffer.Length)) > 0)
+                {
+                    var writeTask = context.Response.Body.WriteAsync(buffer, 0, received);
+                    _writeStartedTcs.TrySetResult();
+
+                    await writeTask;
+                }
+            }
+        }
+
         [Fact]
         [QuarantinedTest("https://github.com/dotnet/aspnetcore/issues/21520")]
         public async Task DATA_Sent_TooSlowlyDueToSocketBackPressureOnSmallWrite_AbortsConnectionAfterGracePeriod()
@@ -347,7 +371,8 @@ public async Task DATA_Sent_TooSlowlyDueToSocketBackPressureOnSmallWrite_AbortsC
 
             _timeoutControl.Initialize(mockSystemClock.UtcNow.Ticks);
 
-            await InitializeConnectionAsync(_echoApplication);
+            var app = new EchoAppWithNotification();
+            await InitializeConnectionAsync(app.RunApp);
 
             await StartStreamAsync(1, _browserRequestHeaders, endStream: false);
             await SendDataAsync(1, _helloWorldBytes, endStream: true);
@@ -357,6 +382,8 @@ await ExpectAsync(Http2FrameType.HEADERS,
                 withFlags: (byte)Http2HeadersFrameFlags.END_HEADERS,
                 withStreamId: 1);
 
+            await app.WriteStartedTask.DefaultTimeout();
+
             // Complete timing of the request body so we don't induce any unexpected request body rate timeouts.
             _timeoutControl.Tick(mockSystemClock.UtcNow);
 
@@ -400,7 +427,8 @@ public async Task DATA_Sent_TooSlowlyDueToSocketBackPressureOnLargeWrite_AbortsC
 
             _timeoutControl.Initialize(mockSystemClock.UtcNow.Ticks);
 
-            await InitializeConnectionAsync(_echoApplication);
+            var app = new EchoAppWithNotification();
+            await InitializeConnectionAsync(app.RunApp);
 
             await StartStreamAsync(1, _browserRequestHeaders, endStream: false);
             await SendDataAsync(1, _maxData, endStream: true);
@@ -410,6 +438,8 @@ await ExpectAsync(Http2FrameType.HEADERS,
                 withFlags: (byte)Http2HeadersFrameFlags.END_HEADERS,
                 withStreamId: 1);
 
+            await app.WriteStartedTask.DefaultTimeout();
+
             // Complete timing of the request body so we don't induce any unexpected request body rate timeouts.
             _timeoutControl.Tick(mockSystemClock.UtcNow);