TOO_MANY_REQUESTS при обновлении токена #15 (#16)

Author: isergeymd

gigachat-java/src/main/java/chat/giga/client/auth/TokenBasedAuthClient.java

@@ -13,6 +13,7 @@
 import java.time.Instant;
 import java.util.Base64;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.concurrent.locks.ReentrantLock;
 
 import static chat.giga.http.client.HttpHeaders.USER_AGENT;
 import static chat.giga.http.client.MediaType.APPLICATION_JSON;
@@ -24,6 +25,7 @@ abstract class TokenBasedAuthClient {
     protected final ObjectMapper objectMapper = JsonUtils.objectMapper();
 
     private final AtomicReference<AccessToken> accessToken = new AtomicReference<>();
+    private final ReentrantLock lock = new ReentrantLock();
 
     protected HttpRequestBuilder getTokenRequest(String url, Scope scope, String user, String password) {
         var credentials = user + ":" + password;
@@ -48,10 +50,20 @@ protected String getBearerAuth() {
     }
 
     protected String retrieveTokenIfExpired() {
-        return accessToken.updateAndGet(t -> {
-            var expiresAt = t == null ? null : t.expiresAt();
-            return expiresAt != null && Instant.now().isBefore(expiresAt) ? t : getToken();
-        }).token();
+        AccessToken token = accessToken.get();
+        if (token == null || Instant.now().isAfter(token.expiresAt())) {
+            lock.lock();
+            try {
+                token = accessToken.get();
+                if (token == null || Instant.now().isAfter(token.expiresAt())) {
+                    token = getToken();
+                    accessToken.set(token);
+                }
+            } finally {
+                lock.unlock();
+            }
+        }
+        return token.token();
     }
 
     protected abstract AccessToken getToken();

gigachat-java/src/test/java/chat/giga/client/auth/OAuthClientTest.java

@@ -22,13 +22,20 @@
 import java.nio.charset.StandardCharsets;
 import java.time.Instant;
 import java.util.List;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.TimeUnit;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
@@ -116,6 +123,54 @@ public void authenticate() throws Exception {
         });
     }
 
+    @Test
+    public void shouldCallGetTokenOnlyOnceUnderContention() throws Exception {
+        OAuthClient authClient = spy(new OAuthClient(null, "client", "secret", null, Scope.GIGACHAT_API_B2B, "url"));
+
+        AccessTokenResponse mockResponse = AccessTokenResponse.builder()
+                .accessToken(token)
+                .expiresAt(Instant.now()
+                        .plusSeconds(1000L).toEpochMilli())
+                .build();
+
+        // Имитируем задержку сети, чтобы увеличить окно конкуренции
+        doAnswer(invocation -> {
+            Thread.sleep(100);
+            return mockResponse;
+        }).when(authClient).oauth();
+
+        int threadCount = 10;
+        ExecutorService executor = Executors.newFixedThreadPool(threadCount);
+        CountDownLatch startLatch = new CountDownLatch(1);
+        CountDownLatch finishLatch = new CountDownLatch(threadCount);
+        HttpRequestBuilder requestBuilder = HttpRequest.builder()
+                .url("test.ru/models");
+
+        for (int i = 0; i < threadCount; i++) {
+            executor.submit(() -> {
+                try {
+                    startLatch.await(); // Ждем сигнала к одновременному старту
+                    authClient.authenticate(requestBuilder);
+                } catch (InterruptedException e) {
+                    Thread.currentThread().interrupt();
+                } finally {
+                    finishLatch.countDown();
+                }
+            });
+        }
+
+        startLatch.countDown();
+
+        // Ждем завершения всех потоков
+        boolean completed = finishLatch.await(2, TimeUnit.SECONDS);
+        assertTrue(completed, "Threads did not finish in time");
+
+        // Проверяем, что сетевой запрос был выполнен ровно 1 раз
+        verify(authClient, times(1)).getToken();
+
+        executor.shutdown();
+    }
+
     @Test
     public void authenticateByKey() throws Exception {
         AuthClient authClient = AuthClientBuilder.builder()