Skip to content

Commit 9d12974

Browse files
prasanna164-codeprasanna164-code
authored
[Core] Don't leave callers waiting forever when a token refresh is cancelled (#50856)
* Don't leave callers waiting forever when a token refresh is cancelled AccessTokenCache and SimpleTokenCache let one caller refresh the token while concurrent callers wait on a Sinks.One. When the refreshing caller was cancelled, the Mono.using cleanup only reset wip and never completed the sink, so every waiting caller stayed pending indefinitely. The cleanup now also completes the sink empty, a no-op when the refresh already completed it. Released callers use the current cached token if it is still valid and otherwise acquire a new one. The branch for an already finished refresh uses the same fallback instead of returning a captured token that could be null or expired. Fixes #50855 * Release waiting token callers only after clearing the in-progress refresh When the token request completed empty, the NO REFRESH branch released waiting callers while wip still held the finished refresh. In SimpleTokenCache a released waiter then re-entered the acquisition, found the same completed sink, and recursed until StackOverflowError. Waiting callers are now released only by the Mono.using cleanup, after wip has been cleared, so a retrying waiter always starts a new acquisition. SimpleTokenCache retries through a private acquireToken() instead of the overridable getToken().
1 parent fcd438d commit 9d12974

5 files changed

Lines changed: 166 additions & 8 deletions

File tree

‎sdk/core/azure-core/CHANGELOG.md‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,10 @@
88

99
### Bugs Fixed
1010

11+
- Fixed `AccessTokenCache` and `SimpleTokenCache` leaving concurrent callers waiting forever when the caller that was
12+
refreshing the token was cancelled, for example by a timeout. Waiting callers now retry the token acquisition when
13+
the refresh they waited on ends without a token, including when `SimpleTokenCache`'s token supplier completes empty.
14+
1115
### Other Changes
1216

1317
#### Dependency Updates

‎sdk/core/azure-core/src/main/java/com/azure/core/credential/AccessTokenCache.java‎

Lines changed: 21 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,13 @@ private Supplier<Mono<? extends AccessToken>> retrieveToken(TokenRequestContext
208208
ignored -> tokenRefresh.materialize()
209209
.flatMap(processTokenRefreshResult(sinksOne, now, fallback))
210210
.doOnError(sinksOne::tryEmitError),
211-
w -> w.set(null));
211+
w -> {
212+
w.set(null);
213+
// If the refresh ended without a token, because it was cancelled or completed empty,
214+
// release the callers waiting on it. This happens after wip is cleared so that they can
215+
// start a new acquisition. No-op if the refresh already completed the sink.
216+
sinksOne.tryEmitEmpty();
217+
});
212218
} else if (cachedToken != null && !cachedToken.isExpired() && !refreshOnContextChange) {
213219
// another thread might be refreshing the token proactively, but the current token is still valid
214220
return Mono.just(cachedToken);
@@ -221,10 +227,10 @@ private Supplier<Mono<? extends AccessToken>> retrieveToken(TokenRequestContext
221227
Sinks.One<AccessToken> sinksOne = wip.get();
222228
if (sinksOne == null) {
223229
// the refreshing thread has finished
224-
return Mono.just(cachedToken);
230+
return validCachedToken();
225231
} else {
226232
// wait for refreshing thread to finish but defer to updated cache in case just missed onNext()
227-
return sinksOne.asMono().switchIfEmpty(Mono.fromSupplier(() -> cachedToken));
233+
return sinksOne.asMono().switchIfEmpty(validCachedToken());
228234
}
229235
}
230236
} catch (Exception ex) {
@@ -308,6 +314,17 @@ private Supplier<AccessToken> retrieveTokenSync(TokenRequestContext tokenRequest
308314
};
309315
}
310316

