Skip to content

Commit 3ab5cbf

Browse files
joviegasjoviegas
committed
S3CrossRegionSyncClient Redirect implementation
1 parent a87b09a commit 3ab5cbf

File tree

6 files changed

+393
-59
lines changed

6 files changed

+393
-59
lines changed

‎services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/crossregion/S3CrossRegionAsyncClient.java

Lines changed: 3 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -20,11 +20,10 @@
2020
import java.util.function.Function;
2121
import software.amazon.awssdk.annotations.SdkInternalApi;
2222
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
23-
import software.amazon.awssdk.endpoints.Endpoint;
2423
import software.amazon.awssdk.services.s3.DelegatingS3AsyncClient;
2524
import software.amazon.awssdk.services.s3.S3AsyncClient;
26-
import software.amazon.awssdk.services.s3.endpoints.S3EndpointParams;
2725
import software.amazon.awssdk.services.s3.endpoints.S3EndpointProvider;
26+
import software.amazon.awssdk.services.s3.internal.crossregion.endpointprovider.BucketEndpointProvider;
2827
import software.amazon.awssdk.services.s3.model.S3Request;
2928

3029
@SdkInternalApi
@@ -67,28 +66,12 @@ private AwsRequestOverrideConfiguration getOrCreateConfigWithEndpointProvider(S3
6766
S3EndpointProvider delegateEndpointProvider = (S3EndpointProvider)
6867
requestOverrideConfig.endpointProvider().orElseGet(() -> serviceClientConfiguration().endpointProvider().get());
6968

69+
// TODO : separate PR to provide supplier for Async client
7070
return requestOverrideConfig.toBuilder()
71-
.endpointProvider(BucketEndpointProvider.create(delegateEndpointProvider, bucket))
71+
.endpointProvider(BucketEndpointProvider.create(delegateEndpointProvider, null))
7272
.build();
7373
}
7474

7575
//TODO: add cross region logic
76-
static final class BucketEndpointProvider implements S3EndpointProvider {
77-
private final S3EndpointProvider delegate;
78-
private final String bucket;
7976

80-
private BucketEndpointProvider(S3EndpointProvider delegate, String bucket) {
81-
this.delegate = delegate;
82-
this.bucket = bucket;
83-
}
84-
85-
public static BucketEndpointProvider create(S3EndpointProvider delegate, String bucket) {
86-
return new BucketEndpointProvider(delegate, bucket);
87-
}
88-
89-
@Override
90-
public CompletableFuture<Endpoint> resolveEndpoint(S3EndpointParams endpointParams) {
91-
return delegate.resolveEndpoint(endpointParams);
92-
}
93-
}
9477
}

‎services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/crossregion/S3CrossRegionSyncClient.java

Lines changed: 65 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -15,80 +15,110 @@
1515

1616
package software.amazon.awssdk.services.s3.internal.crossregion;
1717

18+
import java.util.Map;
1819
import java.util.Optional;
19-
import java.util.concurrent.CompletableFuture;
20+
import java.util.concurrent.ConcurrentHashMap;
2021
import java.util.function.Function;
22+
import java.util.function.Supplier;
2123
import software.amazon.awssdk.annotations.SdkInternalApi;
2224
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
23-
import software.amazon.awssdk.endpoints.Endpoint;
25+
import software.amazon.awssdk.regions.Region;
2426
import software.amazon.awssdk.services.s3.DelegatingS3Client;
2527
import software.amazon.awssdk.services.s3.S3Client;
26-
import software.amazon.awssdk.services.s3.endpoints.S3EndpointParams;
2728
import software.amazon.awssdk.services.s3.endpoints.S3EndpointProvider;
29+
import software.amazon.awssdk.services.s3.internal.crossregion.endpointprovider.BucketEndpointProvider;
30+
import software.amazon.awssdk.services.s3.model.HeadBucketRequest;
31+
import software.amazon.awssdk.services.s3.model.S3Exception;
2832
import software.amazon.awssdk.services.s3.model.S3Request;
2933

