From b2a4fda8a80571a6a63f301ff67c63ecc0380cf7 Mon Sep 17 00:00:00 2001
From: Amogh Jahagirdar <amogh@tabular.io>
Date: Wed, 10 Apr 2024 00:25:35 -0600
Subject: [PATCH] Core: Mark 502 and 504 failures as retryable to the
 exponential retry strategy (#10113)

---
 .../rest/ExponentialHttpRequestRetryStrategy.java    |  8 +++++++-
 .../TestExponentialHttpRequestRetryStrategy.java     | 12 ++++++++++++
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/core/src/main/java/org/apache/iceberg/rest/ExponentialHttpRequestRetryStrategy.java b/core/src/main/java/org/apache/iceberg/rest/ExponentialHttpRequestRetryStrategy.java
index 9d8f5424f532..aadb97bc7112 100644
--- a/core/src/main/java/org/apache/iceberg/rest/ExponentialHttpRequestRetryStrategy.java
+++ b/core/src/main/java/org/apache/iceberg/rest/ExponentialHttpRequestRetryStrategy.java
@@ -60,7 +60,9 @@
  *
  * <ul>
  *   <li>SC_TOO_MANY_REQUESTS (429)
+ *   <li>SC_BAD_GATEWAY (502)
  *   <li>SC_SERVICE_UNAVAILABLE (503)
+ *   <li>SC_GATEWAY_TIMEOUT (504)
  * </ul>
  *
  * Most code and behavior is taken from {@link
@@ -77,7 +79,11 @@ class ExponentialHttpRequestRetryStrategy implements HttpRequestRetryStrategy {
         maximumRetries > 0, "Cannot set retries to %s, the value must be positive", maximumRetries);
     this.maxRetries = maximumRetries;
     this.retriableCodes =
-        ImmutableSet.of(HttpStatus.SC_TOO_MANY_REQUESTS, HttpStatus.SC_SERVICE_UNAVAILABLE);
+        ImmutableSet.of(
+            HttpStatus.SC_TOO_MANY_REQUESTS,
+            HttpStatus.SC_SERVICE_UNAVAILABLE,
+            HttpStatus.SC_BAD_GATEWAY,
+            HttpStatus.SC_GATEWAY_TIMEOUT);
     this.nonRetriableExceptions =
         ImmutableSet.of(
             InterruptedIOException.class,
diff --git a/core/src/test/java/org/apache/iceberg/rest/TestExponentialHttpRequestRetryStrategy.java b/core/src/test/java/org/apache/iceberg/rest/TestExponentialHttpRequestRetryStrategy.java
index e63bdfd06758..7d8c58701a1e 100644
--- a/core/src/test/java/org/apache/iceberg/rest/TestExponentialHttpRequestRetryStrategy.java
+++ b/core/src/test/java/org/apache/iceberg/rest/TestExponentialHttpRequestRetryStrategy.java
@@ -196,4 +196,16 @@ public void invalidRetryAfterHeader() {
     assertThat(retryStrategy.getRetryInterval(response, 3, null).toMilliseconds())
         .isBetween(4000L, 5000L);
   }
+
+  @Test
+  public void testRetryBadGateway() {
+    BasicHttpResponse response502 = new BasicHttpResponse(502, "Bad gateway failure");
+    assertThat(retryStrategy.retryRequest(response502, 3, null)).isTrue();
+  }
+
+  @Test
+  public void testRetryGatewayTimeout() {
+    BasicHttpResponse response504 = new BasicHttpResponse(504, "Gateway timeout");
+    assertThat(retryStrategy.retryRequest(response504, 3, null)).isTrue();
+  }
 }