Skip to content

Commit 12b080a

Browse files
joviegasjoviegas
committed
Handled Anna-karin's comment and also rebased from project branch
2 parents d086ece + f7e4188 commit 12b080a

File tree

6 files changed

+113
-51
lines changed

6 files changed

+113
-51
lines changed

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

Lines changed: 43 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -18,13 +18,16 @@
1818
import static software.amazon.awssdk.services.s3.internal.crossregion.utils.CrossRegionUtils.getBucketRegionFromException;
1919
import static software.amazon.awssdk.services.s3.internal.crossregion.utils.CrossRegionUtils.isS3RedirectException;
2020
import static software.amazon.awssdk.services.s3.internal.crossregion.utils.CrossRegionUtils.requestWithDecoratedEndpointProvider;
21+
import static software.amazon.awssdk.services.s3.internal.crossregion.utils.CrossRegionUtils.updateUserAgentInConfig;
2122

2223
import java.util.Map;
2324
import java.util.Optional;
2425
import java.util.concurrent.CompletableFuture;
2526
import java.util.concurrent.ConcurrentHashMap;
27+
import java.util.function.BiConsumer;
2628
import java.util.function.Function;
2729
import software.amazon.awssdk.annotations.SdkInternalApi;
30+
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
2831
import software.amazon.awssdk.regions.Region;
2932
import software.amazon.awssdk.services.s3.DelegatingS3AsyncClient;
3033
import software.amazon.awssdk.services.s3.S3AsyncClient;
@@ -47,44 +50,60 @@ protected <T extends S3Request, ReturnT> CompletableFuture<ReturnT> invokeOperat
4750

4851
Optional<String> bucket = request.getValueForField("Bucket", String.class);
4952

53+
AwsRequestOverrideConfiguration overrideConfiguration = updateUserAgentInConfig(request);
54+
T userAgentUpdatedRequest = (T) request.toBuilder().overrideConfiguration(overrideConfiguration).build();
55+
5056
if (!bucket.isPresent()) {
51-
return operation.apply(request);
57+
return operation.apply(userAgentUpdatedRequest);
5258
}
5359
String bucketName = bucket.get();
5460

55-
if (bucketToRegionCache.containsKey(bucketName)) {
56-
return operation.apply(requestWithDecoratedEndpointProvider(request,
57-
() -> bucketToRegionCache.get(bucketName),
58-
serviceClientConfiguration().endpointProvider().get()));
59-
}
60-
6161
CompletableFuture<ReturnT> returnFuture = new CompletableFuture<>();
62-
operation.apply(request)
63-
.whenComplete((r, t) -> {
64-
if (t != null) {
65-
if (isS3RedirectException(t)) {
66-
bucketToRegionCache.remove(bucketName);
67-
requestWithCrossRegion(request, operation, bucketName, returnFuture, t);
68-
return;
69-
}
70-
returnFuture.completeExceptionally(t);
71-
return;
72-
}
73-
returnFuture.complete(r);
74-
});
62+
CompletableFuture<ReturnT> apiOperationFuture = bucketToRegionCache.containsKey(bucketName) ?
63+
operation.apply(
64+
requestWithDecoratedEndpointProvider(
65+
userAgentUpdatedRequest,
66+
() -> bucketToRegionCache.get(bucketName),
67+
serviceClientConfiguration().endpointProvider().get()
68+
)
69+
) :
70+
operation.apply(userAgentUpdatedRequest);
71+
72+
apiOperationFuture.whenComplete(redirectToCrossRegionIfRedirectException(operation,
73+
userAgentUpdatedRequest,
74+
bucketName,
75+
returnFuture));
7576
return returnFuture;
7677
}
7778

