Skip to content

Commit e78e977

Browse files
author
Maddie Clayton
authored
Merge pull request Azure#6245 from markcowl/task-cancel
Handle TaskCanceledExceptions for all clients
2 parents 64c9451 + c5bc1ce commit e78e977

File tree

6 files changed

+208
-26
lines changed

6 files changed

+208
-26
lines changed

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

Lines changed: 88 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -24,6 +24,10 @@
2424
using Microsoft.WindowsAzure.Commands.ScenarioTest;
2525
using Xunit;
2626
using Microsoft.Azure.Commands.Common.Authentication.Abstractions;
27+
using System.Linq;
28+
using Microsoft.Azure.Management.Internal.Resources;
29+
using System.Threading;
30+
using System.Threading.Tasks;
2731

2832
namespace Common.Authentication.Test
2933
{
@@ -51,7 +55,7 @@ public void DelegatingHandlersAreCloned()
5155
sub.SetTenant("common");
5256
AzureContext context = new AzureContext
5357
(
54-
sub,
58+
sub,
5559
account,
5660
AzureEnvironment.PublicEnvironments["AzureCloud"]
5761
);
@@ -65,7 +69,66 @@ public void DelegatingHandlersAreCloned()
6569
client = factory.CreateClient<StorageManagementClient>(context, AzureEnvironment.Endpoint.ServiceManagement);
6670
client = factory.CreateClient<StorageManagementClient>(context, AzureEnvironment.Endpoint.ServiceManagement);
6771
client = factory.CreateClient<StorageManagementClient>(context, AzureEnvironment.Endpoint.ServiceManagement);
68-
Assert.Equal(5, MockDelegatingHandler.cloneCount);
72+
Assert.Equal(5, MockDelegatingHandler.cloneCount);
73+
}
74+
75+
[Fact]
76+
[Trait(Category.AcceptanceType, Category.CheckIn)]
77+
public void AddsAppropriateRetryPolicy()
78+
{
79+
AzureSessionInitializer.InitializeAzureSession();
80+
string userAccount = "user@contoso.com";
81+
Guid subscriptionId = Guid.NewGuid();
82+
var account = new AzureAccount()
83+
{
84+
Id = userAccount,
85+
Type = AzureAccount.AccountType.User,
86+
};
87+
account.SetTenants("common");
88+
var sub = new AzureSubscription()
89+
{
90+
Id = subscriptionId.ToString(),
91+
};
92+
sub.SetAccount(userAccount);
93+
sub.SetEnvironment("AzureCloud");
94+
sub.SetTenant("common");
95+
AzureContext context = new AzureContext
96+
(
97+
sub,
98+
account,
99+
AzureEnvironment.PublicEnvironments["AzureCloud"]
100+
);
101+
AzureSession.Instance.AuthenticationFactory = new MockTokenAuthenticationFactory(userAccount, Guid.NewGuid().ToString());
102+
var factory = new ClientFactory();
103+
factory.AddHandler(new RetryTestHandler());
104+
var client = factory.CreateClient<StorageManagementClient>(context, AzureEnvironment.Endpoint.ServiceManagement);
105+
var hyakHandler = EnsureHyakRetryPolicy(client);
106+
hyakHandler.MaxTries = 2;
107+
Assert.Throws<InvalidOperationException>(() => client.StorageAccounts.List());
108+
hyakHandler.MaxTries = 0;
109+
Assert.Throws<TaskCanceledException>(() => client.StorageAccounts.List());
110+
var autorestClient = factory.CreateArmClient<ResourceManagementClient>(context, AzureEnvironment.Endpoint.ResourceManager);
111+
var autoRestHandler = EnsureAutoRestRetryPolicy(autorestClient);
112+
autoRestHandler.MaxTries = 2;
113+
var task = autorestClient.ResourceGroups.ListWithHttpMessagesAsync();
114+
Assert.Throws<InvalidOperationException>(() => task.ConfigureAwait(false).GetAwaiter().GetResult());
115+
autoRestHandler.MaxTries = 0;
116+
task = autorestClient.ResourceGroups.ListWithHttpMessagesAsync();
117+
Assert.Throws<TaskCanceledException>(() => task.ConfigureAwait(false).GetAwaiter().GetResult());
118+
}
119+
120+
private CancelRetryHandler EnsureHyakRetryPolicy<T>( Hyak.Common.ServiceClient<T> client) where T: Hyak.Common.ServiceClient<T>
121+
{
122+
var handler = client.GetHttpPipeline().First(h => h.GetType() == typeof(CancelRetryHandler)) as CancelRetryHandler;
123+
handler.WaitInterval = TimeSpan.Zero;
124+
return handler;
125+
}
126+
127+
private CancelRetryHandler EnsureAutoRestRetryPolicy<T>( Microsoft.Rest.ServiceClient<T> client) where T: Microsoft.Rest.ServiceClient<T>
128+
{
129+
var handler = client.HttpMessageHandlers.First(h => h is CancelRetryHandler)as CancelRetryHandler;
130+
handler.WaitInterval = TimeSpan.Zero;
131+
return handler;
69132
}
70133

