Add response read timeouts

sdk/core/Azure.Core/CHANGELOG.md

@@ -1,8 +1,10 @@
 # Release History
 
 ## 1.1.0-preview.1 (Unreleased)
+
 ### Fixes and improvements
 - Add OPTIONS and TRACE HTTP request methods.
+- Add `TryTimeout` property to `RetryOptions` and apply if to stream reads.
 
 ## 1.0.2 (2020-01-10)
 

sdk/core/Azure.Core/src/Pipeline/HttpClientTransport.cs

@@ -81,7 +81,10 @@ private static HttpClient CreateDefaultClient()
                 httpClientHandler.Proxy = webProxy;
             }
 
-            return new HttpClient(httpClientHandler);
+            return new HttpClient(httpClientHandler)
+            {
+                Timeout = Timeout.InfiniteTimeSpan
+            };
         }
 
         private static HttpRequestMessage BuildRequestMessage(HttpMessage message)

sdk/core/Azure.Core/src/Pipeline/HttpPipelineBuilder.cs

@@ -63,6 +63,8 @@ public static HttpPipeline Build(ClientOptions options, HttpPipelinePolicy[] per
             RetryOptions retryOptions = options.Retry;
             policies.Add(new RetryPolicy(retryOptions.Mode, retryOptions.Delay, retryOptions.MaxDelay, retryOptions.MaxRetries));
 
+            policies.Add(new OperationTimeoutPolicy(options.Retry.TryTimeout));
+
             policies.AddRange(perRetryPolicies);
 
             policies.AddRange(options.PerRetryPolicies);
@@ -73,7 +75,7 @@ public static HttpPipeline Build(ClientOptions options, HttpPipelinePolicy[] per
                     diagnostics.LoggedHeaderNames.ToArray(), diagnostics.LoggedQueryParameters.ToArray()));
             }
 
-            policies.Add(BufferResponsePolicy.Shared);
+            policies.Add(new ResponseBodyPolicy(options.Retry.TryTimeout));
 
             policies.Add(new RequestActivityPolicy(isDistributedTracingEnabled, ClientDiagnostics.GetResourceProviderNamespace(options.GetType().Assembly)));
 

sdk/core/Azure.Core/src/Pipeline/Internal/ReadTimeoutStream.cs

@@ -0,0 +1,123 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Azure.Core.Pipeline
+{
+    internal class ReadTimeoutStream : ReadOnlyStream
+    {
+        private readonly Stream _stream;
+        private TimeSpan _readTimeout;
+        private CancellationTokenSource _cancellationTokenSource;
+
+        public ReadTimeoutStream(Stream stream, TimeSpan readTimeout)
+        {
+            _stream = stream;
+            _readTimeout = readTimeout;
+            UpdateReadTimeout();
+            _cancellationTokenSource = new CancellationTokenSource();
+        }
+
+        public override int Read(byte[] buffer, int offset, int count)
+        {
+            return _stream.Read(buffer, offset, count);
+        }
+
+        public override async Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)
+        {
+            var source = StartTimeout(cancellationToken, out bool dispose);
+            try
+            {
+                source.Token.Register(() =>
+                {
+                    Console.WriteLine("Cancelled");
+                });
+                return await _stream.ReadAsync(buffer, offset, count, source.Token).ConfigureAwait(false);
+            }
+            finally
+            {
+                StopTimeout(source, dispose);
+            }
+        }
+
+        private CancellationTokenSource StartTimeout(CancellationToken additionalToken, out bool dispose)
+        {
+            if (_cancellationTokenSource.IsCancellationRequested)
+            {
+                _cancellationTokenSource = new CancellationTokenSource();
+            }
+
+            CancellationTokenSource source;
+            if (additionalToken.CanBeCanceled)
+            {
+                source = CancellationTokenSource.CreateLinkedTokenSource(additionalToken, _cancellationTokenSource.Token);
+                dispose = true;
+            }
+            else
+            {
+                source = _cancellationTokenSource;
+                dispose = false;
+            }
+
+            _cancellationTokenSource.CancelAfter(_readTimeout);
+
+            return source;
+        }
+
+        private void StopTimeout(CancellationTokenSource source, bool dispose)
+        {
+            _cancellationTokenSource.CancelAfter(Timeout.InfiniteTimeSpan);
+            if (dispose)
+            {
+                source.Dispose();
+            }
+        }
+
+        public override long Seek(long offset, SeekOrigin origin)
+        {
+            return _stream.Seek(offset, origin);
+        }
+
+        public override bool CanRead => _stream.CanRead;
+        public override bool CanSeek => _stream.CanSeek;
+        public override long Length => _stream.Length;
+
+        public override long Position
+        {
+            get => _stream.Position;
+            set => _stream.Position = value;
+        }
+
+        public override int ReadTimeout
+        {
+            get => _readTimeout.Milliseconds;
+            set
+            {
+                _readTimeout = TimeSpan.FromMilliseconds(value);
+                UpdateReadTimeout();
+            }
+        }
+
+        private void UpdateReadTimeout()
+        {
+            try
+            {
+                _stream.ReadTimeout = _readTimeout.Milliseconds;
+            }
+            catch
+            {
+                // ignore
+            }
+        }
+
+        protected override void Dispose(bool disposing)
+        {
+            base.Dispose(disposing);
+            _stream.Dispose();
+        }
+    }
+}