79+
private <T extends S3Request, ReturnT> BiConsumer<ReturnT, Throwable> redirectToCrossRegionIfRedirectException(
80+
Function<T, CompletableFuture<ReturnT>> operation,
81+
T userAgentUpdatedRequest, String bucketName,
82+
CompletableFuture<ReturnT> returnFuture) {
83+
84+
return (response, throwable) -> {
85+
if (throwable != null) {
86+
if (isS3RedirectException(throwable)) {
87+
bucketToRegionCache.remove(bucketName);
88+
requestWithCrossRegion(userAgentUpdatedRequest, operation, bucketName, returnFuture, throwable);
89+
} else {
90+
returnFuture.completeExceptionally(throwable);
91+
}
92+
} else {
93+
returnFuture.complete(response);
94+
}
95+
};
96+
}
97+
7898
private <T extends S3Request, ReturnT> void requestWithCrossRegion(T request,
7999
Function<T, CompletableFuture<ReturnT>> operation,
80100
String bucketName,
81101
CompletableFuture<ReturnT> returnFuture,
82-
Throwable t) {
102+
Throwable throwable) {
83103

84-
Optional<String> bucketRegionFromException = getBucketRegionFromException((S3Exception) t.getCause());
104+
Optional<String> bucketRegionFromException = getBucketRegionFromException((S3Exception) throwable.getCause());
85105
if (bucketRegionFromException.isPresent()) {
86-
sendRequestWithRightRegion(request, operation, bucketName, returnFuture,
87-
bucketRegionFromException);
106+
sendRequestWithRightRegion(request, operation, bucketName, returnFuture, bucketRegionFromException);
88107
} else {
89108
fetchRegionAndSendRequest(request, operation, bucketName, returnFuture);
90109
}
@@ -94,28 +113,21 @@ private <T extends S3Request, ReturnT> void fetchRegionAndSendRequest(T request,
94113
Function<T, CompletableFuture<ReturnT>> operation,
95114
String bucketName,
96115
CompletableFuture<ReturnT> returnFuture) {
97-
98116
// // TODO: will fix the casts with separate PR
99117
((S3AsyncClient) delegate()).headBucket(b -> b.bucket(bucketName)).whenComplete((response,
100118
throwable) -> {
101119
if (throwable != null) {
102120
if (isS3RedirectException(throwable)) {
103121
bucketToRegionCache.remove(bucketName);
104122
Optional<String> bucketRegion = getBucketRegionFromException((S3Exception) throwable.getCause());
105-
106123
if (bucketRegion.isPresent()) {
107-
bucketToRegionCache.put(bucketName, Region.of(bucketRegion.get()));
108124
sendRequestWithRightRegion(request, operation, bucketName, returnFuture, bucketRegion);
109125
} else {
110126
returnFuture.completeExceptionally(throwable);
111127
}
112128
} else {
113129
returnFuture.completeExceptionally(throwable);
114130
}
115-
} else {
116-
CompletableFuture<ReturnT> newFuture = operation.apply(request);
117-
CompletableFutureUtils.forwardResultTo(newFuture, returnFuture);
118-
CompletableFutureUtils.forwardExceptionTo(returnFuture, newFuture);
119131
}
120132
});
121133
}
@@ -127,19 +139,11 @@ private <T extends S3Request, ReturnT> void sendRequestWithRightRegion(T request
127139
Optional<String> bucketRegionFromException) {
128140
String region = bucketRegionFromException.get();
129141
bucketToRegionCache.put(bucketName, Region.of(region));
130-
doSendRequestWithRightRegion(request, operation, returnFuture, region);
131-
}
132-
133-
private <T extends S3Request, ReturnT> void doSendRequestWithRightRegion(T request,
134-
Function<T, CompletableFuture<ReturnT>> operation,
135-
CompletableFuture<ReturnT> returnFuture,
136-
String region) {
137142
CompletableFuture<ReturnT> newFuture = operation.apply(
138143
requestWithDecoratedEndpointProvider(request,
139144
() -> Region.of(region),
140145
serviceClientConfiguration().endpointProvider().get()));
141146
CompletableFutureUtils.forwardResultTo(newFuture, returnFuture);
142-
// forward exception
143147
CompletableFutureUtils.forwardExceptionTo(returnFuture, newFuture);
144148
}
145149
}

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

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -18,12 +18,14 @@
1818
import static software.amazon.awssdk.services.s3.internal.crossregion.utils.CrossRegionUtils.getBucketRegionFromException;
1919
import static software.amazon.awssdk.services.s3.internal.crossregion.utils.CrossRegionUtils.isS3RedirectException;
2020
import static software.amazon.awssdk.services.s3.internal.crossregion.utils.CrossRegionUtils.requestWithDecoratedEndpointProvider;
21+
import static software.amazon.awssdk.services.s3.internal.crossregion.utils.CrossRegionUtils.updateUserAgentInConfig;
2122