317+
/**
318+
* Emits the currently cached token if it hasn't expired, otherwise completes empty so that
319+
* {@link #getToken(TokenRequestContext, boolean)} retries the acquisition.
320+
*/
321+
private Mono<AccessToken> validCachedToken() {
322+
return Mono.defer(() -> {
323+
AccessToken latest = cacheInfo.get().getCachedAccessToken();
324+
return (latest == null || latest.isExpired()) ? Mono.empty() : Mono.just(latest);
325+
});
326+
}
327+
311328
private boolean checkIfForceRefreshRequired(TokenRequestContext tokenRequestContext) {
312329
return !(this.tokenRequestContext != null
313330
&& (this.tokenRequestContext.getClaims() == null
@@ -339,7 +356,7 @@ private boolean checkIfForceRefreshRequired(TokenRequestContext tokenRequestCont
339356
cacheInfo.set(new AccessTokenCacheInfo(cache, nextTokenRefresh));
340357
return fallback.switchIfEmpty(Mono.error(error));
341358
} else { // NO REFRESH
342-
sinksOne.tryEmitEmpty();
359+
// Waiting callers are released by the Mono.using cleanup, once wip has been cleared.
343360
return fallback;
344361
}
345362
};

‎sdk/core/azure-core/src/main/java/com/azure/core/credential/SimpleTokenCache.java‎

Lines changed: 25 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,10 @@ public SimpleTokenCache(Supplier<Mono<AccessToken>> tokenSupplier) {
105105
* @return a Publisher that emits an AccessToken
106106
*/
107107
public Mono<AccessToken> getToken() {
108+
return acquireToken();
109+
}
110+
111+
private Mono<AccessToken> acquireToken() {
108112
return Mono.defer(() -> {
109113
try {
110114
if (wip.compareAndSet(null, Sinks.one())) {
@@ -156,10 +160,16 @@ public Mono<AccessToken> getToken() {
156160
nextTokenRefresh = OffsetDateTime.now().plus(refreshDelay);
157161
return fallback.switchIfEmpty(Mono.error(() -> error));
158162
} else { // NO REFRESH
159-
sinksOne.tryEmitEmpty();
163+
// Waiting callers are released by the cleanup below, once wip has been cleared.
160164
return fallback;
161165
}
162-
}).doOnError(sinksOne::tryEmitError), w -> w.set(null));
166+
}).doOnError(sinksOne::tryEmitError), w -> {
167+
w.set(null);
168+
// If the refresh ended without a token, because it was cancelled or completed empty, release the
169+
// callers waiting on it. This happens after wip is cleared so that they can start a new
170+
// acquisition. No-op if the refresh already completed the sink with a token or an error.
171+
sinksOne.tryEmitEmpty();
172+
});
163173
} else if (cache != null && !cache.isExpired()) {
164174
// another thread might be refreshing the token proactively, but the current token is still valid
165175
return Mono.just(cache);
@@ -168,10 +178,10 @@ public Mono<AccessToken> getToken() {
168178
Sinks.One<AccessToken> sinksOne = wip.get();
169179
if (sinksOne == null) {
170180
// the refreshing thread has finished
171-
return Mono.just(cache);
181+
return validCachedTokenOrRefresh();
172182
} else {
173183
// wait for refreshing thread to finish but defer to updated cache in case just missed onNext()
174-
return sinksOne.asMono().switchIfEmpty(Mono.fromSupplier(() -> cache));
184+
return sinksOne.asMono().switchIfEmpty(validCachedTokenOrRefresh());
175185
}
176186
}
177187
} catch (Exception t) {
@@ -180,6 +190,17 @@ public Mono<AccessToken> getToken() {
180190
});
181191
}
182192

193+
/**
194+
* Emits the currently cached token if it hasn't expired, otherwise starts a new token acquisition. This covers
195+
* waiting callers whose refresh ended without a token, for example because the refreshing caller was cancelled.
196+
*/
197+
private Mono<AccessToken> validCachedTokenOrRefresh() {
198+
return Mono.defer(() -> {
199+
AccessToken latest = cache;
200+
return (latest != null && !latest.isExpired()) ? Mono.just(latest) : acquireToken();
201+
});
202+
}
203+
183204
Sinks.One<AccessToken> getWipValue() {
184205
return wip.get();
185206
}

