Fix race condition in RetryInterceptor

Signed-off-by: andreadimaio <[email protected]>

Author: andreadimaio

modules/watsonx-ai-core/src/main/java/com/ibm/watsonx/ai/core/http/interceptors/RetryInterceptor.java

@@ -4,6 +4,7 @@
  */
 package com.ibm.watsonx.ai.core.http.interceptors;
 
+import static com.ibm.watsonx.ai.core.http.BaseHttpClient.REQUEST_ID_HEADER;
 import static java.util.Objects.isNull;
 import static java.util.Objects.requireNonNull;
 import static java.util.Objects.requireNonNullElse;
@@ -34,11 +35,15 @@ public final class RetryInterceptor implements SyncHttpInterceptor, AsyncHttpInt
 
     private static final Logger logger = LoggerFactory.getLogger(RetryInterceptor.class);
 
+    private record RetryOn(Class<? extends Throwable> clazz, Optional<Predicate<Throwable>> predicate) {}
+
+    private final Duration retryInterval;
+    private final List<RetryOn> retryOn;
+    private final boolean exponentialBackoff;
+    private final Integer maxRetries;
+
     /**
      * Checks whether a {@link WatsonxException} is retryable due to an expired authentication token.
-     * <p>
-     * This condition is met if the HTTP status code is 401 and at least one error in the exception's details has the code
-     * {@code AUTHENTICATION_TOKEN_EXPIRED}.
      *
      * @param maxRetries the maximum number of retry attempts
      * @return a configured {@link RetryInterceptor} instance that handles token expiration retries
@@ -59,28 +64,19 @@ public static RetryInterceptor onTokenExpired(int maxRetries) {
             ).build();
     }
 
-    private record RetryOn(Class<? extends Throwable> clazz, Optional<Predicate<Throwable>> predicate) {}
-
-    private final Duration retryInterval;
-    private final List<RetryOn> retryOn;
-    private final boolean exponentialBackoff;
-    private Integer maxRetries;
-    private Duration timeout;
-
     /**
      * Creates a new {@code RetryInterceptor} using the provided builder.
      *
      * @param builder the builder instance
      */
     public RetryInterceptor(Builder builder) {
         requireNonNull(builder);
-        this.retryInterval = requireNonNullElse(builder.retryInterval, Duration.ofMillis(0));
-        this.timeout = this.retryInterval;
-        this.maxRetries = requireNonNullElse(builder.maxRetries, 1);
-        this.retryOn = builder.retryOn;
-        this.exponentialBackoff = builder.exponentialBackoff;
-        if (isNull(retryOn) || retryOn.isEmpty())
-            throw new RuntimeException("At least one exception must be specified");
+        retryInterval = requireNonNullElse(builder.retryInterval, Duration.ofMillis(0));
+        maxRetries = requireNonNullElse(builder.maxRetries, 1);
+        retryOn = requireNonNull(builder.retryOn, "At least one exception must be specified");
+        exponentialBackoff = builder.exponentialBackoff;
+        if (exponentialBackoff && !retryInterval.isPositive())
+            throw new IllegalArgumentException("Retry interval must be positive when exponential backoff is enabled");
     }
 
     @Override