71134
private class MockDelegatingHandler : DelegatingHandler, ICloneable
@@ -78,5 +141,28 @@ public object Clone()
78141
return this;
79142
}
80143
}
144+
145+
146+
private class RetryTestHandler : DelegatingHandler, ICloneable
147+
{
148+
private static int times = -1;
149+
150+
protected override Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
151+
{
152+
switch (Interlocked.Increment(ref times) % 3)
153+
{
154+
case 0:
155+
case 1:
156+
throw new TaskCanceledException();
157+
default:
158+
throw new InvalidOperationException();
159+
}
160+
}
161+
162+
public object Clone()
163+
{
164+
return new RetryTestHandler();
165+
}
166+
}
81167
}
82168
}

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

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
<?xml version="1.0" encoding="utf-8"?>
22
<Project ToolsVersion="12.0" DefaultTargets="Build" xmlns="http://schemas.microsoft.com/developer/msbuild/2003">
33
<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')" />
4-
<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')"/>
4+
<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')" />
55
<Import Project="$(MSBuildExtensionsPath)\$(MSBuildToolsVersion)\Microsoft.Common.props" Condition="Exists('$(MSBuildExtensionsPath)\$(MSBuildToolsVersion)\Microsoft.Common.props')" />
66
<PropertyGroup>
77
<Configuration Condition=" '$(Configuration)' == '' ">Debug</Configuration>
@@ -136,6 +136,10 @@
136136
<EmbeddedResource Include="Resources\ResourceLocator.cs" />
137137
</ItemGroup>
138138
<ItemGroup>
139+
<ProjectReference Include="..\..\ResourceManager\Common\Commands.ResourceManager.Common\Commands.ResourceManager.Common.csproj">
140+
<Project>{3819d8a7-c62c-4c47-8ddd-0332d9ce1252}</Project>
141+
<Name>Commands.ResourceManager.Common</Name>
142+
</ProjectReference>
139143
<ProjectReference Include="..\Commands.Common.Authentication.Abstractions\Commands.Common.Authentication.Abstractions.csproj">
140144
<Project>{70527617-7598-4aef-b5bd-db9186b8184b}</Project>
141145
<Name>Commands.Common.Authentication.Abstractions</Name>

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,6 +137,7 @@
137137
<Compile Include="ContextAutosaveSettings.cs" />
138138
<Compile Include="Extensions\CloudExceptionExtensions.cs" />
139139
<Compile Include="Factories\AuthenticationFactory.cs" />
140+
<Compile Include="Factories\CancelRetryHandler.cs" />
140141
<Compile Include="Factories\ClientFactory.cs" />
141142
<Compile Include="HttpClientOperationsFactory.cs" />
142143
<Compile Include="Utilities\HttpClientWithRetry.cs" />

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

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
// ----------------------------------------------------------------------------------
2+
//
3+
// Copyright Microsoft Corporation
4+
// Licensed under the Apache License, Version 2.0 (the "License");
5+
// you may not use this file except in compliance with the License.
6+
// You may obtain a copy of the License at
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
// Unless required by applicable law or agreed to in writing, software
9+
// distributed under the License is distributed on an "AS IS" BASIS,
10+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
11+
// See the License for the specific language governing permissions and
12+
// limitations under the License.
13+
// ----------------------------------------------------------------------------------
14+
15+
using System;
16+
using System.Linq;
17+
using System.Net.Http;
18+
using System.Threading;
19+
using System.Threading.Tasks;
20+
21+
namespace Microsoft.Azure.Commands.Common.Authentication.Factories
22+
{
23+
/// <summary>
24+
/// Class to enable automatic retry of TaskCanceledExsceptions.
25+
/// Note that this follows the calling pattern in PowerShell, which does not pass cancellation tokens from
26+
/// the base cmdlet
27+
/// </summary>
28+
public class CancelRetryHandler : DelegatingHandler, ICloneable
29+
{
30+
static readonly string[] TestModeLiveValues = { "Record", "None", "Live" };
31+
const string TestModeVariableKey = "Azure_Test_Mode";
32+
public CancelRetryHandler()
33+
{
34+
WaitInterval = TimeSpan.Zero;
35+
}
36+
public CancelRetryHandler(TimeSpan waitInterval, int maxTries)
37+
{
38+
WaitInterval = waitInterval;
39+
MaxTries = maxTries;
40+
}
41+
42+
public TimeSpan WaitInterval { get; set; } = TimeSpan.Zero;
43+
44+
public int MaxTries { get; set; } = 3;
45+
46+
protected override async Task<HttpResponseMessage> SendAsync(HttpRequestMessage request, CancellationToken cancellationToken)
47+
{
48+
int tries = 0;
49+
do
50+
{
51+
using (var source = new CancellationTokenSource())
52+
{
53+
try
54+
{
55+
return await base.SendAsync(request, source.Token).ConfigureAwait(false);
56+
}
57+
catch (TaskCanceledException) when (tries++ < MaxTries)
58+
{
59+
SleepIfRunningLive(WaitInterval);
60+
}
61+
}
62+
}
63+
while (true);
64+
}
65+
66+
void SleepIfRunningLive(TimeSpan timeSpan)
67+
{
68+
var testModeSetting = Environment.GetEnvironmentVariable(TestModeVariableKey);
69+
if (string.IsNullOrWhiteSpace(testModeSetting) || TestModeLiveValues.Any( v => string.Equals(v, testModeSetting, StringComparison.OrdinalIgnoreCase)))
70+
{
71+
Thread.Sleep(timeSpan);
72+
}
73+
}
74+
75+
public object Clone()
76+
{
77+
return new CancelRetryHandler(WaitInterval, MaxTries);
78+
}
79+
}
80+
}

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

