fix: #330 address check-then-act race for access token using single-flight pattern in OAuth2Client

Author: cportcvent

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

@@ -11,13 +11,15 @@
 import java.time.Instant;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.atomic.AtomicReference;
 import java.util.concurrent.CompletableFuture;
 
 public class OAuth2Client {
     private static final String DEFAULT_API_TOKEN_ISSUER_PATH = "/oauth/token";
 
     private final ApiClient apiClient;
-    private final AccessToken token = new AccessToken();
+    private final AtomicReference<TokenSnapshot> snapshot = new AtomicReference<>(TokenSnapshot.EMPTY);
+    private final AtomicReference<CompletableFuture<String>> inFlight = new AtomicReference<>();
     private final CredentialsFlowRequest authRequest;
     private final Configuration config;
     private final Telemetry telemetry;
@@ -45,26 +47,54 @@ public OAuth2Client(Configuration configuration, ApiClient apiClient) throws Fga
     }
 
     /**
-     * Gets an access token, handling exchange when necessary. The access token is naively cached in memory until it
-     * expires.
+     * Gets an access token, handling exchange when necessary. The token is cached as an immutable
+     * snapshot until it expires. Concurrent calls are deduplicated: only one exchange is in flight
+     * at a time; other callers join the same future rather than issuing redundant requests.
      *
      * @return An access token in a {@link CompletableFuture}
      */
     public CompletableFuture<String> getAccessToken() throws FgaInvalidParameterException, ApiException {
-        if (!token.isValid()) {
-            return exchangeToken().thenCompose(response -> {
-                token.setToken(response.getAccessToken());
-                token.setExpiresAt(Instant.now().plusSeconds(response.getExpiresInSeconds()));
-
-                Map<Attribute, String> attributesMap = new HashMap<>();
+        TokenSnapshot current = snapshot.get();
+        if (current.isValid()) {
+            return CompletableFuture.completedFuture(current.token());
+        }
 
-                telemetry.metrics().credentialsRequest(1L, attributesMap);
+        CompletableFuture<String> promise = new CompletableFuture<>();
+        if (!inFlight.compareAndSet(null, promise)) {
+            // Another thread won the race — join its exchange rather than starting a new one.
+            CompletableFuture<String> existing = inFlight.get();
+            return existing != null ? existing : getAccessToken();
+        }
 
-                return CompletableFuture.completedFuture(token.getToken());
+        // This thread owns the exchange. Start it, wiring completion back to `promise`.
+        try {
+            exchangeToken().whenComplete((response, ex) -> {
+                if (ex != null) {
+                    inFlight.set(null);
+                    promise.completeExceptionally(ex);
+                } else {
+                    String token = response.getAccessToken();
+                    // Write snapshot before clearing the gate so any new caller that arrives
+                    // after inFlight becomes null immediately sees a valid token.
+                    snapshot.set(
+                        new TokenSnapshot(
+                            token,
+                            Instant.now().plusSeconds(response.getExpiresInSeconds())));
+
+                    telemetry.metrics().credentialsRequest(1L, new HashMap<>());
+
+                    // Clear before completing
+                    inFlight.set(null);
+                    promise.complete(token);
+                }
             });
+        } catch (Exception e) {
+            inFlight.set(null);
+            promise.completeExceptionally(e);
+            throw e;
         }
 
-        return CompletableFuture.completedFuture(token.getToken());
+        return promise;
     }
 
     /**

src/main/java/dev/openfga/sdk/api/auth/TokenSnapshot.java

@@ -0,0 +1,39 @@
+package dev.openfga.sdk.api.auth;
+
+import static dev.openfga.sdk.util.StringUtil.isNullOrWhitespace;
+
+import dev.openfga.sdk.constants.FgaConstants;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.concurrent.ThreadLocalRandom;
+
+/**
+ * Immutable snapshot of an access token and its expiry time. The snapshot is valid if the token is non-empty
+ * and the current time is before the expiry time minus a buffer to ensure that callers receive a valid token
+ * even if there is some clock skew or delay between retrieval and use.
+ */
+record TokenSnapshot(String token, Instant expiresAt) {
+    private static final int EXPIRY_BUFFER_SECS = FgaConstants.TOKEN_EXPIRY_THRESHOLD_BUFFER_IN_SEC;
+    private static final int EXPIRY_JITTER_SECS = FgaConstants.TOKEN_EXPIRY_JITTER_IN_SEC;
+
+    static final TokenSnapshot EMPTY = new TokenSnapshot(null, null);
+
+    TokenSnapshot {
+        expiresAt = expiresAt != null ? expiresAt.truncatedTo(ChronoUnit.SECONDS) : null;
+    }
+
+    boolean isValid() {
+        if (isNullOrWhitespace(token)) {
+            return false;
+        }
+        if (expiresAt == null) {
+            return true;
+        }
+        Instant expiresWithLeeway = expiresAt
+                .minusSeconds(EXPIRY_BUFFER_SECS)
+                .minusSeconds(ThreadLocalRandom.current().nextInt(EXPIRY_JITTER_SECS))
+                .truncatedTo(ChronoUnit.SECONDS);
+
+        return Instant.now().truncatedTo(ChronoUnit.SECONDS).isBefore(expiresWithLeeway);
+    }
+}

src/test/java/dev/openfga/sdk/api/auth/AccessTokenTest.java

@@ -38,10 +38,8 @@ private static Stream<Arguments> expTimeAndResults() {
 
     @MethodSource("expTimeAndResults")
     @ParameterizedTest(name = "{0}")
-    public void testTokenValid(String name, Instant exp, boolean valid) {
-        AccessToken accessToken = new AccessToken();
-        accessToken.setToken("token");
-        accessToken.setExpiresAt(exp);
-        assertEquals(valid, accessToken.isValid());
+    void testTokenValid(String name, Instant exp, boolean valid) {
+        TokenSnapshot snapshot = new TokenSnapshot("token", exp);
+        assertEquals(valid, snapshot.isValid());
     }
 }

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

@@ -18,6 +18,11 @@
 import java.net.URLEncoder;
 import java.nio.charset.StandardCharsets;
 import java.time.Duration;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import java.util.stream.Stream;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.params.ParameterizedTest;
@@ -166,6 +171,43 @@ public void exchangeOAuth2TokenWithRetriesFailure(WireMockRuntimeInfo wm) throws
         verify(3, postRequestedFor(urlEqualTo("/oauth/token")));
     }
 
+    @Test
+    void exchangeOAuth2Token_concurrentRequests_singleExchange(WireMockRuntimeInfo wm) throws Exception {
+        // Stub with a delay so concurrent threads pile up before the first exchange completes.
+        stubFor(post(urlEqualTo("/oauth/token"))
+                .willReturn(ok(String.format("{\"access_token\":\"%s\",\"expires_in\":3600}", ACCESS_TOKEN))
+                        .withFixedDelay(100)));
+
+        OAuth2Client client = newOAuth2Client(wm.getHttpBaseUrl(), false);
+
+        int threadCount = 5;
+        CountDownLatch startGate = new CountDownLatch(1);
+        CountDownLatch done = new CountDownLatch(threadCount);
+        List<String> tokens = Collections.synchronizedList(new ArrayList<>());
+        List<Throwable> failures = Collections.synchronizedList(new ArrayList<>());
+
+        for (int i = 0; i < threadCount; i++) {
+            new Thread(() -> {
+                try {
+                    startGate.await();
+                    tokens.add(client.getAccessToken().get());
+                } catch (Exception e) {
+                    failures.add(e);
+                } finally {
+                    done.countDown();
+                }
+            }).start();
+        }
+
+        startGate.countDown();
+        assertTrue(done.await(3, TimeUnit.SECONDS), "threads did not complete in time");
+
+        assertEquals(List.of(), failures, "no thread should have thrown");
+        assertEquals(threadCount, tokens.size(), "all threads should have received a token");
+        assertTrue(tokens.stream().allMatch(ACCESS_TOKEN::equals), "all threads should have received the same token");
+        verify(1, postRequestedFor(urlEqualTo("/oauth/token")));
+    }
+
     @Test
     public void apiTokenIssuer_invalidScheme() {
         // When