How I prevented duplicate token request in an identity server using semaphore technique

Scenario:

I wrote a service where it requests bearer tokens in an identity server for each company separately. Suppose we have 8 companies, Then the service will check if there is any available nonexpired token for that specific company. If there is an available token, the service will return it immediately or it will send a token request to an external server.

The flow looked like this:

Assume we have a BearerTokenService which is a singleton service and will retrieve and cache tokens for each company in it. So the code should look like this.

public class BearerToken
{
    public Guid CompanyId { get; set; }
    public string Token { get; set; }
    public DateTime ExpiresIn { get; set; }
}
public class BearerTokenService 
{
    private readonly IHttpClientFactory _httpClientFactory;
    private Dictionary<Guid, BearerToken> Tokens { get; set; } = new();

    public BearerTokenService(IHttpClientFactory httpClientFactory)
    {
        _httpClientFactory = httpClientFactory;
    }

    public async Task<string> GetTokenAsync(Guid companyId)
    {
        if (Tokens.TryGetValue(companyId, out var token) && token.ExpiresIn > DateTime.Now)
            return token.Token;

        var client = await _httpClientFactory.CreateClient();
        var tokenStr = await client.PostAsync<string>(config.Url, requestPayload);

        token = new BearerToken
        {
            CompanyId = companyId,
            Token = tokenStr
        };

        Tokens[companyId] = token;

        return token.Token;
    }
}

At first glance, this piece of code will look fine except we consider the class calling this service is also async and each company request for GetTokenAsync function is also async.

Suppose 4 token requests for company x arrived at the same time. Then GetTokenAsync will also execute at the same time for each request asynchronously. After entering the method "GetTokenAsync" each execution thread will see that there is no token available for that specific company. Then it will initiate 4 HTTP post requests for getting the token for the same company. Which is unnecessary and inefficient in our case. So we want the first request to retrieve the token and hold the other 3 requests until the first HTTP request finishes and the token is added to the "Tokens" variable in BearerTokenService class. After the first token is saved, the last three requests will see that the token is already available and will return from the "Tokens" immediately.

To fulfill our goal we have to somehow lock our GetTokenAsync method for the other three requests until the first request is complete. That's where semaphore will come to save us. Some of you may refer to the "Lock" mechanism of .NET ecosystem but I am assuring you, that will not work in the asynchronous code block.

Let's see how the codeblock above will change after we add a semaphore to it.

public class BearerToken
{
    public Guid CompanyId { get; set; }
    public string Token { get; set; }
    public DateTime ExpiresIn { get; set; }
}
public class BearerTokenService 
{
    private readonly IHttpClientFactory _httpClientFactory;
    // first argument is initialCount and second argument is maxCount
    // So we are specifying that only one execution can enter the codeblock 
    // at a time
    private readonly SemaphoreSlim _semaphore = new(1, 1);
    private readonly Dictionary<Guid, BearerToken> _tokens = new();

    public BearerTokenService(IHttpClientFactory httpClientFactory)
    {
        _httpClientFactory = httpClientFactory;
    }

    public async Task<string> GetTokenAsync(Guid companyId)
    {
        // from this point only one parralal execution will happen and rest will wait until semaphore lock is released
        await _semaphore.WaitAsync();
        try
        {
            // if token is not null and not expired then the token is valid
            if (_tokens.TryGetValue(companyId, out var token) && token.ExpiresIn > DateTime.Now)
            {
                return token.Token;
            }
            // if execution reach this far that means either our token is expired or this company never requested a token thus token is null, so we have to request a new token.
            var client = await _httpClientFactory.CreateClient();
            var tokenStr = await client.PostAsync<string>(config.Url, requestPayload);

            // cache the token for future checking
            if (_tokens.ContainsKey(companyId))
            {
                _tokens[companyId].Token = tokenStr;
            }
            else
            {
                _tokens[companyId] = new BearerToken
                {
                    CompanyId = companyId,
                    Token = tokenStr
                };
            }

            return tokenStr;
        }
        catch (Exception ex)
        {
            // Handle the exception appropriately.
        }
        finally
        {
            // No matter what is condition, even if the try blocks fail
            // semaphore will be released. 
            // because not releasing semaphore will results in deadlock with other threads
            _semaphore.Release();
        }
    }
}

Yay, Now only one request at a time can send HTTP request to retrieve the token and the other will wait and the problem solved!!!

You may find this official link to SemaphoreSlim class useful. Happy Coding!!!

How I prevented duplicate token request in an identity server using semaphore technique

Comments

More from this blog

A Docker Adventure: Setting Up Your Own Private Docker Registry Inside A Docker Container

Scenario:

Command Palette

Comments

More from this blog

Scenario: