Skip to content

Commit 93de0d7

Browse files
cportcventcportcvent
committed
fix: #330 address possible redundant token exchange with refactor to strict mutex; includes regression test
1 parent f879fde commit 93de0d7

2 files changed

Lines changed: 169 additions & 30 deletions

File tree

‎src/main/java/dev/openfga/sdk/api/auth/OAuth2Client.java‎

Lines changed: 40 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -52,43 +52,53 @@ public OAuth2Client(Configuration configuration, ApiClient apiClient) throws Fga
5252
* @return An access token in a {@link CompletableFuture}
5353
*/
5454
public CompletableFuture<String> getAccessToken() throws FgaInvalidParameterException, ApiException {
55+
// Fast path (lock-free): return cached token if still valid.
5556
AccessToken current = snapshot.get();
5657
if (current.isValid()) {
5758
return CompletableFuture.completedFuture(current.token());
5859
}
5960

60-
CompletableFuture<String> promise = new CompletableFuture<>();
61-
if (!inFlight.compareAndSet(null, promise)) {
62-
// Another thread won the race — join its exchange rather than starting a new one.
63-
CompletableFuture<String> existing = inFlight.get();
64-
return existing != null ? existing : getAccessToken();
65-
}
61+
// Slow path: decide under the lock who starts the exchange.
62+
synchronized (this) {
63+
// Double-check: another thread may have refreshed while we waited.
64+
AccessToken rechecked = snapshot.get();
65+
if (rechecked.isValid()) {
66+
return CompletableFuture.completedFuture(rechecked.token());
67+
}
6668

67-
// This thread owns the exchange. Start it, wiring completion back to `promise`.
68-
try {
69-
exchangeToken().whenComplete((response, ex) -> {
70-
if (ex != null) {
71-
inFlight.set(null);
72-
promise.completeExceptionally(ex);
73-
} else {
74-
String token = response.getAccessToken();
75-
// Write snapshot before clearing the gate so any new caller that arrives
76-
// after inFlight becomes null immediately sees a valid token.
77-
snapshot.set(new AccessToken(token, Instant.now().plusSeconds(response.getExpiresInSeconds())));
78-
79-
// Clear before completing
80-
inFlight.set(null);
81-
promise.complete(token);
82-
telemetry.metrics().credentialsRequest(1L, new HashMap<>());
83-
}
84-
});
85-
} catch (Exception e) {
86-
inFlight.set(null);
87-
promise.completeExceptionally(e);
88-
throw e;
69+
// Join an existing in-flight exchange.
70+
CompletableFuture<String> existing = inFlight.get();
71+
if (existing != null) {
72+
return existing;
73+
}
74+
75+
// Start a new exchange and publish the future so other callers join it.
76+
CompletableFuture<String> promise = new CompletableFuture<>();
77+
inFlight.set(promise);
78+
79+
try {
80+
exchangeToken().whenComplete((response, ex) -> {
81+
if (ex != null) {
82+
inFlight.set(null);
83+
promise.completeExceptionally(ex);
84+
} else {
85+
String token = response.getAccessToken();
86+
// Write snapshot before clearing the gate so any new caller that arrives
87+
// after inFlight becomes null immediately sees a valid token.
88+
snapshot.set(new AccessToken(token, Instant.now().plusSeconds(response.getExpiresInSeconds())));
89+
inFlight.set(null);
90+
promise.complete(token);
91+
telemetry.metrics().credentialsRequest(1L, new HashMap<>());
92+
}
93+
});
94+
} catch (Exception e) {
95+
inFlight.set(null);
96+
promise.completeExceptionally(e);
97+
throw e;
98+
}
99+
100+
return promise;
89101
}
90-
91-
return promise;
92102
}
93103

94104
/**

‎src/test/java/dev/openfga/sdk/api/auth/OAuth2ClientTest.java‎

Lines changed: 129 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -209,6 +209,135 @@ void exchangeOAuth2Token_concurrentRequests_singleExchange(WireMockRuntimeInfo w
209209
verify(1, postRequestedFor(urlEqualTo("/oauth/token")));
210210
}
211211

212+
/**
213+
* Regression: after a successful exchange, a second wave of concurrent callers must see the
214+
* cached snapshot and NOT trigger a redundant exchange.
215+
* This covers a race where thread A reads an invalid snapshot, thread B completes the exchange,
216+
* and thread A then enters the slow path synchronized block - but still returns the cached
217+
* token without issuing another request.
218+
*/
219+
@Test
220+
void getAccessToken_cachedTokenHit(WireMockRuntimeInfo wm) throws Exception {
221+
stubFor(post(urlEqualTo("/oauth/token"))
222+
.willReturn(ok(String.format("{\"access_token\":\"%s\",\"expires_in\":3600}", ACCESS_TOKEN))
223+
.withFixedDelay(100)));
224+
225+
OAuth2Client client = newOAuth2Client(wm.getHttpBaseUrl(), false);
226+
227+
// Wave 1 — triggers the only exchange.
228+
int wave1Count = 3;
229+
CountDownLatch wave1Start = new CountDownLatch(1);
230+
CountDownLatch wave1Done = new CountDownLatch(wave1Count);
231+
List<String> wave1Tokens = Collections.synchronizedList(new ArrayList<>());
232+
List<Throwable> failures = Collections.synchronizedList(new ArrayList<>());
233+
234+
for (int i = 0; i < wave1Count; i++) {
235+
new Thread(() -> {
236+
try {
237+
wave1Start.await();
238+
wave1Tokens.add(client.getAccessToken().get());
239+
} catch (Exception e) {
240+
failures.add(e);
241+
} finally {
242+
wave1Done.countDown();
243+
}
244+
})
245+
.start();
246+
}
247+
wave1Start.countDown();
248+
assertTrue(wave1Done.await(2, TimeUnit.SECONDS), "wave 1 did not complete in time");
249+
assertTrue(failures.isEmpty(), "wave 1 should not have failures");
250+
251+
// Wave 2 — arrives after the exchange has completed; must use the cached token.
252+
int wave2Count = 5;
253+
CountDownLatch wave2Start = new CountDownLatch(1);
254+
CountDownLatch wave2Done = new CountDownLatch(wave2Count);
255+
List<String> wave2Tokens = Collections.synchronizedList(new ArrayList<>());
256+
257+
for (int i = 0; i < wave2Count; i++) {
258+
new Thread(() -> {
259+
try {
260+
wave2Start.await();
261+
wave2Tokens.add(client.getAccessToken().get());
262+
} catch (Exception e) {
263+
failures.add(e);
264+
} finally {
265+
wave2Done.countDown();
266+
}
267+
})
268+
.start();
269+
}
270+
wave2Start.countDown();
271+
assertTrue(wave2Done.await(2, TimeUnit.SECONDS), "wave 2 did not complete in time");
272+
273+
assertEquals(List.of(), failures, "no thread should have thrown");
274+
assertEquals(wave1Count, wave1Tokens.size());
275+
assertEquals(wave2Count, wave2Tokens.size());
276+
assertTrue(wave1Tokens.stream().allMatch(ACCESS_TOKEN::equals));
277+
assertTrue(wave2Tokens.stream().allMatch(ACCESS_TOKEN::equals));
278+
279+
// Only one exchange ever happened across both waves.
280+
verify(1, postRequestedFor(urlEqualTo("/oauth/token")));
281+
}
282+
283+
/**
284+
* Regression: a late wave of callers that arrives while the exchange is still in-flight
285+
* must join the existing future rather than starting a second exchange.
286+
*/
287+
@Test
288+
void getAccessToken_joinInFlightExchange(WireMockRuntimeInfo wm) throws Exception {
289+
stubFor(post(urlEqualTo("/oauth/token"))
290+
.willReturn(ok(String.format("{\"access_token\":\"%s\",\"expires_in\":3600}", ACCESS_TOKEN))
291+
.withFixedDelay(300)));
292+
293+
OAuth2Client client = newOAuth2Client(wm.getHttpBaseUrl(), false);
294+
295+
List<String> allTokens = Collections.synchronizedList(new ArrayList<>());
296+
List<Throwable> failures = Collections.synchronizedList(new ArrayList<>());
297+
298+
// Wave 1 — triggers the exchange (300 ms delay).
299+
int wave1Count = 2;
300+
CountDownLatch wave1Start = new CountDownLatch(1);
301+
CountDownLatch allDone = new CountDownLatch(wave1Count + 3);
302+
303+
for (int i = 0; i < wave1Count; i++) {
304+
new Thread(() -> {
305+
try {
306+
wave1Start.await();
307+
allTokens.add(client.getAccessToken().get());
308+
} catch (Exception e) {
309+
failures.add(e);
310+
} finally {
311+
allDone.countDown();
312+
}
313+
})
314+
.start();
315+
}
316+
wave1Start.countDown();
317+
318+
// Wave 2 — arrives 50 ms later while exchange is still in-flight.
319+
Thread.sleep(50);
320+
int wave2Count = 3;
321+
for (int i = 0; i < wave2Count; i++) {
322+
new Thread(() -> {
323+
try {
324+
allTokens.add(client.getAccessToken().get());
325+
} catch (Exception e) {
326+
failures.add(e);
327+
} finally {
328+
allDone.countDown();
329+
}
330+
})
331+
.start();
332+
}
333+
334+
assertTrue(allDone.await(5, TimeUnit.SECONDS), "threads did not complete in time");
335+
assertEquals(List.of(), failures, "no thread should have thrown");
336+
assertEquals(wave1Count + wave2Count, allTokens.size());
337+
assertTrue(allTokens.stream().allMatch(ACCESS_TOKEN::equals));
338+
verify(1, postRequestedFor(urlEqualTo("/oauth/token")));
339+
}
340+
212341
@Test
213342
public void apiTokenIssuer_invalidScheme() {
214343
// When

0 commit comments

Comments
 (0)