From c85ddaeb9c62eeece282337a705143b526821bc3 Mon Sep 17 00:00:00 2001
From: Jon Chambers <63609320+jon-signal@users.noreply.github.com>
Date: Thu, 16 Nov 2023 12:29:43 -0500
Subject: [PATCH] Make DynamoDB client connection pool sizes configurable

---
 service/pom.xml                               |  8 ++++
 .../DynamoDbClientConfiguration.java          | 44 +++++++------------
 .../util/DynamoDbFromConfig.java              | 18 +++++---
 3 files changed, 37 insertions(+), 33 deletions(-)

diff --git a/service/pom.xml b/service/pom.xml
index a5e3c6247..a91bd8683 100644
--- a/service/pom.xml
+++ b/service/pom.xml
@@ -297,6 +297,14 @@
       <artifactId>reactor-grpc-stub</artifactId>
     </dependency>
 
+    <dependency>
+      <groupId>software.amazon.awssdk</groupId>
+      <artifactId>apache-client</artifactId>
+    </dependency>
+    <dependency>
+      <groupId>software.amazon.awssdk</groupId>
+      <artifactId>netty-nio-client</artifactId>
+    </dependency>
     <dependency>
       <groupId>software.amazon.awssdk</groupId>
       <artifactId>sts</artifactId>
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/DynamoDbClientConfiguration.java b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/DynamoDbClientConfiguration.java
index 3a7c84a7c..804e52504 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/configuration/DynamoDbClientConfiguration.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/configuration/DynamoDbClientConfiguration.java
@@ -5,37 +5,27 @@
 
 package org.whispersystems.textsecuregcm.configuration;
 
-import com.fasterxml.jackson.annotation.JsonCreator;
-import com.fasterxml.jackson.annotation.JsonProperty;
 import java.time.Duration;
-import javax.validation.constraints.NotEmpty;
+import javax.validation.constraints.NotBlank;
+import javax.validation.constraints.NotNull;
+import javax.validation.constraints.Positive;
 
-public class DynamoDbClientConfiguration {
+public record DynamoDbClientConfiguration(@NotBlank String region,
+                                          @NotNull Duration clientExecutionTimeout,
+                                          @NotNull Duration clientRequestTimeout,
+                                          @Positive int maxConnections) {
 
-  private final String region;
-  private final Duration clientExecutionTimeout;
-  private final Duration clientRequestTimeout;
+  public DynamoDbClientConfiguration {
+    if (clientExecutionTimeout == null) {
+      clientExecutionTimeout = Duration.ofSeconds(30);
+    }
 
-  @JsonCreator
-  public DynamoDbClientConfiguration(
-      @JsonProperty("region") final String region,
-      @JsonProperty("clientExcecutionTimeout") final Duration clientExecutionTimeout,
-      @JsonProperty("clientRequestTimeout") final Duration clientRequestTimeout) {
-    this.region = region;
-    this.clientExecutionTimeout = clientExecutionTimeout != null ? clientExecutionTimeout : Duration.ofSeconds(30);
-    this.clientRequestTimeout = clientRequestTimeout != null ? clientRequestTimeout : Duration.ofSeconds(10);
-  }
+    if (clientRequestTimeout == null) {
+      clientRequestTimeout = Duration.ofSeconds(10);
+    }
 
-  @NotEmpty
-  public String getRegion() {
-    return region;
-  }
-
-  public Duration getClientExecutionTimeout() {
-    return clientExecutionTimeout;
-  }
-
-  public Duration getClientRequestTimeout() {
-    return clientRequestTimeout;
+    if (maxConnections == 0) {
+      maxConnections = 50;
+    }
   }
 }
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/util/DynamoDbFromConfig.java b/service/src/main/java/org/whispersystems/textsecuregcm/util/DynamoDbFromConfig.java
index 559baf3d2..79110de16 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/util/DynamoDbFromConfig.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/DynamoDbFromConfig.java
@@ -3,6 +3,8 @@ package org.whispersystems.textsecuregcm.util;
 import org.whispersystems.textsecuregcm.configuration.DynamoDbClientConfiguration;
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
+import software.amazon.awssdk.http.apache.ApacheHttpClient;
+import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@@ -11,12 +13,14 @@ public class DynamoDbFromConfig {
 
   public static DynamoDbClient client(DynamoDbClientConfiguration config, AwsCredentialsProvider credentialsProvider) {
     return DynamoDbClient.builder()
-        .region(Region.of(config.getRegion()))
+        .region(Region.of(config.region()))
         .credentialsProvider(credentialsProvider)
         .overrideConfiguration(ClientOverrideConfiguration.builder()
-            .apiCallTimeout(config.getClientExecutionTimeout())
-            .apiCallAttemptTimeout(config.getClientRequestTimeout())
+            .apiCallTimeout(config.clientExecutionTimeout())
+            .apiCallAttemptTimeout(config.clientRequestTimeout())
             .build())
+        .httpClientBuilder(ApacheHttpClient.builder()
+            .maxConnections(config.maxConnections()))
         .build();
   }
 
@@ -24,12 +28,14 @@ public class DynamoDbFromConfig {
       DynamoDbClientConfiguration config,
       AwsCredentialsProvider credentialsProvider) {
     return DynamoDbAsyncClient.builder()
-        .region(Region.of(config.getRegion()))
+        .region(Region.of(config.region()))
         .credentialsProvider(credentialsProvider)
         .overrideConfiguration(ClientOverrideConfiguration.builder()
-            .apiCallTimeout(config.getClientExecutionTimeout())
-            .apiCallAttemptTimeout(config.getClientRequestTimeout())
+            .apiCallTimeout(config.clientExecutionTimeout())
+            .apiCallAttemptTimeout(config.clientRequestTimeout())
             .build())
+        .httpClientBuilder(NettyNioAsyncHttpClient.builder()
+            .maxConcurrency(config.maxConnections()))
         .build();
   }
 }