‎sdk/core/azure-core/src/test/java/com/azure/core/credential/AccessTokenCacheTests.java‎

Lines changed: 52 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@
44
package com.azure.core.credential;
55

66
import org.junit.jupiter.api.Test;
7+
import reactor.core.Disposable;
78
import reactor.core.publisher.Mono;
9+
import reactor.core.publisher.Sinks;
810
import reactor.test.StepVerifier;
911

1012
import java.time.Duration;
@@ -289,4 +291,54 @@ public void testEmptyStringTenantIdTreatedAsDifferentFromNull() {
289291

290292
assertEquals(2, refreshCount.get());
291293
}
294+
295+
@Test
296+
public void waitingCallerGetsTokenWhenRefreshingCallerIsCancelled() {
297+
// The first token request never completes, any later one succeeds.
298+
AtomicInteger refreshCount = new AtomicInteger(0);
299+
TokenCredential credential = request -> refreshCount.incrementAndGet() == 1
300+
? Mono.never()
301+
: Mono.just(new AccessToken("token-" + refreshCount.get(), OffsetDateTime.now().plusHours(1)));
302+
AccessTokenCache cache = new AccessTokenCache(credential);
303+
TokenRequestContext context = new TokenRequestContext().addScopes(SCOPE);
304+
305+
// The first caller starts the token refresh, the second caller waits on that refresh.
306+
Disposable refreshingCaller = cache.getToken(context, false).subscribe();
307+
assertEquals(1, refreshCount.get());
308+
Mono<AccessToken> waitingCaller = cache.getToken(context, false);
309+
310+
StepVerifier.create(waitingCaller)
311+
.then(refreshingCaller::dispose)
312+
.assertNext(token -> assertEquals("token-2", token.getToken()))
313+
.expectComplete()
314+
.verify(Duration.ofSeconds(10));
315+
316+
assertEquals(2, refreshCount.get());
317+
}
318+
319+
@Test
320+
public void waitingCallerGetsTokenWhenRefreshCompletesEmpty() {
321+
// The first token request completes empty when the test says so, any later one succeeds.
322+
Sinks.One<AccessToken> firstRequest = Sinks.one();
323+
AtomicInteger refreshCount = new AtomicInteger(0);
324+
TokenCredential credential = request -> refreshCount.incrementAndGet() == 1
325+
? firstRequest.asMono()
326+
: Mono.just(new AccessToken("token-" + refreshCount.get(), OffsetDateTime.now().plusHours(1)));
327+
AccessTokenCache cache = new AccessTokenCache(credential);
328+
TokenRequestContext context = new TokenRequestContext().addScopes(SCOPE);
329+
330+
// The first caller starts the token refresh, the second caller waits on that refresh.
331+
Disposable refreshingCaller = cache.getToken(context, false).subscribe();
332+
assertEquals(1, refreshCount.get());
333+
Mono<AccessToken> waitingCaller = cache.getToken(context, false);
334+
335+
StepVerifier.create(waitingCaller)
336+
.then(firstRequest::tryEmitEmpty)
337+
.assertNext(token -> assertEquals("token-2", token.getToken()))
338+
.expectComplete()
339+
.verify(Duration.ofSeconds(10));
340+
341+
assertEquals(2, refreshCount.get());
342+
refreshingCaller.dispose();
343+
}
292344
}

‎sdk/core/azure-core/src/test/java/com/azure/core/credential/SimpleTokenCacheTests.java‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,16 +3,22 @@
33
package com.azure.core.credential;
44

55
import org.junit.jupiter.api.Test;
6+
import reactor.core.Disposable;
67
import reactor.core.publisher.Mono;
8+
import reactor.core.publisher.Sinks;
79
import reactor.test.StepVerifier;
810