@@ -89,15 +85,23 @@ public <T> HttpResponse<T> intercept(HttpRequest request, BodyHandler<T> bodyHan
 
         Throwable exception = null;
 
+        String requestId = request.headers()
+            .firstValue(REQUEST_ID_HEADER)
+            .orElseThrow(); // This should never happen. The SyncHttpClient and AsyncHttpClient add this header if it is not present.
+
+        Duration timeout = Duration.from(retryInterval);
+
         for (int attempt = 0; attempt <= maxRetries; attempt++) {
 
             try {
 
-                if (attempt > 0)
-                    Thread.sleep(timeout.toMillis());
+                if (attempt > 0 && timeout.isPositive()) {
+                    logger.debug("Retry request \"{}\" after {} ms", requestId, timeout.toMillis());
+                    Thread.sleep(timeout);
+                }
 
                 var res = chain.proceed(request, bodyHandler);
-                this.timeout = this.retryInterval;
+                timeout = Duration.from(retryInterval);
                 return res;
 
             } catch (Exception e) {
@@ -117,54 +121,39 @@ public <T> HttpResponse<T> intercept(HttpRequest request, BodyHandler<T> bodyHan
                         timeout = timeout.multipliedBy(2);
                     }
                     if (attempt > 0) {
-                        logger.debug("Retrying request ({}/{}) after failure: {}", attempt, maxRetries,
+                        logger.debug("Retrying request \"{}\" ({}/{}) after failure: {}", requestId, attempt, maxRetries,
                             exception.getMessage());
                     }
                     chain.resetToIndex(index + 1);
                     continue;
                 }
 
-                this.timeout = this.retryInterval;
+                timeout = Duration.from(retryInterval);
                 throw e;
             }
         }
 
-        this.timeout = this.retryInterval;
-        logger.debug("Max retries reached");
-
-        throw new RuntimeException("Max retries reached", isNull(exception) ? new Exception() : exception);
+        timeout = Duration.from(retryInterval);
+        throw new RuntimeException("Max retries reached for request [%s]".formatted(requestId), isNull(exception) ? new Exception() : exception);
     }
 
     @Override
     public <T> CompletableFuture<HttpResponse<T>> intercept(HttpRequest request, BodyHandler<T> bodyHandler,
         Executor executor, int index, AsyncChain chain) {
-        return executeWithRetry(request, bodyHandler, executor, index, 0, chain);
+        return executeWithRetry(request, bodyHandler, executor, index, 0, Duration.from(retryInterval), chain);
     }
 
