Skip to content

Commit e1483c7

Browse files
ashrao94Aiswarya Sadananda Rao
ashrao94
and
Aiswarya Sadananda Rao
authored
Add optional timeout configurations for AWS Lambda plugin (#6413)
* Add optional timeout configurations for AWS Lambda plugin - Add api_call_attempt_timeout configuration for per-attempt timeouts - Make read_timeout optional (only applied when specified) - Add comprehensive unit tests for timeout configurations - Maintain backward compatibility with existing configurations - Follow AWS SDK timeout hierarchy best practices Both timeout parameters are now optional and only configured when explicitly set, allowing users to fine-tune timeout behavior for their Lambda functions. Signed-off-by: Aiswarya Sadananda Rao <aiswarao@amazon.com> * Add license header to ClientOptionsTest.java - Add required OpenSearch Contributors license header - Fixes license header violation in new test file Signed-off-by: Aiswarya Sadananda Rao <aiswarao@amazon.com> * Remove unused Duration import from ClientOptionsTest Signed-off-by: Aiswarya Sadananda Rao <aiswarao@amazon.com> --------- Signed-off-by: Aiswarya Sadananda Rao <aiswarao@amazon.com> Co-authored-by: Aiswarya Sadananda Rao <aiswarao@amazon.com>
1 parent 72ad012 commit e1483c7

4 files changed

Lines changed: 132 additions & 9 deletions

File tree

‎data-prepper-plugins/aws-lambda/src/main/java/org/opensearch/dataprepper/plugins/lambda/common/client/LambdaClientFactory.java‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -27,15 +27,20 @@ public static LambdaAsyncClient createAsyncLambdaClient(
2727
awsCredentialsOptions);
2828
final PluginMetrics awsSdkMetrics = PluginMetrics.fromNames("sdk", "aws");
2929

30+
NettyNioAsyncHttpClient.Builder httpClientBuilder = NettyNioAsyncHttpClient.builder()
31+
.maxConcurrency(clientOptions.getMaxConcurrency())
32+
.connectionTimeout(clientOptions.getConnectionTimeout());
33+
34+
if (clientOptions.getReadTimeout() != null) {
35+
httpClientBuilder.readTimeout(clientOptions.getReadTimeout());
36+
}
37+
3038
return LambdaAsyncClient.builder()
3139
.region(awsAuthenticationOptions.getAwsRegion())
3240
.credentialsProvider(awsCredentialsProvider)
3341
.overrideConfiguration(
3442
createOverrideConfiguration(clientOptions, awsSdkMetrics))
35-
.httpClient(NettyNioAsyncHttpClient.builder()
36-
.maxConcurrency(clientOptions.getMaxConcurrency())
37-
.connectionTimeout(clientOptions.getConnectionTimeout())
38-
.readTimeout(clientOptions.getReadTimeout()).build())
43+
.httpClient(httpClientBuilder.build())
3944
.build();
4045
}
4146

@@ -56,11 +61,16 @@ private static ClientOverrideConfiguration createOverrideConfiguration(
5661
.backoffStrategy(backoffStrategy)
5762
.build();
5863

59-
return ClientOverrideConfiguration.builder()
64+
ClientOverrideConfiguration.Builder configBuilder = ClientOverrideConfiguration.builder()
6065
.retryPolicy(customRetryPolicy)
6166
.addMetricPublisher(new MicrometerMetricPublisher(awsSdkMetrics))
62-
.apiCallTimeout(clientOptions.getApiCallTimeout())
63-
.build();
67+
.apiCallTimeout(clientOptions.getApiCallTimeout());
68+
69+
if (clientOptions.getApiCallAttemptTimeout() != null) {
70+
configBuilder.apiCallAttemptTimeout(clientOptions.getApiCallAttemptTimeout());
71+
}
72+
73+
return configBuilder.build();
6474
}
6575

