Core: Retry requests with an Idempotency-Key on retriable errors

core/src/main/java/org/apache/iceberg/rest/ExponentialHttpRequestRetryStrategy.java

@@ -74,7 +74,10 @@
  *   <li>SC_REQUEST_TIMEOUT (408)
  * </ul>
  *
- * Most code and behavior is taken from {@link
+ * Requests carrying an {@code Idempotency-Key} header are treated as retry-safe for these codes
+ * even when the HTTP method is not idempotent.
+ *
+ * <p>Most code and behavior is taken from {@link
  * org.apache.hc.client5.http.impl.DefaultHttpRequestRetryStrategy}, with minor modifications to
  * {@link #getRetryInterval(HttpResponse, int, HttpContext)} to achieve exponential backoff.
  */
@@ -130,8 +133,10 @@ public boolean retryRequest(
       return false;
     }
 
-    // Retry if the request is considered idempotent
-    return Method.isIdempotent(request.getMethod());
+    // Retry if the request is idempotent, or carries an Idempotency-Key (server guarantees safe
+    // retry)
+    return Method.isIdempotent(request.getMethod())
+        || request.containsHeader(RESTUtil.IDEMPOTENCY_KEY_HEADER);
   }
 
   @Override
@@ -189,8 +194,11 @@ private boolean shouldRetryIdempotent(HttpRequest request, int responseCode) {
       return false;
     }
 
-    // Check if the request is idempotent
-    return Method.isIdempotent(request.getMethod())
-        && idempotentRetriableCodes.contains(responseCode);
+    // A request is retry-safe if its HTTP method is idempotent or it carries an Idempotency-Key
+    // header (which lets the server replay a finalized result on retry).
+    boolean retrySafe =
+        Method.isIdempotent(request.getMethod())
+            || request.containsHeader(RESTUtil.IDEMPOTENCY_KEY_HEADER);
+    return retrySafe && idempotentRetriableCodes.contains(responseCode);
   }
 }

core/src/test/java/org/apache/iceberg/rest/TestExponentialHttpRequestRetryStrategy.java

@@ -229,4 +229,24 @@ public void testRetryHappensWithIdempotentMethods(int statusCode) {
     context.setRequest(new BasicHttpRequest("GET", "/"));
     assertThat(retryStrategy.retryRequest(response, 3, context)).isTrue();
   }
+
+  @ParameterizedTest
+  @ValueSource(ints = {429, 503, 500, 502, 504, 408})
+  public void testRetryHappensForNonIdempotentMethodWithIdempotencyKey(int statusCode) {
+    BasicHttpResponse response = new BasicHttpResponse(statusCode, String.valueOf(statusCode));
+    HttpClientContext context = HttpClientContext.create();
+    BasicHttpRequest request = new BasicHttpRequest("POST", "/");
+    request.addHeader(RESTUtil.IDEMPOTENCY_KEY_HEADER, "017f22e2-79b0-7cc3-98c4-dc0c0c07398f");
+    context.setRequest(request);
+    assertThat(retryStrategy.retryRequest(response, 3, context)).isTrue();
+  }
+
+  @ParameterizedTest
+  @ValueSource(ints = {503, 500, 502, 504, 408})
+  public void testRetryDoesNotHappenForNonIdempotentMethodWithoutIdempotencyKey(int statusCode) {
+    BasicHttpResponse response = new BasicHttpResponse(statusCode, String.valueOf(statusCode));
+    HttpClientContext context = HttpClientContext.create();
+    context.setRequest(new BasicHttpRequest("POST", "/"));
+    assertThat(retryStrategy.retryRequest(response, 3, context)).isFalse();
+  }
 }

core/src/test/java/org/apache/iceberg/rest/TestRESTCatalog.java

@@ -48,6 +48,7 @@
 import java.util.Optional;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicReference;
@@ -166,6 +167,8 @@ private static class HeaderValidatingAdapter extends RESTCatalogAdapter {
     private final HTTPHeaders contextHeaders;
     private final java.util.concurrent.ConcurrentMap<String, RuntimeException>
         simulateFailureOnFirstSuccessByKey = new java.util.concurrent.ConcurrentHashMap<>();
+    // Records the Idempotency-Key value seen on every mutation request, in arrival order.
+    private final List<String> observedMutationIdempotencyKeys = new CopyOnWriteArrayList<>();
 
     HeaderValidatingAdapter(
         Catalog catalog, HTTPHeaders catalogHeaders, HTTPHeaders contextHeaders) {
@@ -194,6 +197,11 @@ public void simulate503OnFirstSuccessForKey(String key) {
               new RuntimeException("simulated transient 503 after success")));
     }
 
