Reduce allocations in Quic/HTTP3 (dotnet#104394)

* Reduce allocations in Quic/HTTP3 * Also wait for WritesClosed when there's no content
mikem8361 · Jul 8, 2024 · 2d92034 · 2d92034
1 parent 26b6161
commit 2d92034
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 24 deletions.
diff --git a/src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/Http3RequestStream.cs b/src/libraries/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/Http3RequestStream.cs
@@ -172,15 +172,9 @@ public async Task<HttpResponseMessage> SendAsync(CancellationToken cancellationT
                     await FlushSendBufferAsync(endStream: _request.Content == null, _requestBodyCancellationSource.Token).ConfigureAwait(false);
                 }
 
-                Task sendRequestTask;
-                if (_request.Content != null)
-                {
-                    sendRequestTask = SendContentAsync(_request.Content!, _requestBodyCancellationSource.Token);
-                }
-                else
-                {
-                    sendRequestTask = Task.CompletedTask;
-                }
+                Task sendRequestTask = _request.Content != null
+                    ? SendContentAsync(_request.Content, _requestBodyCancellationSource.Token)
+                    : Task.CompletedTask;
 
                 // In parallel, send content and read response.
                 // Depending on Expect 100 Continue usage, one will depend on the other making progress.
@@ -219,6 +213,23 @@ await Task.WhenAny(sendRequestTask, readResponseTask).ConfigureAwait(false) == s
                 // Wait for the response headers to be read.
                 await readResponseTask.ConfigureAwait(false);
 
+                // If we've sent a body, wait for the writes to be closed (most likely already done).
+                // If sendRequestTask hasn't completed yet, we're doing duplex content transfers and can't wait for writes to be closed yet.
+                if (sendRequestTask.IsCompletedSuccessfully &&
+                    _stream.WritesClosed is { IsCompletedSuccessfully: false } writesClosed)
+                {
+                    try
+                    {
+                        await writesClosed.WaitAsync(_requestBodyCancellationSource.Token).ConfigureAwait(false);
+                    }
+                    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);
                 // Set our content stream.
                 var responseContent = (HttpConnectionResponseContent)_response.Content;
@@ -460,7 +471,6 @@ private async Task SendContentAsync(HttpContent content, CancellationToken cance
                 else
                 {
                     _stream.CompleteWrites();
-                    await _stream.WritesClosed.ConfigureAwait(false);
                 }
 
                 if (HttpTelemetry.Log.IsEnabled()) HttpTelemetry.Log.RequestContentStop(bytesWritten);
@@ -523,17 +533,11 @@ private async ValueTask WriteRequestContentAsync(ReadOnlyMemory<byte> buffer, Ca
             }
         }
 
-        private async ValueTask FlushSendBufferAsync(bool endStream, CancellationToken cancellationToken)
+        private ValueTask FlushSendBufferAsync(bool endStream, CancellationToken cancellationToken)
         {
-            await _stream.WriteAsync(_sendBuffer.ActiveMemory, endStream, cancellationToken).ConfigureAwait(false);
-            _sendBuffer.Discard(_sendBuffer.ActiveLength);
-
-            await _stream.FlushAsync(cancellationToken).ConfigureAwait(false);
-
-            if (endStream)
-            {
-                await _stream.WritesClosed.ConfigureAwait(false);
-            }
+            ReadOnlyMemory<byte> toSend = _sendBuffer.ActiveMemory;
+            _sendBuffer.Discard(toSend.Length);
+            return _stream.WriteAsync(toSend, endStream, cancellationToken);
         }
 
         private async ValueTask DrainContentLength0Frames(CancellationToken cancellationToken)

diff --git a/src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ResettableValueTaskSource.cs b/src/libraries/System.Net.Quic/src/System/Net/Quic/Internal/ResettableValueTaskSource.cs
@@ -310,6 +310,13 @@ public Task GetTask(object? keepAlive)
         {
             if (_finalTaskSource is null)
             {
+                if (_isSignaled)
+                {
+                    return _exception is null
+                        ? Task.CompletedTask
+                        : Task.FromException(_exception);
+                }
+
                 _finalTaskSource = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
                 if (!_isCompleted)
                 {
@@ -319,10 +326,6 @@ public Task GetTask(object? keepAlive)
                         ((GCHandle)state!).Free();
                     }, handle, CancellationToken.None, TaskContinuationOptions.ExecuteSynchronously, TaskScheduler.Default);
                 }
-                if (_isSignaled)
-                {
-                    TrySignal(out _);
-                }
             }
             return _finalTaskSource.Task;
         }