From 6ae364d5e16de4eacc4fa62f36d054f78bbcbc29 Mon Sep 17 00:00:00 2001
From: Anton Firszov <antonfir@gmail.com>
Date: Tue, 18 Jul 2023 19:10:00 +0200
Subject: [PATCH] Fix
 FailedRequests_ConnectionClosedWhileReceivingHeaders_Recorded (#89047)

---
 .../tests/FunctionalTests/MetricsTest.cs           | 26 ++++++++++++++--------
 1 file changed, 17 insertions(+), 9 deletions(-)

diff --git a/src/libraries/System.Net.Http/tests/FunctionalTests/MetricsTest.cs b/src/libraries/System.Net.Http/tests/FunctionalTests/MetricsTest.cs
index 29694a9..d57eec5 100644
--- a/src/libraries/System.Net.Http/tests/FunctionalTests/MetricsTest.cs
+++ b/src/libraries/System.Net.Http/tests/FunctionalTests/MetricsTest.cs
@@ -530,10 +530,10 @@ namespace System.Net.Http.Functional.Tests
             base.Dispose(disposing);
         }
 
-        protected Task<HttpResponseMessage> SendAsync(HttpMessageInvoker invoker, HttpRequestMessage request) =>
+        protected Task<HttpResponseMessage> SendAsync(HttpMessageInvoker invoker, HttpRequestMessage request, CancellationToken cancellationToken = default) =>
             TestHttpMessageInvoker ?
-            invoker.SendAsync(request, default) :
-            ((HttpClient)invoker).SendAsync(TestAsync, request);
+            invoker.SendAsync(request, cancellationToken) :
+            ((HttpClient)invoker).SendAsync(TestAsync, request, cancellationToken);
 
         protected HttpMessageInvoker CreateHttpMessageInvoker(HttpMessageHandler? handler = null) =>
             TestHttpMessageInvoker ?
@@ -617,24 +617,32 @@ namespace System.Net.Http.Functional.Tests
         [Fact]
         public async Task FailedRequests_ConnectionClosedWhileReceivingHeaders_Recorded()
         {
-            TimeSpan timeout = TimeSpan.FromSeconds(30);
+            using CancellationTokenSource cancelServerCts = new CancellationTokenSource();
             await LoopbackServerFactory.CreateClientAndServerAsync(async uri =>
             {
                 using HttpMessageInvoker client = CreateHttpMessageInvoker();
                 using InstrumentRecorder<long> recorder = SetupInstrumentRecorder<long>(InstrumentNames.FailedRequests);
                 using HttpRequestMessage request = new(HttpMethod.Get, uri) { Version = UseVersion };
 
-                await Assert.ThrowsAsync<HttpRequestException>(async () =>
+                Exception ex = await Assert.ThrowsAnyAsync<Exception>(async () =>
                 {
-                    using HttpResponseMessage response = await SendAsync(client, request);
-                }).WaitAsync(timeout);
+                    // Getting a cancellation is also good if we are unable to detect the peer shutdown.
+                    using CancellationTokenSource cts = new CancellationTokenSource(10_000);
+                    using HttpResponseMessage response = await SendAsync(client, request, cts.Token);
+                });
+                cancelServerCts.Cancel();
+                Assert.True(ex is HttpRequestException or TaskCanceledException);
 
                 Measurement<long> m = recorder.GetMeasurements().Single();
                 VerifyFailedRequests(m, 1, uri, null, null);
             }, async server =>
             {
-                var connection = (LoopbackServer.Connection)await server.EstablishGenericConnectionAsync().WaitAsync(timeout);
-                connection.Socket.Close();
+                try
+                {
+                    var connection = (LoopbackServer.Connection)await server.EstablishGenericConnectionAsync().WaitAsync(cancelServerCts.Token);
+                    connection.Socket.Close();
+                }
+                catch { }
             });
         }
     }
-- 
2.7.4