Improve exception handling on download/upload

src/main/java/software/amazon/nio/spi/s3/S3TransferUtil.java

@@ -7,18 +7,24 @@
 
 import java.io.IOException;
 import java.nio.file.Files;
+import java.nio.file.NoSuchFileException;
+import java.nio.file.OpenOption;
 import java.nio.file.Path;
+import java.nio.file.StandardOpenOption;
+import java.util.Set;
+import java.util.concurrent.CompletionException;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
+import java.util.stream.Stream;
 import software.amazon.awssdk.core.FileTransformerConfiguration;
+import software.amazon.awssdk.core.async.AsyncRequestBody;
 import software.amazon.awssdk.core.async.AsyncResponseTransformer;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
 import software.amazon.awssdk.services.s3.model.GetObjectRequest;
 import software.amazon.awssdk.services.s3.model.GetObjectResponse;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
-import software.amazon.awssdk.transfer.s3.S3TransferManager;
-import software.amazon.awssdk.transfer.s3.model.UploadFileRequest;
+import software.amazon.awssdk.services.s3.model.S3Exception;
 
 final class S3TransferUtil {
     private final S3ObjectIntegrityCheck integrityCheck;
@@ -33,39 +39,59 @@ final class S3TransferUtil {
         this.integrityCheck = integrityCheck;
     }
 
-    void downloadToLocalFile(S3Path path, Path destination, S3OpenOption... options)
-            throws InterruptedException, ExecutionException, TimeoutException {
-        var getObjectRequest = GetObjectRequest.builder()
-            .bucket(path.bucketName())
-            .key(path.getKey());
-        for (var option : options) {
-            option.apply(getObjectRequest);
-        }
-        var transformerConfig = FileTransformerConfiguration.defaultCreateOrReplaceExisting();
-        var responseTransformer = AsyncResponseTransformer.<GetObjectResponse>toFile(destination, transformerConfig);
-        var downloadCompletableFuture = client.getObject(getObjectRequest.build(), responseTransformer)
-            .toCompletableFuture();
+    void downloadToLocalFile(S3Path path, Path destination, Set<? extends OpenOption> options) throws IOException {
+        var s3OpenOptions = options.stream()
+            .flatMap(o -> o instanceof S3OpenOption
+                ? Stream.of((S3OpenOption) o)
+                : Stream.empty())
+            .toArray(S3OpenOption[]::new);
+        try {
+            var getObjectRequest = GetObjectRequest.builder()
+                .bucket(path.bucketName())
+                .key(path.getKey());
+            for (var option : s3OpenOptions) {
+                option.apply(getObjectRequest);
+            }
+            var transformerConfig = FileTransformerConfiguration.defaultCreateOrReplaceExisting();
+            var responseTransformer = AsyncResponseTransformer.<GetObjectResponse>toFile(destination, transformerConfig);
+            var downloadCompletableFuture = client.getObject(getObjectRequest.build(), responseTransformer);
 
-        if (timeout != null && timeUnit != null) {
-            downloadCompletableFuture.get(timeout, timeUnit);
-        } else {
-            downloadCompletableFuture.join();
+            if (timeout != null && timeUnit != null) {
+                downloadCompletableFuture.get(timeout, timeUnit);
+            } else {
+                downloadCompletableFuture.join();
+            }
+        } catch (InterruptedException e) {
+            Thread.currentThread().interrupt();
+            throw new IOException("Could not read from path: " + path, e);
+        } catch (TimeoutException | ExecutionException e) {
+            throw new IOException("Could not read from path: " + path, e);
+        } catch (CompletionException e) {
+            // This complicated download handling is the result of omitting an existence check
+            // with a head object request, instead we look for a 404 status code if available.
+            var cause = e.getCause();
+            if (!(cause instanceof S3Exception)) {
+                throw new IOException("Could not read from path: " + path, e);
+            }
+            var s3e = (S3Exception) cause;
+            if (s3e.statusCode() != 404) {
+                throw new IOException("Could not read from path: " + path, e);
+            }
+            if (!options.contains(StandardOpenOption.CREATE)) {
+                throw new NoSuchFileException(path.toString());
+            }
+            // gracefully handle the file creation
         }
     }
 
     void uploadLocalFile(S3Path path, Path localFile) throws IOException {
-        try (var s3TransferManager = S3TransferManager.builder().s3Client(client).build()) {
+        try {
             var putObjectRequest = PutObjectRequest.builder()
                 .bucket(path.bucketName())
                 .key(path.getKey())
                 .contentType(Files.probeContentType(localFile));
             integrityCheck.addChecksumToRequest(localFile, putObjectRequest);
-            var uploadCompletableFuture = s3TransferManager.uploadFile(
-                UploadFileRequest.builder()
-                    .putObjectRequest(putObjectRequest.build())
-                    .source(localFile)
-                    .build()
-            ).completionFuture();
+            var uploadCompletableFuture = client.putObject(putObjectRequest.build(), AsyncRequestBody.fromFile(localFile));
 
             if (timeout != null && timeUnit != null) {
                 uploadCompletableFuture.get(timeout, timeUnit);
@@ -74,9 +100,11 @@ void uploadLocalFile(S3Path path, Path localFile) throws IOException {
             }
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
-            throw new IOException("Could not write to path:" + path, e);
+            throw new IOException("Could not write to path: " + path, e);
         } catch (TimeoutException | ExecutionException e) {
-            throw new IOException("Could not write to path:" + path, e);
+            throw new IOException("Could not write to path: " + path, e);
+        } catch (CompletionException e) {
+            throw new IOException("Could not write to path: " + path, e);
         }
     }
 }

src/main/java/software/amazon/nio/spi/s3/S3WritableByteChannel.java

@@ -11,20 +11,15 @@
 import java.nio.channels.SeekableByteChannel;
 import java.nio.file.FileAlreadyExistsException;
 import java.nio.file.Files;
-import java.nio.file.NoSuchFileException;
 import java.nio.file.OpenOption;
 import java.nio.file.Path;
 import java.nio.file.StandardOpenOption;
 import java.util.Objects;
 import java.util.Set;
-import java.util.concurrent.CompletionException;
-import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
-import java.util.stream.Stream;
 import org.checkerframework.checker.nullness.qual.NonNull;
 import software.amazon.awssdk.services.s3.S3AsyncClient;
-import software.amazon.awssdk.services.s3.model.S3Exception;
 
 class S3WritableByteChannel implements SeekableByteChannel {
     private final S3Path path;
@@ -54,46 +49,19 @@ class S3WritableByteChannel implements SeekableByteChannel {
 
             tempFile = path.getFileSystem().createTempFile(path);
             if (!options.contains(StandardOpenOption.CREATE_NEW)) {
-                downloadToLocalFile(path, s3TransferUtil, options);
+                s3TransferUtil.downloadToLocalFile(path, tempFile, options);
             }
 
             channel = Files.newByteChannel(this.tempFile, removeCreateNew(options));
         } catch (InterruptedException e) {
             Thread.currentThread().interrupt();
             throw new IOException("Could not open the path:" + path, e);
-        } catch (TimeoutException | ExecutionException e) {
+        } catch (TimeoutException e) {
             throw new IOException("Could not open the path:" + path, e);
         }
         this.open = true;
     }
 
-    private void downloadToLocalFile(S3Path path, S3TransferUtil s3TransferUtil, Set<? extends OpenOption> options)
-            throws InterruptedException, ExecutionException, TimeoutException, NoSuchFileException {
-        // this complicated download handling is the result of
-        // avoiding an existence check with a head-object request
-        try {
-            var s3OpenOptions = options.stream()
-                .flatMap(o -> o instanceof S3OpenOption
-                    ? Stream.of((S3OpenOption) o)
-                    : Stream.empty())
-                .toArray(S3OpenOption[]::new);
-            s3TransferUtil.downloadToLocalFile(path, tempFile, s3OpenOptions);
-        } catch (CompletionException e) {
-            var cause = e.getCause();
-            if (!(cause instanceof S3Exception)) {
-                throw e;
-            }
-            var s3e = (S3Exception) cause;
-            if (s3e.statusCode() != 404) {
-                throw e;
-            }
-            if (!options.contains(StandardOpenOption.CREATE)) {
-                throw new NoSuchFileException("File at path " + path + " does not exist yet");
-            }
-            // gracefully handle the file creation
-        }
-    }
-
     private @NonNull Set<? extends OpenOption> removeCreateNew(Set<? extends OpenOption> options) {
         return options.stream()
             .filter(o -> o != StandardOpenOption.CREATE_NEW)