Lines changed: 33 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,7 @@ namespace Microsoft.Azure.Commands.Common.Authentication.Factories
3232
public class ClientFactory : IClientFactory
3333
{
3434
private static readonly char[] uriPathSeparator = { '/' };
35+
private static readonly CancelRetryHandler DefaultCancelRetryHandler = new CancelRetryHandler(TimeSpan.FromSeconds(10), 3);
3536

3637
private Dictionary<Type, IClientAction> _actions;
3738
private OrderedDictionary _handlers;
@@ -56,11 +57,8 @@ public virtual TClient CreateArmClient<TClient>(IAzureContext context, string en
5657
}
5758

5859
var creds = AzureSession.Instance.AuthenticationFactory.GetServiceClientCredentials(context, endpoint);
59-
var newHandlers = GetCustomHandlers();
60-
TClient client = (newHandlers == null || newHandlers.Length == 0)
61-
? CreateCustomArmClient<TClient>(context.Environment.GetEndpointAsUri(endpoint), creds)
62-
: CreateCustomArmClient<TClient>(context.Environment.GetEndpointAsUri(endpoint), creds, GetCustomHandlers());
63-
60+
var baseUri = context.Environment.GetEndpointAsUri(endpoint);
61+
TClient client = CreateCustomArmClient<TClient>(baseUri, creds);
6462
var subscriptionId = typeof(TClient).GetProperty("SubscriptionId");
6563
if (subscriptionId != null && context.Subscription != null)
6664
{
@@ -73,23 +71,35 @@ public virtual TClient CreateArmClient<TClient>(IAzureContext context, string en
7371
public virtual TClient CreateCustomArmClient<TClient>(params object[] parameters) where TClient : Microsoft.Rest.ServiceClient<TClient>
7472
{
7573
List<Type> types = new List<Type>();
76-
var containsDelegatingHandler = false;
74+
List<object> parameterList = new List<object>();
75+
List<DelegatingHandler> handlerList = new List<DelegatingHandler> { DefaultCancelRetryHandler.Clone() as CancelRetryHandler};
76+
var customHandlers = GetCustomHandlers();
77+
if (customHandlers != null && customHandlers.Any())
78+
{
79+
handlerList.AddRange(customHandlers);
80+
}
81+
82+
bool hasHandlers = false;
7783
foreach (object obj in parameters)
7884
{
79-
if (obj.GetType().Name.Equals("DelegatingHandler[]"))
85+
var paramType = obj.GetType();
86+
types.Add(paramType);
87+
if (paramType == typeof(DelegatingHandler[]))
8088
{
81-
containsDelegatingHandler = true;
89+
handlerList.AddRange(obj as DelegatingHandler[]);
90+
parameterList.Add(handlerList.ToArray());
91+
hasHandlers = true;
92+
}
93+
else
94+
{
95+
parameterList.Add(obj);
8296
}
83-
types.Add(obj.GetType());
8497
}
8598

86-
var copiedParameters = parameters;
87-
if (!containsDelegatingHandler)
99+
if (!hasHandlers)
88100
{
89101
types.Add((new DelegatingHandler[0]).GetType());
90-
var parameterList = copiedParameters.ToList();
91-
parameterList.Add(new DelegatingHandler[0]);
92-
copiedParameters = parameterList.ToArray();
102+
parameterList.Add(handlerList.ToArray());
93103
}
94104

95105
var constructor = typeof(TClient).GetConstructor(types.ToArray());
@@ -99,7 +109,7 @@ public virtual TClient CreateCustomArmClient<TClient>(params object[] parameters
99109
throw new InvalidOperationException(string.Format(Resources.InvalidManagementClientType, typeof(TClient).Name));
100110
}
101111

102-
TClient client = (TClient)constructor.Invoke(copiedParameters);
112+
TClient client = (TClient)constructor.Invoke(parameterList.ToArray());
103113

104114
foreach (ProductInfoHeaderValue userAgent in _userAgents.Keys)
105115
{
@@ -120,13 +130,7 @@ public virtual TClient CreateClient<TClient>(IAzureContext context, string endpo
120130
}
121131

122132
SubscriptionCloudCredentials creds = AzureSession.Instance.AuthenticationFactory.GetSubscriptionCloudCredentials(context, endpoint);
123-
TClient client = CreateCustomClient<TClient>(creds, context.Environment.GetEndpointAsUri(endpoint));
124-
foreach (DelegatingHandler handler in GetCustomHandlers())
125-
{
126-
client.AddHandlerToPipeline(handler);
127-
}
128-
129-
return client;
133+
return CreateCustomClient<TClient>(creds, context.Environment.GetEndpointAsUri(endpoint));
130134
}
131135

132136
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
177181

178182
AzureContext context = new AzureContext(subscription, account, environment);
179183

180-
TClient client = CreateClient<TClient>(context, endpoint);
184+
var client = CreateClient<TClient>(context, endpoint);
181185

182186
foreach (IClientAction action in GetActions())
183187
{
@@ -209,6 +213,12 @@ public virtual TClient CreateCustomClient<TClient>(params object[] parameters) w
209213
client.UserAgent.Add(userAgent);
210214
}
211215

216+
foreach (DelegatingHandler handler in GetCustomHandlers())
217+
{
218+
client.AddHandlerToPipeline(handler);
219+
}
220+
221+
client.AddHandlerToPipeline(DefaultCancelRetryHandler.Clone() as CancelRetryHandler);
212222
return client;
213223
}
214224

‎src/ResourceManager/Profile/ChangeLog.md

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,7 @@
2626
- ParameterSetName
2727
- InvocationName
2828
* Fix issue where version 10.0.3 of Newtonsoft.Json wasn't being loaded on module import
29+
* Retry TaskCanceledException appropriately when connections are left open.
2930

3031
## Version 5.1.0
3132
* 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

0 commit comments

Comments
 (0)