34+
/**
35+
* Decorator S3 Sync client that will fetch the region name whenever there is Redirect 301 error due to cross region bucket
36+
* access.
37+
*/
3038
@SdkInternalApi
3139
public final class S3CrossRegionSyncClient extends DelegatingS3Client {
40+
41+
private static final String AMZ_BUCKET_REGION_HEADER = "x-amz-bucket-region";
42+
public static final int REDIRECT_STATUS_CODE = 301;
43+
private final Map<String, Region> bucketToRegionCache = new ConcurrentHashMap<>();
44+
3245
public S3CrossRegionSyncClient(S3Client s3Client) {
3346
super(s3Client);
3447
}
3548

49+
private static <T extends S3Request> Optional<String> bucketNameFromRequest(T request) {
50+
return request.getValueForField("Bucket", String.class);
51+
}
52+
3653
@Override
3754
protected <T extends S3Request, ReturnT> ReturnT invokeOperation(T request, Function<T, ReturnT> operation) {
3855

39-
Optional<String> bucket = request.getValueForField("Bucket", String.class);
40-
41-
if (bucket.isPresent()) {
42-
try {
43-
return operation.apply(requestWithDecoratedEndpointProvider(request, bucket.get()));
44-
} catch (Exception e) {
45-
handleOperationFailure(e, bucket.get());
56+
Optional<String> bucketRequest = bucketNameFromRequest(request);
57+
if (!bucketRequest.isPresent()) {
58+
return operation.apply(request);
59+
}
60+
String bucketName = bucketRequest.get();
61+
try {
62+
if (bucketToRegionCache.containsKey(bucketName)) {
63+
return operation.apply(requestWithDecoratedEndpointProvider(request, regionSupplier(bucketName)));
64+
}
65+
return operation.apply(request);
66+
} catch (S3Exception exception) {
67+
if (exception.statusCode() == REDIRECT_STATUS_CODE) {
68+
updateCacheFromRedirectException(exception, bucketName);
69+
return operation.apply(requestWithDecoratedEndpointProvider(request, regionSupplier(bucketName)));
4670
}
71+
throw exception;
4772
}
73+
}
74+
75+
private String updateCacheFromRedirectException(S3Exception exception, String bucketName) {
76+
Optional<String> regionStr = getBucketRegionFromException(exception);
77+
// If redirected, clear previous values due to region change. bucketToRegionCache.remove(bucketName);
78+
regionStr.ifPresent(region -> bucketToRegionCache.put(bucketName, Region.of(region)));
79+
return regionStr.orElse(null);
80+
}
4881

49-
return operation.apply(request);
82+
private Supplier<Region> regionSupplier(String bucket) {
83+
return () -> bucketToRegionCache.computeIfAbsent(bucket, this::fetchBucketRegion);
5084
}
5185

52-
private void handleOperationFailure(Throwable t, String bucket) {
53-
//TODO: handle failure case
86+
private Region fetchBucketRegion(String bucketName) {
87+
try {
88+
((S3Client) delegate()).headBucket(HeadBucketRequest.builder().bucket(bucketName).build());
89+
} catch (S3Exception exception) {
90+
if (exception.statusCode() == REDIRECT_STATUS_CODE) {
91+
return Region.of(getBucketRegionFromException(exception).orElseThrow(() -> exception));
92+
}
93+
throw exception;
94+
}
95+
return null;
5496
}
5597

5698
@SuppressWarnings("unchecked")
57-
private <T extends S3Request> T requestWithDecoratedEndpointProvider(T request, String bucket) {
99+
private <T extends S3Request> T requestWithDecoratedEndpointProvider(T request, Supplier<Region> regionSupplier) {
58100
return (T) request.toBuilder()
59-
.overrideConfiguration(getOrCreateConfigWithEndpointProvider(request, bucket))
101+
.overrideConfiguration(getOrCreateConfigWithEndpointProvider(request, regionSupplier))
60102
.build();
61103
}
62104

63-
//TODO: optimize shared sync/async code
64-
private AwsRequestOverrideConfiguration getOrCreateConfigWithEndpointProvider(S3Request request, String bucket) {
105+
private AwsRequestOverrideConfiguration getOrCreateConfigWithEndpointProvider(S3Request request,
106+
Supplier<Region> regionSupplier) {
65107
AwsRequestOverrideConfiguration requestOverrideConfig =
66108
request.overrideConfiguration().orElseGet(() -> AwsRequestOverrideConfiguration.builder().build());
67109

68110
S3EndpointProvider delegateEndpointProvider = (S3EndpointProvider)
69111
requestOverrideConfig.endpointProvider().orElseGet(() -> serviceClientConfiguration().endpointProvider().get());
70112

71113
return requestOverrideConfig.toBuilder()
72-
.endpointProvider(BucketEndpointProvider.create(delegateEndpointProvider, bucket))
114+
.endpointProvider(BucketEndpointProvider.create(delegateEndpointProvider, regionSupplier))
73115
.build();
74116
}
75117

76-
static final class BucketEndpointProvider implements S3EndpointProvider {
77-
private final S3EndpointProvider delegate;
78-
private final String bucket;
79-
80-
private BucketEndpointProvider(S3EndpointProvider delegate, String bucket) {
81-
this.delegate = delegate;
82-
this.bucket = bucket;
83-
}
84-
85-
public static BucketEndpointProvider create(S3EndpointProvider delegate, String bucket) {
86-
return new BucketEndpointProvider(delegate, bucket);
87-
}
88-
89-
@Override
90-
public CompletableFuture<Endpoint> resolveEndpoint(S3EndpointParams endpointParams) {
91-
return delegate.resolveEndpoint(endpointParams);
92-
}
118+
private Optional<String> getBucketRegionFromException(S3Exception exception) {
119+
return exception.awsErrorDetails()
120+
.sdkHttpResponse()
121+
.firstMatchingHeader(AMZ_BUCKET_REGION_HEADER);
93122
}
123+
94124
}

‎services/s3/src/main/java/software/amazon/awssdk/services/s3/internal/crossregion/endpointprovider/BucketEndpointProvider.java

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.services.s3.internal.crossregion.endpointprovider;
17+
18+
import java.util.concurrent.CompletableFuture;
19+
import java.util.function.Supplier;
20+
import software.amazon.awssdk.annotations.SdkInternalApi;
21+
import software.amazon.awssdk.endpoints.Endpoint;
22+
import software.amazon.awssdk.regions.Region;
23+
import software.amazon.awssdk.services.s3.endpoints.S3EndpointParams;
24+
import software.amazon.awssdk.services.s3.endpoints.S3EndpointProvider;
25+
26+
/**
27+
* Decorator S3EndpointProvider which updates the region with the one that is supplied during its instantiation.
28+
*/
29+
@SdkInternalApi
30+
public class BucketEndpointProvider implements S3EndpointProvider {
31+
private final S3EndpointProvider delegateEndPointProvider;
32+
private final Supplier<Region> regionSupplier;
33+
34+
private BucketEndpointProvider(S3EndpointProvider delegateEndPointProvider, Supplier<Region> regionSupplier) {
35+
this.delegateEndPointProvider = delegateEndPointProvider;
36+
this.regionSupplier = regionSupplier;
37+
}
38+
39+
public static BucketEndpointProvider create(S3EndpointProvider delegateEndPointProvider, Supplier<Region> regionSupplier) {
40+
return new BucketEndpointProvider(delegateEndPointProvider, regionSupplier);
41+
}
42+
43+
@Override
44+
public CompletableFuture<Endpoint> resolveEndpoint(S3EndpointParams endpointParams) {
45+
Region crossRegion = regionSupplier.get();
46+
return delegateEndPointProvider.resolveEndpoint(
47+
crossRegion != null ? endpointParams.toBuilder().region(crossRegion).build() : endpointParams);
48+
}
49+
}
50+

‎services/s3/src/test/java/software/amazon/awssdk/services/s3/internal/crossregion/S3CrossRegionAsyncClientTest.java

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,7 @@
3333
import software.amazon.awssdk.services.s3.S3AsyncClient;
3434
import software.amazon.awssdk.services.s3.S3Client;
3535
import software.amazon.awssdk.services.s3.endpoints.internal.DefaultS3EndpointProvider;
36+
import software.amazon.awssdk.services.s3.internal.crossregion.endpointprovider.BucketEndpointProvider;
3637
import software.amazon.awssdk.services.s3.model.GetObjectRequest;
3738
import software.amazon.awssdk.services.s3.model.ListObjectsV2Response;
3839
import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Iterable;
@@ -71,7 +72,7 @@ public void before() {
7172
public void standardOp_crossRegionClient_noOverrideConfig_SuccessfullyIntercepts() {
7273
S3AsyncClient crossRegionClient = new S3CrossRegionAsyncClient(s3Client);
7374
crossRegionClient.getObject(r -> r.bucket(BUCKET).key(KEY), AsyncResponseTransformer.toBytes());
74-
assertThat(captureInterceptor.endpointProvider).isInstanceOf(S3CrossRegionAsyncClient.BucketEndpointProvider.class);
75+
assertThat(captureInterceptor.endpointProvider).isInstanceOf(BucketEndpointProvider.class);
7576
}
7677

7778
@Test
@@ -83,7 +84,7 @@ public void standardOp_crossRegionClient_existingOverrideConfig_SuccessfullyInte
8384
.overrideConfiguration(o -> o.putHeader("someheader", "somevalue"))
8485
.build();
8586
crossRegionClient.getObject(request, AsyncResponseTransformer.toBytes());
86-
assertThat(captureInterceptor.endpointProvider).isInstanceOf(S3CrossRegionAsyncClient.BucketEndpointProvider.class);
87+
assertThat(captureInterceptor.endpointProvider).isInstanceOf(BucketEndpointProvider.class);
8788
assertThat(mockAsyncHttpClient.getLastRequest().headers().get("someheader")).isNotNull();
8889
}
8990

0 commit comments

Comments
 (0)