TestServer returns error with response when server ends with p… (#19249)

src/Hosting/TestHost/src/HttpContextBuilder.cs

@@ -111,9 +111,14 @@ namespace Microsoft.AspNetCore.TestHost
                 {
                     await _application.ProcessRequestAsync(_testContext);
 
+                    // Determine whether request body was complete when the delegate exited.
+                    // This could throw an error if there was a pending server read. Needs to
+                    // happen before completing the response so the response returns the error.
+                    var requestBodyInProgress = RequestBodyReadInProgress();
+
                     // Matches Kestrel server: response is completed before request is drained
                     await CompleteResponseAsync();
-                    await CompleteRequestAsync();
+                    await CompleteRequestAsync(requestBodyInProgress);
                     _application.DisposeContext(_testContext, exception: null);
                 }
                 catch (Exception ex)
@@ -160,18 +165,8 @@ namespace Microsoft.AspNetCore.TestHost
             CancelRequestBody();
         }
 
-        private async Task CompleteRequestAsync()
+        private async Task CompleteRequestAsync(bool requestBodyInProgress)
         {
-            bool requestBodyInProgress;
-            try
-            {
-                requestBodyInProgress = !_requestPipe.Reader.TryRead(out var result) || !result.IsCompleted;
-            }
-            catch (Exception ex)
-            {
-                throw new InvalidOperationException("An error occurred when completing the request. Request delegate may have finished while there is a pending read of the request body.", ex);
-            }
-
             if (requestBodyInProgress)
             {
                 // If request is still in progress then abort it.
@@ -188,6 +183,18 @@ namespace Microsoft.AspNetCore.TestHost
             // Potential future improvement: add logging that the request timed out
         }
 
+        private bool RequestBodyReadInProgress()
+        {
+            try
+            {
+                return !_requestPipe.Reader.TryRead(out var result) || !result.IsCompleted;
+            }
+            catch (Exception ex)
+            {
+                throw new InvalidOperationException("An error occurred when completing the request. Request delegate may have finished while there is a pending read of the request body.", ex);
+            }
+        }
+
         internal async Task CompleteResponseAsync()
         {
             _pipelineFinished = true;

src/Hosting/TestHost/src/ResponseBodyReaderStream.cs

@@ -20,6 +20,7 @@ namespace Microsoft.AspNetCore.TestHost
         private bool _aborted;
         private Exception _abortException;
 
+        private readonly object _abortLock = new object();
         private readonly Action _abortRequest;
         private readonly Action _readComplete;
         private readonly Pipe _pipe;
@@ -124,16 +125,24 @@ namespace Microsoft.AspNetCore.TestHost
         internal void Abort(Exception innerException)
         {
             Contract.Requires(innerException != null);
-            _aborted = true;
+
-            _abortException = innerException;
+            lock (_abortLock)
+            {
+                _abortException = innerException;
+                _aborted = true;
+            }
+
             _pipe.Reader.CancelPendingRead();
         }
 
         private void CheckAborted()
         {
-            if (_aborted)
+            lock (_abortLock)
             {
-                throw new IOException(string.Empty, _abortException);
+                if (_aborted)
+                {
+                    throw new IOException(string.Empty, _abortException);
+                }
             }
         }
 

src/Hosting/TestHost/test/TestClientTests.cs

@@ -2,7 +2,6 @@
 // Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
 
 using System;
-using System.Globalization;
 using System.IO;
 using System.Linq;
 using System.Net;
@@ -386,7 +385,6 @@ namespace Microsoft.AspNetCore.TestHost
         }
 
         [Fact]
-        [Flaky("<No longer used; tracked in Kusto>", FlakyOn.All)]
         public async Task ClientStreaming_ResponseCompletesWithPendingRead_ThrowError()
         {
             // Arrange