911
import java.time.Duration;
1012
import java.time.OffsetDateTime;
13+
import java.util.concurrent.CountDownLatch;
14+
import java.util.concurrent.TimeUnit;
15+
import java.util.concurrent.atomic.AtomicInteger;
1116
import java.util.concurrent.atomic.AtomicLong;
1217

1318
import static org.junit.jupiter.api.Assertions.assertEquals;
1419
import static org.junit.jupiter.api.Assertions.assertNotNull;
1520
import static org.junit.jupiter.api.Assertions.assertNull;
21+
import static org.junit.jupiter.api.Assertions.assertTrue;
1622

1723
/**
1824
* Tests {@link SimpleTokenCache}.
@@ -35,6 +41,64 @@ public void wipResetsOnCancel() {
3541
assertNull(simpleTokenCache.getWipValue());
3642
}
3743

44+
@Test
45+
public void waitingCallerGetsTokenWhenRefreshingCallerIsCancelled() throws InterruptedException {
46+
// The first token request never completes, any later one succeeds.
47+
AtomicInteger refreshCount = new AtomicInteger(0);
48+
CountDownLatch firstRequestStarted = new CountDownLatch(1);
49+
SimpleTokenCache simpleTokenCache = new SimpleTokenCache(() -> {
50+
if (refreshCount.incrementAndGet() == 1) {
51+
firstRequestStarted.countDown();
52+
return Mono.never();
53+
}
54+
return Mono.just(new AccessToken("token-" + refreshCount.get(), OffsetDateTime.now().plusHours(1)));
55+
}, Duration.ZERO);
56+
57+
// The first caller starts the token refresh, the second caller waits on that refresh.
58+
Disposable refreshingCaller = simpleTokenCache.getToken().subscribe();
59+
// A new cache starts its first token request asynchronously.
60+
assertTrue(firstRequestStarted.await(10, TimeUnit.SECONDS));
61+
Mono<AccessToken> waitingCaller = simpleTokenCache.getToken();
62+
63+
StepVerifier.create(waitingCaller)
64+
.then(refreshingCaller::dispose)
65+
.assertNext(token -> assertEquals("token-2", token.getToken()))
66+
.expectComplete()
67+
.verify(Duration.ofSeconds(10));
68+
69+
assertEquals(2, refreshCount.get());
70+
}
71+
72+
@Test
73+
public void waitingCallerGetsTokenWhenRefreshCompletesEmpty() throws InterruptedException {
74+
// The first token request completes empty when the test says so, any later one succeeds.
75+
Sinks.One<AccessToken> firstRequest = Sinks.one();
76+
AtomicInteger refreshCount = new AtomicInteger(0);
77+
CountDownLatch firstRequestStarted = new CountDownLatch(1);
78+
SimpleTokenCache simpleTokenCache = new SimpleTokenCache(() -> {
79+
if (refreshCount.incrementAndGet() == 1) {
80+
firstRequestStarted.countDown();
81+
return firstRequest.asMono();
82+
}
83+
return Mono.just(new AccessToken("token-" + refreshCount.get(), OffsetDateTime.now().plusHours(1)));
84+
}, Duration.ZERO);
85+
86+
// The first caller starts the token refresh, the second caller waits on that refresh.
87+
Disposable refreshingCaller = simpleTokenCache.getToken().subscribe();
88+
// A new cache starts its first token request asynchronously.
89+
assertTrue(firstRequestStarted.await(10, TimeUnit.SECONDS));
90+
Mono<AccessToken> waitingCaller = simpleTokenCache.getToken();
91+
92+
StepVerifier.create(waitingCaller)
93+
.then(firstRequest::tryEmitEmpty)
94+
.assertNext(token -> assertEquals("token-2", token.getToken()))
95+
.expectComplete()
96+
.verify(Duration.ofSeconds(10));
97+
98+
assertEquals(2, refreshCount.get());
99+
refreshingCaller.dispose();
100+
}
101+
38102
@Test
39103
public void testRefreshOnFlow() throws InterruptedException {
40104
AtomicLong refreshes = new AtomicLong(0);

0 commit comments

Comments
 (0)