[H/3] Fix disposing of buffers in parallel with running tasks

src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/Http3RequestStream.cs

@@ -26,7 +26,9 @@ internal sealed class Http3RequestStream : IHttpStreamHeadersHandler, IAsyncDisp
         private Http3Connection _connection;
         private long _streamId = -1; // A stream does not have an ID until the first I/O against it. This gets set almost immediately following construction.
         private readonly QuicStream _stream;
+        private volatile bool _finishingBackgroundWrite;
         private ArrayBuffer _sendBuffer;
+        private volatile bool _finishingBackgroundRead;
         private ArrayBuffer _recvBuffer;
         private TaskCompletionSource<bool>? _expect100ContinueCompletionSource; // True indicates we should send content (e.g. received 100 Continue).
         private bool _disposed;
@@ -130,8 +132,18 @@ private void DisposeSyncHelper()
         {
             _connection.RemoveStream(_stream);
 
-            _sendBuffer.Dispose();
-            _recvBuffer.Dispose();
+            // If the request sending was offloaded to be done concurrently and not awaited within SendAsync (by calling connection.LogException),
+            // the _sendBuffer disposal is the responsibility of that offloaded task to prevent returning the buffer to the pool while it still might be in use.
+            if (!_finishingBackgroundWrite)
+            {
+                _sendBuffer.Dispose();
+            }
+            // If the response receiving was offloaded to be done concurrently and not awaited within SendAsync (by calling connection.LogException),
+            // the _recvBuffer disposal is the responsibility of that offloaded task to prevent returning the buffer to the pool while it still might be in use.
+            if (!_finishingBackgroundRead)
+            {
+                _recvBuffer.Dispose();
+            }
         }
 
         public void GoAway()
@@ -196,6 +208,11 @@ await Task.WhenAny(sendRequestTask, readResponseTask).ConfigureAwait(false) == s
                     }
                     catch
                     {
+                        // This is a best effort attempt to transfer the responsibility of disposing _recvBuffer to ReadResponseAsync.
+                        // The task might be past checking the variable or already finished, in which case the buffer won't be returned to the pool.
+                        // Not returning the buffer to the pool is an acceptable trade-off for making sure that the buffer is not used after it's been returned.
+                        _finishingBackgroundRead = true;
+
                         // Exceptions will be bubbled up from sendRequestTask here,
                         // which means the result of readResponseTask won't be observed directly:
                         // Do a background await to log any exceptions.
@@ -205,6 +222,11 @@ await Task.WhenAny(sendRequestTask, readResponseTask).ConfigureAwait(false) == s
                 }
                 else
                 {
+                    // This is a best effort attempt to transfer the responsibility of disposing _sendBuffer to SendContentAsync.
+                    // The task might be past checking the variable or already finished, in which case the buffer won't be returned to the pool.
+                    // Not returning the buffer to the pool is an acceptable trade-off for making sure that the buffer is not used after it's been returned.
+                    _finishingBackgroundWrite = true;
+
                     // Duplex is being used, so we can't wait for content to finish sending.
                     // Do a background await to log any exceptions.
                     _connection.LogExceptions(sendRequestTask);
@@ -226,12 +248,6 @@ await Task.WhenAny(sendRequestTask, readResponseTask).ConfigureAwait(false) == s
                     {
                         // The server doesn't need the whole request to respond so it's aborting its reading side gracefully, see https://datatracker.ietf.org/doc/html/rfc9114#section-4.1-15.
                     }
-                    catch (OperationCanceledException)
-                    {
-                        // If the request got cancelled before WritesClosed completed, avoid leaking an unobserved task exception.
-                        _connection.LogExceptions(writesClosed);
-                        throw;
-                    }
                 }
 
                 Debug.Assert(_response != null && _response.Content != null);
@@ -387,32 +403,44 @@ await Task.WhenAny(sendRequestTask, readResponseTask).ConfigureAwait(false) == s
         /// </summary>
         private async Task ReadResponseAsync(CancellationToken cancellationToken)
         {
-            if (HttpTelemetry.Log.IsEnabled()) HttpTelemetry.Log.ResponseHeadersStart();
-
-            Debug.Assert(_response == null);
-            do
+            try
             {
-                _headerState = HeaderState.StatusHeader;
+                if (HttpTelemetry.Log.IsEnabled()) HttpTelemetry.Log.ResponseHeadersStart();
 
-                (Http3FrameType? frameType, long payloadLength) = await ReadFrameEnvelopeAsync(cancellationToken).ConfigureAwait(false);
-
-                if (frameType != Http3FrameType.Headers)
+                Debug.Assert(_response == null);
+                do
                 {
-                    if (NetEventSource.Log.IsEnabled())
+                    _headerState = HeaderState.StatusHeader;
+
+                    (Http3FrameType? frameType, long payloadLength) = await ReadFrameEnvelopeAsync(cancellationToken).ConfigureAwait(false);
+
+                    if (frameType != Http3FrameType.Headers)
                     {
-                        Trace($"Expected HEADERS as first response frame; received {frameType}.");
+                        if (NetEventSource.Log.IsEnabled())
+                        {
+                            Trace($"Expected HEADERS as first response frame; received {frameType}.");
+                        }
+                        throw new HttpIOException(HttpRequestError.InvalidResponse, SR.net_http_invalid_response);
                     }
-                    throw new HttpIOException(HttpRequestError.InvalidResponse, SR.net_http_invalid_response);
-                }
 
-                await ReadHeadersAsync(payloadLength, cancellationToken).ConfigureAwait(false);
-                Debug.Assert(_response != null);
-            }
-            while ((int)_response.StatusCode < 200);
+                    await ReadHeadersAsync(payloadLength, cancellationToken).ConfigureAwait(false);
+                    Debug.Assert(_response != null);
+                }
+                while ((int)_response.StatusCode < 200);
 
-            _headerState = HeaderState.TrailingHeaders;
+                _headerState = HeaderState.TrailingHeaders;
 
-            if (HttpTelemetry.Log.IsEnabled()) HttpTelemetry.Log.ResponseHeadersStop((int)_response.StatusCode);
+                if (HttpTelemetry.Log.IsEnabled()) HttpTelemetry.Log.ResponseHeadersStop((int)_response.StatusCode);
+            }
+            finally
+            {
+                // Note that we might still observe false here even if we're responsible for the _recvBuffer disposal.
+                // But in that case, we just don't return the rented buffer to the pool, which is lesser evil than writing to returned one.
+                if (_finishingBackgroundRead)
+                {
+                    _recvBuffer.Dispose();
+                }
+            }
         }
 
         private async Task SendContentAsync(HttpContent content, CancellationToken cancellationToken)
@@ -491,6 +519,12 @@ private async Task SendContentAsync(HttpContent content, CancellationToken cance
             }
             finally
             {
+                // Note that we might still observe false here even if we're responsible for the _sendBuffer disposal.
+                // But in that case, we just don't return the rented buffer to the pool, which is lesser evil than writing to returned one.
+                if (_finishingBackgroundWrite)
+                {
+                    _sendBuffer.Dispose();
+                }
                 _requestSendCompleted = true;
                 RemoveFromConnectionIfDone();
             }