Repository: apache/iceberg (language: python)

Issue:
AWS: Add retry logic for S3InputStream and S3OutputStream
@danielcweeks @jackye1995 @amogh-jahagirdar @nastra 
Add retry for both S3InputStream and S3OutputStream so that when we encounter network failures (mostly SSLException for server side connection reset and SocketTimoutException for client side connection reset) or any other retriable S3 error (like throttling error), we can retry at operation level without failing the entire query.


Discussion / hints:
It feels like this adds a lot of baggage/complexity due to requiring quite a few configuration options. Also I was wondering whether we could leverage the retry behavior of the underlying S3 client by properly configuring it. The default retry behavior is shown in https://github.com/aws/aws-sdk-java-v2/blob/2.20.18/core/sdk-core/src/main/java/software/amazon/awssdk/core/internal/retry/SdkDefaultRetrySetting.java#L72-L84 and a configuration could look similar to what has been proposed in https://github.com/apache/iceberg/pull/8043
Thanks. @nastra I see what you are proposing but in our case the SDK level retry wouldn't help. Because we are using this method 
https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/services/s3/S3Client.html#getObject(software.amazon.awssdk.services.s3.model.GetObjectRequest,software.amazon.awssdk.core.sync.ResponseTransformer) which returns us a stream we will read from. The SDK level retry doesn't give us coverage beyond this call. In our case, the exceptions are thrown post this call when we are reading from the stream. That is why we need to add retries here.
@danielcweeks, what are your thoughts on retries in the S3 layer?
This pull request has been marked as stale due to 30 days of inactivity. It will be closed in 1 week if no further activity occurs. If you think that’s incorrect or this pull request requires a review, please simply write any comment. If closed, you can revive the PR at any time and @mention a reviewer or discuss it on the dev@iceberg.apache.org list. Thank you for your contributions.
This pull request has been closed due to lack of activity. This is not a judgement on the merit of the PR in any way. It is just a way of keeping the PR queue manageable. If you think that is incorrect, or the pull request requires review, you can revive the PR at any time.

Fix patch:
diff --git a/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java b/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java
index d9ea511f9bcd..4c3c305d4bae 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java
@@ -47,6 +47,7 @@ public S3Client s3() {
         .applyMutation(s3FileIOProperties::applyEndpointConfigurations)
         .applyMutation(s3FileIOProperties::applyServiceConfigurations)
         .applyMutation(s3FileIOProperties::applySignerConfiguration)
+        .applyMutation(s3FileIOProperties::applyRetryConfigurations)
         .build();
   }
 
diff --git a/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java b/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java
index 81c7bd6b4bab..5974e21209e3 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java
@@ -114,6 +114,7 @@ public S3Client s3() {
           .applyMutation(s3FileIOProperties::applySignerConfiguration)
           .applyMutation(s3FileIOProperties::applyS3AccessGrantsConfigurations)
           .applyMutation(s3FileIOProperties::applyUserAgentConfigurations)
+          .applyMutation(s3FileIOProperties::applyRetryConfigurations)
           .build();
     }
 
diff --git a/aws/src/main/java/org/apache/iceberg/aws/lakeformation/LakeFormationAwsClientFactory.java b/aws/src/main/java/org/apache/iceberg/aws/lakeformation/LakeFormationAwsClientFactory.java
index 552da4bc949c..5d37470066d2 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/lakeformation/LakeFormationAwsClientFactory.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/lakeformation/LakeFormationAwsClientFactory.java
@@ -81,6 +81,7 @@ public S3Client s3() {
           .applyMutation(httpClientProperties()::applyHttpClientConfigurations)
           .applyMutation(s3FileIOProperties()::applyEndpointConfigurations)
           .applyMutation(s3FileIOProperties()::applyServiceConfigurations)
+          .applyMutation(s3FileIOProperties()::applyRetryConfigurations)
           .credentialsProvider(
               new LakeFormationCredentialsProvider(lakeFormation(), buildTableArn()))
           .region(Region.of(region()))
diff --git a/aws/src/main/java/org/apache/iceberg/aws/s3/DefaultS3FileIOAwsClientFactory.java b/aws/src/main/java/org/apache/iceberg/aws/s3/DefaultS3FileIOAwsClientFactory.java
index 18b40000a91a..8687d737a5d7 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/s3/DefaultS3FileIOAwsClientFactory.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/s3/DefaultS3FileIOAwsClientFactory.java
@@ -55,6 +55,7 @@ public S3Client s3() {
         .applyMutation(s3FileIOProperties::applySignerConfiguration)
         .applyMutation(s3FileIOProperties::applyS3AccessGrantsConfigurations)
         .applyMutation(s3FileIOProperties::applyUserAgentConfigurations)
+        .applyMutation(s3FileIOProperties::applyRetryConfigurations)
         .build();
   }
 }
