AWS: Add AWS crt client support

apache · Apr 24, 2024 · c8ec7d1 · c8ec7d1
1 parent bfe0daa
commit c8ec7d1
Show file tree

Hide file tree

Showing 8 changed files with 174 additions and 10 deletions.
diff --git a/aws-bundle/build.gradle b/aws-bundle/build.gradle
@@ -27,6 +27,7 @@ project(":iceberg-aws-bundle") {
     implementation platform(libs.awssdk.bom)
     implementation "software.amazon.awssdk:apache-client"
     implementation "software.amazon.awssdk:auth"
+    implementation "software.amazon.awssdk:aws-crt-client"
     implementation "software.amazon.awssdk:iam"
     implementation "software.amazon.awssdk:sso"
     implementation "software.amazon.awssdk:s3"

diff --git a/aws/src/main/java/org/apache/iceberg/aws/AwsCrtHttpClientConfigurations.java b/aws/src/main/java/org/apache/iceberg/aws/AwsCrtHttpClientConfigurations.java
@@ -0,0 +1,71 @@
+/*
+ * 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;
+
+import java.time.Duration;
+import java.util.Map;
+import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
+import org.apache.iceberg.util.PropertyUtil;
+import software.amazon.awssdk.awscore.client.builder.AwsSyncClientBuilder;
+import software.amazon.awssdk.http.crt.AwsCrtHttpClient;
+
+class AwsCrtHttpClientConfigurations {
+  private Long connectionTimeoutMs;
+  private Long connectionMaxIdleTimeMs;
+  private Integer maxConcurrency;
+
+  private AwsCrtHttpClientConfigurations() {}
+
+  public <T extends AwsSyncClientBuilder> void configureHttpClientBuilder(T awsClientBuilder) {
+    AwsCrtHttpClient.Builder httpClientBuilder = AwsCrtHttpClient.builder();
+    configureAwsCrtHttpClientBuilder(httpClientBuilder);
+    awsClientBuilder.httpClientBuilder(httpClientBuilder);
+  }
+
+  private void initialize(Map<String, String> httpClientProperties) {
+    this.connectionTimeoutMs =
+        PropertyUtil.propertyAsNullableLong(
+            httpClientProperties, HttpClientProperties.AWS_CRT_CONNECTION_TIMEOUT_MS);
+    this.connectionMaxIdleTimeMs =
+        PropertyUtil.propertyAsNullableLong(
+            httpClientProperties, HttpClientProperties.AWS_CRT_CONNECTION_MAX_IDLE_TIME_MS);
+    this.maxConcurrency =
+        PropertyUtil.propertyAsNullableInt(
+            httpClientProperties, HttpClientProperties.AWS_CRT_MAX_CONCURRENCY);
+  }
+
+  @VisibleForTesting
+  void configureAwsCrtHttpClientBuilder(AwsCrtHttpClient.Builder httpClientBuilder) {
+    if (connectionTimeoutMs != null) {
+      httpClientBuilder.connectionTimeout(Duration.ofMillis(connectionTimeoutMs));
+    }
+    if (connectionMaxIdleTimeMs != null) {
+      httpClientBuilder.connectionMaxIdleTime(Duration.ofMillis(connectionMaxIdleTimeMs));
+    }
+    if (maxConcurrency != null) {
+      httpClientBuilder.maxConcurrency(maxConcurrency);
+    }
+  }
+
+  public static AwsCrtHttpClientConfigurations create(Map<String, String> properties) {
+    AwsCrtHttpClientConfigurations configurations = new AwsCrtHttpClientConfigurations();
+    configurations.initialize(properties);
+    return configurations;
+  }
+}
diff --git a/aws/src/main/java/org/apache/iceberg/aws/HttpClientProperties.java b/aws/src/main/java/org/apache/iceberg/aws/HttpClientProperties.java
@@ -51,6 +51,8 @@ public class HttpClientProperties implements Serializable {
    */
   public static final String CLIENT_TYPE_URLCONNECTION = "urlconnection";
 
+  public static final String CLIENT_TYPE_AWS_CRT = "aws-crt";
+
   public static final String CLIENT_TYPE_DEFAULT = CLIENT_TYPE_APACHE;
   /**
    * Used to configure the connection timeout in milliseconds for {@link
@@ -167,6 +169,27 @@ public class HttpClientProperties implements Serializable {
   public static final String APACHE_USE_IDLE_CONNECTION_REAPER_ENABLED =
       "http-client.apache.use-idle-connection-reaper-enabled";
 
+  /**
+   * Used to configure the connection timeout in milliseconds for {@link
+   * software.amazon.awssdk.http.crt.AwsCrtHttpClient.Builder}. This flag only works when {@link
+   * #CLIENT_TYPE} is set to {@link #CLIENT_TYPE_AWS_CRT}
+   */
+  public static final String AWS_CRT_CONNECTION_TIMEOUT_MS =
+      "http-client.aws-crt.connection-timeout-ms";
+  /**
+   * Used to configure the connection max idle time in milliseconds for {@link
+   * software.amazon.awssdk.http.crt.AwsCrtHttpClient.Builder}. This flag only works when {@link
+   * #CLIENT_TYPE} is set to {@link #CLIENT_TYPE_AWS_CRT}
+   */
+  public static final String AWS_CRT_CONNECTION_MAX_IDLE_TIME_MS =
+      "http-client.aws-crt.connection-max-idle-time-ms";
+  /**
+   * Used to configure the max concurrency number for {@link
+   * software.amazon.awssdk.http.crt.AwsCrtHttpClient.Builder}. This flag only works when {@link
+   * #CLIENT_TYPE} is set to {@link #CLIENT_TYPE_AWS_CRT}
+   */
+  public static final String AWS_CRT_MAX_CONCURRENCY = "http-client.aws-crt.max-concurrency";
+
   private String httpClientType;
   private final Map<String, String> httpClientProperties;
 
@@ -183,8 +206,8 @@ public HttpClientProperties(Map<String, String> properties) {
   }
 
   /**
-   * Configure the httpClient for a client according to the HttpClientType. The two supported
-   * HttpClientTypes are urlconnection and apache
+   * Configure the httpClient for a client according to the HttpClientType. The three supported
+   * HttpClientTypes are: urlconnection, apache, and aws-crt
    *
    * <p>Sample usage:
    *
@@ -208,6 +231,11 @@ public <T extends AwsSyncClientBuilder> void applyHttpClientConfigurations(T bui
             loadHttpClientConfigurations(ApacheHttpClientConfigurations.class.getName());
         apacheHttpClientConfigurations.configureHttpClientBuilder(builder);
         break;
+      case CLIENT_TYPE_AWS_CRT:
+        AwsCrtHttpClientConfigurations awsCrtHttpClientConfigurations =
+            loadHttpClientConfigurations(AwsCrtHttpClientConfigurations.class.getName());
+        awsCrtHttpClientConfigurations.configureHttpClientBuilder(builder);
+        break;
       default:
         throw new IllegalArgumentException("Unrecognized HTTP client type " + httpClientType);
     }

diff --git a/aws/src/test/java/org/apache/iceberg/aws/TestAwsClientFactories.java b/aws/src/test/java/org/apache/iceberg/aws/TestAwsClientFactories.java
@@ -30,6 +30,8 @@
 import org.assertj.core.api.Assertions;
 import org.assertj.core.api.ThrowableAssert;
 import org.junit.jupiter.api.Test;
+import org.junit.jupiter.params.ParameterizedTest;
+import org.junit.jupiter.params.provider.MethodSource;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.auth.credentials.AwsCredentials;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
@@ -135,10 +137,16 @@ public void testLakeFormationAwsClientFactorySerializable() throws IOException {
         .isInstanceOf(LakeFormationAwsClientFactory.class);
   }
 
-  @Test
-  public void testWithDummyValidCredentialsProvider() {
+  private static String[] httpClientTypes() {
+    return new String[] {"apache", "aws-crt", "urlconnection", null};
+  }
+
+  @ParameterizedTest
+  @MethodSource("httpClientTypes")
+  public void testWithDummyValidCredentialsProvider(String httpClientType) {
     AwsClientFactory defaultAwsClientFactory =
-        getAwsClientFactoryByCredentialsProvider(DummyValidProvider.class.getName());
+        getAwsClientFactoryByCredentialsProvider(
+            DummyValidProvider.class.getName(), httpClientType);
     assertDefaultAwsClientFactory(defaultAwsClientFactory);
     assertClientObjectsNotNull(defaultAwsClientFactory);
     // Ensuring S3Exception thrown instead exception thrown by resolveCredentials() implemented by
@@ -187,7 +195,7 @@ public void testWithClassDoesNotImplementCredentialsProvider() {
 
   private void testProviderAndAssertThrownBy(String providerClassName, String containsMessage) {
     AwsClientFactory defaultAwsClientFactory =
-        getAwsClientFactoryByCredentialsProvider(providerClassName);
+        getAwsClientFactoryByCredentialsProvider(providerClassName, null);
     assertDefaultAwsClientFactory(defaultAwsClientFactory);
     assertAllClientObjectsThrownBy(defaultAwsClientFactory, containsMessage);
   }
@@ -222,8 +230,12 @@ private void assertDefaultAwsClientFactory(AwsClientFactory awsClientFactory) {
         .isInstanceOf(AwsClientFactories.DefaultAwsClientFactory.class);
   }
 
-  private AwsClientFactory getAwsClientFactoryByCredentialsProvider(String providerClass) {
+  private AwsClientFactory getAwsClientFactoryByCredentialsProvider(
+      String providerClass, String httpClientType) {
     Map<String, String> properties = getDefaultClientFactoryProperties(providerClass);
+    if (httpClientType != null) {
+      properties.put("http-client.type", httpClientType);
+    }
     AwsClientFactory defaultAwsClientFactory = AwsClientFactories.from(properties);
     return defaultAwsClientFactory;
   }

diff --git a/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientConfigurations.java b/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientConfigurations.java
@@ -24,6 +24,7 @@
 import org.junit.jupiter.api.Test;
 import org.mockito.Mockito;
 import software.amazon.awssdk.http.apache.ApacheHttpClient;
+import software.amazon.awssdk.http.crt.AwsCrtHttpClient;
 import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
 
 public class TestHttpClientConfigurations {
@@ -124,4 +125,36 @@ public void testApacheDefaultConfigurations() {
     Mockito.verify(spyApacheHttpClientBuilder, Mockito.never())
         .useIdleConnectionReaper(Mockito.anyBoolean());
   }
+
+  @Test
+  public void testAwsCrtOverrideConfigurations() {
+    Map<String, String> properties = Maps.newHashMap();
+    properties.put(HttpClientProperties.URLCONNECTION_SOCKET_TIMEOUT_MS, "90");
+    properties.put(HttpClientProperties.URLCONNECTION_CONNECTION_TIMEOUT_MS, "80");
+    properties.put(HttpClientProperties.AWS_CRT_CONNECTION_TIMEOUT_MS, "200");
+    properties.put(HttpClientProperties.AWS_CRT_CONNECTION_MAX_IDLE_TIME_MS, "102");
+    properties.put(HttpClientProperties.AWS_CRT_MAX_CONCURRENCY, "104");
+    AwsCrtHttpClientConfigurations awsCrtHttpClientConfigurations =
+        AwsCrtHttpClientConfigurations.create(properties);
+    AwsCrtHttpClient.Builder awsCrtHttpClientBuilder = AwsCrtHttpClient.builder();
+    AwsCrtHttpClient.Builder spyAwsCrtHttpClientBuilder = Mockito.spy(awsCrtHttpClientBuilder);
+
+    awsCrtHttpClientConfigurations.configureAwsCrtHttpClientBuilder(spyAwsCrtHttpClientBuilder);
+
+    Mockito.verify(spyAwsCrtHttpClientBuilder).connectionTimeout(Duration.ofMillis(200));
+    Mockito.verify(spyAwsCrtHttpClientBuilder).connectionMaxIdleTime(Duration.ofMillis(102));
+    Mockito.verify(spyAwsCrtHttpClientBuilder).maxConcurrency(104);
+  }
+
+  @Test
+  public void testAwsCrtDefaultConfigurations() {
+    AwsCrtHttpClientConfigurations awsCrtHttpClientConfigurations =
+        AwsCrtHttpClientConfigurations.create(Maps.newHashMap());
+    AwsCrtHttpClient.Builder awsCrtHttpClientBuilder = AwsCrtHttpClient.builder();
+    AwsCrtHttpClient.Builder spyAwsCrtHttpClientBuilder = Mockito.spy(awsCrtHttpClientBuilder);
+
+    awsCrtHttpClientConfigurations.configureAwsCrtHttpClientBuilder(spyAwsCrtHttpClientBuilder);
+
+    Mockito.verifyNoInteractions(spyAwsCrtHttpClientBuilder);
+  }
 }
diff --git a/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientProperties.java b/aws/src/test/java/org/apache/iceberg/aws/TestHttpClientProperties.java
@@ -26,6 +26,7 @@
 import org.mockito.Mockito;
 import software.amazon.awssdk.http.SdkHttpClient;
 import software.amazon.awssdk.http.apache.ApacheHttpClient;
+import software.amazon.awssdk.http.crt.AwsCrtHttpClient;
 import software.amazon.awssdk.http.urlconnection.UrlConnectionHttpClient;
 import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.s3.S3ClientBuilder;
@@ -67,6 +68,23 @@ public void testApacheHttpClientConfiguration() {
         .isInstanceOf(ApacheHttpClient.Builder.class);
   }
 
+  @Test
+  public void testAwsCrtHttpClientConfiguration() {
+    Map<String, String> properties = Maps.newHashMap();
+    properties.put(HttpClientProperties.CLIENT_TYPE, "aws-crt");
+    HttpClientProperties httpProperties = new HttpClientProperties(properties);
+    S3ClientBuilder mockS3ClientBuilder = Mockito.mock(S3ClientBuilder.class);
+    ArgumentCaptor<SdkHttpClient.Builder> httpClientBuilderCaptor =
+        ArgumentCaptor.forClass(SdkHttpClient.Builder.class);
+
+    httpProperties.applyHttpClientConfigurations(mockS3ClientBuilder);
+    Mockito.verify(mockS3ClientBuilder).httpClientBuilder(httpClientBuilderCaptor.capture());
+    SdkHttpClient.Builder capturedHttpClientBuilder = httpClientBuilderCaptor.getValue();
+    Assertions.assertThat(capturedHttpClientBuilder)
+        .as("Should use aws crt http client")
+        .isInstanceOf(AwsCrtHttpClient.Builder.class);
+  }
+
   @Test
   public void testInvalidHttpClientType() {
     Map<String, String> properties = Maps.newHashMap();

diff --git a/build.gradle b/build.gradle
@@ -476,6 +476,7 @@ project(':iceberg-aws') {
     compileOnly(libs.awssdk.s3accessgrants)
     compileOnly("software.amazon.awssdk:url-connection-client")
     compileOnly("software.amazon.awssdk:apache-client")
+    compileOnly("software.amazon.awssdk:aws-crt-client")
     compileOnly("software.amazon.awssdk:auth")
     compileOnly("software.amazon.awssdk:s3")
     compileOnly("software.amazon.awssdk:kms")

diff --git a/docs/docs/aws.md b/docs/docs/aws.md
@@ -580,9 +580,9 @@ For more details of configuration, see sections [URL Connection HTTP Client Conf
 
 Configure the following property to set the type of HTTP client:
 
-| Property         | Default | Description                                                                                                |
-|------------------|---------|------------------------------------------------------------------------------------------------------------|
-| http-client.type | apache  | Types of HTTP Client. <br/> `urlconnection`: URL Connection HTTP Client <br/> `apache`: Apache HTTP Client |
+| Property         | Default | Description                                                                                                                                 |
+|------------------|---------|---------------------------------------------------------------------------------------------------------------------------------------------|
+| http-client.type | apache  | Types of HTTP Client. <br/> `urlconnection`: URL Connection HTTP Client <br/> `apache`: Apache HTTP Client  <br/> `aws-crt`: AWS CRT client |
 
 #### URL Connection HTTP Client Configurations