Skip to content

Commit 0f28069

Browse files
bgavrilMSbgavrilMS
authored
Expose access token cache count (#5330)
Co-authored by: Robbie Ginsburg <[email protected]>
1 parent bbaa1bf commit 0f28069

File tree

18 files changed

+135
-32
lines changed

18 files changed

+135
-32
lines changed

‎src/client/Microsoft.Identity.Client/AuthenticationResultMetadata.cs

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -97,5 +97,10 @@ public AuthenticationResultMetadata(TokenSource tokenSource)
9797
/// See https://aka.ms/msal-net-logging for more details about logging.
9898
/// </remarks>
9999
public string Telemetry { get; set; }
100+
101+
/// <summary>
102+
/// The number of access tokens in the cache.
103+
/// </summary>
104+
public int CachedAccessTokenCount { get; set; }
100105
}
101106
}

‎src/client/Microsoft.Identity.Client/Cache/CacheSessionManager.cs

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,9 +48,11 @@ public async Task<MsalAccessTokenCacheItem> FindAccessTokenAsync()
4848
return await TokenCacheInternal.FindAccessTokenAsync(_requestParams).ConfigureAwait(false);
4949
}
5050

51-
public Task<Tuple<MsalAccessTokenCacheItem, MsalIdTokenCacheItem, Account>> SaveTokenResponseAsync(MsalTokenResponse tokenResponse)
51+
public async Task<Tuple<MsalAccessTokenCacheItem, MsalIdTokenCacheItem, Account>> SaveTokenResponseAsync(MsalTokenResponse tokenResponse)
5252
{
53-
return TokenCacheInternal.SaveTokenResponseAsync(_requestParams, tokenResponse);
53+
var result = await TokenCacheInternal.SaveTokenResponseAsync(_requestParams, tokenResponse).ConfigureAwait(false);
54+
RequestContext.ApiEvent.CachedAccessTokenCount = TokenCacheInternal.Accessor.EntryCount;
55+
return result;
5456
}
5557

5658
public async Task<Account> GetAccountAssociatedWithAccessTokenAsync(MsalAccessTokenCacheItem msalAccessTokenCacheItem)
@@ -181,6 +183,8 @@ private async Task RefreshCacheForReadOperationsAsync()
181183
{
182184
RequestContext.ApiEvent.CacheLevel = CacheLevel.L1Cache;
183185
}
186+
187+
RequestContext.ApiEvent.CachedAccessTokenCount = TokenCacheInternal.Accessor.EntryCount;
184188
}
185189
}
186190
}

