Handle TaskCanceledExceptions for all clients

src/Common/Commands.Common.Authentication.Test/ClientFactoryHandlerTests.cs

@@ -24,6 +24,10 @@
 using Microsoft.WindowsAzure.Commands.ScenarioTest;
 using Xunit;
 using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
+using System.Linq;
+using Microsoft.Azure.Management.Internal.Resources;
+using System.Threading;
+using System.Threading.Tasks;
 
 namespace Common.Authentication.Test
 {
@@ -51,7 +55,7 @@ public void DelegatingHandlersAreCloned()
             sub.SetTenant("common");
             AzureContext context = new AzureContext
             (
-                sub, 
+                sub,
                 account,
                 AzureEnvironment.PublicEnvironments["AzureCloud"]
             );
@@ -65,7 +69,66 @@ public void DelegatingHandlersAreCloned()
             client = factory.CreateClient<StorageManagementClient>(context, AzureEnvironment.Endpoint.ServiceManagement);
             client = factory.CreateClient<StorageManagementClient>(context, AzureEnvironment.Endpoint.ServiceManagement);
             client = factory.CreateClient<StorageManagementClient>(context, AzureEnvironment.Endpoint.ServiceManagement);
-            Assert.Equal(5, MockDelegatingHandler.cloneCount); 
+            Assert.Equal(5, MockDelegatingHandler.cloneCount);
+        }
+
+        [Fact]
+        [Trait(Category.AcceptanceType, Category.CheckIn)]
+        public void AddsAppropriateRetryPolicy()
+        {
+            AzureSessionInitializer.InitializeAzureSession();
+            string userAccount = "user@contoso.com";
+            Guid subscriptionId = Guid.NewGuid();
+            var account = new AzureAccount()
+            {
+                Id = userAccount,
+                Type = AzureAccount.AccountType.User,
+            };
+            account.SetTenants("common");
+            var sub = new AzureSubscription()
+            {
+                Id = subscriptionId.ToString(),
+            };
+            sub.SetAccount(userAccount);
+            sub.SetEnvironment("AzureCloud");
+            sub.SetTenant("common");
+            AzureContext context = new AzureContext
+            (
+                sub,
+                account,
+                AzureEnvironment.PublicEnvironments["AzureCloud"]
+            );
+            AzureSession.Instance.AuthenticationFactory = new MockTokenAuthenticationFactory(userAccount, Guid.NewGuid().ToString());
+            var factory = new ClientFactory();
+            factory.AddHandler(new RetryTestHandler());
+            var client = factory.CreateClient<StorageManagementClient>(context, AzureEnvironment.Endpoint.ServiceManagement);
+            var hyakHandler = EnsureHyakRetryPolicy(client);
+            hyakHandler.MaxTries = 2;
+            Assert.Throws<InvalidOperationException>(() => client.StorageAccounts.List());
+            hyakHandler.MaxTries = 0;
+            Assert.Throws<TaskCanceledException>(() => client.StorageAccounts.List());
+            var autorestClient = factory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager);
+            var autoRestHandler = EnsureAutoRestRetryPolicy(autorestClient);
+            autoRestHandler.MaxTries = 2;
+            var task = autorestClient.ResourceGroups.ListWithHttpMessagesAsync();
+            Assert.Throws<InvalidOperationException>(() => task.ConfigureAwait(false).GetAwaiter().GetResult());
+            autoRestHandler.MaxTries = 0;
+            task = autorestClient.ResourceGroups.ListWithHttpMessagesAsync();
+            Assert.Throws<TaskCanceledException>(() => task.ConfigureAwait(false).GetAwaiter().GetResult());
+        }
+
+        private CancelRetryHandler EnsureHyakRetryPolicy<T>( Hyak.Common.ServiceClient<T> client) where T: Hyak.Common.ServiceClient<T>
+        {
+            var handler = client.GetHttpPipeline().First(h => h.GetType() == typeof(CancelRetryHandler)) as CancelRetryHandler;
+            handler.WaitInterval = TimeSpan.Zero;
+            return handler;
+        }
+
+        private CancelRetryHandler EnsureAutoRestRetryPolicy<T>( Microsoft.Rest.ServiceClient<T> client) where T: Microsoft.Rest.ServiceClient<T>
+        {
+            var handler = client.HttpMessageHandlers.First(h => h is CancelRetryHandler)as CancelRetryHandler;
+            handler.WaitInterval = TimeSpan.Zero;
+            return handler;
         }
 
         private class MockDelegatingHandler : DelegatingHandler, ICloneable
