|
| 1 | +// Copyright (c) 2026 dexpace and Omar Aljarrah. |
| 2 | +// Licensed under the MIT License. See LICENSE in the repository root for details. |
| 3 | + |
| 4 | +namespace Dexpace.Sdk.Core.Auth; |
| 5 | + |
| 6 | +/// <summary> |
| 7 | +/// An in-memory token cache that wraps a <see cref="TokenCredential"/> and provides |
| 8 | +/// proactive refresh, expiry-based invalidation, and single-flight protection against |
| 9 | +/// concurrent refresh stampedes. |
| 10 | +/// </summary> |
| 11 | +/// <remarks> |
| 12 | +/// <para> |
| 13 | +/// Tokens are cached keyed by the <see cref="TokenRequestContext.CacheKey"/>. A cached token |
| 14 | +/// is served without calling the underlying credential as long as both: |
| 15 | +/// <list type="bullet"> |
| 16 | +/// <item><c>now < ExpiresOn</c></item> |
| 17 | +/// <item><c>RefreshOn</c> is <see langword="null"/> OR <c>now < RefreshOn</c></item> |
| 18 | +/// </list> |
| 19 | +/// </para> |
| 20 | +/// <para> |
| 21 | +/// Once either condition fails, a single caller acquires the per-key semaphore and calls |
| 22 | +/// the credential. All concurrent callers wait on the semaphore and perform a double-checked |
| 23 | +/// read after acquiring it. Only one network round-trip fires per key per refresh cycle. |
| 24 | +/// </para> |
| 25 | +/// <para> |
| 26 | +/// <strong>Failure while valid:</strong> if the credential throws but a token remains valid |
| 27 | +/// (<c>now < ExpiresOn</c>), the cached token is returned silently. If no valid token |
| 28 | +/// exists the exception propagates. |
| 29 | +/// </para> |
| 30 | +/// <para> |
| 31 | +/// This class is thread-safe. Inject a custom <see cref="TimeProvider"/> for deterministic |
| 32 | +/// unit testing. |
| 33 | +/// </para> |
| 34 | +/// </remarks> |
| 35 | +public sealed class AccessTokenCache |
| 36 | +{ |
| 37 | + private readonly TokenCredential _credential; |
| 38 | + private readonly TimeProvider _time; |
| 39 | + |
| 40 | + // Per-key state: the current cached token (null = never fetched) + a semaphore for |
| 41 | + // single-flight. One CacheEntry per unique TokenRequestContext.CacheKey. |
| 42 | + private readonly System.Collections.Concurrent.ConcurrentDictionary<string, CacheEntry> _entries = new(); |
| 43 | + |
| 44 | + /// <summary> |
| 45 | + /// Initializes an <see cref="AccessTokenCache"/> backed by the given credential. |
| 46 | + /// </summary> |
| 47 | + /// <param name="credential">The underlying credential to call when a token is needed.</param> |
| 48 | + /// <param name="timeProvider"> |
| 49 | + /// A <see cref="TimeProvider"/> used to determine "now". Defaults to |
| 50 | + /// <see cref="TimeProvider.System"/> when <see langword="null"/>. |
| 51 | + /// </param> |
| 52 | + /// <exception cref="ArgumentNullException"><paramref name="credential"/> is <see langword="null"/>.</exception> |
| 53 | + public AccessTokenCache(TokenCredential credential, TimeProvider? timeProvider = null) |
| 54 | + { |
| 55 | + ArgumentNullException.ThrowIfNull(credential); |
| 56 | + _credential = credential; |
| 57 | + _time = timeProvider ?? TimeProvider.System; |
| 58 | + } |
| 59 | + |
| 60 | + /// <summary> |
| 61 | + /// Returns a valid <see cref="AccessToken"/> for the given <paramref name="context"/>, |
| 62 | + /// fetching and caching one if necessary. |
| 63 | + /// </summary> |
| 64 | + /// <param name="context">The token request context identifying the scopes and claims.</param> |
| 65 | + /// <param name="ct">A token to cancel an in-progress credential call.</param> |
| 66 | + /// <returns> |
| 67 | + /// A <see cref="ValueTask{AccessToken}"/> resolving to a valid access token. |
| 68 | + /// </returns> |
| 69 | + public async ValueTask<AccessToken> GetAsync(TokenRequestContext context, CancellationToken ct = default) |
| 70 | + { |
| 71 | + var entry = _entries.GetOrAdd(context.CacheKey, static _ => new CacheEntry()); |
| 72 | + var now = _time.GetUtcNow(); |
| 73 | + |
| 74 | + // Fast path: read the holder once (volatile acquire) and return immediately when the |
| 75 | + // token is still valid and does not need a proactive refresh. The volatile read ensures |
| 76 | + // that any holder published by a slow-path writer is fully visible to this thread. |
| 77 | + var fastHolder = entry.Holder; |
| 78 | + if (fastHolder is not null && IsValid(fastHolder.Token, now) && !NeedsRefresh(fastHolder.Token, now)) |
| 79 | + { |
| 80 | + return fastHolder.Token; |
| 81 | + } |
| 82 | + |
| 83 | + // Slow path: acquire the semaphore for this key. |
| 84 | + await entry.Semaphore.WaitAsync(ct).ConfigureAwait(false); |
| 85 | + try |
| 86 | + { |
| 87 | + // Double-check after acquiring. |
| 88 | + now = _time.GetUtcNow(); |
| 89 | + var recheckedHolder = entry.Holder; |
| 90 | + if (recheckedHolder is not null && IsValid(recheckedHolder.Token, now) && !NeedsRefresh(recheckedHolder.Token, now)) |
| 91 | + { |
| 92 | + return recheckedHolder.Token; |
| 93 | + } |
| 94 | + |
| 95 | + // Try to refresh from the credential. |
| 96 | + try |
| 97 | + { |
| 98 | + var fresh = await _credential.GetTokenAsync(context, ct).ConfigureAwait(false); |
| 99 | + // Volatile write releases the fully-constructed holder to all threads. |
| 100 | + entry.Holder = new TokenHolder(fresh); |
| 101 | + return fresh; |
| 102 | + } |
| 103 | + catch |
| 104 | + { |
| 105 | + // Re-sample time: GetTokenAsync may have taken a long time. Validate the |
| 106 | + // cached token against the clock *after* the (slow) failing call, so a token |
| 107 | + // that expired during the attempt is not returned as valid. |
| 108 | + var nowAfterFailure = _time.GetUtcNow(); |
| 109 | + var fallbackHolder = entry.Holder; |
| 110 | + if (fallbackHolder is not null && IsValid(fallbackHolder.Token, nowAfterFailure)) |
| 111 | + { |
| 112 | + return fallbackHolder.Token; |
| 113 | + } |
| 114 | + |
| 115 | + throw; |
| 116 | + } |
| 117 | + } |
| 118 | + finally |
| 119 | + { |
| 120 | + entry.Semaphore.Release(); |
| 121 | + } |
| 122 | + } |
| 123 | + |
| 124 | + // A token is "valid" while the clock hasn't yet reached ExpiresOn. |
| 125 | + private static bool IsValid(AccessToken token, DateTimeOffset now) => |
| 126 | + now < token.ExpiresOn; |
| 127 | + |
| 128 | + // A token "needs refresh" once RefreshOn has been reached (proactive hint). |
| 129 | + private static bool NeedsRefresh(AccessToken token, DateTimeOffset now) => |
| 130 | + token.RefreshOn is { } refreshOn && now >= refreshOn; |
| 131 | + |
| 132 | + // Immutable wrapper so the token can be published through a single volatile reference, |
| 133 | + // giving acquire/release ordering on all platforms. AccessToken is a multi-field struct; |
| 134 | + // publishing it directly would allow fast-path readers to observe a torn or partially |
| 135 | + // written value. Boxing it inside a reference type (TokenHolder) reduces the published |
| 136 | + // value to a single pointer-width write, which the .NET memory model guarantees to be |
| 137 | + // atomic and not torn. |
| 138 | + private sealed class TokenHolder |
| 139 | + { |
| 140 | + /// <summary>Initializes a <see cref="TokenHolder"/> wrapping the given token.</summary> |
| 141 | + /// <param name="token">The token to publish atomically.</param> |
| 142 | + public TokenHolder(AccessToken token) => Token = token; |
| 143 | + |
| 144 | + /// <summary>The wrapped access token.</summary> |
| 145 | + public AccessToken Token { get; } |
| 146 | + } |
| 147 | + |
| 148 | + private sealed class CacheEntry |
| 149 | + { |
| 150 | + // volatile ensures that a fast-path reader always observes the most recently published |
| 151 | + // holder (acquire semantics on read, release semantics on write). Combined with the |
| 152 | + // pointer-width atomicity of reference reads/writes, this is safe without a lock on |
| 153 | + // the fast path. |
| 154 | + private volatile TokenHolder? _holder; |
| 155 | + |
| 156 | + /// <summary> |
| 157 | + /// The most recently published token holder, or <see langword="null"/> if no token |
| 158 | + /// has been fetched yet. Reads and writes are reference-atomic and carry |
| 159 | + /// acquire/release ordering via the <see langword="volatile"/> modifier. |
| 160 | + /// </summary> |
| 161 | + public TokenHolder? Holder |
| 162 | + { |
| 163 | + get => _holder; |
| 164 | + set => _holder = value; |
| 165 | + } |
| 166 | + |
| 167 | + /// <summary> |
| 168 | + /// A semaphore (initial count = 1) that serializes refresh calls for this key. |
| 169 | + /// </summary> |
| 170 | + public SemaphoreSlim Semaphore { get; } = new SemaphoreSlim(1, 1); |
| 171 | + } |
| 172 | +} |
0 commit comments