6676
public static AwsCredentialsOptions convertToCredentialsOptions(

‎data-prepper-plugins/aws-lambda/src/main/java/org/opensearch/dataprepper/plugins/lambda/common/config/ClientOptions.java‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ public class ClientOptions {
1111
public static final int DEFAULT_CONNECTION_RETRIES = 3;
1212
public static final int DEFAULT_MAXIMUM_CONCURRENCY = 200;
1313
public static final Duration DEFAULT_CONNECTION_TIMEOUT = Duration.ofSeconds(60);
14-
public static final Duration DEFAULT_READ_TIMEOUT = Duration.ofSeconds(60);
14+
1515
public static final Duration DEFAULT_API_TIMEOUT = Duration.ofSeconds(60);
1616
public static final Duration DEFAULT_BASE_DELAY = Duration.ofMillis(100);
1717
public static final Duration DEFAULT_MAX_BACKOFF = Duration.ofSeconds(20);
@@ -24,13 +24,17 @@ public class ClientOptions {
2424
@JsonProperty("api_call_timeout")
2525
private Duration apiCallTimeout = DEFAULT_API_TIMEOUT;
2626

27+
@JsonPropertyDescription("api call attempt timeout defines the time sdk waits for a single attempt before timing out")
28+
@JsonProperty("api_call_attempt_timeout")
29+
private Duration apiCallAttemptTimeout;
30+
2731
@JsonPropertyDescription("sdk timeout defines the time sdk maintains the connection to the client before timing out")
2832
@JsonProperty("connection_timeout")
2933
private Duration connectionTimeout = DEFAULT_CONNECTION_TIMEOUT;
3034

3135
@JsonPropertyDescription("read timeout defines the time sdk waits for data to be read from an established connection")
3236
@JsonProperty("read_timeout")
33-
private Duration readTimeout = DEFAULT_READ_TIMEOUT;
37+
private Duration readTimeout;
3438

3539
@JsonPropertyDescription("max concurrency defined from the client side")
3640
@JsonProperty("max_concurrency")

‎data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/common/client/LambdaClientFactoryTest.java‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,8 @@
3333
import static org.mockito.Mockito.mock;
3434
import static org.mockito.Mockito.when;
3535

36+
import java.time.Duration;
37+
3638
@ExtendWith(MockitoExtension.class)
3739
@MockitoSettings(strictness = Strictness.LENIENT)
3840
class LambdaClientFactoryTest {
@@ -86,6 +88,58 @@ void testCreateAsyncLambdaClientOverrideConfiguration() {
8688
assertNotNull(overrideConfig.retryPolicy());
8789
assertNotNull(overrideConfig.metricPublishers());
8890
assertFalse(overrideConfig.metricPublishers().isEmpty());
91+
// apiCallAttemptTimeout should not be set when null
92+
assertFalse(overrideConfig.apiCallAttemptTimeout().isPresent());
93+
}
94+
95+
@Test
96+
void testCreateAsyncLambdaClientWithApiCallAttemptTimeout() {
97+
// Arrange
98+
ClientOptions clientOptions = mock(ClientOptions.class);
99+
when(clientOptions.getMaxConcurrency()).thenReturn(200);
100+
when(clientOptions.getConnectionTimeout()).thenReturn(Duration.ofSeconds(60));
101+
when(clientOptions.getReadTimeout()).thenReturn(Duration.ofSeconds(60));
102+
when(clientOptions.getApiCallTimeout()).thenReturn(Duration.ofSeconds(60));
103+
when(clientOptions.getApiCallAttemptTimeout()).thenReturn(Duration.ofSeconds(30));
104+
when(clientOptions.getMaxConnectionRetries()).thenReturn(3);
105+
when(clientOptions.getBaseDelay()).thenReturn(Duration.ofMillis(100));
106+
when(clientOptions.getMaxBackoff()).thenReturn(Duration.ofSeconds(20));
107+
108+
// Act
109+
LambdaAsyncClient client = LambdaClientFactory.createAsyncLambdaClient(
110+
awsAuthenticationOptions,
111+
awsCredentialsSupplier,
112+
clientOptions
113+
);
114+
115+
// Assert
116+
assertNotNull(client);
117+
ClientOverrideConfiguration overrideConfig = client.serviceClientConfiguration().overrideConfiguration();
118+
assertEquals(Duration.ofSeconds(30), overrideConfig.apiCallAttemptTimeout().get());
119+
}
120+
121+
@Test
122+
void testCreateAsyncLambdaClientWithoutReadTimeout() {
123+
// Arrange
124+
ClientOptions clientOptions = mock(ClientOptions.class);
125+
when(clientOptions.getMaxConcurrency()).thenReturn(200);
126+
when(clientOptions.getConnectionTimeout()).thenReturn(Duration.ofSeconds(60));
127+
when(clientOptions.getReadTimeout()).thenReturn(null); // No read timeout
128+
when(clientOptions.getApiCallTimeout()).thenReturn(Duration.ofSeconds(60));
129+
when(clientOptions.getApiCallAttemptTimeout()).thenReturn(null); // No attempt timeout
130+
when(clientOptions.getMaxConnectionRetries()).thenReturn(3);
131+
when(clientOptions.getBaseDelay()).thenReturn(Duration.ofMillis(100));
132+
when(clientOptions.getMaxBackoff()).thenReturn(Duration.ofSeconds(20));
133+
134+
// Act
135+
LambdaAsyncClient client = LambdaClientFactory.createAsyncLambdaClient(
136+
awsAuthenticationOptions,
137+
awsCredentialsSupplier,
138+
clientOptions
139+
);
140+
141+
// Assert - should not throw exception when readTimeout is null
142+
assertNotNull(client);
89143
}
90144

91145
@Test
@@ -184,4 +238,28 @@ void testRetryConditionFirstFailsAndThenSucceeds() {
184238
assertTrue(successReached, "Should have reached successful completion");
185239
}
186240

241+
@Test
242+
void testClientUsesConfiguredReadTimeout() {
243+
ClientOptions clientOptions = new ClientOptions();
244+
Duration customReadTimeout = Duration.ofSeconds(30);
245+
246+
// Use reflection to set the readTimeout since there's no setter
247+
try {
248+
java.lang.reflect.Field readTimeoutField = ClientOptions.class.getDeclaredField("readTimeout");
249+
readTimeoutField.setAccessible(true);
250+
readTimeoutField.set(clientOptions, customReadTimeout);
251+
} catch (Exception e) {
252+
throw new RuntimeException("Failed to set readTimeout", e);
253+
}
254+
255+
LambdaAsyncClient client = LambdaClientFactory.createAsyncLambdaClient(
256+
awsAuthenticationOptions,
257+
awsCredentialsSupplier,
258+
clientOptions
259+
);
260+
261+
assertNotNull(client);
262+
assertEquals(customReadTimeout, clientOptions.getReadTimeout());
263+
}
264+
187265
}

‎data-prepper-plugins/aws-lambda/src/test/java/org/opensearch/dataprepper/plugins/lambda/common/config/ClientOptionsTest.java‎

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*
5+
* The OpenSearch Contributors require contributions made to
6+
* this file be licensed under the Apache-2.0 license or a
7+
* compatible open source license.
8+
*
9+
*/
10+
11+
package org.opensearch.dataprepper.plugins.lambda.common.config;
12+
13+
import org.junit.jupiter.api.Test;
14+
15+
16+
import static org.junit.jupiter.api.Assertions.assertEquals;
17+
18+
class ClientOptionsTest {
19+
20+
@Test
21+
void testDefaultReadTimeout() {
22+
ClientOptions clientOptions = new ClientOptions();
23+
assertEquals(null, clientOptions.getReadTimeout());
24+
}
25+
26+
@Test
27+
void testDefaultApiCallAttemptTimeout() {
28+
ClientOptions clientOptions = new ClientOptions();
29+
assertEquals(null, clientOptions.getApiCallAttemptTimeout());
30+
}
31+
}

0 commit comments

Comments
 (0)