sdk/core/Azure.Core/src/Pipeline/Internal/ResponseBodyPolicy.cs

@@ -0,0 +1,68 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.IO;
+using System.Threading.Tasks;
+
+namespace Azure.Core.Pipeline
+{
+    internal class ResponseBodyPolicy : HttpPipelinePolicy
+    {
+        private readonly TimeSpan _readTimeout;
+
+        public ResponseBodyPolicy(TimeSpan readTimeout)
+        {
+            _readTimeout = readTimeout;
+        }
+
+        public override async ValueTask ProcessAsync(HttpMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline)
+        {
+            await ProcessAsync(message, pipeline, true).ConfigureAwait(false);
+        }
+
+        public override void Process(HttpMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline)
+        {
+            ProcessAsync(message, pipeline, false).EnsureCompleted();
+        }
+
+        private async ValueTask ProcessAsync(HttpMessage message, ReadOnlyMemory<HttpPipelinePolicy> pipeline, bool async)
+        {
+            if (async)
+            {
+                await ProcessNextAsync(message, pipeline).ConfigureAwait(false);
+            }
+            else
+            {
+                ProcessNext(message, pipeline);
+            }
+
+            Stream? responseContentStream = message.Response.ContentStream;
+            if (responseContentStream == null || responseContentStream.CanSeek)
+            {
+                return;
+            }
+
+            if (message.BufferResponse)
+            {
+                var bufferedStream = new MemoryStream();
+                if (async)
+                {
+                    await responseContentStream.CopyToAsync(bufferedStream, message.CancellationToken).ConfigureAwait(false);
+                }
+                else
+                {
+                    responseContentStream.CopyTo(bufferedStream, message.CancellationToken);
+                }
+
+                responseContentStream.Dispose();
+                bufferedStream.Position = 0;
+                message.Response.ContentStream = bufferedStream;
+            }
+            else
+            {
+                message.Response.ContentStream = new ReadTimeoutStream(responseContentStream, _readTimeout);
+            }
+        }
+    }
+}

sdk/core/Azure.Core/src/Pipeline/Internal/StreamExtensions.cs

@@ -0,0 +1,53 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+using System;
+using System.Buffers;
+using System.IO;
+using System.Threading;
+using System.Threading.Tasks;
+using Azure.Core.Buffers;
+
+namespace Azure.Core.Pipeline
+{
+    internal static class StreamExtensions
+    {
+        internal const int DefaultCopyBufferSize = 81920;
+
+        public static async Task CopyToAsync(this Stream source, Stream destination, CancellationToken cancellationToken)
+        {
+            byte[] buffer = ArrayPool<byte>.Shared.Rent(DefaultCopyBufferSize);
+            try
+            {
+                while (true)
+                {
+                    int bytesRead = await source.ReadAsync(buffer, 0, buffer.Length, cancellationToken).ConfigureAwait(false);
+                    if (bytesRead == 0) break;
+                    await destination.WriteAsync(new ReadOnlyMemory<byte>(buffer, 0, bytesRead), cancellationToken).ConfigureAwait(false);
+                }
+            }
+            finally
+            {
+                ArrayPool<byte>.Shared.Return(buffer);
+            }
+        }
+
+        public static void CopyTo(this Stream source, Stream destination, CancellationToken cancellationToken)
+        {
+            byte[] buffer = ArrayPool<byte>.Shared.Rent(DefaultCopyBufferSize);
+            try
+            {
+                int read;
+                while ((read = source.Read(buffer, 0, buffer.Length)) != 0)
+                {
+                    cancellationToken.ThrowIfCancellationRequested();
+                    destination.Write(buffer, 0, read);
+                }
+            }
+            finally
+            {
+                ArrayPool<byte>.Shared.Return(buffer);
+            }
+        }
+    }
+}