feat: add the V2 Async S3 Client (#4396)

* feat: add the V2 Async S3 Client

* apply formatting

Author: jkwatson-verta

backend/common/pom.xml

@@ -198,6 +198,10 @@
             <groupId>com.amazonaws</groupId>
             <artifactId>aws-java-sdk-sts</artifactId>
         </dependency>
+        <dependency>
+            <groupId>software.amazon.awssdk</groupId>
+            <artifactId>s3</artifactId>
+        </dependency>
         <!-- https://mvnrepository.com/artifact/com.google.cloud/google-cloud-storage -->
         <dependency>
             <groupId>com.google.cloud</groupId>

backend/common/src/main/java/ai/verta/modeldb/common/artifactStore/storageservice/s3/RefCountedS3Client.java

@@ -5,24 +5,35 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
 
 public class RefCountedS3Client implements AutoCloseable {
   private static final Logger LOGGER = LogManager.getLogger(RefCountedS3Client.class);
   private final AWSCredentials credentials;
   private final AmazonS3 s3Client;
+  private final S3AsyncClient asyncClient;
   private final AtomicInteger referenceCounter;
 
-  RefCountedS3Client(AWSCredentials credentials, AmazonS3 client, AtomicInteger counter) {
+  RefCountedS3Client(
+      AWSCredentials credentials,
+      AmazonS3 client,
+      S3AsyncClient asyncClient,
+      AtomicInteger counter) {
     this.credentials = credentials;
-    s3Client = client;
-    referenceCounter = counter;
-    referenceCounter.incrementAndGet();
+    this.s3Client = client;
+    this.asyncClient = asyncClient;
+    this.referenceCounter = counter;
+    this.referenceCounter.incrementAndGet();
   }
 
   public AmazonS3 getClient() {
     return s3Client;
   }
 
+  public S3AsyncClient getAsyncClient() {
+    return asyncClient;
+  }
+
   public AWSCredentials getCredentials() {
     return credentials;
   }
@@ -34,6 +45,9 @@ public void close() {
         LOGGER.debug("shutting client down");
         s3Client.shutdown();
       }
+      if (asyncClient != null) {
+        asyncClient.close();
+      }
     }
   }
 }

backend/common/src/main/java/ai/verta/modeldb/common/artifactStore/storageservice/s3/RefreshS3ClientCron.java

