AWS: Enable RetryMode for AWS KMS client (#11420)

aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java

@@ -38,6 +38,7 @@ public class AssumeRoleAwsClientFactory implements AwsClientFactory {
   private HttpClientProperties httpClientProperties;
   private S3FileIOProperties s3FileIOProperties;
   private String roleSessionName;
+  private AwsClientProperties awsClientProperties;
 
   @Override
   public S3Client s3() {
@@ -64,6 +65,7 @@ public KmsClient kms() {
     return KmsClient.builder()
         .applyMutation(this::applyAssumeRoleConfigurations)
         .applyMutation(httpClientProperties::applyHttpClientConfigurations)
+        .applyMutation(awsClientProperties::applyRetryConfigurations)
         .build();
   }
 
@@ -81,6 +83,7 @@ public void initialize(Map<String, String> properties) {
     this.awsProperties = new AwsProperties(properties);
     this.s3FileIOProperties = new S3FileIOProperties(properties);
     this.httpClientProperties = new HttpClientProperties(properties);
+    this.awsClientProperties = new AwsClientProperties(properties);
     this.roleSessionName = genSessionName();
     Preconditions.checkNotNull(
         awsProperties.clientAssumeRoleArn(),
@@ -126,6 +129,10 @@ protected S3FileIOProperties s3FileIOProperties() {
     return s3FileIOProperties;
   }
 
+  protected AwsClientProperties awsClientProperties() {
+    return awsClientProperties;
+  }
+
   private StsClient sts() {
     return StsClient.builder()
         .applyMutation(httpClientProperties::applyHttpClientConfigurations)

aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java

@@ -134,6 +134,7 @@ public KmsClient kms() {
           .applyMutation(awsClientProperties::applyClientRegionConfiguration)
           .applyMutation(httpClientProperties::applyHttpClientConfigurations)
           .applyMutation(awsClientProperties::applyClientCredentialConfigurations)
+          .applyMutation(awsClientProperties::applyRetryConfigurations)
           .build();
     }
 

aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java

@@ -32,6 +32,8 @@
 import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
 import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder;
+import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
+import software.amazon.awssdk.core.retry.RetryMode;
 import software.amazon.awssdk.regions.Region;
 
 public class AwsClientProperties implements Serializable {
@@ -178,6 +180,26 @@ public AwsCredentialsProvider credentialsProvider(
     return DefaultCredentialsProvider.builder().build();
   }
 
+  /**
+   * Configure <a
+   * href="https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/core/retry/RetryMode.html">RetryMode</a>
+   * to ADAPTIVE_V2 for AWS clients
+   *
+   * <p>Sample usage:
+   *
+   * <pre>
+   *   KmsClient.builder().applyMutation(awsClientProperties::applyRetryConfigurations)
+   * </pre>
+   */
+  public <T extends AwsClientBuilder> void applyRetryConfigurations(T builder) {
+    ClientOverrideConfiguration.Builder configBuilder =
+        null != builder.overrideConfiguration()
+            ? builder.overrideConfiguration().toBuilder()
+            : ClientOverrideConfiguration.builder();
+
+    builder.overrideConfiguration(configBuilder.retryStrategy(RetryMode.ADAPTIVE_V2).build());
+  }
+
   private AwsCredentialsProvider credentialsProvider(String credentialsProviderClass) {
     Class<?> providerClass;
     try {

aws/src/main/java/org/apache/iceberg/aws/lakeformation/LakeFormationAwsClientFactory.java

@@ -96,6 +96,7 @@ public KmsClient kms() {
     if (isTableRegisteredWithLakeFormation()) {
       return KmsClient.builder()
           .applyMutation(httpClientProperties()::applyHttpClientConfigurations)
+          .applyMutation(awsClientProperties()::applyRetryConfigurations)
           .credentialsProvider(
               new LakeFormationCredentialsProvider(lakeFormation(), buildTableArn()))
           .region(Region.of(region()))

aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java

@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.iceberg.aws.kms;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+import org.apache.iceberg.aws.AwsClientProperties;
+import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.core.retry.RetryMode;
+import software.amazon.awssdk.services.kms.KmsClient;
+import software.amazon.awssdk.services.kms.KmsClientBuilder;
+
+public class TestKmsClientProperties {
+  @Test
+  public void testApplyRetryConfiguration() {
+    AwsClientProperties awsClientProperties = new AwsClientProperties();
+
+    KmsClientBuilder builder = KmsClient.builder();
+    awsClientProperties.applyRetryConfigurations(builder);
+    RetryMode retryPolicy = builder.overrideConfiguration().retryMode().get();
+
+    assertThat(retryPolicy).as("retry mode should be ADAPTIVE_V2").isEqualTo(RetryMode.ADAPTIVE_V2);
+  }
+}