-    /**
-     * The current timeout interval.
-     *
-     * @return the timeout duration
-     */
-    public Duration getTimeout() {
-        return timeout;
-    }
-
-    /**
-     * Returns a new {@link Builder} instance.
-     *
-     * @return {@link Builder} instance.
-     */
-    public static Builder builder() {
-        return new Builder();
-    }
 
     private <T> CompletableFuture<HttpResponse<T>> executeWithRetry(HttpRequest request, BodyHandler<T> bodyHandler,
-        Executor executor, int index, int attempt, AsyncChain chain) {
+        Executor executor, int index, int attempt, Duration timeout, AsyncChain chain) {
 
         return chain.proceed(request, bodyHandler, executor)
             .exceptionallyCompose(throwable -> {
 
+                String requestId = request.headers()
+                    .firstValue(REQUEST_ID_HEADER)
+                    .orElseThrow(); // This should never happen. The SyncHttpClient and AsyncHttpClient add this header if it is not present.
+
                 Throwable cause = throwable.getCause() != null ? throwable.getCause() : throwable;
 
                 var shouldRetry =
@@ -178,27 +167,36 @@ private <T> CompletableFuture<HttpResponse<T>> executeWithRetry(HttpRequest requ
 
                 if (!shouldRetry || attempt >= maxRetries) {
                     CompletableFuture<HttpResponse<T>> failed = new CompletableFuture<>();
-                    logger.debug("Max retries reached");
+                    logger.debug("Max retries reached for request \"{}\"", requestId);
                     failed.completeExceptionally(cause);
-                    this.timeout = this.retryInterval;
                     return failed;
                 }
 
-                if (exponentialBackoff && attempt > 0)
-                    timeout = timeout.multipliedBy(2);
+                Duration nextTimeout = exponentialBackoff ? timeout.multipliedBy(2) : timeout;
 
-                logger.debug("Retrying request ({}/{}) after failure: {}", attempt + 1, maxRetries, cause.getMessage());
+                if (timeout.isPositive())
+                    logger.debug("Retry request \"{}\" after {} ms", requestId, nextTimeout.toMillis());
 
                 return CompletableFuture.supplyAsync(
                     () -> {
+                        logger.debug("Retrying request \"{}\" ({}/{}) after failure: {}", requestId, attempt + 1, maxRetries, cause.getMessage());
                         chain.resetToIndex(index + 1);
-                        return executeWithRetry(request, bodyHandler, executor, index, attempt + 1, chain);
+                        return executeWithRetry(request, bodyHandler, executor, index, attempt + 1, nextTimeout, chain);
                     },
-                    CompletableFuture.delayedExecutor(timeout.toMillis(), TimeUnit.MILLISECONDS, executor)
+                    CompletableFuture.delayedExecutor(nextTimeout.toMillis(), TimeUnit.MILLISECONDS, executor)
                 ).thenCompose(Function.identity());
             });
     }
 
+    /**
+     * Returns a new {@link Builder} instance.
+     *
+     * @return {@link Builder} instance.
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
     /**
      * Builder for {@link RetryInterceptor}.
      */

modules/watsonx-ai-core/src/test/java/com/ibm/watsonx/ai/core/RetryInterceptorTest.java

@@ -78,6 +78,15 @@ void setup() {
         when(httpRequest.headers()).thenReturn(HttpHeaders.of(Map.of(), (k, v) -> true));
     }
 
+    @Test
+    void test_exponential_backoff_without_retry_interval() {
+
+        assertThrows(IllegalArgumentException.class, () -> RetryInterceptor.builder()
+            .exponentialBackoff(true)
+            .retryOn(NullPointerException.class)
+            .build(), "Retry interval must be positive when exponential backoff is enabled");
+    }
+
     @Nested
     public class Sync {
 
@@ -318,44 +327,7 @@ void retry_with_tool_exception() throws Exception {
             verify(mockInterceptor, times(2)).intercept(any(), eq(bodyHandler), anyInt(), any());
         }
 
-        @Test
-        void retry_with_exponential_backoff_fail_retries() throws Exception {
-            Duration timeout = Duration.ofMillis(10);
-            RetryInterceptor retryInterceptor = RetryInterceptor.builder()
-                .maxRetries(3)
-                .retryInterval(timeout)
-                .exponentialBackoff(true)
-                .retryOn(NullPointerException.class)
-                .build();
-
-            SyncHttpInterceptor mockInterceptor = new SyncHttpInterceptor() {
-                int numCalled = 0;
 
-                @Override
-                public <T> HttpResponse<T> intercept(HttpRequest request, BodyHandler<T> bodyHandler, int index, Chain chain)
-                    throws WatsonxException, IOException, InterruptedException {
-                    if (this.numCalled > 0)
-                        assertEquals((long) (timeout.toMillis() * Math.pow(2, this.numCalled - 1)),
-                            retryInterceptor.getTimeout().toMillis());
-                    this.numCalled++;
-                    return chain.proceed(request, bodyHandler);
-                }
-            };
-
-            SyncHttpClient client = SyncHttpClient.builder()
-                .httpClient(httpClient)
-                .interceptor(retryInterceptor)
-                .interceptor(mockInterceptor)
-                .build();
-
-            when(httpClient.send(any(), eq(bodyHandler)))
-                .thenThrow(NullPointerException.class);
-
-            var ex = assertThrows(RuntimeException.class, () -> client.send(httpRequest, bodyHandler));
-            assertEquals(NullPointerException.class, ex.getCause().getClass());
-            verify(httpClient, times(4)).send(any(), eq(bodyHandler));
-            assertEquals(retryInterceptor.getTimeout().toMillis(), timeout.toMillis());
-        }
 
         @Test
         void retry_with_exponential_backoff_succeed_after_retry() throws Exception {
@@ -368,16 +340,9 @@ void retry_with_exponential_backoff_succeed_after_retry() throws Exception {
                 .build();
 
             SyncHttpInterceptor mockInterceptor = new SyncHttpInterceptor() {
-                int numCalled = 0;
-
                 @Override
-                public <T> HttpResponse<T> intercept(HttpRequest request, BodyHandler<T> bodyHandler, int index,
-                    Chain chain)
+                public <T> HttpResponse<T> intercept(HttpRequest request, BodyHandler<T> bodyHandler, int index, Chain chain)
                     throws WatsonxException, IOException, InterruptedException {
-                    if (this.numCalled > 0)
-                        assertEquals((long) (timeout.toMillis() * Math.pow(2, this.numCalled - 1)),
-                            retryInterceptor.getTimeout().toMillis());
-                    this.numCalled++;
                     return chain.proceed(request, bodyHandler);
                 }
             };
@@ -389,15 +354,16 @@ public <T> HttpResponse<T> intercept(HttpRequest request, BodyHandler<T> bodyHan
                 .build();
 
             when(httpClient.send(any(), eq(bodyHandler)))
+                .thenThrow(NullPointerException.class)
+                .thenThrow(NullPointerException.class)
                 .thenThrow(NullPointerException.class)
                 .thenReturn(httpResponse);
 
             when(httpResponse.statusCode())
                 .thenReturn(200);
 
             client.send(httpRequest, bodyHandler);
-            verify(httpClient, times(2)).send(any(), eq(bodyHandler));
-            assertEquals(retryInterceptor.getTimeout().toMillis(), timeout.toMillis());
+            verify(httpClient, times(4)).send(any(), eq(bodyHandler));
         }
     }
 
@@ -586,46 +552,6 @@ void retry_with_watsonx_exception() throws Exception {
             verify(mockInterceptor, times(2)).intercept(any(), eq(bodyHandler), any(), anyInt(), any());
         }
 
-        @Test
-        @SuppressWarnings("unchecked")
-        void retry_with_exponential_backoff_fail_retries() throws Exception {
-            Duration timeout = Duration.ofMillis(10);
-            RetryInterceptor retryInterceptor = RetryInterceptor.builder()
-                .maxRetries(3)
-                .retryInterval(timeout)
-                .exponentialBackoff(true)
-                .retryOn(NullPointerException.class)
-                .build();
-
-            AsyncHttpInterceptor mockInterceptor = new AsyncHttpInterceptor() {
-                int numCalled = 0;
-
-                @Override
-                public <T> CompletableFuture<HttpResponse<T>> intercept(HttpRequest request, BodyHandler<T> bodyHandler,
-                    Executor executor, int index, AsyncChain chain) {
-                    if (this.numCalled > 0)
-                        assertEquals((long) (timeout.toMillis() * Math.pow(2, this.numCalled - 1)),
-                            retryInterceptor.getTimeout().toMillis());
-                    this.numCalled++;
-                    return chain.proceed(request, bodyHandler, executor);
-                }
-            };
-
-            AsyncHttpClient client = AsyncHttpClient.builder()
-                .httpClient(httpClient)
-                .interceptor(retryInterceptor)
-                .interceptor(mockInterceptor)
-                .build();
-
-            when(httpClient.sendAsync(any(), any(BodyHandler.class)))
-                .thenReturn(CompletableFuture.failedFuture(new NullPointerException()));
-
-            var ex = assertThrows(RuntimeException.class, () -> client.send(httpRequest, bodyHandler).join());
-            assertEquals(NullPointerException.class, ex.getCause().getClass());
-            verify(httpClient, times(4)).sendAsync(any(), any(BodyHandler.class));
-            assertEquals(retryInterceptor.getTimeout().toMillis(), timeout.toMillis());
-        }
-
         @Test
         @SuppressWarnings("unchecked")
         void retry_with_exponential_backoff_succeed_after_retry() throws Exception {
@@ -638,15 +564,9 @@ void retry_with_exponential_backoff_succeed_after_retry() throws Exception {
                 .build();
 
             AsyncHttpInterceptor mockInterceptor = new AsyncHttpInterceptor() {
-                int numCalled = 0;
-
                 @Override
                 public <T> CompletableFuture<HttpResponse<T>> intercept(HttpRequest request, BodyHandler<T> bodyHandler,
                     Executor executor, int index, AsyncChain chain) {
-                    if (this.numCalled > 0)
-                        assertEquals((long) (timeout.toMillis() * Math.pow(2, this.numCalled - 1)),
-                            retryInterceptor.getTimeout().toMillis());
-                    this.numCalled++;
                     return chain.proceed(request, bodyHandler, executor);
                 }
             };
@@ -658,12 +578,13 @@ public <T> CompletableFuture<HttpResponse<T>> intercept(HttpRequest request, Bod
                 .build();
 
             when(httpClient.sendAsync(any(), any(BodyHandler.class)))
+                .thenReturn(failedFuture(new NullPointerException()))
+                .thenReturn(failedFuture(new NullPointerException()))
                 .thenReturn(failedFuture(new NullPointerException()))
                 .thenReturn(completedFuture(httpResponse));
 
             client.send(httpRequest, bodyHandler).join();
-            verify(httpClient, times(2)).sendAsync(any(), any(BodyHandler.class));
-            assertEquals(retryInterceptor.getTimeout().toMillis(), timeout.toMillis());
+            verify(httpClient, times(4)).sendAsync(any(), any(BodyHandler.class));
         }
     }