Skip to content

Commit bf77785

Browse files
arouelarouel
authored
Avoid unnecessary head-object request (#640)
1 parent 9fd29c5 commit bf77785

File tree

3 files changed

+86
-15
lines changed

3 files changed

+86
-15
lines changed

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

+22-7
Original file line numberDiff line numberDiff line change
@@ -18,10 +18,12 @@
1818
import java.util.HashSet;
1919
import java.util.Objects;
2020
import java.util.Set;
21+
import java.util.concurrent.CompletionException;
2122
import java.util.concurrent.ExecutionException;
2223
import java.util.concurrent.TimeoutException;
2324
import org.checkerframework.checker.nullness.qual.NonNull;
2425
import software.amazon.awssdk.services.s3.S3AsyncClient;
26+
import software.amazon.awssdk.services.s3.model.S3Exception;
2527

2628
class S3WritableByteChannel implements SeekableByteChannel {
2729
private final S3Path path;
@@ -44,18 +46,31 @@ class S3WritableByteChannel implements SeekableByteChannel {
4446

4547
try {
4648
var fileSystemProvider = (S3FileSystemProvider) path.getFileSystem().provider();
47-
var exists = fileSystemProvider.exists(client, path);
4849

49-
if (exists && options.contains(StandardOpenOption.CREATE_NEW)) {
50+
if (options.contains(StandardOpenOption.CREATE_NEW) && fileSystemProvider.exists(client, path)) {
5051
throw new FileAlreadyExistsException("File at path:" + path + " already exists");
5152
}
52-
if (!exists && !options.contains(StandardOpenOption.CREATE_NEW) && !options.contains(StandardOpenOption.CREATE)) {
53-
throw new NoSuchFileException("File at path:" + path + " does not exist yet");
54-
}
5553

5654
tempFile = Files.createTempFile("aws-s3-nio-", ".tmp");
57-
if (exists) {
58-
s3TransferUtil.downloadToLocalFile(path, tempFile);
55+
// this complicated download handling is the result of
56+
// avoiding an existence check with a head-object request
57+
if (!options.contains(StandardOpenOption.CREATE_NEW)) {
58+
try {
59+
s3TransferUtil.downloadToLocalFile(path, tempFile);
60+
} catch (CompletionException e) {
61+
var cause = e.getCause();
62+
if (!(cause instanceof S3Exception)) {
63+
throw e;
64+
}
65+
var s3e = (S3Exception) cause;
66+
if (s3e.statusCode() != 404) {
67+
throw e;
68+
}
69+
if (!options.contains(StandardOpenOption.CREATE)) {
70+
throw new NoSuchFileException("File at path " + path + " does not exist yet");
71+
}
72+
// gracefully handle the file creation
73+
}
5974
}
6075

6176
channel = Files.newByteChannel(this.tempFile, removeCreateNew(options));

‎src/test/java/software/amazon/nio/spi/s3/S3SeekableByteChannelTest.java

+3-2
Original file line numberDiff line numberDiff line change
@@ -42,9 +42,9 @@
4242
import software.amazon.awssdk.services.s3.model.GetObjectResponse;
4343
import software.amazon.awssdk.services.s3.model.HeadObjectRequest;
4444
import software.amazon.awssdk.services.s3.model.HeadObjectResponse;
45-
import software.amazon.awssdk.services.s3.model.NoSuchKeyException;
4645
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
4746
import software.amazon.awssdk.services.s3.model.PutObjectResponse;
47+
import software.amazon.awssdk.services.s3.model.S3Exception;
4848

4949
@ExtendWith(MockitoExtension.class)
5050
@SuppressWarnings("unchecked")
@@ -102,7 +102,8 @@ public void read() throws IOException {
102102

103103
@Test
104104
public void write() throws IOException {
105-
when(mockClient.headObject(any(HeadObjectRequest.class))).thenThrow(NoSuchKeyException.class);
105+
var exception = S3Exception.builder().statusCode(404).build();
106+
when(mockClient.getObject(any(GetObjectRequest.class), any(AsyncResponseTransformer.class))).thenThrow(exception);
106107
when(mockClient.putObject(any(PutObjectRequest.class), any(AsyncRequestBody.class))).thenReturn(CompletableFuture.supplyAsync(() ->
107108
PutObjectResponse.builder().build()));
108109
try(var channel = new S3SeekableByteChannel(path, mockClient, Set.<OpenOption>of(CREATE, WRITE), DisabledFileIntegrityCheck.INSTANCE)){

‎src/test/java/software/amazon/nio/spi/s3/S3WritableByteChannelTest.java

+61-6
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,8 @@
2121
import java.nio.file.Path;
2222
import java.nio.file.StandardOpenOption;
2323
import java.util.Set;
24+
import java.util.concurrent.CompletionException;
25+
import java.util.concurrent.ExecutionException;
2426
import java.util.concurrent.TimeoutException;
2527
import java.util.stream.Stream;
2628

@@ -32,11 +34,13 @@
3234
import static org.assertj.core.api.Assertions.assertThat;
3335
import static org.assertj.core.api.Assertions.assertThatThrownBy;
3436
import static org.mockito.ArgumentMatchers.any;
37+
import static org.mockito.Mockito.doThrow;
3538
import static org.mockito.Mockito.mock;
3639
import static org.mockito.Mockito.times;
3740
import static org.mockito.Mockito.verify;
3841
import static org.mockito.Mockito.when;
3942
import software.amazon.awssdk.services.s3.S3AsyncClient;
43+
import software.amazon.awssdk.services.s3.model.S3Exception;
4044

4145
@TestInstance(TestInstance.Lifecycle.PER_CLASS)
4246
class S3WritableByteChannelTest {
@@ -55,18 +59,69 @@ void whenFileExistsAndCreateNewShouldThrowFileAlreadyExistsException() throws In
5559
.isInstanceOf(FileAlreadyExistsException.class);
5660
}
5761

62+
@Test
63+
@DisplayName("when file does not exist and constructor is invoked the option `CREATE` option")
64+
void whenFileDoesNotExistsAndCreateOptionIsPresent() throws Exception {
65+
var provider = mock(S3FileSystemProvider.class);
66+
var fs = mock(S3FileSystem.class);
67+
when(fs.provider()).thenReturn(provider);
68+
var file = S3Path.getPath(fs, "somefile");
69+
var s3Client = mock(S3AsyncClient.class);
70+
var transferManager = mock(S3TransferUtil.class);
71+
var exception = new CompletionException(S3Exception.builder().statusCode(404).build());
72+
doThrow(exception).when(transferManager).downloadToLocalFile(any(S3Path.class), any(Path.class));
73+
74+
try (var channel = new S3WritableByteChannel(file, s3Client, transferManager, Set.of(CREATE))) {
75+
assertThat(channel.position()).isZero();
76+
}
77+
}
78+
5879
@Test
5980
@DisplayName("when file does not exist and constructor is invoked without option `CREATE_NEW` nor `CREATE` should throw NoSuchFileException")
60-
void whenFileDoesNotExistsAndNoCreateNewShouldThrowNoSuchFileException() throws InterruptedException, TimeoutException {
61-
S3FileSystemProvider provider = mock();
62-
when(provider.exists(any(S3AsyncClient.class), any())).thenReturn(false);
81+
void whenFileDoesNotExistsAndNoCreateShouldThrowNoSuchFileException() throws Exception {
82+
var provider = mock(S3FileSystemProvider.class);
83+
var fs = mock(S3FileSystem.class);
84+
when(fs.provider()).thenReturn(provider);
85+
var file = S3Path.getPath(fs, "somefile");
86+
var s3Client = mock(S3AsyncClient.class);
87+
var transferManager = mock(S3TransferUtil.class);
88+
var exception = new CompletionException(S3Exception.builder().statusCode(404).build());
89+
doThrow(exception).when(transferManager).downloadToLocalFile(any(S3Path.class), any(Path.class));
6390

64-
S3FileSystem fs = mock();
91+
assertThatThrownBy(() -> new S3WritableByteChannel(file, s3Client, transferManager, emptySet()))
92+
.isInstanceOf(NoSuchFileException.class);
93+
}
94+
95+
@Test
96+
@DisplayName("when file the download fails due to an invalid request")
97+
void whenFileDownloadFailsDueToInvalidRequestTheExceptionShouldBePropagated() throws InterruptedException, TimeoutException, ExecutionException {
98+
var provider = mock(S3FileSystemProvider.class);
99+
var fs = mock(S3FileSystem.class);
65100
when(fs.provider()).thenReturn(provider);
101+
var file = S3Path.getPath(fs, "somefile");
102+
var s3Client = mock(S3AsyncClient.class);
103+
var transferManager = mock(S3TransferUtil.class);
104+
var exception = new CompletionException(S3Exception.builder().statusCode(400).message("Invalid Request").build());
105+
doThrow(exception).when(transferManager).downloadToLocalFile(any(S3Path.class), any(Path.class));
66106

107+
assertThatThrownBy(() -> new S3WritableByteChannel(file, s3Client, transferManager, emptySet()))
108+
.isEqualTo(exception);
109+
}
110+
111+
@Test
112+
@DisplayName("when file the download fails for an unknown reason")
113+
void whenFileDownloadFailsForUnknownReasonTheExceptionShouldBePropagated() throws InterruptedException, TimeoutException, ExecutionException {
114+
var provider = mock(S3FileSystemProvider.class);
115+
var fs = mock(S3FileSystem.class);
116+
when(fs.provider()).thenReturn(provider);
67117
var file = S3Path.getPath(fs, "somefile");
68-
assertThatThrownBy(() -> new S3WritableByteChannel(file, mock(), mock(), emptySet()))
69-
.isInstanceOf(NoSuchFileException.class);
118+
var s3Client = mock(S3AsyncClient.class);
119+
var transferManager = mock(S3TransferUtil.class);
120+
var exception = new CompletionException(new RuntimeException("unknown error"));
121+
doThrow(exception).when(transferManager).downloadToLocalFile(any(S3Path.class), any(Path.class));
122+
123+
assertThatThrownBy(() -> new S3WritableByteChannel(file, s3Client, transferManager, emptySet()))
124+
.isEqualTo(exception);
70125
}
71126

72127
@Test

0 commit comments

Comments
 (0)