Can I perform an atomic try-get-or-add operation with a ConcurrentDictionary? - asp.net-core

Question

In my ASP.NET Core MVC web-application (that uses OIDC) I have a class that automatically refreshes the access_token stored in the visitor's authentication cookie before it expires.
It is based on AutomaticTokenManagementCookieEvents from the IdentityServer4 samples. Available here: https://github.com/IdentityServer/IdentityServer4/blob/0155beb2cea850144b6407684a2eda22e4eea3db/samples/Clients/src/MvcHybridAutomaticRefresh/AutomaticTokenManagement/AutomaticTokenManagementCookieEvents.cs
static readonly ConcurrentDictionary<String,Object> _pendingRefreshes = new ConcurrentDictionary<String,Object>();
public override async Task ValidatePrincipal( CookieValidatePrincipalContext context )
{
DateTime accessTokenExpiresAt = GetAccessTokenExpiry( context ); // gets the 'expires_at' value from `context.Properties.GetTokens();`
String refreshToken = GetRefreshToken( context ); // Gets the 'refresh_token' value from `context.Properties.GetTokens();`
Boolean isExpired = DateTime.UtcNow > accessTokenExpiresAt;
Boolean willExpireSoon = DateTime.UtcNow > accessTokenExpiresAt.Subtract( TimeSpan.FromSeconds( 60 ) );
if( isExpired || willExpireSoon )
{
Boolean canRefresh = _pendingRefreshes.TryAdd( refreshToken, null );
if( canRefresh )
{
try
{
await RefreshAccessTokenAsync( context, refreshToken );
}
finally
{
_pendingRefreshes.TryRemove( refreshToken );
}
}
else
{
// TODO: What should happen here?
}
}
}
private async Task RefreshAccessTokenAsync( CookieValidatePrincipalContext context, String refreshToken )
{
// using IdentityModel.Client.HttpClientTokenRequestExtensions.RequestRefreshTokenAsync
TokenResponse response = await this.httpClient.RefreshTokenAsync( refreshToken );
if( response.IsError )
{
// (Error logging code here)
if( response.Error == "invalid_grant" )
{
// Usually invalid_grant errors happen if the user's refresh_token has been revoked or expired
// refresh_token expiry is separate from access_token expiry.
// If a refresh_token has expired or been revoked the only thing to do is force the user to login again. `RejectPrincipal()` will send the user to the OIDC OP login page - though this will cause the user to lose their data if this is a POST request.
context.RejectPrincipal();
}
else
{
// Something else bad happened. Don't invalidate the user's credentials unless they're actually expired, though.
throw new Exception( "Unexpected error." );
}
}
else
{
context.Properties.UpdateTokenValue( "access_token" , response.AccessToken );
context.Properties.UpdateTokenValue( "refresh_token", response.RefreshToken );
DateTime newExpiresAt = DateTime.UtcNow + TimeSpan.FromSeconds( response.ExpiresIn );
context.Properties.UpdateTokenValue( "expires_at", newExpiresAt.ToString( "o", CultureInfo.InvariantCulture ) );
await context.HttpContext.SignInAsync( context.Principal, context.Properties );
}
}
The problem with this code is that if the user's browser makes two simultaneous requests after their access_token has already expired then the user will get an error message if code later-on in the ASP.NET Core pipeline for the second simultaneous request uses the now-expired access_token.
...how can I get it so that the second concurrent request using an expired access_token will await the same Task (from RefreshAccessTokenAsync)?
My thought is something like this:
Change _pendingRefreshes to ConcurrentDictionary<String,Task<String>>.
Change Boolean canRefresh = _pendingRefreshes.TryAdd( refreshToken, null ); to something like this (using a hypothetical TryGetOrAdd method):
Boolean addedNewTask = _pendingRefreshes
.TryGetOrAdd(
key: refreshToken,
valueFactory: rt => this.RefreshTokenAsync( context, rt ),
value: out Task task
);
if( addedNewTask )
{
// wait for the new access_token to be saved before continuing.
await task;
}
else
{
if( isExpired )
{
// If the current request's access_token is already expired and its refresh_token is currently being refrehsed, then wait for it to finish as well, then update the access_token but only for this request's lifetime (i.e. don't call `ReplacePrincipal` or `SignInAsync`.
await task;
}
}
The problem is ConcurrentDictionary<TKey,TValue> does not have a TryGetOrAdd method I can use to atomically get-existing-or-add-new items.
AddOrUpdate - Doesn't return any existing items. Doesn't indicate if the value returned was an existing item or not.
GetOrAdd - Doesn't indicate if the value returned was an existing item or not.
TryAdd - Doesn't let you atomically get any existing value with the same key.
TryGetValue - Doesn't let you atomically add a new item if there isn't a value for the given key.
TryRemove - Doesn't let you atomically add a new item.
TryUpdate - Doesn't let you add new items.
This is fixable using a lock but that negates the advantages of using ConcurrentDictionary. Something like this:
Task<String> task;
Boolean addedNewTask;
lock( _pendingRefreshes )
{
Boolean taskExists = _pendingRefreshes.TryGetValue( refreshToken, out task );
if( taskExists )
{
addedNewTask = false;
}
else
{
task = RefreshAccessTokenAsync( context, refreshToken );
if( !_pendingRefreshes.TryAdd( refreshToken, task ) )
{
throw new InvalidOperationException( "Could not add the Task." ); // This should never happen.
}
addedNewTask = true;
}
}
if( addedNewTask || isExpired )
{
String newAccessToken = await task;
if( isExpired )
{
context.Properties.UpdateTokenValue( "access_token", newAccessToken );
}
}
...or is this the correct use of ConcurrentDictionary for this scenario?

Answer 1

I learned that ConcurrentDictionary does not guarantee that valueFactory won't be invoked exactly once per key, but may be invoked multiple times - which I can't live with because my valueFactory is an expensive operation with side-effects that may adversely negatively affect the user-experience (e.g. invalidating a token stored inside the user's cookies unnecessarily).
I settled on an approach using lock and a Dictionary<String,Task<TokenResponse>> instead. I do wonder how it could be improved and I'd really appreciate any feedback:
private static readonly Dictionary<String,Task<TokenResponse>> _currentlyRunningRefreshOperations = new Dictionary<String,Task<TokenResponse>>();
/// <summary>Returns <c>true</c> if a new Task was returned. Returns <c>false</c> if an existing Task was returned.</summary>
private Boolean GetOrCreateRefreshTokenTaskAsync( CookieValidatePrincipalContext context, String refreshToken, out Task<TokenResponse> task )
{
lock( _currentlyRunningRefreshOperations )
{
if( _currentlyRunningRefreshOperations.TryGetValue( refreshToken, out task ) )
{
return false;
}
else
{
task = this.AttemptRefreshTokensAsync( context, refreshToken );
_currentlyRunningRefreshOperations.Add( refreshToken, task );
return true;
}
}
}
private async Task<TokenResponse> AttemptRefreshTokensAsync( CookieValidatePrincipalContext context, String refreshToken )
{
try
{
TokenResponse response = await this.service.RefreshTokenAsync( refreshToken );
if( response.IsError )
{
this.logger.LogWarning( "Encountered " + nameof(this.service.RefreshTokenAsync) + " error: {error}. Type: {errorType}. Description: {errorDesc}. refresh_token: {refreshToken}.", response.Error, response.ErrorType, response.ErrorDescription, refreshToken );
return response;
}
}
finally
{
lock( _currentlyRunningRefreshOperations )
{
_currentlyRunningRefreshOperations.Remove( refreshToken );
}
}
}