From c24511034f8409c77230bcf43837e8c16e3d300f Mon Sep 17 00:00:00 2001
From: ManickaP <mapichov@microsoft.com>
Date: Wed, 16 Aug 2023 10:46:33 +0200
Subject: [PATCH] Release the stream once the response is sent

---
 .../Net/Http/Http3LoopbackConnection.cs       | 38 +++++++++++++------
 1 file changed, 26 insertions(+), 12 deletions(-)

diff --git a/src/libraries/Common/tests/System/Net/Http/Http3LoopbackConnection.cs b/src/libraries/Common/tests/System/Net/Http/Http3LoopbackConnection.cs
index ec5c7a022c65a..9d6fef5fb3a72 100644
--- a/src/libraries/Common/tests/System/Net/Http/Http3LoopbackConnection.cs
+++ b/src/libraries/Common/tests/System/Net/Http/Http3LoopbackConnection.cs
@@ -108,11 +108,6 @@ public static int GetRequestId(QuicStream stream)
             return checked((int)stream.Id + 1);
         }
 
-        public Http3LoopbackStream GetOpenRequest(int requestId = 0)
-        {
-            return requestId == 0 ? _currentStream : _openStreams[requestId - 1];
-        }
-
         public override Task InitializeConnectionAsync()
         {
             throw new NotImplementedException();
@@ -195,6 +190,17 @@ public async Task EstablishControlStreamAsync(SettingsEntry[] settingsEntries)
             await _outboundControlStream.SendSettingsFrameAsync(settingsEntries);
         }
 
+        public async Task DisposeCurrentStream()
+        {
+            Assert.NotNull(_currentStream);
+            Assert.True(_currentStreamId >= 0);
+
+            await _currentStream.DisposeAsync().ConfigureAwait(false);
+            _openStreams.Remove((int)_currentStreamId);
+            _currentStream = null;
+            _currentStreamId = -4;
+        }
+
         public override async Task<byte[]> ReadRequestBodyAsync()
         {
             return await _currentStream.ReadRequestBodyAsync().ConfigureAwait(false);
@@ -206,24 +212,32 @@ public override async Task<HttpRequestData> ReadRequestDataAsync(bool readBody =
             return await stream.ReadRequestDataAsync(readBody).ConfigureAwait(false);
         }
 
-        public override Task SendResponseAsync(HttpStatusCode statusCode = HttpStatusCode.OK, IList<HttpHeaderData> headers = null, string content = "", bool isFinal = true)
+        public override async Task SendResponseAsync(HttpStatusCode statusCode = HttpStatusCode.OK, IList<HttpHeaderData> headers = null, string content = "", bool isFinal = true)
         {
-            return GetOpenRequest().SendResponseAsync(statusCode, headers, content, isFinal);
+            await _currentStream.SendResponseAsync(statusCode, headers, content, isFinal);
+            if (isFinal)
+            {
+                await DisposeCurrentStream().ConfigureAwait(false);
+            }
         }
 
-        public override Task SendResponseBodyAsync(byte[] content, bool isFinal = true)
+        public override async Task SendResponseBodyAsync(byte[] content, bool isFinal = true)
         {
-            return GetOpenRequest().SendResponseBodyAsync(content, isFinal);
+            await _currentStream.SendResponseBodyAsync(content, isFinal);
+            if (isFinal)
+            {
+                await DisposeCurrentStream().ConfigureAwait(false);
+            }
         }
 
         public override Task SendResponseHeadersAsync(HttpStatusCode statusCode = HttpStatusCode.OK, IList<HttpHeaderData> headers = null)
         {
-            return GetOpenRequest().SendResponseHeadersAsync(statusCode, headers);
+            return _currentStream.SendResponseHeadersAsync(statusCode, headers);
         }
 
         public override Task SendPartialResponseHeadersAsync(HttpStatusCode statusCode = HttpStatusCode.OK, IList<HttpHeaderData> headers = null)
         {
-            return GetOpenRequest().SendPartialResponseHeadersAsync(statusCode, headers);
+            return _currentStream.SendPartialResponseHeadersAsync(statusCode, headers);
         }
 
         public override async Task<HttpRequestData> HandleRequestAsync(HttpStatusCode statusCode = HttpStatusCode.OK, IList<HttpHeaderData> headers = null, string content = "")
@@ -310,7 +324,7 @@ public async Task WaitForClientDisconnectAsync(bool refuseNewRequests = true)
 
         public override async Task WaitForCancellationAsync(bool ignoreIncomingData = true)
         {
-            await GetOpenRequest().WaitForCancellationAsync(ignoreIncomingData).ConfigureAwait(false);
+            await _currentStream.WaitForCancellationAsync(ignoreIncomingData).ConfigureAwait(false);
         }
 
         public override Task WaitForCloseAsync(CancellationToken cancellationToken)