‎src/client/Microsoft.Identity.Client/Cache/ITokenCacheAccessor.cs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,8 @@ namespace Microsoft.Identity.Client.Cache
1010
{
1111
internal interface ITokenCacheAccessor
1212
{
13+
int EntryCount { get; }
14+
1315
void SaveAccessToken(MsalAccessTokenCacheItem item);
1416

1517
void SaveRefreshToken(MsalRefreshTokenCacheItem item);

‎src/client/Microsoft.Identity.Client/Internal/Requests/RequestBase.cs

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -223,6 +223,7 @@ private void UpdateTelemetry(long elapsedMilliseconds, ApiEvent apiEvent, Authen
223223
authenticationResult.AuthenticationResultMetadata.CacheLevel = GetCacheLevel(authenticationResult);
224224
authenticationResult.AuthenticationResultMetadata.Telemetry = apiEvent.MsalRuntimeTelemetry;
225225
authenticationResult.AuthenticationResultMetadata.RegionDetails = CreateRegionDetails(apiEvent);
226+
authenticationResult.AuthenticationResultMetadata.CachedAccessTokenCount = apiEvent.CachedAccessTokenCount;
226227

227228
Metrics.IncrementTotalDurationInMs(authenticationResult.AuthenticationResultMetadata.DurationTotalInMs);
228229
}

‎src/client/Microsoft.Identity.Client/Platforms/Android/AndroidTokenCacheAccessor.cs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -161,6 +161,8 @@ public List<MsalAccountCacheItem> GetAllAccounts(string optionalPartitionKey = n
161161
}
162162
#endregion
163163

164+
public int EntryCount { get; } = 0; // not implemented for Android
165+
164166
public MsalAccountCacheItem GetAccount(MsalAccountCacheItem accountCacheItem)
165167
{
166168
return MsalAccountCacheItem.FromJsonString(_accountSharedPreference.GetString(accountCacheItem.CacheKey, null));

‎src/client/Microsoft.Identity.Client/Platforms/iOS/iOSTokenCacheAccessor.cs

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -167,7 +167,10 @@ public List<MsalAccountCacheItem> GetAllAccounts(string optionalPartitionKey = n
167167
.Select(x => MsalAccountCacheItem.FromJsonString(x))
168168
.ToList();
169169
}
170-
#endregion
170+
#endregion
171+
172+
public int EntryCount { get; } = 0; // not implemented for iOS
173+
171174

172175
internal SecStatusCode TryGetBrokerApplicationToken(string clientId, out string appToken)
173176
{

‎src/client/Microsoft.Identity.Client/PlatformsCommon/Shared/InMemoryPartitionedAppTokenCacheAccessor.cs

Lines changed: 37 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
using System.Collections.Concurrent;
66
using System.Collections.Generic;
77
using System.Linq;
8+
using System.Threading;
89
using Microsoft.Identity.Client.Cache;
910
using Microsoft.Identity.Client.Cache.Items;
1011
using Microsoft.Identity.Client.Core;
@@ -34,6 +35,11 @@ internal class InMemoryPartitionedAppTokenCacheAccessor : ITokenCacheAccessor
3435
protected readonly ILoggerAdapter _logger;
3536
private readonly CacheOptions _tokenCacheAccessorOptions;
3637

38+
private int _entryCount = 0;
39+
private static int s_entryCount = 0;
40+
41+
public int EntryCount => GetEntryCountRef();
42+
3743
public InMemoryPartitionedAppTokenCacheAccessor(
3844
ILoggerAdapter logger,
3945
CacheOptions tokenCacheAccessorOptions)
@@ -59,9 +65,18 @@ public void SaveAccessToken(MsalAccessTokenCacheItem item)
5965
string itemKey = item.CacheKey;
6066
string partitionKey = CacheKeyFactory.GetAppTokenCacheItemKey(item.ClientId, item.TenantId, item.KeyId, item.AdditionalCacheKeyComponents);
6167

62-
// if a conflict occurs, pick the latest value
63-
AccessTokenCacheDictionary
64-
.GetOrAdd(partitionKey, new ConcurrentDictionary<string, MsalAccessTokenCacheItem>())[itemKey] = item;
68+
var partition = AccessTokenCacheDictionary.GetOrAdd(partitionKey, _ => new ConcurrentDictionary<string, MsalAccessTokenCacheItem>());
69+
bool added = partition.TryAdd(itemKey, item);
70+
71+
// only increment the entry count if the item was added, not updated
72+
if (added)
73+
{
74+
Interlocked.Increment(ref GetEntryCountRef());
75+
}
76+
else
77+
{
78+
partition[itemKey] = item;
79+
}
6580
}
6681

6782
/// <summary>
@@ -129,12 +144,19 @@ public void DeleteAccessToken(MsalAccessTokenCacheItem item)
129144
{
130145
var partitionKey = CacheKeyFactory.GetAppTokenCacheItemKey(item.ClientId, item.TenantId, item.KeyId);
131146

132-
AccessTokenCacheDictionary.TryGetValue(partitionKey, out var partition);
133-
if (partition == null || !partition.TryRemove(item.CacheKey, out _))
147+
if (AccessTokenCacheDictionary.TryGetValue(partitionKey, out var partition))
134148
{
135-
_logger.InfoPii(
136-
() => $"[Internal cache] Cannot delete access token because it was not found in the cache. Key {item.CacheKey}.",
137-
() => "[Internal cache] Cannot delete access token because it was not found in the cache.");
149+
bool removed = partition.TryRemove(item.CacheKey, out _);
150+
if (removed)
151+
{
152+
Interlocked.Decrement(ref GetEntryCountRef());
153+
}
154+
else
155+
{
156+
_logger.InfoPii(
157+
() => $"[Internal cache] Cannot delete access token because it was not found in the cache. Key {item.CacheKey}.",
158+
() => "[Internal cache] Cannot delete access token because it was not found in the cache.");
159+
}
138160
}
139161
}
140162

@@ -219,6 +241,7 @@ public virtual void Clear(ILoggerAdapter requestlogger = null)
219241
{
220242
var logger = requestlogger ?? _logger;
221243
AccessTokenCacheDictionary.Clear();
244+
Interlocked.Exchange(ref GetEntryCountRef(), 0);
222245
logger.Always("[Internal cache] Clearing app token cache accessor.");
223246
// app metadata isn't removable
224247
}
@@ -227,5 +250,11 @@ public virtual bool HasAccessOrRefreshTokens()
227250
{
228251
return AccessTokenCacheDictionary.Any(partition => partition.Value.Any(token => !token.Value.IsExpiredWithBuffer()));
229252
}
253+
254+
private ref int GetEntryCountRef()
255+
{
256+
return ref _tokenCacheAccessorOptions.UseSharedCache ? ref s_entryCount : ref _entryCount;
257+
}
258+
230259
}
231260
}

‎src/client/Microsoft.Identity.Client/PlatformsCommon/Shared/InMemoryPartitionedUserTokenCacheAccessor.cs

Lines changed: 38 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -41,9 +41,15 @@ internal class InMemoryPartitionedUserTokenCacheAccessor : ITokenCacheAccessor
4141
private static readonly ConcurrentDictionary<string, MsalAppMetadataCacheItem> s_appMetadataDictionary =
4242
new ConcurrentDictionary<string, MsalAppMetadataCacheItem>();
4343

44+
private static int s_entryCount = 0;
45+
4446
protected readonly ILoggerAdapter _logger;
4547
private readonly CacheOptions _tokenCacheAccessorOptions;
4648

49+
private int _entryCount = 0;
50+
51+
public int EntryCount => GetEntryCountRef();
52+
4753
public InMemoryPartitionedUserTokenCacheAccessor(ILoggerAdapter logger, CacheOptions tokenCacheAccessorOptions)
4854
{
4955
_logger = logger ?? throw new ArgumentNullException(nameof(logger));
@@ -73,8 +79,17 @@ public void SaveAccessToken(MsalAccessTokenCacheItem item)
7379
string itemKey = item.CacheKey;
7480
string partitionKey = CacheKeyFactory.GetKeyFromCachedItem(item);
7581

76-
AccessTokenCacheDictionary
77-
.GetOrAdd(partitionKey, new ConcurrentDictionary<string, MsalAccessTokenCacheItem>())[itemKey] = item; // if a conflict occurs, pick the latest value
82+
var partition = AccessTokenCacheDictionary.GetOrAdd(partitionKey, _ => new ConcurrentDictionary<string, MsalAccessTokenCacheItem>());
83+
bool added = partition.TryAdd(itemKey, item);
84+
// only increment the entry count if this is a new item
85+
if (added)
86+
{
87+
System.Threading.Interlocked.Increment(ref GetEntryCountRef());
88+
}
89+
else
90+
{
91+
partition[itemKey] = item;
92+
}
7893
}
7994

8095
public void SaveRefreshToken(MsalRefreshTokenCacheItem item)
@@ -151,12 +166,19 @@ public void DeleteAccessToken(MsalAccessTokenCacheItem item)
151166
{
152167
string partitionKey = CacheKeyFactory.GetKeyFromCachedItem(item);
153168

154-
AccessTokenCacheDictionary.TryGetValue(partitionKey, out var partition);
155-
if (partition == null || !partition.TryRemove(item.CacheKey, out _))
169+
if (AccessTokenCacheDictionary.TryGetValue(partitionKey, out var partition))
156170
{
157-
_logger.InfoPii(
158-
() => $"[Internal cache] Cannot delete access token because it was not found in the cache. Key {item.CacheKey}.",
159-
() => "[Internal cache] Cannot delete access token because it was not found in the cache.");
171+
bool removed = partition.TryRemove(item.CacheKey, out _);
172+
if (removed)
173+
{
174+
System.Threading.Interlocked.Decrement(ref GetEntryCountRef());
175+
}
176+
else
177+
{
178+
_logger.InfoPii(
179+
() => $"[Internal cache] Cannot delete access token because it was not found in the cache. Key {item.CacheKey}.",
180+
() => "[Internal cache] Cannot delete access token because it was not found in the cache.");
181+
}
160182
}
161183
}
162184

@@ -246,7 +268,7 @@ public virtual List<MsalRefreshTokenCacheItem> GetAllRefreshTokens(string partit
246268
if (RefreshTokenCacheDictionary.Count == 1 && result.Count == 0)
247269
{
248270
logger.VerbosePii(
249-
() => $"[Internal cache] 0 RTs and 1 partition. Partition in cache is {RefreshTokenCacheDictionary.Keys.First()}",
271+
() => $"[Internal cache] 0 RTs and 1 partition. Partition in cache is {RefreshTokenCacheDictionary.Keys.First()}",
250272
() => "[Internal cache] 0 RTs and 1 partition] 0 RTs and 1 partition.");
251273
}
252274
}
@@ -290,7 +312,7 @@ public virtual List<MsalAccountCacheItem> GetAllAccounts(string partitionKey = n
290312
{
291313
AccountCacheDictionary.TryGetValue(partitionKey, out ConcurrentDictionary<string, MsalAccountCacheItem> partition);
292314
result = partition?.Select(kv => kv.Value)?.ToList() ?? CollectionHelpers.GetEmptyList<MsalAccountCacheItem>();
293-
315+
294316
if (logger.IsLoggingEnabled(LogLevel.Verbose))
295317
{
296318
logger.Verbose(() => $"[Internal cache] GetAllAccounts (with partition - exists? {partition != null}) found {result.Count} accounts.");
@@ -327,6 +349,8 @@ public virtual void Clear(ILoggerAdapter requestlogger = null)
327349
IdTokenCacheDictionary.Clear();
328350
AccountCacheDictionary.Clear();
329351
// app metadata isn't removable
352+
System.Threading.Interlocked.Exchange(ref GetEntryCountRef(), 0);
353+
330354
}
331355

332356
/// WARNING: this API is slow as it loads all tokens, not just from 1 partition.
@@ -336,5 +360,10 @@ public virtual bool HasAccessOrRefreshTokens()
336360
return RefreshTokenCacheDictionary.Any(partition => partition.Value.Count > 0) ||
337361
AccessTokenCacheDictionary.Any(partition => partition.Value.Any(token => !token.Value.IsExpiredWithBuffer()));
338362
}
363+
364+
private ref int GetEntryCountRef()
365+
{
366+
return ref _tokenCacheAccessorOptions.UseSharedCache ? ref s_entryCount : ref _entryCount;
367+
}
339368
}
340369
}

