update

Azure · Dec 10, 2024 · 1dd8e0d · 1dd8e0d
1 parent 8eda0f9
commit 1dd8e0d
Showing 1 changed file with 8 additions and 7 deletions.
diff --git a/src/Microsoft.Azure.SignalR.Common/Auth/MicrosoftEntra/MicrosoftEntraAccessKey.cs b/src/Microsoft.Azure.SignalR.Common/Auth/MicrosoftEntra/MicrosoftEntraAccessKey.cs
@@ -120,8 +120,7 @@ public async Task<string> GenerateAccessTokenAsync(string audience,
         var updateTask = Task.CompletedTask;
         if (NeedRefresh)
         {
-            var source = new CancellationTokenSource(Constants.Periods.DefaultUpdateAccessKeyTimeout);
-            updateTask = UpdateAccessKeyAsync(source.Token).OrCancelAsync(ctoken);
+            updateTask = UpdateAccessKeyAsync().OrCancelAsync(ctoken);
         }
 
         if (!Available)
@@ -147,21 +146,23 @@ internal void UpdateAccessKey(string kid, string keyStr)
 
     internal async Task<bool> UpdateAccessKeyAsync(CancellationToken ctoken = default)
     {
-        TaskCompletionSource<bool> source;
+        TaskCompletionSource<bool> tcs;
         lock (_lock)
         {
             if (!_updateTaskSource.Task.IsCompleted)
             {
-                source = _updateTaskSource;
+                tcs = _updateTaskSource;
             }
             else
             {
                 _updateTaskSource = new TaskCompletionSource<bool>(TaskCreationOptions.RunContinuationsAsynchronously);
-                _ = UpdateAccessKeyInternalAsync(_updateTaskSource, ctoken);
-                source = _updateTaskSource;
+                var cts = new CancellationTokenSource(Constants.Periods.DefaultUpdateAccessKeyTimeout);
+                var linkedSource = CancellationTokenSource.CreateLinkedTokenSource(ctoken, cts.Token);
+                _ = UpdateAccessKeyInternalAsync(_updateTaskSource, linkedSource.Token);
+                tcs = _updateTaskSource;
             }
         }
-        return await source.Task;
+        return await tcs.Task;
     }
 
     internal async Task UpdateAccessKeyInternalAsync(TaskCompletionSource<bool> tcs, CancellationToken ctoken)