[aws-lambda] Add retryable_status_codes option to retry transient HTTP status codes

data-prepper-plugins/aws-lambda/README.md

@@ -21,6 +21,7 @@ lambda-pipeline:
             read_timeout: 60s
             api_call_timeout: 60s
             max_retries: 3
+            retryable_status_codes: [500, 502, 503, 504]
         batch:
             key_name: "osi_key"
             threshold:
@@ -31,6 +32,8 @@ lambda-pipeline:
 
 `invocation_type` as request-response is used when the response from aws lambda comes back to dataprepper.
 
+`client.retryable_status_codes` is an optional set of HTTP status codes (for example 500, 502, 503, 504) that the Lambda client retries in addition to the retryable exception types. It is empty by default. Set it only for idempotent functions, since retrying a 5xx may re-run an invocation that already executed on the server.
+
 In batch options, an implicit batch threshold option is that if events size is 3mb, we flush it.
 `payload_model` this is used to define how the payload should be constructed from a dataprepper event by converting it to corresponding json.
 `payload_model` as batch_event is used when the output needs to be formed as a batch of multiple events, and a key(key_name) will be associated with the set of events.
@@ -106,6 +109,7 @@ lambda-pipeline:
             read_timeout: 60s
             api_call_timeout: 60s
             max_retries: 3
+            retryable_status_codes: [500, 502, 503, 504]
         batch:
             key_name: "osi_key"
             threshold:

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

@@ -12,9 +12,14 @@
 import software.amazon.awssdk.core.retry.RetryPolicy;
 import software.amazon.awssdk.core.retry.backoff.BackoffStrategy;
 import software.amazon.awssdk.core.retry.backoff.EqualJitterBackoffStrategy;
+import software.amazon.awssdk.core.retry.conditions.OrRetryCondition;
+import software.amazon.awssdk.core.retry.conditions.RetryCondition;
+import software.amazon.awssdk.core.retry.conditions.RetryOnStatusCodeCondition;
 import software.amazon.awssdk.http.nio.netty.NettyNioAsyncHttpClient;
 import software.amazon.awssdk.services.lambda.LambdaAsyncClient;
 
