Updating Cancel handler to create new cancellation tokens

Author: markcowl

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

@@ -19,18 +19,24 @@
 
 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
     {
         public CancelRetryHandler()
         {
             WaitInterval = TimeSpan.Zero;
         }
-        public CancelRetryHandler(TimeSpan waitInterval)
+        public CancelRetryHandler(TimeSpan waitInterval, int maxTries)
         {
             WaitInterval = waitInterval;
+            MaxTries = maxTries;
         }
 
-        public TimeSpan WaitInterval { get; set; }
+        public TimeSpan WaitInterval { get; set; } = TimeSpan.Zero;
 
         public int MaxTries { get; set; } = 3;
 
@@ -39,21 +45,24 @@ protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage
             int tries = 0;
             do
             {
-                try
+                using (var source = new CancellationTokenSource())
                 {
-                    return await base.SendAsync(request, cancellationToken).ConfigureAwait(false);
-                }
-                catch (TaskCanceledException) when (tries++ < MaxTries)
-                {
-                    Thread.Sleep(WaitInterval);
+                    try
+                    {
+                        return await base.SendAsync(request, source.Token).ConfigureAwait(false);
+                    }
+                    catch (TaskCanceledException) when (tries++ < MaxTries)
+                    {
+                        Thread.Sleep(WaitInterval);
+                    }
                 }
             }
             while (true);
         }
 
         public object Clone()
         {
-            return new CancelRetryHandler(WaitInterval);
+            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;
@@ -71,7 +72,7 @@ public virtual TClient CreateCustomArmClient<TClient>(params object[] parameters
         {
             List<Type> types = new List<Type>();
             List<object> parameterList = new List<object>();
-            List<DelegatingHandler> handlerList = new List<DelegatingHandler> { new CancelRetryHandler(TimeSpan.FromSeconds(10))};
+            List<DelegatingHandler> handlerList = new List<DelegatingHandler> { DefaultCancelRetryHandler.Clone() as CancelRetryHandler};
             var customHandlers = GetCustomHandlers();
             if (customHandlers != null && customHandlers.Any())
             {
@@ -218,7 +219,7 @@ public virtual TClient CreateCustomClient<TClient>(params object[] parameters) w
                 client.AddHandlerToPipeline(handler);
             }
 
-            client.AddHandlerToPipeline(new CancelRetryHandler(TimeSpan.FromSeconds(10)));
+            client.AddHandlerToPipeline(DefaultCancelRetryHandler.Clone() as CancelRetryHandler);
             return client;
         }