From 10d70c5c9eb1c132c70976cd93640d996a1550ae Mon Sep 17 00:00:00 2001
From: Justin Kotalik <jukotali@microsoft.com>
Date: Thu, 18 Apr 2019 12:44:54 -0700
Subject: [PATCH] Fix
 AbortedStream_ResetsAndDrainsRequest_RefusesFramesAfterCooldownExpires
 (#9487)

---
 .../Http2/Http2TestBase.cs                    |  9 +++++++-
 .../Http2/Http2TimeoutTests.cs                | 21 ++++++++++---------
 2 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2TestBase.cs b/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2TestBase.cs
index 6752e74e3b..3e024a9bce 100644
--- a/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2TestBase.cs
+++ b/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2TestBase.cs
@@ -1163,6 +1163,13 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
 
         internal async Task WaitForConnectionErrorAsync<TException>(bool ignoreNonGoAwayFrames, int expectedLastStreamId, Http2ErrorCode expectedErrorCode, params string[] expectedErrorMessage)
             where TException : Exception
+        {
+            await WaitForConnectionErrorAsyncDoNotCloseTransport<TException>(ignoreNonGoAwayFrames, expectedLastStreamId, expectedErrorCode, expectedErrorMessage);
+            _pair.Application.Output.Complete();
+        }
+
+        internal async Task WaitForConnectionErrorAsyncDoNotCloseTransport<TException>(bool ignoreNonGoAwayFrames, int expectedLastStreamId, Http2ErrorCode expectedErrorCode, params string[] expectedErrorMessage)
+            where TException : Exception
         {
             var frame = await ReceiveFrameAsync();
 
@@ -1184,7 +1191,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
             }
 
             await _connectionTask.DefaultTimeout();
-            _pair.Application.Output.Complete();
+            TestApplicationErrorLogger.LogInformation("Stopping Connection From ConnectionErrorAsync");
         }
 
         internal async Task WaitForStreamErrorAsync(int expectedStreamId, Http2ErrorCode expectedErrorCode, string expectedErrorMessage)
diff --git a/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2TimeoutTests.cs b/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2TimeoutTests.cs
index 029e8b5e15..812c533fa4 100644
--- a/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2TimeoutTests.cs
+++ b/src/Servers/Kestrel/test/InMemory.FunctionalTests/Http2/Http2TimeoutTests.cs
@@ -11,6 +11,7 @@ using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http2;
 using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure;
 using Microsoft.AspNetCore.Testing;
 using Microsoft.AspNetCore.Testing.xunit;
+using Microsoft.Extensions.Logging.Testing;
 using Microsoft.Net.Http.Headers;
 using Moq;
 using Xunit;
@@ -193,11 +194,13 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
         }
 
         [Theory]
-        [Flaky("https://github.com/aspnet/AspNetCore-Internal/issues/1879", FlakyOn.All)]
+        [Repeat(20)]
         [InlineData((int)Http2FrameType.DATA)]
         [InlineData((int)Http2FrameType.CONTINUATION)]
         public async Task AbortedStream_ResetsAndDrainsRequest_RefusesFramesAfterCooldownExpires(int intFinalFrameType)
         {
+            var closeLock = new object();
+            var closed = false;
             var finalFrameType = (Http2FrameType)intFinalFrameType;
             // Remove callback that completes _pair.Application.Output on abort.
             _mockConnectionContext.Reset();
@@ -216,8 +219,6 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
 
             await WaitForStreamErrorAsync(1, Http2ErrorCode.INTERNAL_ERROR, "The connection was aborted by the application.");
 
-            var cts = new CancellationTokenSource();
-
             async Task AdvanceClockAndSendFrames()
             {
                 if (finalFrameType == Http2FrameType.CONTINUATION)
@@ -227,7 +228,7 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
                 }
 
                 // There's a race when the appfunc is exiting about how soon it unregisters the stream, so retry until success.
-                while (!cts.Token.IsCancellationRequested)
+                while (!closed)
                 {
                     // Just past the timeout
                     mockSystemClock.UtcNow += Constants.RequestBodyDrainTimeout + TimeSpan.FromTicks(1);
@@ -247,24 +248,24 @@ namespace Microsoft.AspNetCore.Server.Kestrel.Core.Tests
                             throw new NotImplementedException(finalFrameType.ToString());
                     }
 
-                    if (!cts.Token.IsCancellationRequested)
-                    {
-                        await Task.Delay(10);
-                    }
+                    // TODO how do I force a function to go async?
+                    await Task.Delay(1);
                 }
             }
 
             var sendTask = AdvanceClockAndSendFrames();
 
-            await WaitForConnectionErrorAsync<Http2ConnectionErrorException>(
+            await WaitForConnectionErrorAsyncDoNotCloseTransport<Http2ConnectionErrorException>(
                 ignoreNonGoAwayFrames: false,
                 expectedLastStreamId: 1,
                 Http2ErrorCode.STREAM_CLOSED,
                 CoreStrings.FormatHttp2ErrorStreamClosed(finalFrameType, 1));
 
-            cts.Cancel();
+            closed = true;
 
             await sendTask.DefaultTimeout();
+
+            _pair.Application.Output.Complete();
         }
 
         [Fact]