@@ -78,5 +141,28 @@ public object Clone()
                 return this;
             }
         }
+
+
+        private class RetryTestHandler : DelegatingHandler, ICloneable
+        {
+            private static int times = -1;
+
+            protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+            {
+                switch (Interlocked.Increment(ref times) % 3)
+                {
+                    case 0:
+                    case 1:
+                        throw new TaskCanceledException();
+                    default:
+                        throw new InvalidOperationException();
+                }
+            }
+
+            public object Clone()
+            {
+                return new RetryTestHandler();
+            }
+        }
     }
 }

src/Common/Commands.Common.Authentication.Test/Commands.Common.Authentication.Test.csproj

@@ -1,7 +1,7 @@
 <?xml version="1.0" encoding="utf-8"?>
 <Project ToolsVersion="12.0" DefaultTargets="Build" xmlns="http://schemas.microsoft.com/developer/msbuild/2003">
   <Import Project="..\..\packages\xunit.runner.visualstudio.2.1.0\build\net20\xunit.runner.visualstudio.props" Condition="Exists('..\..\packages\xunit.runner.visualstudio.2.1.0\build\net20\xunit.runner.visualstudio.props')" />
-  <Import Project="..\..\packages\xunit.core.2.1.0\build\portable-net45+win8+wp8+wpa81\xunit.core.props"  Condition="Exists('..\..\packages\xunit.core.2.1.0\build\portable-net45+win8+wp8+wpa81\xunit.core.props')"/>
+  <Import Project="..\..\packages\xunit.core.2.1.0\build\portable-net45+win8+wp8+wpa81\xunit.core.props" Condition="Exists('..\..\packages\xunit.core.2.1.0\build\portable-net45+win8+wp8+wpa81\xunit.core.props')" />
   <Import Project="$(MSBuildExtensionsPath)\$(MSBuildToolsVersion)\Microsoft.Common.props" Condition="Exists('$(MSBuildExtensionsPath)\$(MSBuildToolsVersion)\Microsoft.Common.props')" />
   <PropertyGroup>
     <Configuration Condition=" '$(Configuration)' == '' ">Debug</Configuration>
@@ -136,6 +136,10 @@
     <EmbeddedResource Include="Resources\ResourceLocator.cs" />
   </ItemGroup>
   <ItemGroup>
+    <ProjectReference Include="..\..\ResourceManager\Common\Commands.ResourceManager.Common\Commands.ResourceManager.Common.csproj">
+      <Project>{3819d8a7-c62c-4c47-8ddd-0332d9ce1252}</Project>
+      <Name>Commands.ResourceManager.Common</Name>
+    </ProjectReference>
     <ProjectReference Include="..\Commands.Common.Authentication.Abstractions\Commands.Common.Authentication.Abstractions.csproj">
       <Project>{70527617-7598-4aef-b5bd-db9186b8184b}</Project>
       <Name>Commands.Common.Authentication.Abstractions</Name>

src/Common/Commands.Common.Authentication/Commands.Common.Authentication.csproj

@@ -137,6 +137,7 @@
     <Compile Include="ContextAutosaveSettings.cs" />
     <Compile Include="Extensions\CloudExceptionExtensions.cs" />
     <Compile Include="Factories\AuthenticationFactory.cs" />
+    <Compile Include="Factories\CancelRetryHandler.cs" />
     <Compile Include="Factories\ClientFactory.cs" />
     <Compile Include="HttpClientOperationsFactory.cs" />
     <Compile Include="Utilities\HttpClientWithRetry.cs" />

src/Common/Commands.Common.Authentication/Factories/CancelRetryHandler.cs