+    /** Returns all Idempotency-Key values observed on mutation requests, in arrival order. */
+    public List<String> observedMutationIdempotencyKeys() {
+      return observedMutationIdempotencyKeys;
+    }
+
     @Override
     public <T extends RESTResponse> T execute(
         HTTPRequest request,
@@ -233,13 +241,16 @@ protected <T extends RESTResponse> T execute(
               handleRequest(
                   routeAndVars.first(), vars.build(), request, responseType, responseHeaders);
 
-          // For tests: simulate a transient 503 after the first successful mutation for a key.
+          // For tests: record observed keys and simulate transient failures for keyed mutations.
           Optional<HTTPHeaders.HTTPHeader> keyHeader =
               request.headers().firstEntry(RESTUtil.IDEMPOTENCY_KEY_HEADER);
           boolean isMutation =
               request.method() == HTTPMethod.POST || request.method() == HTTPMethod.DELETE;
           if (isMutation && keyHeader.isPresent()) {
             String key = keyHeader.get().value();
+            // Record every Idempotency-Key seen on a mutation (including retries) so tests can
+            // assert that all transport attempts carry the identical key.
+            observedMutationIdempotencyKeys.add(key);
             RuntimeException failure = simulateFailureOnFirstSuccessByKey.remove(key);
             if (failure != null) {
               throw failure;
@@ -3477,30 +3488,48 @@ public void testIdempotentCreateReplayAfterSimulated503() {
     IdempotentEnv env = idempotentEnv(key, ns, "t_idemp");
     CreateTableRequest req = createReq(env.ident);
 
-    // First attempt: server finalizes success but responds 503
-    assertThatThrownBy(
-            () ->
-                env.http.post(
-                    ResourcePaths.forCatalogProperties(ImmutableMap.of()).tables(ns),
-                    req,
-                    LoadTableResponse.class,
-                    env.headers,
-                    ErrorHandlers.tableErrorHandler()))
-        .isInstanceOf(RuntimeException.class)
-        .hasMessageContaining("simulated transient 503");
+    // The client auto-retries the keyed POST on 503; the server replays the finalized 200, so
+    // the call succeeds transparently without the caller needing to retry manually.
+    LoadTableResponse response =
+        env.http.post(
+            ResourcePaths.forCatalogProperties(ImmutableMap.of()).tables(ns),
+            req,
+            LoadTableResponse.class,
+            env.headers,
+            ErrorHandlers.tableErrorHandler());
+    assertThat(response).isNotNull();
 
     // Verify request shape (method, path, headers including Idempotency-Key)
     verifyCreatePost(ns, env.headers);
+  }
+
+  @Test
+  public void testIdempotentCreateRetryCarriesSameKey() {
+    // Pin the invariant: when the client auto-retries a keyed POST (503-then-200), every transport
+    // attempt must carry the identical Idempotency-Key so the server can replay the cached result.
+    String key = "idemp-same-key-retry";
+    adapterForRESTServer.simulate503OnFirstSuccessForKey(key);
+    Namespace ns = Namespace.of("ns_samekey");
+    IdempotentEnv env = idempotentEnv(key, ns, "t_samekey");
+    CreateTableRequest req = createReq(env.ident);
 
-    // Retry with same key: server should replay 200 OK
-    LoadTableResponse replay =
+    // Trigger the 503-then-200 retry cycle; the call must succeed transparently.
+    LoadTableResponse response =
         env.http.post(
             ResourcePaths.forCatalogProperties(ImmutableMap.of()).tables(ns),
             req,
             LoadTableResponse.class,
             env.headers,
             ErrorHandlers.tableErrorHandler());
-    assertThat(replay).isNotNull();
+    assertThat(response).isNotNull();
+
+    // The adapter must have observed the Idempotency-Key on exactly 2 transport attempts
+    // (initial attempt + one auto-retry) and both values must be identical.
+    List<String> observedKeys =
+        adapterForRESTServer.observedMutationIdempotencyKeys().stream()
+            .filter(k -> k.equals(key))
+            .collect(Collectors.toList());
+    assertThat(observedKeys).hasSize(2);
   }
 
   @Test