Skip to content

Commit

Permalink
Fix #15013: Intermittent Exceptions when calling DefaultAzureCredenti…
Browse files Browse the repository at this point in the history
…al (#15025)
  • Loading branch information
AlexanderSher committed Sep 9, 2020
1 parent 57b3ebb commit 1618e22
Show file tree
Hide file tree
Showing 4 changed files with 139 additions and 129 deletions.
4 changes: 2 additions & 2 deletions sdk/identity/Azure.Identity/src/CredentialDiagnosticScope.cs
Original file line number Diff line number Diff line change
Expand Up @@ -16,10 +16,10 @@ namespace Azure.Identity
private readonly TokenRequestContext _context;
private readonly IScopeHandler _scopeHandler;

public CredentialDiagnosticScope(string name, TokenRequestContext context, IScopeHandler scopeHandler)
public CredentialDiagnosticScope(ClientDiagnostics diagnostics, string name, TokenRequestContext context, IScopeHandler scopeHandler)
{
_name = name;
_scope = scopeHandler.CreateScope(name);
_scope = scopeHandler.CreateScope(diagnostics, name);
_context = context;
_scopeHandler = scopeHandler;
}
Expand Down
133 changes: 7 additions & 126 deletions sdk/identity/Azure.Identity/src/CredentialPipeline.cs
Original file line number Diff line number Diff line change
Expand Up @@ -2,22 +2,18 @@
// Licensed under the MIT License.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using Azure.Core;
using Azure.Core.Diagnostics;
using Azure.Core.Pipeline;
using Microsoft.Identity.Client;

namespace Azure.Identity
{
internal class CredentialPipeline
internal class CredentialPipeline
{
private static readonly Lazy<CredentialPipeline> s_singleton = new Lazy<CredentialPipeline>(() => new CredentialPipeline(new TokenCredentialOptions()));

private readonly IScopeHandler _defaultScopeHandler;
private IScopeHandler _groupScopeHandler;
private static readonly IScopeHandler _defaultScopeHandler = new ScopeHandler();

private CredentialPipeline(TokenCredentialOptions options)
{
Expand All @@ -26,8 +22,6 @@ private CredentialPipeline(TokenCredentialOptions options)
HttpPipeline = HttpPipelineBuilder.Build(options, Array.Empty<HttpPipelinePolicy>(), Array.Empty<HttpPipelinePolicy>(), new CredentialResponseClassifier());

Diagnostics = new ClientDiagnostics(options);

_defaultScopeHandler = new ScopeHandler(Diagnostics);
}

public static CredentialPipeline GetInstance(TokenCredentialOptions options)
Expand All @@ -48,18 +42,18 @@ public IConfidentialClientApplication CreateMsalConfidentialClient(string tenant

public CredentialDiagnosticScope StartGetTokenScope(string fullyQualifiedMethod, TokenRequestContext context)
{
IScopeHandler scopeHandler = _groupScopeHandler ?? _defaultScopeHandler;
IScopeHandler scopeHandler = ScopeGroupHandler.Current ?? _defaultScopeHandler;

CredentialDiagnosticScope scope = new CredentialDiagnosticScope(fullyQualifiedMethod, context, scopeHandler);
CredentialDiagnosticScope scope = new CredentialDiagnosticScope(Diagnostics, fullyQualifiedMethod, context, scopeHandler);
scope.Start();
return scope;
}

public CredentialDiagnosticScope StartGetTokenScopeGroup(string fullyQualifiedMethod, TokenRequestContext context)
{
var scopeHandler = new ScopeGroupHandler(this, fullyQualifiedMethod);
var scopeHandler = new ScopeGroupHandler(fullyQualifiedMethod);

CredentialDiagnosticScope scope = new CredentialDiagnosticScope(fullyQualifiedMethod, context, scopeHandler);
CredentialDiagnosticScope scope = new CredentialDiagnosticScope(Diagnostics, fullyQualifiedMethod, context, scopeHandler);
scope.Start();
return scope;
}
Expand All @@ -74,123 +68,10 @@ public override bool IsRetriableResponse(HttpMessage message)

private class ScopeHandler : IScopeHandler
{
private readonly ClientDiagnostics _diagnostics;

public ScopeHandler(ClientDiagnostics diagnostics)
{
_diagnostics = diagnostics;
}

public DiagnosticScope CreateScope(string name) => _diagnostics.CreateScope(name);
public DiagnosticScope CreateScope(ClientDiagnostics diagnostics, string name) => diagnostics.CreateScope(name);
public void Start(string name, in DiagnosticScope scope) => scope.Start();
public void Dispose(string name, in DiagnosticScope scope) => scope.Dispose();
public void Fail(string name, in DiagnosticScope scope, Exception exception) => scope.Failed(exception);
}

private class ScopeGroupHandler : IScopeHandler
{
private readonly CredentialPipeline _pipeline;
private readonly string _groupName;
private Dictionary<string, (DateTime StartDateTime, Exception Exception)> _childScopes;

public ScopeGroupHandler(CredentialPipeline pipeline, string groupName)
{
_pipeline = pipeline;
_groupName = groupName;
}

public DiagnosticScope CreateScope(string name)
{
if (IsGroup(name))
{
_pipeline._groupScopeHandler = this;
return _pipeline.Diagnostics.CreateScope(name);
}

_childScopes ??= new Dictionary<string, (DateTime startDateTime, Exception exception)>();
_childScopes[name] = default;
return default;
}

public void Start(string name, in DiagnosticScope scope)
{
if (IsGroup(name))
{
scope.Start();
}
else
{
_childScopes[name] = (DateTime.UtcNow, default);
}
}

public void Dispose(string name, in DiagnosticScope scope)
{
if (!IsGroup(name))
{
return;
}

if (_childScopes != null)
{
var succeededScope = _childScopes.LastOrDefault(kvp => kvp.Value.Exception == default);
if (succeededScope.Key != default)
{
SucceedChildScope(succeededScope.Key, succeededScope.Value.StartDateTime);
}
}

scope.Dispose();
_pipeline._groupScopeHandler = default;
}

public void Fail(string name, in DiagnosticScope scope, Exception exception)
{
if (_childScopes == default)
{
scope.Failed(exception);
return;
}

if (IsGroup(name))
{
if (exception is OperationCanceledException)
{
var canceledScope = _childScopes.Last(kvp => kvp.Value.Exception == exception);
FailChildScope(canceledScope.Key, canceledScope.Value.StartDateTime, canceledScope.Value.Exception);
}
else
{
foreach (var childScope in _childScopes)
{
FailChildScope(childScope.Key, childScope.Value.StartDateTime, childScope.Value.Exception);
}
}

scope.Failed(exception);
}
else
{
_childScopes[name] = (_childScopes[name].StartDateTime, exception);
}
}

private void SucceedChildScope(string name, DateTime dateTime)
{
using DiagnosticScope scope = _pipeline.Diagnostics.CreateScope(name);
scope.SetStartTime(dateTime);
scope.Start();
}

private void FailChildScope(string name, DateTime dateTime, Exception exception)
{
using DiagnosticScope scope = _pipeline.Diagnostics.CreateScope(name);
scope.SetStartTime(dateTime);
scope.Start();
scope.Failed(exception);
}

private bool IsGroup(string name) => string.Equals(name, _groupName, StringComparison.Ordinal);
}
}
}
2 changes: 1 addition & 1 deletion sdk/identity/Azure.Identity/src/IScopeHandler.cs
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ namespace Azure.Identity
{
internal interface IScopeHandler
{
DiagnosticScope CreateScope(string name);
DiagnosticScope CreateScope(ClientDiagnostics diagnostics, string name);
void Start(string name, in DiagnosticScope scope);
void Dispose(string name, in DiagnosticScope scope);
void Fail(string name, in DiagnosticScope scope, Exception exception);
Expand Down
129 changes: 129 additions & 0 deletions sdk/identity/Azure.Identity/src/ScopeGroupHandler.cs
Original file line number Diff line number Diff line change
@@ -0,0 +1,129 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.

using System;
using System.Collections.Generic;
using System.Linq;
using System.Threading;
using Azure.Core.Pipeline;

namespace Azure.Identity
{
internal sealed class ScopeGroupHandler : IScopeHandler
{
private static readonly AsyncLocal<IScopeHandler> _currentAsyncLocal = new AsyncLocal<IScopeHandler>();
private readonly string _groupName;
private Dictionary<string, ChildScopeInfo> _childScopes;

public static IScopeHandler Current => _currentAsyncLocal.Value;

public ScopeGroupHandler(string groupName)
{
_groupName = groupName;
_currentAsyncLocal.Value = this;
}

public DiagnosticScope CreateScope(ClientDiagnostics diagnostics, string name)
{
if (IsGroup(name))
{
return diagnostics.CreateScope(name);
}

_childScopes ??= new Dictionary<string, ChildScopeInfo>();
_childScopes[name] = new ChildScopeInfo(diagnostics, name);
return default;
}

public void Start(string name, in DiagnosticScope scope)
{
if (IsGroup(name))
{
scope.Start();
}
else
{
_childScopes[name].StartDateTime = DateTime.UtcNow;
}
}

public void Dispose(string name, in DiagnosticScope scope)
{
if (!IsGroup(name))
{
return;
}

var succeededScope = _childScopes?.Values.LastOrDefault(i => i.Exception == default);
if (succeededScope != null)
{
SucceedChildScope(succeededScope);
}

scope.Dispose();
_currentAsyncLocal.Value = default;
}

public void Fail(string name, in DiagnosticScope scope, Exception exception)
{
if (_childScopes == default)
{
scope.Failed(exception);
return;
}

if (IsGroup(name))
{
if (exception is OperationCanceledException)
{
var canceledScope = _childScopes.Values.Last(i => i.Exception == exception);
FailChildScope(canceledScope);
}
else
{
foreach (var childScope in _childScopes.Values)
{
FailChildScope(childScope);
}
}

scope.Failed(exception);
}
else
{
_childScopes[name].Exception = exception;
}
}

private static void SucceedChildScope(ChildScopeInfo scopeInfo)
{
using DiagnosticScope scope = scopeInfo.Diagnostics.CreateScope(scopeInfo.Name);
scope.SetStartTime(scopeInfo.StartDateTime);
scope.Start();
}

private static void FailChildScope(ChildScopeInfo scopeInfo)
{
using DiagnosticScope scope = scopeInfo.Diagnostics.CreateScope(scopeInfo.Name);
scope.SetStartTime(scopeInfo.StartDateTime);
scope.Start();
scope.Failed(scopeInfo.Exception);
}

private bool IsGroup(string name) => string.Equals(name, _groupName, StringComparison.Ordinal);

private class ChildScopeInfo
{
public ClientDiagnostics Diagnostics { get; }
public string Name { get; }
public DateTime StartDateTime { get; set; }
public Exception Exception { get; set; }

public ChildScopeInfo(ClientDiagnostics diagnostics, string name)
{
Diagnostics = diagnostics;
Name = name;
}
}
}
}

0 comments on commit 1618e22

Please sign in to comment.