@@ -0,0 +1,80 @@
+// ----------------------------------------------------------------------------------
+//
+// Copyright Microsoft Corporation
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+// http://www.apache.org/licenses/LICENSE-2.0
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// ----------------------------------------------------------------------------------
+
+using System;
+using System.Linq;
+using System.Net.Http;
+using System.Threading;
+using System.Threading.Tasks;
+
+namespace Microsoft.Azure.Commands.Common.Authentication.Factories
+{
+    /// <summary>
+    /// Class to enable automatic retry of TaskCanceledExsceptions.
+    /// Note that this follows the calling pattern in PowerShell, which does not pass cancellation tokens from
+    /// the base cmdlet
+    /// </summary>
+    public class CancelRetryHandler : DelegatingHandler, ICloneable
+    {
+        static readonly string[] TestModeLiveValues = { "Record", "None", "Live" };
+        const string TestModeVariableKey = "Azure_Test_Mode";
+        public CancelRetryHandler()
+        {
+            WaitInterval = TimeSpan.Zero;
+        }
+        public CancelRetryHandler(TimeSpan waitInterval, int maxTries)
+        {
+            WaitInterval = waitInterval;
+            MaxTries = maxTries;
+        }
+
+        public TimeSpan WaitInterval { get; set; } = TimeSpan.Zero;
+
+        public int MaxTries { get; set; } = 3;
+
+        protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
+        {
+            int tries = 0;
+            do
+            {
+                using (var source = new CancellationTokenSource())
+                {
+                    try
+                    {
+                        return await base.SendAsync(request, source.Token).ConfigureAwait(false);
+                    }
+                    catch (TaskCanceledException) when (tries++ < MaxTries)
+                    {
+                        SleepIfRunningLive(WaitInterval);
+                    }
+                }
+            }
+            while (true);
+        }
+
+        void SleepIfRunningLive(TimeSpan timeSpan)
+        {
+            var testModeSetting = Environment.GetEnvironmentVariable(TestModeVariableKey);
+            if (string.IsNullOrWhiteSpace(testModeSetting) || TestModeLiveValues.Any( v => string.Equals(v, testModeSetting, StringComparison.OrdinalIgnoreCase)))
+            {
+                Thread.Sleep(timeSpan);
+            }
+        }
+
+        public object Clone()
+        {
+            return new CancelRetryHandler(WaitInterval, MaxTries);
+        }
+    }
+}

src/Common/Commands.Common.Authentication/Factories/ClientFactory.cs

