[Performance] Low hanging fruit in Tokens.JWT

[BrennanConroy]

Recently while comparing a benchmark that didn't use auth and one that used jwt auth I noticed the allocations were extremely high when using jwt. Digging into it I found many places where allocations could be improved in System.IdentityModel.Tokens.Jwt, I'll focus on the largest allocation I saw for this issue.

Our code first checks if we the token can be read via

azure-activedirectory-identitymodel-extensions-for-dotnet/src/System.IdentityModel.Tokens.Jwt/JwtSecurityTokenHandler.cs

Line 278 in b7a4fdd

public override bool CanReadToken(string token)

Then it proceeds to call ValidateToken

azure-activedirectory-identitymodel-extensions-for-dotnet/src/System.IdentityModel.Tokens.Jwt/JwtSecurityTokenHandler.cs

Line 835 in b7a4fdd

public override ClaimsPrincipal ValidateToken(string token, TokenValidationParameters validationParameters, out SecurityToken validatedToken)

Both these methods call string.Split on the provided token, check how many values are in the array, and then throw away the results.

azure-activedirectory-identitymodel-extensions-for-dotnet/src/System.IdentityModel.Tokens.Jwt/JwtSecurityTokenHandler.cs

Lines 292 to 293 in b7a4fdd

	string[] tokenParts = token.Split(new char[] { '.' }, JwtConstants.MaxJwtSegmentCount + 1);
	if (tokenParts.Length == JwtConstants.JwsSegmentCount)

Trying out a quick test to count the '.' characters in a string (then adding 1 to count the segments) shows we can do it allocation free and almost 6 times faster:

Benchmark code

[MemoryDiagnoser]
public class SplitCount
{
    private string JwtToken = "eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJodHRwOi8vc2NoZW1hcy54bWxzb2FwLm9yZy93cy8yMDA1LzA1L2lkZW50aXR5L2NsYWltcy9uYW1laWRlbnRpZmllciI6InRlc3QiLCJodHRwOi8vc2NoZW1hcy54bWxzb2FwLm9yZy93cy8yMDA1LzA1L2lkZW50aXR5L2NsYWltcy9lbWFpbGFkZHJlc3MiOiJ0ZXN0QHRlc3QuY29tIiwic3ViIjoidGVzdEB0ZXN0LmNvbSIsImV4cCI6MjAxMzUzMjk3NiwiaXNzIjoiVGVzdCIsImF1ZCI6InRlc3QifQ.tQ-s9uAokahLfeKY-SKnvYQWDU4HrFCRV-0Ae7LAabY";

    [Benchmark]
    public int Split()
    {
        var tokenParts = JwtToken.Split(new char[] { '.' }, 6);
        return tokenParts.Length;
    }

    [Benchmark]
    public int ManualCount()
    {
        var count = 1;
        var span = JwtToken.AsSpan();
        while (span.Length > 0)
        {
            var index = span.IndexOf('.');
            if (index < 0)
            {
                break;
            }
            count++;
            span = span.Slice(index);
            if (count == 6)
            {
                break;
            }
        }
        return count;
    }
}

Method	Mean	Error	StdDev	Gen0	Allocated
Split	193.57 ns	1.845 ns	1.725 ns	0.1183	928 B
ManualCount	35.83 ns	0.155 ns	0.130 ns	-	-

And since there are at least 2 calls to string.Split per request in the scenario I'm looking at, it would be even bigger savings!

Proposal

1. Create an internal helper method to count the segments in a token instead of using string.Split, could even try to vectorize 😄
2. Then call this method from everywhere that string.Split is used just for the length of the output array
3. ???
4. Profit!

[brentschmaltz]

@BrennanConroy thanks for taking time to investigate.
We are thinking along the same lines.
Have a look at this PR where we made some of the changes you suggested, the performance was about 30%.
#1934

We are debating if we should make the simple fixes for the easy performance gains, it may make sense as the gains for everyone could be significant.

[BrennanConroy]

Have a look at this PR where we made some of the changes you suggested, the performance was about 30%.

Nice! Do you know if most of the savings were from switching to System.Text.Json or from the other changes (like removing string.Split)? I just tried with the 6.25.0 packages, sadly the scenario I'm looking at doesn't hit most of the code paths in that PR.

[BrennanConroy]

Done via #1964