Skip to content

Commit b3a72e4

Browse files
sandeepvinayaksandeepvinayak
authored
blobstore: fix the retry config for the bucketclient (#135)
1 parent 575f364 commit b3a72e4

File tree

10 files changed

+256
-46
lines changed

10 files changed

+256
-46
lines changed

‎blob/blob-aws/src/main/java/com/salesforce/multicloudj/blob/aws/AwsBlobClient.java‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
import software.amazon.awssdk.services.s3.S3Client;
1818
import software.amazon.awssdk.services.s3.S3ClientBuilder;
1919

20+
import java.time.Duration;
2021
import java.util.stream.Collectors;
2122

2223
/**
@@ -88,6 +89,20 @@ private static S3Client buildS3Client(Builder builder) {
8889
.proxyConfiguration(proxyConfig)
8990
.build());
9091
}
92+
if (builder.getRetryConfig() != null) {
93+
// Create a temporary transformer instance for retry strategy conversion
94+
AwsTransformer transformer = new AwsTransformer(null);
95+
b.overrideConfiguration(config -> {
96+
config.retryStrategy(transformer.toAwsRetryStrategy(builder.getRetryConfig()));
97+
// Set API call timeouts if provided
98+
if (builder.getRetryConfig().getAttemptTimeout() != null) {
99+
config.apiCallAttemptTimeout(Duration.ofMillis(builder.getRetryConfig().getAttemptTimeout()));
100+
}
101+
if (builder.getRetryConfig().getTotalTimeout() != null) {
102+
config.apiCallTimeout(Duration.ofMillis(builder.getRetryConfig().getTotalTimeout()));
103+
}
104+
});
105+
}
91106

92107
return b.build();
93108
}

‎blob/blob-aws/src/test/java/com/salesforce/multicloudj/blob/aws/AwsBlobClientTest.java‎

Lines changed: 99 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33

44
import com.salesforce.multicloudj.common.exceptions.InvalidArgumentException;
55
import com.salesforce.multicloudj.common.exceptions.UnknownException;
6+
import com.salesforce.multicloudj.common.retries.RetryConfig;
67
import com.salesforce.multicloudj.sts.model.CredentialsOverrider;
78
import com.salesforce.multicloudj.sts.model.CredentialsType;
89
import com.salesforce.multicloudj.sts.model.StsCredentials;
@@ -12,6 +13,7 @@
1213
import org.mockito.MockedStatic;
1314
import software.amazon.awssdk.awscore.exception.AwsErrorDetails;
1415
import software.amazon.awssdk.awscore.exception.AwsServiceException;
16+
import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
1517
import software.amazon.awssdk.core.exception.SdkClientException;
1618
import software.amazon.awssdk.services.s3.S3Client;
1719
import software.amazon.awssdk.services.s3.S3ClientBuilder;
@@ -20,11 +22,14 @@
2022

2123
import java.io.IOException;
2224
import java.net.URI;
25+
import java.time.Duration;
2326
import java.util.List;
27+
import java.util.function.Consumer;
2428

2529
import static org.junit.jupiter.api.Assertions.assertEquals;
2630
import static org.junit.jupiter.api.Assertions.assertNotNull;
2731
import static org.mockito.ArgumentMatchers.any;
32+
import static org.mockito.Mockito.doAnswer;
2833
import static org.mockito.Mockito.mock;
2934
import static org.mockito.Mockito.mockStatic;
3035
import static org.mockito.Mockito.verify;
@@ -38,9 +43,20 @@ public class AwsBlobClientTest {
3843

3944
@BeforeEach
4045
void setup() {
41-
var mockBuilder = mock(S3ClientBuilder.class);
46+
S3ClientBuilder mockBuilder = mock(S3ClientBuilder.class);
4247
when(mockBuilder.region(any())).thenReturn(mockBuilder);
4348

49+
// Execute the consumer lambda to cover retry config lines 96-104
50+
doAnswer(invocation -> {
51+
Consumer<ClientOverrideConfiguration.Builder> consumer = invocation.getArgument(0);
52+
ClientOverrideConfiguration.Builder configBuilder = mock(ClientOverrideConfiguration.Builder.class);
53+
when(configBuilder.retryStrategy(any(software.amazon.awssdk.retries.api.RetryStrategy.class))).thenReturn(configBuilder);
54+
when(configBuilder.apiCallAttemptTimeout(any(Duration.class))).thenReturn(configBuilder);
55+
when(configBuilder.apiCallTimeout(any(Duration.class))).thenReturn(configBuilder);
56+
consumer.accept(configBuilder);
57+
return mockBuilder;
58+
}).when(mockBuilder).overrideConfiguration(any(Consumer.class));
59+
4460
s3Client = mockStatic(S3Client.class);
4561
s3Client.when(S3Client::builder).thenReturn(mockBuilder);
4662

@@ -122,4 +138,86 @@ void testExceptionHandling() {
122138
cls = aws.getException(new IOException("Channel is closed"));
123139
assertEquals(cls, UnknownException.class);
124140
}
141+
142+
@Test
143+
void testBuildS3ClientWithExponentialRetryConfig() {
144+
// Test with exponential retry config including timeouts
145+
RetryConfig exponentialConfig = RetryConfig.builder()
146+
.mode(RetryConfig.Mode.EXPONENTIAL)
147+
.maxAttempts(5)
148+
.initialDelayMillis(100L)
149+
.multiplier(2.0)
150+
.maxDelayMillis(5000L)
151+
.attemptTimeout(30000L)
152+
.totalTimeout(120000L)
153+
.build();
154+
155+
var client = new AwsBlobClient.Builder()
156+
.withRegion("us-east-2")
157+
.withRetryConfig(exponentialConfig)
158+
.build();
159+
160+
assertNotNull(client);
161+
assertEquals("aws", client.getProviderId());
162+
}
163+
164+
@Test
165+
void testBuildS3ClientWithFixedRetryConfig() {
166+
// Test with fixed retry config
167+
RetryConfig fixedConfig = RetryConfig.builder()
168+
.mode(RetryConfig.Mode.FIXED)
169+
.maxAttempts(3)
170+
.fixedDelayMillis(500L)
171+
.build();
172+
173+
var client = new AwsBlobClient.Builder()
174+
.withRegion("us-east-2")
175+
.withRetryConfig(fixedConfig)
176+
.build();
177+
178+
assertNotNull(client);
179+
assertEquals("aws", client.getProviderId());
180+
}
181+
182+
@Test
183+
void testBuildS3ClientWithRetryConfigWithAttemptTimeout() {
184+
// Test with attempt timeout only
185+
RetryConfig config = RetryConfig.builder()
186+
.mode(RetryConfig.Mode.EXPONENTIAL)
187+
.maxAttempts(3)
188+
.initialDelayMillis(100L)
189+
.multiplier(2.0)
190+
.maxDelayMillis(2000L)
191+
.attemptTimeout(10000L)
192+
.build();
193+
194+
var client = new AwsBlobClient.Builder()
195+
.withRegion("us-east-2")
196+
.withRetryConfig(config)
197+
.build();
198+
199+
assertNotNull(client);
200+
assertEquals("aws", client.getProviderId());
201+
}
202+
203+
@Test
204+
void testBuildS3ClientWithRetryConfigWithTotalTimeout() {
205+
// Test with total timeout only
206+
RetryConfig config = RetryConfig.builder()
207+
.mode(RetryConfig.Mode.EXPONENTIAL)
208+
.maxAttempts(3)
209+
.initialDelayMillis(100L)
210+
.multiplier(2.0)
211+
.maxDelayMillis(2000L)
212+
.totalTimeout(60000L)
213+
.build();
214+
215+
var client = new AwsBlobClient.Builder()
216+
.withRegion("us-east-2")
217+
.withRetryConfig(config)
218+
.build();
219+
220+
assertNotNull(client);
221+
assertEquals("aws", client.getProviderId());
222+
}
125223
}

‎blob/blob-client/src/main/java/com/salesforce/multicloudj/blob/async/client/AsyncBucketClient.java‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import com.salesforce.multicloudj.blob.driver.UploadResponse;
2929
import com.salesforce.multicloudj.common.exceptions.ExceptionHandler;
3030
import com.salesforce.multicloudj.common.exceptions.SubstrateSdkException;
31+
import com.salesforce.multicloudj.common.retries.RetryConfig;
3132
import com.salesforce.multicloudj.sts.model.CredentialsOverrider;
3233

3334
import java.io.File;
@@ -563,6 +564,17 @@ public Builder withMaxNativeMemoryLimitInBytes(Long maxNativeMemoryLimitInBytes)
563564
return this;
564565
}
565566

567+
/**
568+
* Method to supply retry configuration
569+
* @param retryConfig The retry configuration to use for retrying failed requests
570+
* @return An instance of self
571+
*/
572+
@Override
573+
public Builder withRetryConfig(RetryConfig retryConfig) {
574+
super.withRetryConfig(retryConfig);
575+
return this;
576+
}
577+
566578
/**
567579
* {@inheritDoc}
568580
*/

‎blob/blob-client/src/main/java/com/salesforce/multicloudj/blob/client/BlobClient.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import com.salesforce.multicloudj.blob.driver.ListBucketsResponse;
55
import com.salesforce.multicloudj.common.exceptions.ExceptionHandler;
66
import com.salesforce.multicloudj.common.exceptions.SubstrateSdkException;
7+
import com.salesforce.multicloudj.common.retries.RetryConfig;
78
import com.salesforce.multicloudj.sts.model.CredentialsOverrider;
89

910
import java.net.URI;
@@ -95,6 +96,16 @@ public BlobClientBuilder withCredentialsOverrider(CredentialsOverrider credentia
9596
return this;
9697
}
9798

99+
/**
100+
* Method to supply retry configuration
101+
* @param retryConfig The retry configuration to use for retrying failed requests
102+
* @return An instance of self
103+
*/
104+
public BlobClientBuilder withRetryConfig(RetryConfig retryConfig) {
105+
this.blobClientBuilder.withRetryConfig(retryConfig);
106+
return this;
107+
}
108+
98109
/**
99110
* Builds and returns an instance of BlobClient.
100111
*

‎blob/blob-client/src/main/java/com/salesforce/multicloudj/blob/client/BucketClient.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,7 @@
2222
import com.salesforce.multicloudj.blob.driver.UploadResponse;
2323
import com.salesforce.multicloudj.common.exceptions.ExceptionHandler;
2424
import com.salesforce.multicloudj.common.exceptions.SubstrateSdkException;
25+
import com.salesforce.multicloudj.common.retries.RetryConfig;
2526
import com.salesforce.multicloudj.sts.model.CredentialsOverrider;
2627

2728
import java.io.File;
@@ -555,6 +556,16 @@ public BlobBuilder withCredentialsOverrider(CredentialsOverrider credentialsOver
555556
return this;
556557
}
557558

559+
/**
560+
* Method to supply retry configuration
561+
* @param retryConfig The retry configuration to use for retrying failed requests
562+
* @return An instance of self
563+
*/
564+
public BlobBuilder withRetryConfig(RetryConfig retryConfig) {
565+
this.blobStoreBuilder.withRetryConfig(retryConfig);
566+
return this;
567+
}
568+
558569
/**
559570
* Builds and returns an instance of BucketClient.
560571
* @return An instance of BucketClient.

‎blob/blob-client/src/main/java/com/salesforce/multicloudj/blob/driver/BlobBuilder.java‎

Lines changed: 14 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,10 @@
11
package com.salesforce.multicloudj.blob.driver;
22

33
import com.salesforce.multicloudj.common.provider.SdkProvider;
4+
import com.salesforce.multicloudj.common.retries.RetryConfig;
45
import com.salesforce.multicloudj.common.service.SdkService;
56
import com.salesforce.multicloudj.sts.model.CredentialsOverrider;
7+
import lombok.Getter;
68

79
import java.net.URI;
810
import java.util.Properties;
@@ -16,57 +18,15 @@
1618
*
1719
* @param <T>
1820
*/
21+
@Getter
1922
public abstract class BlobBuilder<T extends SdkService> implements SdkProvider.Builder<T> {
2023
private String providerId;
2124
private String region;
2225
private URI endpoint;
2326
private URI proxyEndpoint;
2427
private CredentialsOverrider credentialsOverrider;
2528
private Properties properties = new Properties();
26-
27-
/**
28-
* Gets the providerId that this builder was built with.
29-
* @return the id of the provider
30-
*/
31-
public String getProviderId() {
32-
return this.providerId;
33-
}
34-
35-
/**
36-
* Gets the region.
37-
* @return The region.
38-
*/
39-
public String getRegion(){
40-
return this.region;
41-
}
42-
43-
/**
44-
* Gets the endpoint.
45-
* @return The endpoint.
46-
*/
47-
public URI getEndpoint(){
48-
return this.endpoint;
49-
}
50-
51-
/**
52-
* Gets the proxy endpoint.
53-
* @return The proxy endpoint.
54-
*/
55-
public URI getProxyEndpoint(){
56-
return this.proxyEndpoint;
57-
}
58-
59-
/**
60-
* Gets the CredentialsOverrider.
61-
* @return The CredentialsOverrider.
62-
*/
63-
public CredentialsOverrider getCredentialsOverrider() {
64-
return this.credentialsOverrider;
65-
}
66-
67-
public Properties getProperties() {
68-
return this.properties;
69-
}
29+
private RetryConfig retryConfig;
7030

7131
public BlobBuilder<T> providerId(String providerId) {
7232
this.providerId = providerId;
@@ -118,4 +78,14 @@ public BlobBuilder<T> withProperties(Properties properties) {
11878
return this;
11979
}
12080

81+
/**
82+
* Method to supply retry configuration
83+
* @param retryConfig The retry configuration to use for retrying failed requests
84+
* @return An instance of self
85+
*/
86+
public BlobBuilder<T> withRetryConfig(RetryConfig retryConfig) {
87+
this.retryConfig = retryConfig;
88+
return this;
89+
}
90+
12191
}

‎blob/blob-client/src/main/java/com/salesforce/multicloudj/blob/driver/BlobClientBuilder.java‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package com.salesforce.multicloudj.blob.driver;
22

3+
import com.salesforce.multicloudj.common.retries.RetryConfig;
34
import com.salesforce.multicloudj.common.service.SdkService;
45
import com.salesforce.multicloudj.sts.model.CredentialsOverrider;
56

@@ -226,6 +227,16 @@ public BlobClientBuilder<C, S> withTransferDirectoryMaxConcurrency(Integer trans
226227
return this;
227228
}
228229

230+
/**
231+
* Method to supply retry configuration
232+
* @param retryConfig The retry configuration to use for retrying failed requests
233+
* @return An instance of self
234+
*/
235+
public BlobClientBuilder<C, S> withRetryConfig(RetryConfig retryConfig) {
236+
this.storeBuilder.withRetryConfig(retryConfig);
237+
return this;
238+
}
239+
229240
/**
230241
* Builds and returns an instance of the target client implementation.
231242
* @return A fully constructed client implementation.

‎blob/blob-client/src/test/java/com/salesforce/multicloudj/blob/async/client/AsyncBucketClientTest.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -27,6 +27,7 @@
2727
import com.salesforce.multicloudj.blob.driver.UploadResponse;
2828
import com.salesforce.multicloudj.common.exceptions.ExceptionHandler;
2929
import com.salesforce.multicloudj.common.exceptions.UnAuthorizedException;
30+
import com.salesforce.multicloudj.common.retries.RetryConfig;
3031
import com.salesforce.multicloudj.sts.model.CredentialsOverrider;
3132
import com.salesforce.multicloudj.sts.model.CredentialsType;
3233
import com.salesforce.multicloudj.sts.model.StsCredentials;
@@ -798,4 +799,31 @@ void testDownloadDirectory_WithNullResponse() throws ExecutionException, Interru
798799
verify(mockBlobStore, times(1)).downloadDirectory(eq(request));
799800
assertNull(actualResponse);
800801
}
802+
803+
@Test
804+
void testAsyncBucketClientBuilderWithRetryConfig() {
805+
RetryConfig retryConfig = RetryConfig.builder()
806+
.maxAttempts(5)
807+
.attemptTimeout(3000L)
808+
.totalTimeout(10000L)
809+
.build();
810+
811+
AsyncBlobStoreProvider.Builder mockBuilder2 = mock(AsyncBlobStoreProvider.Builder.class);
812+
when(mockBuilder2.withBucket(any())).thenReturn(mockBuilder2);
813+
when(mockBuilder2.withRegion(any())).thenReturn(mockBuilder2);
814+
when(mockBuilder2.withRetryConfig(any())).thenReturn(mockBuilder2);
815+
when(mockBuilder2.build()).thenReturn(mockBlobStore);
816+
817+
providerSupplier.when(() -> ProviderSupplier.findAsyncBuilder("test2"))
818+
.thenReturn(mockBuilder2);
819+
820+
AsyncBucketClient testClient = AsyncBucketClient.builder("test2")
821+
.withBucket("test-bucket")
822+
.withRegion("us-east-1")
823+
.withRetryConfig(retryConfig)
824+
.build();
825+
826+
verify(mockBuilder2, times(1)).withRetryConfig(retryConfig);
827+
assertInstanceOf(AsyncBucketClient.class, testClient);
828+
}
801829
}

0 commit comments

Comments
 (0)