Fix cancellation of WinHttpHandler response stream reads

Reading from the response stream using ReadAsync() was ignoring the cancellation token passed in. To fix this, registered a delegate that will be called to dispose the WinHTTP request handle during pending I/O. This will cause the WinHTTP operation to cancel.

Added tests. For now, left them disabled to be resolved in issue dotnet/corefx#8692 next week. But wanted to get this fix in asap since it blocks partners for RTM.

Fixes dotnet/corefx#8662.


Commit migrated from dotnet/corefx@a6e019d

src/libraries/System.Net.Http.WinHttpHandler/src/System/Net/Http/WinHttpResponseStream.cs

@@ -17,6 +17,7 @@ internal class WinHttpResponseStream : Stream
         private volatile bool _disposed;
         private readonly WinHttpRequestState _state;
         private SafeWinHttpHandle _requestHandle;
+        private CancellationTokenRegistration _ctr;
 
         internal WinHttpResponseStream(SafeWinHttpHandle requestHandle, WinHttpRequestState state)
         {
@@ -150,6 +151,7 @@ public override Task<int> ReadAsync(byte[] buffer, int offset, int count, Cancel
 
                         lock (_state.Lock)
                         {
+                            Debug.Assert(!_requestHandle.IsInvalid);
                             if (!Interop.WinHttp.WinHttpReadData(
                                 _requestHandle,
                                 Marshal.UnsafeAddrOfPinnedArrayElement(buffer, offset),
@@ -164,10 +166,20 @@ public override Task<int> ReadAsync(byte[] buffer, int offset, int count, Cancel
                 }, 
                 CancellationToken.None, TaskContinuationOptions.None, TaskScheduler.Default);
 
-            // TODO: Issue #2165. Register callback on cancellation token to cancel WinHTTP operation.
-
+            // Register callback on cancellation token to cancel any pending WinHTTP operation.
+            if (token != CancellationToken.None)
+            {
+                WinHttpTraceHelper.Trace("WinHttpResponseStream.ReadAsync: registering for cancellation token request");
+                _ctr = token.Register(() => CancelPendingResponseStreamReadOperation());
+            }
+            else
+            {
+                WinHttpTraceHelper.Trace("WinHttpResponseStream.ReadAsync: received no cancellation token");
+            }
+
             lock (_state.Lock)
             {
+                Debug.Assert(!_requestHandle.IsInvalid);
                 if (!Interop.WinHttp.WinHttpQueryDataAvailable(_requestHandle, IntPtr.Zero))
                 {
                     _state.TcsReadFromResponseStream.TrySetException(
@@ -227,5 +239,33 @@ private void CheckDisposed()
                 throw new ObjectDisposedException(this.GetType().FullName);
             }
         }
+
+        // The only way to abort pending async operations in WinHTTP is to close the request handle.
+        // This causes WinHTTP to cancel any pending I/O and accelerating its callbacks on the handle.
+        // This causes our related TaskCompletionSource objects to move to a terminal state.
+        //
+        // We only want to dispose the handle if we are actually waiting for a pending WinHTTP I/O to complete,
+        // meaning that we are await'ing for a Task to complete. While we could simply call dispose without
+        // a pending operation, it would cause random failures in the other threads when we expect a valid handle.
+        private void CancelPendingResponseStreamReadOperation()
+        {
+            WinHttpTraceHelper.Trace("WinHttpResponseStream.CancelPendingResponseStreamReadOperation");
+            lock (_state.Lock)
+            {
+                WinHttpTraceHelper.Trace("WinHttpResponseStream.CancelPendingResponseStreamReadOperation: in lock");
+                WinHttpTraceHelper.Trace(
+                    string.Format("WinHttpResponseStream.CancelPendingResponseStreamReadOperation: {0} {1}",
+                    (int)_state.TcsQueryDataAvailable.Task.Status, (int)_state.TcsReadFromResponseStream.Task.Status));
+                if (!_state.TcsQueryDataAvailable.Task.IsCompleted)
+                {
+                    Debug.Assert(_requestHandle != null);
+                    Debug.Assert(!_requestHandle.IsInvalid);
+
+                    WinHttpTraceHelper.Trace("WinHttpResponseStream.CancelPendingResponseStreamReadOperation: before dispose");
+                    _requestHandle.Dispose();
+                    WinHttpTraceHelper.Trace("WinHttpResponseStream.CancelPendingResponseStreamReadOperation: after dispose");
+                }
+            }
+        }        
     }
 }

src/libraries/System.Net.Http/tests/FunctionalTests/CancellationTest.cs

@@ -0,0 +1,125 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Diagnostics;
+using System.IO;
+using System.Net.Test.Common;
+using System.Threading;
+using System.Threading.Tasks;
+
+using Xunit;
+using Xunit.Abstractions;
+
+namespace System.Net.Http.Functional.Tests
+{
+    public class CancellationTest
+    {
+        // TODO: Issue #8692. Move this test server capability to Azure test server or Loopback server.
+        private const string FastHeadersSlowBodyHost = "winrtnet.corp.microsoft.com";
+        private const int FastHeadersSlowBodyPort = 1337;
+        private const int ResponseBodyReadDelayInMilliseconds = 15000; // 15 seconds.
+        private const int ResponseBodyLength = 1024;
+
+        private static Uri s_fastHeadersSlowBodyServer = new Uri(string.Format(
+            "http://{0}:{1}/?slow={2}&length={3}",
+            FastHeadersSlowBodyHost,
+            FastHeadersSlowBodyPort,
+            ResponseBodyReadDelayInMilliseconds,
+            ResponseBodyLength));
+
+        private readonly ITestOutputHelper _output;
+
+        public CancellationTest(ITestOutputHelper output)
+        {
+            _output = output;
+            _output.WriteLine(s_fastHeadersSlowBodyServer.ToString());
+        }
+
+        [ActiveIssue(8663)]
+        [Fact]
+        public async Task GetIncludesReadingResponseBody_CancelUsingTimeout_TaskCanceledQuickly()
+        {
+            var stopwatch = new Stopwatch();
+            stopwatch.Start();
+
+            using (var client = new HttpClient())
+            {
+                client.Timeout = new TimeSpan(0, 0, 1);
+
+                Task<HttpResponseMessage> getResponse =
+                    client.GetAsync(s_fastHeadersSlowBodyServer, HttpCompletionOption.ResponseContentRead);
+
+                stopwatch.Restart();
+                await Assert.ThrowsAsync<TaskCanceledException>(
+                    () => getResponse);
+                stopwatch.Stop();
+                _output.WriteLine("GetAsync() completed at: {0}", stopwatch.Elapsed.ToString());
+                Assert.True(stopwatch.Elapsed < new TimeSpan(0,0,3), "Elapsed time should be short");
+            }
+        }
+
+        [ActiveIssue(8663)]
+        [Fact]
+        public async Task GetIncludesReadingResponseBody_CancelUsingToken_TaskCanceledQuickly()
+        {
+            var stopwatch = new Stopwatch();
+            stopwatch.Start();
+
+            var cts = new CancellationTokenSource();
+            using (var client = new HttpClient())
+            {
+                Task<HttpResponseMessage> getResponse =
+                    client.GetAsync(s_fastHeadersSlowBodyServer, HttpCompletionOption.ResponseContentRead, cts.Token);
+
+                Task ignore = Task.Delay(new TimeSpan(0, 0, 1)).ContinueWith(_ =>
+                {
+                    _output.WriteLine("Calling cts.Cancel() at: {0}", stopwatch.Elapsed.ToString());
+                    cts.Cancel();
+                });
+
+                stopwatch.Restart();
+                await Assert.ThrowsAsync<TaskCanceledException>(
+                    () => getResponse);
+                stopwatch.Stop();
+                _output.WriteLine("GetAsync() completed at: {0}", stopwatch.Elapsed.ToString());
+                Assert.True(stopwatch.Elapsed < new TimeSpan(0,0,3), "Elapsed time should be short");
+            }
+        }
+
+        [ActiveIssue(8692)]
+        [Fact]
+        public async Task ResponseStreamRead_CancelUsingToken_TaskCanceledQuickly()
+        {
+            var stopwatch = new Stopwatch();
+
+            stopwatch.Start();
+            using (var client = new HttpClient())
+            using (HttpResponseMessage response =
+                await client.GetAsync(s_fastHeadersSlowBodyServer, HttpCompletionOption.ResponseHeadersRead))
+            {
+                stopwatch.Stop();
+                _output.WriteLine("Time to get headers: {0}", stopwatch.Elapsed.ToString());
+                Assert.Equal(HttpStatusCode.OK, response.StatusCode);
+
+                var cts = new CancellationTokenSource();
+
+                Stream stream = await response.Content.ReadAsStreamAsync();
+                byte[] buffer = new byte[ResponseBodyLength];
+
+                Task ignore = Task.Delay(new TimeSpan(0,0,1)).ContinueWith(_ =>
+                {
+                    _output.WriteLine("Calling cts.Cancel() at: {0}", stopwatch.Elapsed.ToString());
+                    cts.Cancel();
+                });
+
+                stopwatch.Restart();
+                await Assert.ThrowsAsync<TaskCanceledException>(
+                    () => stream.ReadAsync(buffer, 0, buffer.Length, cts.Token));
+                stopwatch.Stop();
+                _output.WriteLine("ReadAsync() completed at: {0}", stopwatch.Elapsed.ToString());
+                Assert.True(stopwatch.Elapsed < new TimeSpan(0,0,3), "Elapsed time should be short");
+            }
+        }
+    }
+}

src/libraries/System.Net.Http/tests/FunctionalTests/System.Net.Http.Functional.Tests.csproj

@@ -55,6 +55,7 @@
       <Link>Common\System\Net\Http\LoopbackServer.cs</Link>
     </Compile>
     <Compile Include="ByteArrayContentTest.cs" />
+    <Compile Include="CancellationTest.cs" />
     <Compile Include="ChannelBindingAwareContent.cs" />
     <Compile Include="CustomContent.cs" />
     <Compile Include="DelegatingHandlerTest.cs" />