feat: DH-18145: Write timeout configuration with S3 (deephaven#6650)

Author: AlSpinks

extensions/s3/src/main/java/io/deephaven/extensions/s3/S3ClientFactory.java

@@ -17,17 +17,13 @@
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
 import software.amazon.awssdk.core.client.config.SdkAdvancedAsyncClientOption;
 import software.amazon.awssdk.core.exception.SdkClientException;
-import software.amazon.awssdk.core.retry.RetryMode;
 import software.amazon.awssdk.http.SdkHttpClient;
 import software.amazon.awssdk.http.async.SdkAsyncHttpClient;
 import software.amazon.awssdk.http.crt.AwsCrtAsyncHttpClient;
 import software.amazon.awssdk.http.crt.AwsCrtHttpClient;
 import software.amazon.awssdk.regions.Region;
-import software.amazon.awssdk.services.s3.S3AsyncClient;
-import software.amazon.awssdk.services.s3.S3AsyncClientBuilder;
-import software.amazon.awssdk.services.s3.S3BaseClientBuilder;
-import software.amazon.awssdk.services.s3.S3Client;
-import software.amazon.awssdk.services.s3.S3ClientBuilder;
+import software.amazon.awssdk.retries.StandardRetryStrategy;
+import software.amazon.awssdk.services.s3.*;
 import software.amazon.awssdk.utils.ThreadFactoryBuilder;
 
 import java.time.Duration;
@@ -46,6 +42,10 @@ final class S3ClientFactory {
             getOrComputeThreadCountProperty("S3.numFutureCompletionThreads", -1);
     private static final int NUM_SCHEDULED_EXECUTOR_THREADS =
             getOrComputeThreadCountProperty("S3.numScheduledExecutorThreads", 5);
+    // The default retries value of 3 matches the value from the deprecated
+    // software.amazon.awssdk.core.retry.RetryMode#STANDARD
+    // which allowed 2 retries (3 attempts).
+    static final int RETRY_STRATEGY_MAX_ATTEMPTS = 3;
 
     private static final Logger log = LoggerFactory.getLogger(S3ClientFactory.class);
     private static final Map<HttpClientConfig, SdkAsyncHttpClient> httpAsyncClientCache = new ConcurrentHashMap<>();
@@ -150,9 +150,9 @@ private static <Builder extends SdkClientBuilder<Builder, Client>, Client> void
                 // If we find that the STANDARD retry policy does not work well in all situations, we might
                 // try experimenting with ADAPTIVE retry policy, potentially with fast fail.
                 // .retryPolicy(RetryPolicy.builder(RetryMode.ADAPTIVE).fastFailRateLimiting(true).build())
-                .retryPolicy(RetryMode.STANDARD)
-                .apiCallAttemptTimeout(instructions.readTimeout().dividedBy(3))
-                .apiCallTimeout(instructions.readTimeout())
+                .retryStrategy(StandardRetryStrategy.builder()
+                        .maxAttempts(RETRY_STRATEGY_MAX_ATTEMPTS)
+                        .build())
                 // Adding a metrics publisher may be useful for debugging, but it's very verbose.
                 // .addMetricPublisher(LoggingMetricPublisher.create(Level.INFO, Format.PRETTY))
                 .scheduledExecutorService(ensureScheduledExecutor());

extensions/s3/src/main/java/io/deephaven/extensions/s3/S3CompletableOutputStream.java

@@ -7,35 +7,22 @@
 import io.deephaven.util.channel.CompletableOutputStream;
 import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
+import software.amazon.awssdk.awscore.AwsRequest;
 import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.S3Uri;
 import software.amazon.awssdk.services.s3.internal.multipart.SdkPojoConversionUtils;
-import software.amazon.awssdk.services.s3.model.AbortMultipartUploadRequest;
-import software.amazon.awssdk.services.s3.model.AbortMultipartUploadResponse;
-import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest;
-import software.amazon.awssdk.services.s3.model.CompletedMultipartUpload;
-import software.amazon.awssdk.services.s3.model.CompletedPart;
-import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
-import software.amazon.awssdk.services.s3.model.CreateMultipartUploadResponse;
-import software.amazon.awssdk.services.s3.model.UploadPartRequest;
-import software.amazon.awssdk.services.s3.model.UploadPartResponse;
+import software.amazon.awssdk.services.s3.model.*;
 
 import java.io.IOException;
 import java.net.URI;
 import java.nio.ByteBuffer;
-import java.util.ArrayList;
-import java.util.Collections;
-import java.util.Comparator;
-import java.util.List;
-import java.util.Objects;
-import java.util.concurrent.CancellationException;
-import java.util.concurrent.CompletableFuture;
-import java.util.concurrent.CompletionException;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.Semaphore;
+import java.time.Duration;
+import java.util.*;
+import java.util.concurrent.*;
 
 import static io.deephaven.extensions.s3.S3ReadContext.handleS3Exception;
+import static io.deephaven.extensions.s3.S3Utils.addTimeout;
 
 class S3CompletableOutputStream extends CompletableOutputStream {
 
@@ -290,10 +277,10 @@ public void close() throws IOException {
     ////////// Helper methods and classes //////////
 
     private String initiateMultipartUpload() throws IOException {
-        final CreateMultipartUploadRequest createMultipartUploadRequest = CreateMultipartUploadRequest.builder()
+        final CreateMultipartUploadRequest.Builder builder = CreateMultipartUploadRequest.builder()
                 .bucket(uri.bucket().orElseThrow())
-                .key(uri.key().orElseThrow())
-                .build();
+                .key(uri.key().orElseThrow());
+        final CreateMultipartUploadRequest createMultipartUploadRequest = applyOverrideConfiguration(builder).build();
         // Note: We can add support for other parameters like tagging, storage class, encryption, permissions, etc. in
         // future
         final CompletableFuture<CreateMultipartUploadResponse> initiateUploadFuture =
@@ -306,6 +293,7 @@ private String initiateMultipartUpload() throws IOException {
         try {
             response = initiateUploadFuture.get();
         } catch (final InterruptedException | ExecutionException | CancellationException e) {
+            initiateUploadFuture.cancel(true);
             throw handleS3Exception(e, String.format("initiating multipart upload for uri %s", uri), s3Instructions);
         }
         return response.uploadId();
@@ -334,12 +322,12 @@ private void sendPartRequest(boolean onDone) throws IOException {
             return;
         }
 
-        final UploadPartRequest uploadPartRequest = UploadPartRequest.builder()
+        final UploadPartRequest.Builder builder = UploadPartRequest.builder()
                 .bucket(uri.bucket().orElseThrow())
                 .key(uri.key().orElseThrow())
                 .uploadId(uploadId)
-                .partNumber(nextPartNumber)
-                .build();
+                .partNumber(nextPartNumber);
+        final UploadPartRequest uploadPartRequest = applyOverrideConfiguration(builder).build();
 
         final int partNumber = nextPartNumber;
         final CompletableFuture<UploadPartResponse> uploadPartFuture = s3AsyncClient.uploadPart(uploadPartRequest,
@@ -449,17 +437,20 @@ private void completeMultipartUpload() throws IOException {
         // Sort the completed parts by part number, as required by S3
         completedParts.sort(Comparator.comparingInt(CompletedPart::partNumber));
 
-        final CompleteMultipartUploadRequest completeRequest = CompleteMultipartUploadRequest.builder()
+        final CompleteMultipartUploadRequest.Builder builder = CompleteMultipartUploadRequest.builder()
                 .bucket(uri.bucket().orElseThrow())
                 .key(uri.key().orElseThrow())
                 .uploadId(uploadId)
                 .multipartUpload(CompletedMultipartUpload.builder()
                         .parts(completedParts)
-                        .build())
-                .build();
+                        .build());
+        final CompleteMultipartUploadRequest completeRequest = applyOverrideConfiguration(builder).build();
+        final CompletableFuture<CompleteMultipartUploadResponse> uploadFuture =
+                s3AsyncClient.completeMultipartUpload(completeRequest);
         try {
-            s3AsyncClient.completeMultipartUpload(completeRequest).get();
+            uploadFuture.get();
         } catch (final InterruptedException | ExecutionException | CancellationException e) {
+            uploadFuture.cancel(true);
             final IOException ex = handleS3Exception(e,
                     String.format("completing multipart upload for uri %s", uri), s3Instructions);
             failAll(ex);
@@ -503,19 +494,32 @@ private void abortMultipartUpload() throws IOException {
         }
 
         // Initiate the abort request
-        final AbortMultipartUploadRequest abortRequest = AbortMultipartUploadRequest.builder()
+        final AbortMultipartUploadRequest.Builder builder = AbortMultipartUploadRequest.builder()
                 .bucket(uri.bucket().orElseThrow())
                 .key(uri.key().orElseThrow())
-                .uploadId(uploadId)
-                .build();
+                .uploadId(uploadId);
+        final AbortMultipartUploadRequest abortRequest = applyOverrideConfiguration(builder).build();
         final CompletableFuture<AbortMultipartUploadResponse> future = s3AsyncClient.abortMultipartUpload(abortRequest);
 
         // Wait for the abort to complete
         try {
             future.get();
         } catch (final InterruptedException | ExecutionException | CancellationException e) {
+            future.cancel(true);
             throw handleS3Exception(e,
                     String.format("aborting multipart upload for uri %s", uri), s3Instructions);
         }
     }
+
+    /**
+     * Applies the write timeout, then generates a request for the given builder and request class.
+     *
+     * @param builder an instance of a {@link S3Request.Builder} class
+     * @return the builder
+     */
+    private <B extends AwsRequest.Builder> B applyOverrideConfiguration(@NotNull B builder) {
+        final Duration writeTimeout = s3Instructions.writeTimeout();
+        builder.overrideConfiguration(b -> addTimeout(b, writeTimeout));
+        return builder;
+    }
 }

extensions/s3/src/main/java/io/deephaven/extensions/s3/S3Instructions.java

@@ -33,9 +33,13 @@ public abstract class S3Instructions implements LogOutputAppendable {
     private static final int DEFAULT_FRAGMENT_SIZE = 1 << 16; // 64 KiB
     private static final int MIN_FRAGMENT_SIZE = 8 << 10; // 8 KiB
     private static final Duration DEFAULT_CONNECTION_TIMEOUT = Duration.ofSeconds(2);
-    private static final Duration DEFAULT_READ_TIMEOUT = Duration.ofSeconds(2);
+    @VisibleForTesting
+    static final Duration DEFAULT_READ_TIMEOUT = Duration.ofSeconds(2);
+    @VisibleForTesting
+    static final Duration DEFAULT_WRITE_TIMEOUT = Duration.ofSeconds(2);
     private static final int DEFAULT_NUM_CONCURRENT_WRITE_PARTS = 64;
     private static final int MIN_CONCURRENT_WRITE_PARTS = 1;
+    private static final Duration MIN_READ_WRITE_TIMEOUT = Duration.ofMillis(1);
 
     /**
      * We set default part size to 10 MiB. The maximum number of parts allowed is 10,000. This means maximum size of a
@@ -115,6 +119,16 @@ public Duration readTimeout() {
         return DEFAULT_READ_TIMEOUT;
     }
 
+    /**
+     * The amount of time to wait when writing a fragment before giving up and timing out, defaults to 2 seconds. The
+     * implementation may choose to internally retry the request multiple times, so long as the total time does not
+     * exceed this timeout.
+     */
+    @Default
+    public Duration writeTimeout() {
+        return DEFAULT_WRITE_TIMEOUT;
+    }
+
     /**
      * The credentials to use when reading or writing to S3. By default, uses {@link Credentials#resolving()}.
      */
@@ -220,6 +234,8 @@ public interface Builder {
 
         Builder readTimeout(Duration connectionTimeout);
 
+        Builder writeTimeout(Duration connectionTimeout);
+
         Builder credentials(Credentials credentials);
 
         Builder endpointOverride(URI endpointOverride);
@@ -307,6 +323,24 @@ final void boundsCheckMinNumConcurrentWriteParts() {
         }
     }
 
+    @Check
+    final void boundsCheckReadTimeout() {
+        if (MIN_READ_WRITE_TIMEOUT.compareTo(readTimeout()) > 0) {
+            throw new IllegalArgumentException(
+                    "readTimeout(=" + readTimeout() + ") must be >= " +
+                            MIN_READ_WRITE_TIMEOUT);
+        }
+    }
+
+    @Check
+    final void boundsCheckWriteTimeout() {
+        if (MIN_READ_WRITE_TIMEOUT.compareTo(writeTimeout()) > 0) {
+            throw new IllegalArgumentException(
+                    "writeTimeout(=" + writeTimeout() + ") must be >= " +
+                            MIN_READ_WRITE_TIMEOUT);
+        }
+    }
+
     @Check
     final void boundsCheckMaxNumConcurrentWriteParts() {
         if (numConcurrentWriteParts() > maxConcurrentRequests()) {

extensions/s3/src/main/java/io/deephaven/extensions/s3/S3ReadRequest.java

@@ -31,6 +31,8 @@
 import java.util.concurrent.TimeoutException;
 import java.util.function.BiConsumer;
 
+import static io.deephaven.extensions.s3.S3Utils.addTimeout;
+
 /**
  * A request for a single fragment of an S3 object, which can be used concurrently.
  *
@@ -312,11 +314,13 @@ private int requestLength() {
     }
 
     private GetObjectRequest getObjectRequest() {
-        return GetObjectRequest.builder()
+        final GetObjectRequest.Builder requestBuilder = GetObjectRequest.builder()
                 .bucket(s3Uri.bucket().orElseThrow())
                 .key(s3Uri.key().orElseThrow())
-                .range("bytes=" + from + "-" + to)
-                .build();
+                .range("bytes=" + from + "-" + to);
+        final Duration readTimeout = instructions.readTimeout();
+        requestBuilder.overrideConfiguration(b -> addTimeout(b, readTimeout));
+        return requestBuilder.build();
     }
 
     String requestStr() {

extensions/s3/src/main/java/io/deephaven/extensions/s3/S3SeekableChannelProvider.java

@@ -39,6 +39,8 @@
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.time.Duration;
+import java.util.concurrent.*;
 import java.util.concurrent.atomic.AtomicReferenceFieldUpdater;
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
@@ -47,6 +49,7 @@
 import static io.deephaven.base.FileUtils.REPEATED_URI_SEPARATOR_PATTERN;
 import static io.deephaven.base.FileUtils.URI_SEPARATOR;
 import static io.deephaven.extensions.s3.S3ReadContext.handleS3Exception;
+import static io.deephaven.extensions.s3.S3Utils.addTimeout;
 
 /**
  * {@link SeekableChannelsProvider} implementation that is used to fetch objects from an S3-compatible API.
@@ -225,12 +228,17 @@ private void fetchNextBatch() throws IOException {
                     // Add a delimiter to the request if we don't want to fetch all files recursively
                     requestBuilder.delimiter("/");
                 }
-                final long readTimeoutNanos = s3Instructions.readTimeout().toNanos();
+                final Duration readTimeout = s3Instructions.readTimeout();
+                final long readTimeoutNanos = readTimeout.toNanos();
+                requestBuilder.overrideConfiguration(b -> addTimeout(b, readTimeout));
+
                 final ListObjectsV2Request request = requestBuilder.continuationToken(continuationToken).build();
+                final CompletableFuture<ListObjectsV2Response> responseFuture = s3AsyncClient.listObjectsV2(request);
                 final ListObjectsV2Response response;
                 try {
-                    response = s3AsyncClient.listObjectsV2(request).get(readTimeoutNanos, TimeUnit.NANOSECONDS);
+                    response = responseFuture.get(readTimeoutNanos, TimeUnit.NANOSECONDS);
                 } catch (final InterruptedException | ExecutionException | TimeoutException | CancellationException e) {
+                    responseFuture.cancel(true);
                     throw handleS3Exception(e, String.format("fetching list of files in directory %s", directory),
                             s3Instructions);
                 }
@@ -292,14 +300,16 @@ long fetchFileSize(@NotNull final S3Uri s3Uri) throws IOException {
             log.debug().append("Head: ").append(s3Uri.toString()).endl();
         }
         final HeadObjectResponse headObjectResponse;
+        final HeadObjectRequest.Builder requestBuilder = HeadObjectRequest.builder()
+                .bucket(s3Uri.bucket().orElseThrow())
+                .key(s3Uri.key().orElseThrow());
+        final Duration readTimeout = s3Instructions.readTimeout();
+        requestBuilder.overrideConfiguration(b -> addTimeout(b, readTimeout));
+        final CompletableFuture<HeadObjectResponse> responseFuture = s3AsyncClient.headObject(requestBuilder.build());
         try {
-            headObjectResponse = s3AsyncClient
-                    .headObject(HeadObjectRequest.builder()
-                            .bucket(s3Uri.bucket().orElseThrow())
-                            .key(s3Uri.key().orElseThrow())
-                            .build())
-                    .get(s3Instructions.readTimeout().toNanos(), TimeUnit.NANOSECONDS);
+            headObjectResponse = responseFuture.get(readTimeout.toNanos(), TimeUnit.NANOSECONDS);
         } catch (final InterruptedException | ExecutionException | TimeoutException | CancellationException e) {
+            responseFuture.cancel(true);
             throw handleS3Exception(e, String.format("fetching HEAD for file %s", s3Uri), s3Instructions);
         }
         final long fileSize = headObjectResponse.contentLength();

extensions/s3/src/main/java/io/deephaven/extensions/s3/S3Utils.java

@@ -4,12 +4,16 @@
 package io.deephaven.extensions.s3;
 
 import org.jetbrains.annotations.NotNull;
+import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
 import software.amazon.awssdk.profiles.ProfileFile;
 import software.amazon.awssdk.profiles.ProfileFileLocation;
 
 import java.nio.file.Path;
+import java.time.Duration;
 import java.util.Optional;
 
+import static io.deephaven.extensions.s3.S3ClientFactory.RETRY_STRATEGY_MAX_ATTEMPTS;
+
 class S3Utils {
 
     /**
@@ -48,4 +52,15 @@ private static void addProfileFile(
                 .content(path)
                 .build());
     }
+
+    /**
+     * Helper function to add timeout to the builder.
+     *
+     * @param builder the {@link AwsRequestOverrideConfiguration.Builder} to add the timeout to
+     * @param timeout the timeout to add
+     */
+    static void addTimeout(AwsRequestOverrideConfiguration.Builder builder, final Duration timeout) {
+        builder.apiCallAttemptTimeout(timeout.dividedBy(RETRY_STRATEGY_MAX_ATTEMPTS))
+                .apiCallTimeout(timeout);
+    }
 }