Add response drain max time to SocketsHttpHandler (#27708)

* Add response drain max time to SocketsHttpHandler

Adds support for a max time limit on how long we'll try to drain a connection before returning it to the pool, and adds a (currently non-public) knob for controlling that limit.

This also fixes a race condition with regards to cancellation and chunked responses.  If cancellation was requested after returning the connection to the pool but before we disposed of the cancellation registration, someone else could get the pooled connection and then cancellation of one request could end up canceling the other request.

* Address PR feedback

Author: stephentoub

src/Common/src/System/Net/Http/HttpHandlerDefaults.cs

@@ -14,6 +14,7 @@ internal static class HttpHandlerDefaults
         public const int DefaultMaxAutomaticRedirections = 50;
         public const int DefaultMaxConnectionsPerServer = int.MaxValue;
         public const int DefaultMaxResponseDrainSize = 1024 * 1024;
+        public static readonly TimeSpan DefaultMaxResponseDrainTime = TimeSpan.FromSeconds(2);
         public const int DefaultMaxResponseHeadersLength = 64; // Units in K (1024) bytes.
         public const DecompressionMethods DefaultAutomaticDecompression = DecompressionMethods.None;
         public const bool DefaultAutomaticRedirection = true;

src/System.Net.Http/src/ILLinkTrim.xml

@@ -2,5 +2,9 @@
   <assembly fullname="System.Net.Http">
     <!-- Anonymous types are used with DiagnosticSource logging and subscribers reflect over those, calling their public getters. -->
     <type fullname="*f__AnonymousType*" />
+    <type fullname="System.Net.Http.SocketsHttpHandler"> <!-- TODO #27685: Remove once public. -->
+      <method name="get_MaxResponseDrainTime" />
+      <method name="set_MaxResponseDrainTime" />
+    </type>
   </assembly>
 </linker>

src/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/ChunkedEncodingReadStream.cs

@@ -45,7 +45,7 @@ public override ValueTask<int> ReadAsync(Memory<byte> destination, CancellationT
                 }
 
                 // Try to consume from data we already have in the buffer.
-                int bytesRead = ReadChunksFromConnectionBuffer(destination.Span);
+                int bytesRead = ReadChunksFromConnectionBuffer(destination.Span, cancellationRegistration: default);
                 if (bytesRead > 0)
                 {
                     return new ValueTask<int>(bytesRead);
@@ -108,7 +108,7 @@ private async ValueTask<int> ReadAsyncCore(Memory<byte> destination, Cancellatio
 
                         // Now that we have more, see if we can get any response data, and if
                         // we can we're done.
-                        int bytesCopied = ReadChunksFromConnectionBuffer(destination.Span);
+                        int bytesCopied = ReadChunksFromConnectionBuffer(destination.Span, ctr);
                         if (bytesCopied > 0)
                         {
                             return bytesCopied;
@@ -144,7 +144,7 @@ private async Task CopyToAsyncCore(Stream destination, CancellationToken cancell
                     {
                         while (true)
                         {
-                            ReadOnlyMemory<byte> bytesRead = ReadChunkFromConnectionBuffer(int.MaxValue);
+                            ReadOnlyMemory<byte> bytesRead = ReadChunkFromConnectionBuffer(int.MaxValue, ctr);
                             if (bytesRead.Length == 0)
                             {
                                 break;
@@ -171,12 +171,12 @@ private async Task CopyToAsyncCore(Stream destination, CancellationToken cancell
                 }
             }
 
-            private int ReadChunksFromConnectionBuffer(Span<byte> destination)
+            private int ReadChunksFromConnectionBuffer(Span<byte> destination, CancellationTokenRegistration cancellationRegistration)
             {
                 int totalBytesRead = 0;
                 while (destination.Length > 0)
                 {
-                    ReadOnlyMemory<byte> bytesRead = ReadChunkFromConnectionBuffer(destination.Length);
+                    ReadOnlyMemory<byte> bytesRead = ReadChunkFromConnectionBuffer(destination.Length, cancellationRegistration);
                     Debug.Assert(bytesRead.Length <= destination.Length);
                     if (bytesRead.Length == 0)
                     {
@@ -190,7 +190,7 @@ private int ReadChunksFromConnectionBuffer(Span<byte> destination)
                 return totalBytesRead;
             }
 
-            private ReadOnlyMemory<byte> ReadChunkFromConnectionBuffer(int maxBytesToRead)
+            private ReadOnlyMemory<byte> ReadChunkFromConnectionBuffer(int maxBytesToRead, CancellationTokenRegistration cancellationRegistration)
             {
                 Debug.Assert(maxBytesToRead > 0);
 
@@ -287,6 +287,16 @@ private ReadOnlyMemory<byte> ReadChunkFromConnectionBuffer(int maxBytesToRead)
 
                                 if (currentLine.IsEmpty)
                                 {
+                                    // Dispose of the registration and then check whether cancellation has been
+                                    // requested. This is necessary to make determinstic a race condition between
+                                    // cancellation being requested and unregistering from the token.  Otherwise,
+                                    // it's possible cancellation could be requested just before we unregister and
+                                    // we then return a connection to the pool that has been or will be disposed
+                                    // (e.g. if a timer is used and has already queued its callback but the
+                                    // callback hasn't yet run).
+                                    cancellationRegistration.Dispose();
+                                    cancellationRegistration.Token.ThrowIfCancellationRequested();
+
                                     _state = ParsingState.Done;
                                     _connection.CompleteResponse();
                                     _connection = null;
@@ -344,32 +354,52 @@ public override async Task<bool> DrainAsync(int maxDrainBytes)
             {
                 Debug.Assert(_connection != null);
 
-                int drainedBytes = 0;
-                while (true)
+                CancellationTokenSource cts = null;
+                CancellationTokenRegistration ctr = default;
+                try
                 {
-                    drainedBytes += _connection.RemainingBuffer.Length;
+                    int drainedBytes = 0;
                     while (true)
                     {
-                        ReadOnlyMemory<byte> bytesRead = ReadChunkFromConnectionBuffer(int.MaxValue);
-                        if (bytesRead.Length == 0)
+                        drainedBytes += _connection.RemainingBuffer.Length;
+                        while (true)
                         {
-                            break;
+                            ReadOnlyMemory<byte> bytesRead = ReadChunkFromConnectionBuffer(int.MaxValue, ctr);
+                            if (bytesRead.Length == 0)
+                            {
+                                break;
+                            }
                         }
-                    }
 
-                    // When ReadChunkFromConnectionBuffer reads the final chunk, it will clear out _connection
-                    // and return the connection to the pool.
-                    if (_connection == null)
-                    {
-                        return true;
-                    }
+                        // When ReadChunkFromConnectionBuffer reads the final chunk, it will clear out _connection
+                        // and return the connection to the pool.
+                        if (_connection == null)
+                        {
+                            return true;
+                        }
 
-                    if (drainedBytes >= maxDrainBytes)
-                    {
-                        return false;
-                    }
+                        if (drainedBytes >= maxDrainBytes)
+                        {
+                            return false;
+                        }
+
+                        if (cts == null) // only create the drain timer if we have to go async
+                        {
+                            TimeSpan drainTime = _connection._pool.Settings._maxResponseDrainTime;
+                            if (drainTime != Timeout.InfiniteTimeSpan)
+                            {
+                                cts = new CancellationTokenSource((int)drainTime.TotalMilliseconds);
+                                ctr = cts.Token.Register(s => ((HttpConnection)s).Dispose(), _connection);
+                            }
+                        }
 
-                    await _connection.FillAsync().ConfigureAwait(false);
+                        await _connection.FillAsync().ConfigureAwait(false);
+                    }
+                }
+                finally
+                {
+                    ctr.Dispose();
+                    cts?.Dispose();
                 }
             }
         }

src/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/ContentLengthReadStream.cs

@@ -174,16 +174,42 @@ public override async Task<bool> DrainAsync(int maxDrainBytes)
                     return false;
                 }
 
-                while (true)
+                CancellationTokenSource cts = null;
+                CancellationTokenRegistration ctr = default;
+                TimeSpan drainTime = _connection._pool.Settings._maxResponseDrainTime;
+                if (drainTime != Timeout.InfiniteTimeSpan)
                 {
-                    await _connection.FillAsync().ConfigureAwait(false);
-                    ReadFromConnectionBuffer(int.MaxValue);
-                    if (_contentBytesRemaining == 0)
+                    cts = new CancellationTokenSource((int)drainTime.TotalMilliseconds);
+                    ctr = cts.Token.Register(s => ((HttpConnection)s).Dispose(), _connection);
+                }
+                try
+                {
+                    while (true)
                     {
-                        Finish();
-                        return true;
+                        await _connection.FillAsync().ConfigureAwait(false);
+                        ReadFromConnectionBuffer(int.MaxValue);
+                        if (_contentBytesRemaining == 0)
+                        {
+                            // Dispose of the registration and then check whether cancellation has been
+                            // requested. This is necessary to make determinstic a race condition between
+                            // cancellation being requested and unregistering from the token.  Otherwise,
+                            // it's possible cancellation could be requested just before we unregister and
+                            // we then return a connection to the pool that has been or will be disposed
+                            // (e.g. if a timer is used and has already queued its callback but the
+                            // callback hasn't yet run).
+                            ctr.Dispose();
+                            ctr.Token.ThrowIfCancellationRequested();
+
+                            Finish();
+                            return true;
+                        }
                     }
                 }
+                finally
+                {
+                    ctr.Dispose();
+                    cts?.Dispose();
+                }
             }
         }
     }

src/System.Net.Http/src/System/Net/Http/SocketsHttpHandler/HttpConnectionSettings.cs

@@ -27,6 +27,7 @@ internal sealed class HttpConnectionSettings
 
         internal int _maxConnectionsPerServer = HttpHandlerDefaults.DefaultMaxConnectionsPerServer;
         internal int _maxResponseDrainSize = HttpHandlerDefaults.DefaultMaxResponseDrainSize;
+        internal TimeSpan _maxResponseDrainTime = HttpHandlerDefaults.DefaultMaxResponseDrainTime;
         internal int _maxResponseHeadersLength = HttpHandlerDefaults.DefaultMaxResponseHeadersLength;
 
         internal TimeSpan _pooledConnectionLifetime = HttpHandlerDefaults.DefaultPooledConnectionLifetime;
@@ -58,6 +59,7 @@ public HttpConnectionSettings Clone()
                 _maxAutomaticRedirections = _maxAutomaticRedirections,
                 _maxConnectionsPerServer = _maxConnectionsPerServer,
                 _maxResponseDrainSize = _maxResponseDrainSize,
+                _maxResponseDrainTime = _maxResponseDrainTime,
                 _maxResponseHeadersLength = _maxResponseHeadersLength,
                 _pooledConnectionLifetime = _pooledConnectionLifetime,
                 _pooledConnectionIdleTimeout = _pooledConnectionIdleTimeout,

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

@@ -167,6 +167,22 @@ public int MaxResponseDrainSize
             }
         }
 
+        internal TimeSpan MaxResponseDrainTime // TODO #27685: Expose publicly.
+        {
+            get => _settings._maxResponseDrainTime;
+            set
+            {
+                if ((value < TimeSpan.Zero && value != Timeout.InfiniteTimeSpan) ||
+                    (value.TotalMilliseconds > int.MaxValue))
+                {
+                    throw new ArgumentOutOfRangeException(nameof(value));
+                }
+
+                CheckDisposedOrStarted();
+                _settings._maxResponseDrainTime = value;
+            }
+        }
+
         public int MaxResponseHeadersLength
         {
             get => _settings._maxResponseHeadersLength;

src/System.Net.Http/tests/FunctionalTests/HttpClientHandlerTest.ResponseDrain.cs

@@ -5,13 +5,16 @@
 using System.IO;
 using System.Net.Test.Common;
 using System.Text;
+using System.Threading;
 using System.Threading.Tasks;
 using Xunit;
 
 namespace System.Net.Http.Functional.Tests
 {
     public class HttpClientHandler_ResponseDrain_Test : HttpClientTestBase
     {
+        protected virtual void SetMaxResponseDrainTime(HttpClientHandler handler, TimeSpan time) { }
+
         public enum ContentMode
         {
             ContentLength,
@@ -131,6 +134,7 @@ await LoopbackServer.CreateClientAndServerAsync(
                 async url =>
                 {
                     HttpClientHandler handler = CreateHttpClientHandler();
+                    SetMaxResponseDrainTime(handler, Timeout.InfiniteTimeSpan);
 
                     // Set MaxConnectionsPerServer to 1.  This will ensure we will wait for the previous request to drain (or fail to)
                     handler.MaxConnectionsPerServer = 1;
@@ -208,6 +212,7 @@ await LoopbackServer.CreateClientAndServerAsync(
                 async url =>
                 {
                     HttpClientHandler handler = CreateHttpClientHandler();
+                    SetMaxResponseDrainTime(handler, Timeout.InfiniteTimeSpan);
 
                     // Set MaxConnectionsPerServer to 1.  This will ensure we will wait for the previous request to drain (or fail to)
                     handler.MaxConnectionsPerServer = 1;

src/System.Net.Http/tests/FunctionalTests/HttpClientTestBase.cs

@@ -42,21 +42,13 @@ protected static HttpClientHandler CreateHttpClientHandler(bool useSocketsHttpHa
             return handler;
         }
 
-        protected static bool IsSocketsHttpHandler(HttpClientHandler handler)
+        protected static bool IsSocketsHttpHandler(HttpClientHandler handler) =>
+            GetUnderlyingSocketsHttpHandler(handler) != null;
+
+        protected static object GetUnderlyingSocketsHttpHandler(HttpClientHandler handler)
         {
             FieldInfo field = typeof(HttpClientHandler).GetField("_socketsHttpHandler", BindingFlags.Instance | BindingFlags.NonPublic);
-            if (field == null)
-            {
-                return false;
-            }
-
-            object socketsHttpHandler = field.GetValue(handler);
-            if (socketsHttpHandler == null)
-            {
-                return false;
-            }
-
-            return true;
+            return field?.GetValue(handler);
         }
     }
 }