diff --git a/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIOProperties.java b/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIOProperties.java
index 343540d87f95..6813913a4db0 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIOProperties.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIOProperties.java
@@ -20,10 +20,12 @@
 
 import java.io.Serializable;
 import java.net.URI;
+import java.time.Duration;
 import java.util.Collections;
 import java.util.Map;
 import java.util.Set;
 import java.util.stream.Collectors;
+import javax.xml.stream.XMLStreamException;
 import org.apache.iceberg.EnvironmentContext;
 import org.apache.iceberg.aws.AwsClientProperties;
 import org.apache.iceberg.aws.glue.GlueCatalog;
@@ -38,6 +40,14 @@
 import software.amazon.awssdk.auth.credentials.AnonymousCredentialsProvider;
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
 import software.amazon.awssdk.core.client.config.SdkAdvancedClientOption;
+import software.amazon.awssdk.core.exception.SdkServiceException;
+import software.amazon.awssdk.core.retry.RetryMode;
+import software.amazon.awssdk.core.retry.RetryPolicy;
+import software.amazon.awssdk.core.retry.backoff.EqualJitterBackoffStrategy;
+import software.amazon.awssdk.core.retry.conditions.OrRetryCondition;
+import software.amazon.awssdk.core.retry.conditions.RetryCondition;
+import software.amazon.awssdk.core.retry.conditions.RetryOnExceptionsCondition;
+import software.amazon.awssdk.core.retry.conditions.TokenBucketRetryCondition;
 import software.amazon.awssdk.services.s3.S3ClientBuilder;
 import software.amazon.awssdk.services.s3.S3Configuration;
 import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