@@ -32,6 +32,7 @@ namespace Microsoft.Azure.Commands.Common.Authentication.Factories
     public class ClientFactory : IClientFactory
     {
         private static readonly char[] uriPathSeparator = { '/' };
+        private static readonly CancelRetryHandler DefaultCancelRetryHandler  = new CancelRetryHandler(TimeSpan.FromSeconds(10), 3);
 
         private Dictionary<Type, IClientAction> _actions;
         private OrderedDictionary _handlers;
@@ -56,11 +57,8 @@ public virtual TClient CreateArmClient<TClient>(IAzureContext context, string en
             }
 
             var creds = AzureSession.Instance.AuthenticationFactory.GetServiceClientCredentials(context, endpoint);
-            var newHandlers = GetCustomHandlers();
-            TClient client = (newHandlers == null || newHandlers.Length == 0)
-                ? CreateCustomArmClient<TClient>(context.Environment.GetEndpointAsUri(endpoint), creds)
-                : CreateCustomArmClient<TClient>(context.Environment.GetEndpointAsUri(endpoint), creds, GetCustomHandlers());
-
+            var baseUri = context.Environment.GetEndpointAsUri(endpoint);
+            TClient client = CreateCustomArmClient<TClient>(baseUri, creds);
             var subscriptionId = typeof(TClient).GetProperty("SubscriptionId");
             if (subscriptionId != null && context.Subscription != null)
             {
@@ -73,23 +71,35 @@ public virtual TClient CreateArmClient<TClient>(IAzureContext context, string en
         public virtual TClient CreateCustomArmClient<TClient>(params object[] parameters) where TClient : Microsoft.Rest.ServiceClient<TClient>
         {
             List<Type> types = new List<Type>();
-            var containsDelegatingHandler = false;
+            List<object> parameterList = new List<object>();
+            List<DelegatingHandler> handlerList = new List<DelegatingHandler> { DefaultCancelRetryHandler.Clone() as CancelRetryHandler};
+            var customHandlers = GetCustomHandlers();
+            if (customHandlers != null && customHandlers.Any())
+            {
+                handlerList.AddRange(customHandlers);
+            }
+
+            bool hasHandlers = false;
             foreach (object obj in parameters)
             {
-                if (obj.GetType().Name.Equals("DelegatingHandler[]"))
+                var paramType = obj.GetType();
+                types.Add(paramType);
+                if (paramType == typeof(DelegatingHandler[]))
                 {
-                    containsDelegatingHandler = true;
+                    handlerList.AddRange(obj as DelegatingHandler[]);
+                    parameterList.Add(handlerList.ToArray());
+                    hasHandlers = true;
+                }
+                else
+                {
+                    parameterList.Add(obj);
                 }
-                types.Add(obj.GetType());
             }
 
-            var copiedParameters = parameters;
-            if (!containsDelegatingHandler)
+            if (!hasHandlers)
             {
                 types.Add((new DelegatingHandler[0]).GetType());
-                var parameterList = copiedParameters.ToList();
-                parameterList.Add(new DelegatingHandler[0]);
-                copiedParameters = parameterList.ToArray();
+                parameterList.Add(handlerList.ToArray());
             }
 
             var constructor = typeof(TClient).GetConstructor(types.ToArray());
@@ -99,7 +109,7 @@ public virtual TClient CreateCustomArmClient<TClient>(params object[] parameters
                 throw new InvalidOperationException(string.Format(Resources.InvalidManagementClientType, typeof(TClient).Name));
             }
 
-            TClient client = (TClient)constructor.Invoke(copiedParameters);
+            TClient client = (TClient)constructor.Invoke(parameterList.ToArray());
 
             foreach (ProductInfoHeaderValue userAgent in _userAgents.Keys)
             {
@@ -120,13 +130,7 @@ public virtual TClient CreateClient<TClient>(IAzureContext context, string endpo
             }
 
             SubscriptionCloudCredentials creds = AzureSession.Instance.AuthenticationFactory.GetSubscriptionCloudCredentials(context, endpoint);
-            TClient client = CreateCustomClient<TClient>(creds, context.Environment.GetEndpointAsUri(endpoint));
-            foreach (DelegatingHandler handler in GetCustomHandlers())
-            {
-                client.AddHandlerToPipeline(handler);
-            }
-
-            return client;
+            return CreateCustomClient<TClient>(creds, context.Environment.GetEndpointAsUri(endpoint));
         }
 
         public virtual TClient CreateClient<TClient>(IAzureContextContainer container, string endpoint) where TClient : ServiceClient<TClient>
@@ -177,7 +181,7 @@ public virtual TClient CreateClient<TClient>(IAzureContextContainer profile, IAz
 
             AzureContext context = new AzureContext(subscription, account, environment);
 
-            TClient client = CreateClient<TClient>(context, endpoint);
+            var client = CreateClient<TClient>(context, endpoint);
 
             foreach (IClientAction action in GetActions())
             {
@@ -209,6 +213,12 @@ public virtual TClient CreateCustomClient<TClient>(params object[] parameters) w
                 client.UserAgent.Add(userAgent);
             }
 
+            foreach (DelegatingHandler handler in GetCustomHandlers())
+            {
+                client.AddHandlerToPipeline(handler);
+            }
+
+            client.AddHandlerToPipeline(DefaultCancelRetryHandler.Clone() as CancelRetryHandler);
             return client;
         }
 

src/ResourceManager/Profile/ChangeLog.md

@@ -25,6 +25,7 @@
   - ParameterSetName
   - InvocationName
 * Fix issue where version 10.0.3 of Newtonsoft.Json wasn't being loaded on module import
+* Retry TaskCanceledException appropriately when connections are left open.
 
 ## Version 5.1.0
 * Fix issue where running `Clear-AzureRmContext` would keep an empty context with the name of the previous default context, which prevented the user from creating a new context with the old name