@@ -16,6 +16,7 @@
 import java.util.UUID;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import software.amazon.awssdk.auth.credentials.AwsSessionCredentials;
 
 public class RefreshS3ClientCron implements Runnable {
   private static final Logger LOGGER = LogManager.getLogger(RefreshS3ClientCron.class);
@@ -57,7 +58,13 @@ private void createAndRefreshNewClient(BasicSessionCredentials awsCredentials) {
     newS3Client.doesBucketExistV2(bucketName);
     LOGGER.trace("New S3 Client created");
 
-    s3Client.refreshS3Client(awsCredentials, newS3Client);
+    s3Client.refreshS3Client(
+        awsCredentials,
+        newS3Client,
+        AwsSessionCredentials.create(
+            awsCredentials.getAWSAccessKeyId(),
+            awsCredentials.getAWSSecretKey(),
+            awsCredentials.getSessionToken()));
   }
 
   private BasicSessionCredentials getBasicSessionCredentials() throws IOException {

backend/common/src/main/java/ai/verta/modeldb/common/artifactStore/storageservice/s3/S3Client.java

@@ -14,6 +14,10 @@
 import java.util.function.Supplier;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
+import software.amazon.awssdk.auth.credentials.AwsBasicCredentials;
+import software.amazon.awssdk.auth.credentials.AwsCredentials;
+import software.amazon.awssdk.regions.Region;
+import software.amazon.awssdk.services.s3.S3AsyncClient;
 
 // S3Client provides a wrapper to the regular AmazonS3 object. The goal is to ensure that the
 // AmazonS3 object is
@@ -33,8 +37,10 @@ public class S3Client {
   private final S3Config s3Config;
 
   private volatile AmazonS3 s3Client;
+  private volatile S3AsyncClient asyncClient;
   private volatile AtomicInteger referenceCounter;
   private volatile AWSCredentials awsCredentials;
+  private volatile AwsCredentials v2Credentials;
 
   public S3Client(S3Config s3Config) throws ModelDBException {
     this.s3Config = s3Config;
@@ -89,7 +95,7 @@ private AmazonS3 buildEnvironmentClient(Regions awsRegion) {
    */
   private void scheduleRefresh(Supplier<AmazonS3> s3) {
     CommonUtils.scheduleTask(
-        () -> refreshS3Client(awsCredentials, s3.get()),
+        () -> refreshS3Client(awsCredentials, s3.get(), v2Credentials),
         0L,
         s3Config.getRefreshIntervalSeconds(),
         TimeUnit.SECONDS);
@@ -98,6 +104,7 @@ private void scheduleRefresh(Supplier<AmazonS3> s3) {
   private void initializeMinioClient(
       String cloudAccessKey, String cloudSecretKey, Regions awsRegion, String minioEndpoint) {
     this.awsCredentials = new BasicAWSCredentials(cloudAccessKey, cloudSecretKey);
+    this.v2Credentials = AwsBasicCredentials.create(cloudAccessKey, cloudSecretKey);
     var clientConfiguration = new ClientConfiguration(defaultClientConfig);
     clientConfiguration.setSignerOverride("VertaSignOverrideS3Signer");
     SignerFactory.registerSigner("VertaSignOverrideS3Signer", SignOverrideS3Signer.class);
@@ -120,6 +127,7 @@ private AmazonS3 buildMinioClient(
   private void initializeS3ClientWithAccessKey(
       String cloudAccessKey, String cloudSecretKey, Regions awsRegion) {
     this.awsCredentials = new BasicAWSCredentials(cloudAccessKey, cloudSecretKey);
+    this.v2Credentials = AwsBasicCredentials.create(cloudAccessKey, cloudSecretKey);
     this.s3Client = buildAccessKeyClient(awsRegion);
     scheduleRefresh(() -> buildAccessKeyClient(awsRegion));
   }
@@ -133,7 +141,7 @@ private AmazonS3 buildAccessKeyClient(Regions awsRegion) {
   }
 
   public RefCountedS3Client getRefCountedClient() {
-    return new RefCountedS3Client(awsCredentials, s3Client, referenceCounter);
+    return new RefCountedS3Client(awsCredentials, s3Client, asyncClient, referenceCounter);
   }
 
   private void initializeWithWebIdentity(Regions awsRegion) {
@@ -152,7 +160,13 @@ private void initializeWithWebIdentity(Regions awsRegion) {
         TimeUnit.SECONDS);
   }
 
-  void refreshS3Client(AWSCredentials awsCredentials, AmazonS3 s3Client) {
+  void refreshS3Client(
+      AWSCredentials awsCredentials, AmazonS3 s3Client, AwsCredentials v2Credentials) {
+    var s3AsyncClient =
+        S3AsyncClient.builder()
+            .credentialsProvider(() -> v2Credentials)
+            .region(Region.of(s3Config.getAwsRegion()))
+            .build();
     // Once we get to this point, we know that we have a good new s3 client, so it's time to swap
     // it. No fail can happen now
     LOGGER.debug("Replacing S3 Client");
@@ -163,7 +177,9 @@ void refreshS3Client(AWSCredentials awsCredentials, AmazonS3 s3Client) {
       // Swap the references
       this.referenceCounter = new AtomicInteger(1);
       this.awsCredentials = awsCredentials;
+      this.v2Credentials = v2Credentials;
       this.s3Client = s3Client;
+      this.asyncClient = s3AsyncClient;
       LOGGER.debug("S3 Client replaced");
       // At the end of the try, the reference counter will be decremented again and shutdown will
       // be called

pom.xml

@@ -30,6 +30,7 @@
         <prometheus.version>0.16.0</prometheus.version>
         <protobuf.version>3.25.2</protobuf.version>
         <aws.version>1.12.656</aws.version>
+        <aws.java.sdk.v2.version>2.25.1</aws.java.sdk.v2.version>
         <jackson.version>2.16.1</jackson.version>
         <os-maven-plugin.version>1.7.1</os-maven-plugin.version>
         <otel.core.version>1.32.0</otel.core.version>
@@ -258,6 +259,13 @@
                 <artifactId>aws-java-sdk-sts</artifactId>
                 <version>${aws.version}</version>
             </dependency>
+            <dependency>
+                <groupId>software.amazon.awssdk</groupId>
+                <artifactId>bom</artifactId>
+                <version>${aws.java.sdk.v2.version}</version>
+                <type>pom</type>
+                <scope>import</scope>
+            </dependency>
 
 
             <!-- https://mvnrepository.com/artifact/com.google.cloud/google-cloud-storage -->