From 4305db5579dab7c10d531c187032f7d0a4e8f55a Mon Sep 17 00:00:00 2001
From: ravi-signal <99042880+ravi-signal@users.noreply.github.com>
Date: Tue, 30 Jan 2024 12:46:27 -0600
Subject: [PATCH] Update to aws sdk 2.23.8 and use AwsCrtHttpClient

---
 pom.xml                                       |  6 +++-
 .../textsecuregcm/WhisperServerService.java   |  4 +++
 .../util/DynamoDbFromConfig.java              |  7 +++--
 .../textsecuregcm/util/ManagedAwsCrt.java     | 29 +++++++++++++++++++
 4 files changed, 42 insertions(+), 4 deletions(-)
 create mode 100644 service/src/main/java/org/whispersystems/textsecuregcm/util/ManagedAwsCrt.java

diff --git a/pom.xml b/pom.xml
index 101fc4f14..2ba7feaa9 100644
--- a/pom.xml
+++ b/pom.xml
@@ -37,7 +37,7 @@
   </modules>
 
   <properties>
-    <aws.sdk2.version>2.21.5</aws.sdk2.version>
+    <aws.sdk2.version>2.23.8</aws.sdk2.version>
     <braintree.version>3.27.0</braintree.version>
     <commons-csv.version>1.10.0</commons-csv.version>
     <commons-io.version>2.14.0</commons-io.version>
@@ -313,6 +313,10 @@
       <version>1.3</version>
       <scope>test</scope>
     </dependency>
+    <dependency>
+      <groupId>software.amazon.awssdk</groupId>
+      <artifactId>aws-crt-client</artifactId>
+    </dependency>
     <dependency>
       <groupId>org.wiremock</groupId>
       <!-- use standalone until Dropwizard 4 + jakarta.* -->
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
index 44f969218..f7287b7d7 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
@@ -209,6 +209,7 @@ import org.whispersystems.textsecuregcm.subscriptions.BankMandateTranslator;
 import org.whispersystems.textsecuregcm.subscriptions.BraintreeManager;
 import org.whispersystems.textsecuregcm.subscriptions.StripeManager;
 import org.whispersystems.textsecuregcm.util.DynamoDbFromConfig;
+import org.whispersystems.textsecuregcm.util.ManagedAwsCrt;
 import org.whispersystems.textsecuregcm.util.SystemMapper;
 import org.whispersystems.textsecuregcm.util.UsernameHashZkProofVerifier;
 import org.whispersystems.textsecuregcm.util.logging.LoggingUnhandledExceptionMapper;
@@ -237,6 +238,7 @@ import reactor.core.scheduler.Schedulers;
 import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
 import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider;
 import software.amazon.awssdk.auth.credentials.WebIdentityTokenFileCredentialsProvider;
+import software.amazon.awssdk.http.crt.AwsCrtHttpClient;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
 import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
@@ -321,6 +323,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
         headerControlledResourceBundleLookup);
     BankMandateTranslator bankMandateTranslator = new BankMandateTranslator(headerControlledResourceBundleLookup);
 
+    environment.lifecycle().manage(new ManagedAwsCrt());
     DynamoDbAsyncClient dynamoDbAsyncClient = DynamoDbFromConfig.asyncClient(config.getDynamoDbClientConfiguration(),
         AWSSDK_CREDENTIALS_PROVIDER);
 
@@ -646,6 +649,7 @@ public class WhisperServerService extends Application<WhisperServerConfiguration
     S3Client cdnS3Client = S3Client.builder()
         .credentialsProvider(cdnCredentialsProvider)
         .region(Region.of(config.getCdnConfiguration().region()))
+        .httpClientBuilder(AwsCrtHttpClient.builder())
         .build();
     S3AsyncClient asyncCdnS3Client = S3AsyncClient.builder()
         .credentialsProvider(cdnCredentialsProvider)
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 79110de16..804dcd6bf 100644
--- a/service/src/main/java/org/whispersystems/textsecuregcm/util/DynamoDbFromConfig.java
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/DynamoDbFromConfig.java
@@ -3,7 +3,7 @@ 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.crt.AwsCrtHttpClient;
 import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
@@ -19,8 +19,9 @@ public class DynamoDbFromConfig {
             .apiCallTimeout(config.clientExecutionTimeout())
             .apiCallAttemptTimeout(config.clientRequestTimeout())
             .build())
-        .httpClientBuilder(ApacheHttpClient.builder()
-            .maxConnections(config.maxConnections()))
+        .httpClientBuilder(AwsCrtHttpClient
+            .builder()
+            .maxConcurrency(config.maxConnections()))
         .build();
   }
 
diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/util/ManagedAwsCrt.java b/service/src/main/java/org/whispersystems/textsecuregcm/util/ManagedAwsCrt.java
new file mode 100644
index 000000000..25f52ac2e
--- /dev/null
+++ b/service/src/main/java/org/whispersystems/textsecuregcm/util/ManagedAwsCrt.java
@@ -0,0 +1,29 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+package org.whispersystems.textsecuregcm.util;
+
+import io.dropwizard.lifecycle.Managed;
+import software.amazon.awssdk.crt.CRT;
+
+/**
+ * The AWS CRT client registers its own JVM shutdown handler which makes sure asynchronous operations in the CRT don't
+ * call back into the JVM after the JVM shuts down. Unfortunately, this hook kills all outstanding operations using the
+ * CRT, even though we typically orchestrate a graceful shutdown where we give outstanding requests time to complete.
+ *
+ * The CRT lets you take over the shutdown sequencing by incrementing/decrementing a ref count, so we introduce a
+ * lifecycle object that will be shutdown after our graceful shutdown process.
+ */
+public class ManagedAwsCrt implements Managed {
+
+  @Override
+  public void start() {
+    CRT.acquireShutdownRef();
+  }
+
+  @Override
+  public void stop() {
+    CRT.releaseShutdownRef();
+  }
+}