@@ -393,6 +403,21 @@ public class S3FileIOProperties implements Serializable {
    */
   private static final String S3_FILE_IO_USER_AGENT = "s3fileio/" + EnvironmentContext.get();
 
+  /** Number of times to retry S3 operations. */
+  public static final String S3_RETRY_NUM_RETRIES = "s3.retry.num-retries";
+
+  public static final int S3_RETRY_NUM_RETRIES_DEFAULT = 5;
+
+  /** Minimum wait time to retry a S3 operation */
+  public static final String S3_RETRY_MIN_WAIT_MS = "s3.retry.min-wait-ms";
+
+  public static final long S3_RETRY_MIN_WAIT_MS_DEFAULT = 2_000; // 2 seconds
+
+  /** Maximum wait time to retry a S3 read operation */
+  public static final String S3_RETRY_MAX_WAIT_MS = "s3.retry.max-wait-ms";
+
+  public static final long S3_RETRY_MAX_WAIT_MS_DEFAULT = 20_000; // 20 seconds
+
   private String sseType;
   private String sseKey;
   private String sseMd5;
@@ -423,6 +448,9 @@ public class S3FileIOProperties implements Serializable {
   private final String endpoint;
   private final boolean isRemoteSigningEnabled;
   private String writeStorageClass;
+  private int s3RetryNumRetries;
+  private long s3RetryMinWaitMs;
+  private long s3RetryMaxWaitMs;
   private final Map<String, String> allProperties;
 
   public S3FileIOProperties() {
@@ -455,6 +483,9 @@ public S3FileIOProperties() {
     this.isRemoteSigningEnabled = REMOTE_SIGNING_ENABLED_DEFAULT;
     this.isS3AccessGrantsEnabled = S3_ACCESS_GRANTS_ENABLED_DEFAULT;
     this.isS3AccessGrantsFallbackToIamEnabled = S3_ACCESS_GRANTS_FALLBACK_TO_IAM_ENABLED_DEFAULT;
+    this.s3RetryNumRetries = S3_RETRY_NUM_RETRIES_DEFAULT;
+    this.s3RetryMinWaitMs = S3_RETRY_MIN_WAIT_MS_DEFAULT;
+    this.s3RetryMaxWaitMs = S3_RETRY_MAX_WAIT_MS_DEFAULT;
     this.allProperties = Maps.newHashMap();
 
     ValidationException.check(
@@ -553,6 +584,12 @@ public S3FileIOProperties(Map<String, String> properties) {
             properties,
             S3_ACCESS_GRANTS_FALLBACK_TO_IAM_ENABLED,
             S3_ACCESS_GRANTS_FALLBACK_TO_IAM_ENABLED_DEFAULT);
+    this.s3RetryNumRetries =
+        PropertyUtil.propertyAsInt(properties, S3_RETRY_NUM_RETRIES, S3_RETRY_NUM_RETRIES_DEFAULT);
+    this.s3RetryMinWaitMs =
+        PropertyUtil.propertyAsLong(properties, S3_RETRY_MIN_WAIT_MS, S3_RETRY_MIN_WAIT_MS_DEFAULT);
+    this.s3RetryMaxWaitMs =
+        PropertyUtil.propertyAsLong(properties, S3_RETRY_MAX_WAIT_MS, S3_RETRY_MAX_WAIT_MS_DEFAULT);
 
     ValidationException.check(
         keyIdAccessKeyBothConfigured(),
@@ -753,6 +790,34 @@ public void setS3AccessGrantsFallbackToIamEnabled(boolean s3AccessGrantsFallback
     this.isS3AccessGrantsFallbackToIamEnabled = s3AccessGrantsFallbackToIamEnabled;
   }
 
+  public int s3RetryNumRetries() {
+    return s3RetryNumRetries;
+  }
+
+  public void setS3RetryNumRetries(int s3RetryNumRetries) {
+    this.s3RetryNumRetries = s3RetryNumRetries;
+  }
+
+  public long s3RetryMinWaitMs() {
+    return s3RetryMinWaitMs;
+  }
+
+  public void setS3RetryMinWaitMs(long s3RetryMinWaitMs) {
+    this.s3RetryMinWaitMs = s3RetryMinWaitMs;
+  }
+
+  public long s3RetryMaxWaitMs() {
+    return s3RetryMaxWaitMs;
+  }
+
+  public void setS3RetryMaxWaitMs(long s3RetryMaxWaitMs) {
+    this.s3RetryMaxWaitMs = s3RetryMaxWaitMs;
+  }
+
+  public long s3RetryTotalWaitMs() {
+    return (long) s3RetryNumRetries() * s3RetryMaxWaitMs();
+  }
+
   private boolean keyIdAccessKeyBothConfigured() {
     return (accessKeyId == null) == (secretAccessKey == null);
   }
@@ -824,6 +889,65 @@ public <T extends S3ClientBuilder> void applyEndpointConfigurations(T builder) {
     }
   }
 
+  /**
+   * Override the retry configurations for an S3 client.
+   *
+   * <p>Sample usage:
+   *
+   * <pre>
+   *     S3Client.builder().applyMutation(s3FileIOProperties::applyRetryConfigurations)
+   * </pre>
+   */
+  public <T extends S3ClientBuilder> void applyRetryConfigurations(T builder) {
+    builder.overrideConfiguration(
+        config ->
+            config.retryPolicy(
+                // Use a retry strategy which will persistently retry throttled exceptions with
+                // exponential backoff, to give S3 a chance to autoscale.
+                // LEGACY mode works best here, as it will allow throttled exceptions to use all of
+                // the configured retry attempts.
+                RetryPolicy.builder(RetryMode.LEGACY)
+                    .numRetries(s3RetryNumRetries)
+                    .throttlingBackoffStrategy(
+                        EqualJitterBackoffStrategy.builder()
+                            .baseDelay(Duration.ofMillis(s3RetryMinWaitMs))
+                            .maxBackoffTime(Duration.ofMillis(s3RetryMaxWaitMs))
+                            .build())
+
+                    // Workaround: add XMLStreamException as a retryable exception.
+                    // https://github.com/aws/aws-sdk-java-v2/issues/5442
+                    // Without this workaround, we see SDK failures if there's a socket exception
+                    // while parsing an error XML response.
+                    .retryCondition(
+                        OrRetryCondition.create(
+                            RetryCondition.defaultRetryCondition(),
+                            RetryOnExceptionsCondition.create(XMLStreamException.class)))
+
+                    // Workaround: exclude all 503s from consuming retry tokens.
+                    // https://github.com/aws/aws-sdk-java-v2/issues/5414
+                    // Without this workaround, workloads which see 503s from S3 HEAD will fail
+                    // prematurely.
+                    .retryCapacityCondition(
+                        TokenBucketRetryCondition.builder()
+                            .tokenBucketSize(500) // 500 is the SDK default
+                            .exceptionCostFunction(
+                                e -> {
+                                  if (e instanceof SdkServiceException) {
+                                    SdkServiceException sdkServiceException =
+                                        (SdkServiceException) e;
+                                    if (sdkServiceException.isThrottlingException()
+                                        || sdkServiceException.statusCode() == 503) {
+                                      return 0;
+                                    }
+                                  }
+
+                                  // 5 is the SDK default for non-throttling exceptions
+                                  return 5;
+                                })
+                            .build())
+                    .build()));
+  }
+
   /**
    * Add the S3 Access Grants Plugin for an S3 client.
    *
diff --git a/docs/docs/aws.md b/docs/docs/aws.md
index b9638bb50dad..5a166c0c9193 100644
--- a/docs/docs/aws.md
+++ b/docs/docs/aws.md
@@ -378,6 +378,20 @@ However, for the older versions up to 0.12.0, the logic is as follows:
 
 For more details, please refer to the [LocationProvider Configuration](custom-catalog.md#custom-location-provider-implementation) section.  
 
+### S3 Retries
+
+Workloads which encounter S3 throttling should persistently retry, with exponential backoff, to make progress while S3
+automatically scales. We provide the configurations below to adjust S3 retries for this purpose. For workloads that encounter
+throttling and fail due to retry exhaustion, we recommend retry count to set 32 in order allow S3 to auto-scale. Note that
+workloads with exceptionally high throughput against tables that S3 has not yet scaled, it may be necessary to increase the retry count further.
+
+
+| Property             | Default | Description                                                                           |
+|----------------------|---------|---------------------------------------------------------------------------------------|
+| s3.retry.num-retries | 5       | Number of times to retry S3 operations. Recommended 32 for high-throughput workloads. |
+| s3.retry.min-wait-ms | 2s      | Minimum wait time to retry a S3 operation.                                            |
+| s3.retry.max-wait-ms | 20s     | Maximum wait time to retry a S3 read operation.                                       |
+
 ### S3 Strong Consistency
 
 In November 2020, S3 announced [strong consistency](https://aws.amazon.com/s3/consistency/) for all read operations, and Iceberg is updated to fully leverage this feature.


Test patch:
diff --git a/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3FileIOProperties.java b/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3FileIOProperties.java
index f445a2d22422..a61b9efb9fec 100644
--- a/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3FileIOProperties.java
+++ b/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3FileIOProperties.java
@@ -36,6 +36,8 @@
 import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
+import software.amazon.awssdk.core.retry.RetryPolicy;
+import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.s3.S3ClientBuilder;
 import software.amazon.awssdk.services.s3.S3Configuration;
 import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
@@ -491,4 +493,17 @@ public void testApplyUserAgentConfigurations() {
     Mockito.verify(mockS3ClientBuilder)
         .overrideConfiguration(Mockito.any(ClientOverrideConfiguration.class));
   }
+
+  @Test
+  public void testApplyRetryConfiguration() {
+    Map<String, String> properties = Maps.newHashMap();
+    properties.put(S3FileIOProperties.S3_RETRY_NUM_RETRIES, "999");
+    S3FileIOProperties s3FileIOProperties = new S3FileIOProperties(properties);
+
+    S3ClientBuilder builder = S3Client.builder();
+    s3FileIOProperties.applyRetryConfigurations(builder);
+
+    RetryPolicy retryPolicy = builder.overrideConfiguration().retryPolicy().get();
+    assertThat(retryPolicy.numRetries()).as("retries was not set").isEqualTo(999);
+  }
 }