HTTP/2 Request Cancellation (#35118)

HTTP/2 cancellation support, plus improvements to outgoing write buffering.
dotnet · Apr 11, 2019 · d6e36e4 · d6e36e4
1 parent d26c6cd
commit d6e36e4
Show file tree

Hide file tree

Showing 11 changed files with 451 additions and 175 deletions.
diff --git a/src/System.Net.Http/src/System.Net.Http.csproj b/src/System.Net.Http/src/System.Net.Http.csproj
@@ -157,6 +157,7 @@
     <Compile Include="System\Net\Http\SocketsHttpHandler\SocketsHttpHandler.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\RawConnectionStream.cs" />
     <Compile Include="System\Net\Http\SocketsHttpHandler\RedirectHandler.cs" />
+    <Compile Include="System\Net\Http\SocketsHttpHandler\TaskCompletionSourceWithCancellation.cs" />
     <Compile Include="$(CommonPath)\CoreLib\System\Collections\Concurrent\ConcurrentQueueSegment.cs">
       <Link>Common\CoreLib\System\Collections\Concurrent\ConcurrentQueueSegment.cs</Link>
     </Compile>

diff --git a/src/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/ArrayBuffer.cs b/src/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/ArrayBuffer.cs
@@ -59,6 +59,8 @@ public void Dispose()
         public Memory<byte> ActiveMemory => new Memory<byte>(_bytes, _activeStart, _availableStart - _activeStart);
         public Memory<byte> AvailableMemory => new Memory<byte>(_bytes, _availableStart, _bytes.Length - _availableStart);
 
+        public int Capacity => _bytes.Length;
+
         public void Discard(int byteCount)
         {
             Debug.Assert(byteCount <= ActiveSpan.Length, $"Expected {byteCount} <= {ActiveSpan.Length}");

diff --git a/src/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/CreditManager.cs b/src/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/CreditManager.cs
@@ -4,6 +4,7 @@
 
 using System.Diagnostics;
 using System.Collections.Generic;
+using System.Threading;
 using System.Threading.Tasks;
 
 namespace System.Net.Http
@@ -13,7 +14,7 @@ internal sealed class CreditManager : IDisposable
         private struct Waiter
         {
             public int Amount;
-            public TaskCompletionSource<int> TaskCompletionSource;
+            public TaskCompletionSourceWithCancellation<int> TaskCompletionSource;
         }
 
         private int _current;
@@ -37,7 +38,7 @@ private object SyncObject
             }
         }
 
-        public ValueTask<int> RequestCreditAsync(int amount)
+        public ValueTask<int> RequestCreditAsync(int amount, CancellationToken cancellationToken)
         {
             lock (SyncObject)
             {
@@ -55,16 +56,21 @@ public ValueTask<int> RequestCreditAsync(int amount)
                     return new ValueTask<int>(granted);
                 }
 
-                var tcs = new TaskCompletionSource<int>(TaskContinuationOptions.RunContinuationsAsynchronously);
+                // Uses RunContinuationsAsynchronously internally.
+                var tcs = new TaskCompletionSourceWithCancellation<int>();
 
                 if (_waiters == null)
                 {
                     _waiters = new Queue<Waiter>();
                 }
 
-                _waiters.Enqueue(new Waiter { Amount = amount, TaskCompletionSource = tcs });
+                Waiter waiter = new Waiter { Amount = amount, TaskCompletionSource = tcs };
 
-                return new ValueTask<int>(tcs.Task);
+                _waiters.Enqueue(waiter);
+
+                return new ValueTask<int>(cancellationToken.CanBeCanceled ?
+                                          tcs.WaitWithCancellationAsync(cancellationToken) :
+                                          tcs.Task);
             }
         }
 
@@ -92,8 +98,12 @@ public void AdjustCredit(int amount)
                     while (_current > 0 && _waiters.TryDequeue(out Waiter waiter))
                     {
                         int granted = Math.Min(waiter.Amount, _current);
-                        _current -= granted;
-                        waiter.TaskCompletionSource.SetResult(granted);
+
+                        // Ensure that we grant credit only if the task has not been canceled.
+                        if (waiter.TaskCompletionSource.TrySetResult(granted))
+                        {
+                            _current -= granted;
+                        }
                     }
                 }
             }
@@ -114,7 +124,7 @@ public void Dispose()
                 {
                     while (_waiters.TryDequeue(out Waiter waiter))
                     {
-                        waiter.TaskCompletionSource.SetException(new ObjectDisposedException(nameof(CreditManager)));
+                        waiter.TaskCompletionSource.TrySetException(new ObjectDisposedException(nameof(CreditManager)));
                     }
                 }
             }