Skip to content

Commit 4cdb67c

Browse files
committed
fix(auth): resolve flakiness in ClientSideCredentialAccessBoundaryFactoryTest
Synchronize the decision to refresh and task creation to avoid race conditions. Use Double-Checked Locking to minimize lock contention on the hot path. Resolves: #12871
1 parent 058efbb commit 4cdb67c

2 files changed

Lines changed: 50 additions & 30 deletions

File tree

google-auth-library-java/cab-token-generator/java/com/google/auth/credentialaccessboundary/ClientSideCredentialAccessBoundaryFactory.java

Lines changed: 40 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -146,7 +146,7 @@ public class ClientSideCredentialAccessBoundaryFactory {
146146
private final Duration refreshMargin;
147147
private RefreshTask refreshTask;
148148
private final Object refreshLock = new byte[0];
149-
private IntermediateCredentials intermediateCredentials = null;
149+
private volatile IntermediateCredentials intermediateCredentials = null;
150150
private final Clock clock;
151151
private final CelCompiler celCompiler;
152152

@@ -234,8 +234,26 @@ void refreshCredentialsIfRequired() throws IOException {
234234
return;
235235
}
236236

237-
// If a refresh is required, create or retrieve the refresh task.
238-
RefreshTask currentRefreshTask = getOrCreateRefreshTask();
237+
RefreshTask currentRefreshTask;
238+
// Synchronize both the decision to refresh and the task creation to prevent a race
239+
// condition. Without this, multiple threads might concurrently determine a refresh is
240+
// needed (e.g., in ASYNC mode) and return ASYNC. The first thread would start the refresh
241+
// task. If that task completes quickly and clears the `refreshTask` reference, a subsequent
242+
// thread that also determined a refresh was needed would then see `refreshTask` as null
243+
// and incorrectly start a second, redundant refresh task.
244+
synchronized (refreshLock) {
245+
// Re-evaluate the refresh type under the lock, as another thread might have completed
246+
// the refresh while this thread was waiting for the lock.
247+
refreshType = determineRefreshType();
248+
249+
if (refreshType == RefreshType.NONE) {
250+
// No refresh needed, token is still valid.
251+
return;
252+
}
253+
254+
// If a refresh is required, create or retrieve the refresh task.
255+
currentRefreshTask = getOrCreateRefreshTask();
256+
}
239257

240258
// Handle the refresh based on the determined refresh type.
241259
switch (refreshType) {
@@ -283,16 +301,14 @@ void refreshCredentialsIfRequired() throws IOException {
283301
}
284302
}
285303

304+
// Assumes the caller holds refreshLock.
286305
private RefreshType determineRefreshType() {
287-
AccessToken intermediateAccessToken;
288-
synchronized (refreshLock) {
289-
if (intermediateCredentials == null
290-
|| intermediateCredentials.intermediateAccessToken == null) {
291-
// A blocking refresh is needed if the intermediate access token doesn't exist.
292-
return RefreshType.BLOCKING;
293-
}
294-
intermediateAccessToken = intermediateCredentials.intermediateAccessToken;
306+
if (intermediateCredentials == null
307+
|| intermediateCredentials.intermediateAccessToken == null) {
308+
// A blocking refresh is needed if the intermediate access token doesn't exist.
309+
return RefreshType.BLOCKING;
295310
}
311+
AccessToken intermediateAccessToken = intermediateCredentials.intermediateAccessToken;
296312

297313
Date expirationTime = intermediateAccessToken.getExpirationTime();
298314
if (expirationTime == null) {
@@ -322,23 +338,22 @@ private RefreshType determineRefreshType() {
322338
* responsibility of the caller to execute it. The task will clear the single flight slot upon
323339
* completion.
324340
*/
341+
// Assumes the caller holds refreshLock.
325342
private RefreshTask getOrCreateRefreshTask() {
326-
synchronized (refreshLock) {
327-
if (refreshTask != null) {
328-
// An existing refresh task is already in progress. Return a NEW RefreshTask instance with
329-
// the existing task, but set isNew to false. This indicates to the caller that a new
330-
// refresh task was NOT created.
331-
return new RefreshTask(refreshTask.task, false);
332-
}
343+
if (refreshTask != null) {
344+
// An existing refresh task is already in progress. Return a NEW RefreshTask instance with
345+
// the existing task, but set isNew to false. This indicates to the caller that a new
346+
// refresh task was NOT created.
347+
return new RefreshTask(refreshTask.task, false);
348+
}
333349

334-
final ListenableFutureTask<IntermediateCredentials> task =
335-
ListenableFutureTask.create(this::fetchIntermediateCredentials);
350+
final ListenableFutureTask<IntermediateCredentials> task =
351+
ListenableFutureTask.create(this::fetchIntermediateCredentials);
336352

337-
// Store the new refresh task in the refreshTask field before returning. This ensures that
338-
// subsequent calls to this method will return the existing task while it's still in progress.
339-
refreshTask = new RefreshTask(task, true);
340-
return refreshTask;
341-
}
353+
// Store the new refresh task in the refreshTask field before returning. This ensures that
354+
// subsequent calls to this method will return the existing task while it's still in progress.
355+
refreshTask = new RefreshTask(task, true);
356+
return refreshTask;
342357
}
343358

344359
/**

google-auth-library-java/cab-token-generator/javatests/com/google/auth/credentialaccessboundary/ClientSideCredentialAccessBoundaryFactoryTest.java

Lines changed: 10 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -325,7 +325,6 @@ void refreshCredentialsIfRequired_blockingMultiThread() throws IOException, Inte
325325
}
326326

327327
@Test
328-
@Disabled("Flaky test: https://github.com/googleapis/google-cloud-java/issues/12871")
329328
void refreshCredentialsIfRequired_asyncMultiThread() throws IOException, InterruptedException {
330329
final ClientSideCredentialAccessBoundaryFactory factory =
331330
getClientSideCredentialAccessBoundaryFactory(RefreshType.ASYNC);
@@ -623,10 +622,16 @@ private Clock createMockClock(RefreshType refreshType, GoogleCredentials sourceC
623622
// (within the refresh margin).
624623
mockedTimeInMillis = expirationTimeInMillis - refreshMarginInMillis + 60000;
625624
when(mockClock.currentTimeMillis())
626-
.thenReturn(
627-
mockedTimeInMillis, // First call: Stale (triggers the async refresh)
628-
currentTimeInMillis // Subsequent calls: Fresh (skips redundant refreshes)
629-
);
625+
.thenAnswer(
626+
invocation -> {
627+
// If the async refresh has already been triggered (request count >= 2),
628+
// return the fresh time to skip redundant refreshes.
629+
// Note: 1st request was the initial blocking refresh.
630+
if (mockStsTransportFactory.transport.getRequestCount() >= 2) {
631+
return currentTimeInMillis;
632+
}
633+
return mockedTimeInMillis;
634+
});
630635
break;
631636
case BLOCKING:
632637
// Set mocked time so that the token requires immediate refresh (just after the minimum

0 commit comments

Comments
 (0)