2223
import java.util.Map;
2324
import java.util.Optional;
2425
import java.util.concurrent.ConcurrentHashMap;
2526
import java.util.function.Function;
2627
import software.amazon.awssdk.annotations.SdkInternalApi;
28+
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
2729
import software.amazon.awssdk.regions.Region;
2830
import software.amazon.awssdk.services.s3.DelegatingS3Client;
2931
import software.amazon.awssdk.services.s3.S3Client;
@@ -52,24 +54,29 @@ private static <T extends S3Request> Optional<String> bucketNameFromRequest(T re
5254
protected <T extends S3Request, ReturnT> ReturnT invokeOperation(T request, Function<T, ReturnT> operation) {
5355

5456
Optional<String> bucketRequest = bucketNameFromRequest(request);
57+
58+
AwsRequestOverrideConfiguration overrideConfiguration = updateUserAgentInConfig(request);
59+
T userAgentUpdatedRequest = (T) request.toBuilder().overrideConfiguration(overrideConfiguration).build();
60+
61+
5562
if (!bucketRequest.isPresent()) {
56-
return operation.apply(request);
63+
return operation.apply(userAgentUpdatedRequest);
5764
}
5865
String bucketName = bucketRequest.get();
5966
try {
6067
if (bucketToRegionCache.containsKey(bucketName)) {
6168
return operation.apply(
62-
requestWithDecoratedEndpointProvider(request,
69+
requestWithDecoratedEndpointProvider(userAgentUpdatedRequest,
6370
() -> bucketToRegionCache.get(bucketName),
6471
serviceClientConfiguration().endpointProvider().get()));
6572
}
66-
return operation.apply(request);
73+
return operation.apply(userAgentUpdatedRequest);
6774
} catch (S3Exception exception) {
6875
if (isS3RedirectException(exception)) {
6976
updateCacheFromRedirectException(exception, bucketName);
7077
return operation.apply(
7178
requestWithDecoratedEndpointProvider(
72-
request,
79+
userAgentUpdatedRequest,
7380
() -> bucketToRegionCache.computeIfAbsent(bucketName, this::fetchBucketRegion),
7481
serviceClientConfiguration().endpointProvider().get()));
7582
}

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

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,9 +18,11 @@
1818

1919
import java.util.Optional;
2020
import java.util.concurrent.CompletionException;
21+
import java.util.function.Consumer;
2122
import java.util.function.Supplier;
2223
import software.amazon.awssdk.annotations.SdkInternalApi;
2324
import software.amazon.awssdk.awscore.AwsRequestOverrideConfiguration;
25+
import software.amazon.awssdk.core.ApiName;
2426
import software.amazon.awssdk.endpoints.EndpointProvider;
2527
import software.amazon.awssdk.regions.Region;
2628
import software.amazon.awssdk.services.s3.endpoints.S3EndpointProvider;
@@ -32,6 +34,9 @@
3234
public final class CrossRegionUtils {
3335
public static final int REDIRECT_STATUS_CODE = 301;
3436
public static final String AMZ_BUCKET_REGION_HEADER = "x-amz-bucket-region";
37+
private static final ApiName API_NAME = ApiName.builder().version("cross-region").name("hll").build();
38+
private static final Consumer<AwsRequestOverrideConfiguration.Builder> USER_AGENT_APPLIER = b -> b.addApiName(API_NAME);
39+
3540

3641
private CrossRegionUtils() {
3742
}
@@ -49,6 +54,7 @@ public static boolean isS3RedirectException(Throwable exception) {
4954
}
5055

5156

57+
@SuppressWarnings("unchecked")
5258
public static <T extends S3Request> T requestWithDecoratedEndpointProvider(T request, Supplier<Region> regionSupplier,
5359
EndpointProvider clientEndpointProvider) {
5460
AwsRequestOverrideConfiguration requestOverrideConfig =
@@ -64,4 +70,15 @@ public static <T extends S3Request> T requestWithDecoratedEndpointProvider(T req
6470
.build())
6571
.build();
6672
}
73+
74+
public static <T extends S3Request> AwsRequestOverrideConfiguration updateUserAgentInConfig(T request) {
75+
AwsRequestOverrideConfiguration overrideConfiguration =
76+
request.overrideConfiguration().map(c -> c.toBuilder()
77+
.applyMutation(USER_AGENT_APPLIER)
78+
.build())
79+
.orElse(AwsRequestOverrideConfiguration.builder()
80+
.applyMutation(USER_AGENT_APPLIER)
81+
.build());
82+
return overrideConfiguration;
83+
}
6784
}

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -248,6 +248,22 @@ void crossRegionClient_cancelsTheThread_when_futureIsCancelled(){
248248
assertThat(completableFuture.isCancelled()).isTrue();
249249
}
250250

251+
@Test
252+
void standardOp_crossRegionClient_containUserAgent() {
253+
mockAsyncHttpClient.stubResponses(successHttpResponse());
254+
S3AsyncClient crossRegionClient = clientBuilder().serviceConfiguration(c -> c.crossRegionAccessEnabled(true)).build();
255+
crossRegionClient.getObject(r -> r.bucket(BUCKET).key(KEY), AsyncResponseTransformer.toBytes()).join();
256+
assertThat(mockAsyncHttpClient.getLastRequest().firstMatchingHeader("User-Agent").get()).contains("hll/cross-region");
257+
}
258+
259+
@Test
260+
void standardOp_simpleClient_doesNotContainCrossRegionUserAgent() {
261+
mockAsyncHttpClient.stubResponses(successHttpResponse());
262+
S3AsyncClient crossRegionClient = clientBuilder().serviceConfiguration(c -> c.crossRegionAccessEnabled(false)).build();
263+
crossRegionClient.getObject(r -> r.bucket(BUCKET).key(KEY), AsyncResponseTransformer.toBytes()).join();
264+
assertThat(mockAsyncHttpClient.getLastRequest().firstMatchingHeader("User-Agent").get()).doesNotContain("hll/cross-region");
265+
}
266+
251267
private S3AsyncClientBuilder clientBuilder() {
252268
return S3AsyncClient.builder()
253269
.httpClient(mockAsyncHttpClient)

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

Lines changed: 21 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -150,11 +150,11 @@ void crossRegionClient_createdWithWrapping_SuccessfullyIntercepts(Consumer<MockS
150150
@ParameterizedTest
151151
@MethodSource("stubOverriddenEndpointProviderResponses")
152152
void standardOp_crossRegionClient_takesCustomEndpointProviderInRequest(Consumer<MockSyncHttpClient> stubConsumer,
153-
Class<?> endpointProviderType,
153+
Class<?> endpointProviderType,
154154
String region) {
155155
stubConsumer.accept(mockSyncHttpClient);
156156
S3Client crossRegionClient = clientBuilder().serviceConfiguration(c -> c.crossRegionAccessEnabled(true))
157-
.endpointProvider(new TestEndpointProvider())
157+
.endpointProvider(new TestEndpointProvider())
158158
.build();
159159
GetObjectRequest request = GetObjectRequest.builder()
160160
.bucket(BUCKET)
@@ -172,7 +172,7 @@ void standardOp_crossRegionClient_takesCustomEndpointProviderInRequest(Consumer<
172172
@ParameterizedTest
173173
@MethodSource("stubOverriddenEndpointProviderResponses")
174174
void standardOp_crossRegionClient_takesCustomEndpointProviderInClient(Consumer<MockSyncHttpClient> stubConsumer,
175-
Class<?> endpointProviderType,
175+
Class<?> endpointProviderType,
176176
String region) {
177177
stubConsumer.accept(mockSyncHttpClient);
178178
S3Client crossRegionClient = clientBuilder().serviceConfiguration(c -> c.crossRegionAccessEnabled(true))
@@ -277,6 +277,24 @@ void crossRegionClient_CallsHeadObjectWithNoRegion_shouldTerminateHeadBucketAPI(
277277
SdkHttpMethod.HEAD));
278278
}
279279

280+
281+
@Test
282+
void standardOp_crossRegionClient_containUserAgent() {
283+
mockSyncHttpClient.stubResponses(successHttpResponse());
284+
S3Client crossRegionClient = clientBuilder().serviceConfiguration(c -> c.crossRegionAccessEnabled(true)).build();
285+
crossRegionClient.getObject(r -> r.bucket(BUCKET).key(KEY));
286+
assertThat(mockSyncHttpClient.getLastRequest().firstMatchingHeader("User-Agent").get()).contains("hll/cross-region");
287+
}
288+
289+
@Test
290+
void standardOp_simpleClient_doesNotContainCrossRegionUserAgent() {
291+
mockSyncHttpClient.stubResponses(successHttpResponse());
292+
S3Client crossRegionClient = clientBuilder().serviceConfiguration(c -> c.crossRegionAccessEnabled(false)).build();
293+
crossRegionClient.getObject(r -> r.bucket(BUCKET).key(KEY));
294+
assertThat(mockSyncHttpClient.getLastRequest().firstMatchingHeader("User-Agent").get())
295+
.doesNotContain("hll/cross-region");
296+
}
297+
280298
private S3ClientBuilder clientBuilder() {
281299
return S3Client.builder()
282300
.httpClient(mockSyncHttpClient)

0 commit comments

Comments
 (0)