+import java.util.Set;
+
 public final class LambdaClientFactory {
 
   public static LambdaAsyncClient createAsyncLambdaClient(
@@ -55,8 +60,14 @@ private static ClientOverrideConfiguration createOverrideConfiguration(
             .maxBackoffTime(clientOptions.getMaxBackoff())
             .build();
 
+    final Set<Integer> retryableStatusCodes = clientOptions.getRetryableStatusCodes();
+    final RetryCondition retryCondition = retryableStatusCodes.isEmpty()
+            ? new CustomLambdaRetryCondition()
+            : OrRetryCondition.create(new CustomLambdaRetryCondition(),
+                    RetryOnStatusCodeCondition.create(retryableStatusCodes));
+
     final RetryPolicy customRetryPolicy = RetryPolicy.builder()
-            .retryCondition(new CustomLambdaRetryCondition())
+            .retryCondition(retryCondition)
             .numRetries(clientOptions.getMaxConnectionRetries())
             .backoffStrategy(backoffStrategy)
             .build();

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

@@ -5,6 +5,8 @@
 import lombok.Getter;
 
 import java.time.Duration;
+import java.util.Collections;
+import java.util.Set;
 
 @Getter
 public class ClientOptions {
@@ -48,4 +50,8 @@ public class ClientOptions {
     @JsonProperty("max_backoff")
     private Duration maxBackoff = DEFAULT_MAX_BACKOFF;
 
-}
+    @JsonPropertyDescription("HTTP status codes to retry in addition to retryable exceptions. Set only for idempotent functions")
+    @JsonProperty("retryable_status_codes")
+    private Set<Integer> retryableStatusCodes = Collections.emptySet();
+
+}

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

@@ -18,13 +18,15 @@
 import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider;
 import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration;
 import software.amazon.awssdk.core.retry.RetryPolicyContext;
+import software.amazon.awssdk.core.retry.conditions.RetryCondition;
 import software.amazon.awssdk.regions.Region;
 import software.amazon.awssdk.services.lambda.LambdaAsyncClient;
 import software.amazon.awssdk.services.lambda.model.InvokeRequest;
 import software.amazon.awssdk.services.lambda.model.InvokeResponse;
 import software.amazon.awssdk.services.lambda.model.TooManyRequestsException;
 
 import java.util.HashMap;
+import java.util.Set;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
 
@@ -262,4 +264,71 @@ void testClientUsesConfiguredReadTimeout() {
     assertEquals(customReadTimeout, clientOptions.getReadTimeout());
   }
 
+  @Test
+  void testRetryableStatusCodesThreadsThroughToRetryCondition() {
+    final ClientOptions clientOptions = mock(ClientOptions.class);
+    when(clientOptions.getMaxConcurrency()).thenReturn(200);
+    when(clientOptions.getConnectionTimeout()).thenReturn(Duration.ofSeconds(60));
+    when(clientOptions.getReadTimeout()).thenReturn(null);
+    when(clientOptions.getApiCallTimeout()).thenReturn(Duration.ofSeconds(60));
+    when(clientOptions.getApiCallAttemptTimeout()).thenReturn(null);
+    when(clientOptions.getMaxConnectionRetries()).thenReturn(3);
+    when(clientOptions.getBaseDelay()).thenReturn(Duration.ofMillis(100));
+    when(clientOptions.getMaxBackoff()).thenReturn(Duration.ofSeconds(20));
+    when(clientOptions.getRetryableStatusCodes()).thenReturn(Set.of(502, 504));
+
+    final LambdaAsyncClient client = LambdaClientFactory.createAsyncLambdaClient(
+            awsAuthenticationOptions, awsCredentialsSupplier, clientOptions);
+
+    final RetryCondition retryCondition = client.serviceClientConfiguration()
+            .overrideConfiguration().retryPolicy().get().retryCondition();
+    final RetryPolicyContext gatewayError = RetryPolicyContext.builder()
+            .httpStatusCode(502)
+            .retriesAttempted(0)
+            .build();
+    assertTrue(retryCondition.shouldRetry(gatewayError));
+  }
+
+  @Test
+  void testStatusCodePathInactiveByDefault() {
+    final ClientOptions clientOptions = new ClientOptions();
+
+    final LambdaAsyncClient client = LambdaClientFactory.createAsyncLambdaClient(
+            awsAuthenticationOptions, awsCredentialsSupplier, clientOptions);
+
+    final RetryCondition retryCondition = client.serviceClientConfiguration()
+            .overrideConfiguration().retryPolicy().get().retryCondition();
+    final RetryPolicyContext gatewayError = RetryPolicyContext.builder()
+            .httpStatusCode(502)
+            .retriesAttempted(0)
+            .build();
+    assertFalse(retryCondition.shouldRetry(gatewayError));
+  }
+
+  @Test
+  void testComposedConditionStillRetriesRetryableExceptions() {
+    final ClientOptions clientOptions = mock(ClientOptions.class);
+    when(clientOptions.getMaxConcurrency()).thenReturn(200);
+    when(clientOptions.getConnectionTimeout()).thenReturn(Duration.ofSeconds(60));
+    when(clientOptions.getReadTimeout()).thenReturn(null);
+    when(clientOptions.getApiCallTimeout()).thenReturn(Duration.ofSeconds(60));
+    when(clientOptions.getApiCallAttemptTimeout()).thenReturn(null);
+    when(clientOptions.getMaxConnectionRetries()).thenReturn(3);
+    when(clientOptions.getBaseDelay()).thenReturn(Duration.ofMillis(100));
+    when(clientOptions.getMaxBackoff()).thenReturn(Duration.ofSeconds(20));
+    when(clientOptions.getRetryableStatusCodes()).thenReturn(Set.of(502, 504));
+
+    final LambdaAsyncClient client = LambdaClientFactory.createAsyncLambdaClient(
+            awsAuthenticationOptions, awsCredentialsSupplier, clientOptions);
+
+    final RetryCondition retryCondition = client.serviceClientConfiguration()
+            .overrideConfiguration().retryPolicy().get().retryCondition();
+    // No matching status code on the context: the exception path must still trigger a retry.
+    final RetryPolicyContext throttling = RetryPolicyContext.builder()
+            .exception(TooManyRequestsException.builder().build())
+            .retriesAttempted(0)
+            .build();
+    assertTrue(retryCondition.shouldRetry(throttling));
+  }
+
 }

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

@@ -14,6 +14,7 @@
 
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 class ClientOptionsTest {
 
@@ -28,4 +29,10 @@ void testDefaultApiCallAttemptTimeout() {
         ClientOptions clientOptions = new ClientOptions();
         assertEquals(null, clientOptions.getApiCallAttemptTimeout());
     }
+
+    @Test
+    void testDefaultRetryableStatusCodes() {
+        final ClientOptions clientOptions = new ClientOptions();
+        assertTrue(clientOptions.getRetryableStatusCodes().isEmpty());
+    }
 }