some fb

src/Extensions/Azure/AccessTokenCache.cs

@@ -0,0 +1,50 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+using Azure.Core;
+
+namespace Microsoft.DurableTask.Extensions.Azure;
+
+/// <summary>
+/// Caches and manages refresh for Azure access tokens.
+/// </summary>
+internal sealed class AccessTokenCache
+{
+    readonly TokenCredential credential;
+    readonly TokenRequestContext context;
+    readonly TimeSpan margin;
+
+    AccessToken? token;
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="AccessTokenCache"/> class.
+    /// </summary>
+    /// <param name="credential">The token credential to use for authentication.</param>
+    /// <param name="context">The token request context.</param>
+    /// <param name="margin">The time margin to use for token refresh.</param>
+    public AccessTokenCache(TokenCredential credential, TokenRequestContext context, TimeSpan margin)
+    {
+        this.credential = credential;
+        this.context = context;
+        this.margin = margin;
+    }
+
+    /// <summary>
+    /// Gets a token, either from cache or by requesting a new one if needed.
+    /// </summary>
+    /// <param name="cancellationToken">A cancellation token.</param>
+    /// <returns>An access token.</returns>
+    public async Task<AccessToken> GetTokenAsync(CancellationToken cancellationToken)
+    {
+        DateTimeOffset nowWithMargin = DateTimeOffset.UtcNow.Add(this.margin);
+
+        if (this.token is null
+            || this.token.Value.RefreshOn < nowWithMargin
+            || this.token.Value.ExpiresOn < nowWithMargin)
+        {
+            this.token = await this.credential.GetTokenAsync(this.context, cancellationToken);
+        }
+
+        return this.token.Value;
+    }
+}

src/Extensions/Azure/DurableTaskSchedulerOptions.cs

@@ -81,11 +81,11 @@ internal GrpcChannel GetGrpcChannel()
         int processId = Environment.ProcessId;
         string workerId = this.WorkerId ?? $"{Environment.MachineName},{processId},{Guid.NewGuid():N}";
 
-        TokenCache? cache =
+        AccessTokenCache? cache =
             this.Credential is not null
-                ? new(
+                ? new AccessTokenCache(
                     this.Credential,
-                    new(new[] { $"{resourceId}/.default" }),
+                    new TokenRequestContext(new[] { $"{resourceId}/.default" }),
                     TimeSpan.FromMinutes(5))
                 : null;
 
@@ -95,13 +95,12 @@ this.Credential is not null
                 metadata.Add("taskhub", taskHubName);
                 metadata.Add("workerid", workerId);
 
-                if (cache is null)
+                if (cache == null)
                 {
                     return;
                 }
-
+                
                 AccessToken token = await cache.GetTokenAsync(context.CancellationToken);
-
                 metadata.Add("Authorization", $"Bearer {token.Token}");
             });
 
@@ -219,27 +218,4 @@ public static DurableTaskSchedulerOptions FromConnectionString(
         DurableTaskSchedulerOptions options = new(endpointAddress, connectionString.TaskHubName, credential);
         return options;
     }
-
-    sealed class TokenCache(TokenCredential credential, TokenRequestContext context, TimeSpan margin)
-    {
-        readonly TokenCredential credential = credential;
-        readonly TokenRequestContext context = context;
-        readonly TimeSpan margin = margin;
-
-        AccessToken? token;
-
-        public async Task<AccessToken> GetTokenAsync(CancellationToken cancellationToken)
-        {
-            DateTimeOffset nowWithMargin = DateTimeOffset.UtcNow.Add(this.margin);
-
-            if (this.token is null
-                || this.token.Value.RefreshOn < nowWithMargin
-                || this.token.Value.ExpiresOn < nowWithMargin)
-            {
-                this.token = await this.credential.GetTokenAsync(this.context, cancellationToken);
-            }
-
-            return this.token.Value;
-        }
-    }
 }