‎src/client/Microsoft.Identity.Client/PublicApi/net462/PublicAPI.Unshipped.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,3 +5,5 @@ Microsoft.Identity.Client.Utils.MacMainThreadScheduler.IsRunning() -> bool
55
Microsoft.Identity.Client.Utils.MacMainThreadScheduler.Stop() -> void
66
Microsoft.Identity.Client.Utils.MacMainThreadScheduler.RunOnMainThreadAsync(System.Func<System.Threading.Tasks.Task> asyncAction) -> System.Threading.Tasks.Task
77
Microsoft.Identity.Client.Utils.MacMainThreadScheduler.StartMessageLoop() -> void
8+
Microsoft.Identity.Client.AuthenticationResultMetadata.CachedAccessTokenCount.get -> int
9+
Microsoft.Identity.Client.AuthenticationResultMetadata.CachedAccessTokenCount.set -> void

‎src/client/Microsoft.Identity.Client/PublicApi/net472/PublicAPI.Unshipped.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,3 +5,5 @@ Microsoft.Identity.Client.Utils.MacMainThreadScheduler.IsRunning() -> bool
55
Microsoft.Identity.Client.Utils.MacMainThreadScheduler.Stop() -> void
66
Microsoft.Identity.Client.Utils.MacMainThreadScheduler.RunOnMainThreadAsync(System.Func<System.Threading.Tasks.Task> asyncAction) -> System.Threading.Tasks.Task
77
Microsoft.Identity.Client.Utils.MacMainThreadScheduler.StartMessageLoop() -> void
8+
Microsoft.Identity.Client.AuthenticationResultMetadata.CachedAccessTokenCount.get -> int
9+
Microsoft.Identity.Client.AuthenticationResultMetadata.CachedAccessTokenCount.set -> void

0 commit comments

Comments
 (0)