feat(sdk): Classify Durable API errors as retryable or non-retryable (e.g. customer KMS errors)

denzyll · web-flow · commit c1fcea6e0baf · 2026-04-27T15:05:59.000-04:00
feat(sdk): Classify Durable API errors as retryable or non-retryable (e.g. customer KMS errors) The Java SDK currently fails the execution when encountering any error from CheckpointDurableExecution and GetDurableExecutionState API calls, including retryable errors like throttling (429) and service errors (5xx). KMS exceptions (e.g., KMSAccessDeniedException, KMSDisabledException) also arrive as 502 errors from these APIs but, unlike other 5xx errors, would never self-resolve on retry. Add DurableApiErrorClassifier in the util package with a NON_RETRYABLE_ERROR_CODES set and classifyException method. Use it in CheckpointManager to wrap both checkpoint() and getExecutionState() calls. Non-retryable errors (KMS, 4xx non-429) return UnrecoverableDurableExecutionException with retryable=false, terminating the execution immediately with status FAILED. Retryable errors (429, 5xx, invalid checkpoint token) return UnrecoverableDurableExecutionException with retryable=true, integrating with the retryable exception support from #350 so the backend can retry the invocation. The classifier also matches the JS SDK's classifyCheckpointError logic for 4xx and invalid checkpoint token handling. Testing — parameterized unit tests for all four KMS exceptions with realistic error messages, verifying retryable=false. Tests for retryable cases (429, 5xx, invalid checkpoint token, non-KMS 502) verify retryable=true. Additional tests for non-token InvalidParameterValueException, error detail preservation, and two wiring tests verifying CheckpointManager delegates to the classifier for both checkpoint and getExecutionState API paths. By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/CheckpointManager.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/CheckpointManager.java
@@ -14,12 +14,14 @@
 import java.util.function.Consumer;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
+import software.amazon.awssdk.awscore.exception.AwsServiceException;
 import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState;
 import software.amazon.awssdk.services.lambda.model.Operation;
 import software.amazon.awssdk.services.lambda.model.OperationUpdate;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.retry.PollingStrategies;
 import software.amazon.lambda.durable.retry.PollingStrategy;
+import software.amazon.lambda.durable.util.DurableApiErrorClassifier;
 
 /**
  * Package-private checkpoint manager for batching and queueing checkpoint API calls.
@@ -163,14 +165,18 @@ List<Operation> fetchAllPages(CheckpointUpdatedExecutionState checkpointUpdatedE
         var nextMarker = checkpointUpdatedExecutionState.nextMarker();
         while (nextMarker != null && !nextMarker.isEmpty()) {
             var startTime = System.nanoTime();
-            var response = config.getDurableExecutionClient()
-                    .getExecutionState(durableExecutionArn, checkpointToken, nextMarker);
-            logger.debug(
-                    "Durable getExecutionState API called (latency={}ns): {}.",
-                    System.nanoTime() - startTime,
-                    response);
-            operations.addAll(response.operations());
-            nextMarker = response.nextMarker();
+            try {
+                var response = config.getDurableExecutionClient()
+                        .getExecutionState(durableExecutionArn, checkpointToken, nextMarker);
+                logger.debug(
+                        "Durable getExecutionState API called (latency={}ns): {}.",
+                        System.nanoTime() - startTime,
+                        response);
+                operations.addAll(response.operations());
+                nextMarker = response.nextMarker();
+            } catch (AwsServiceException e) {
+                throw DurableApiErrorClassifier.classifyException(e);
+            }
         }
         return operations;
     }
@@ -187,35 +193,41 @@ private void checkpointBatch(List<OperationUpdate> updates) {
 
             var startTime = System.nanoTime();
             logger.debug("Calling durable checkpoint API with {} updates: {}", updates.size(), request);
-            var response = config.getDurableExecutionClient().checkpoint(durableExecutionArn, checkpointToken, request);
-            logger.debug("Durable checkpoint API called (latency={}ns): {}.", System.nanoTime() - startTime, response);
-
-            // Notify callback of completion
-            checkpointToken = response.checkpointToken();
-            if (response.newExecutionState() != null) {
-                // fetch all pages of operations
-                var operations = fetchAllPages(response.newExecutionState());
-
-                var processStartTime = System.nanoTime();
-                int completedFutures = 0;
+            try {
+                var response =
+                        config.getDurableExecutionClient().checkpoint(durableExecutionArn, checkpointToken, request);
                 logger.debug(
-                        "Processing {} operations. ({} pending pollers)", operations.size(), pollingFutures.size());
-                // call the callback
-                callback.accept(operations);
-
-                // complete the registered pollingFutures
-                for (var operation : operations) {
-                    var pollers = pollingFutures.remove(operation.id());
-                    if (pollers != null) {
-                        completedFutures += pollers.size();
-                        pollers.forEach(poller -> poller.complete(operation));
+                        "Durable checkpoint API called (latency={}ns): {}.", System.nanoTime() - startTime, response);
+
+                // Notify callback of completion
+                checkpointToken = response.checkpointToken();
+                if (response.newExecutionState() != null) {
+                    // fetch all pages of operations
+                    var operations = fetchAllPages(response.newExecutionState());
+
+                    var processStartTime = System.nanoTime();
+                    int completedFutures = 0;
+                    logger.debug(
+                            "Processing {} operations. ({} pending pollers)", operations.size(), pollingFutures.size());
+                    // call the callback
+                    callback.accept(operations);
+
+                    // complete the registered pollingFutures
+                    for (var operation : operations) {
+                        var pollers = pollingFutures.remove(operation.id());
+                        if (pollers != null) {
+                            completedFutures += pollers.size();
+                            pollers.forEach(poller -> poller.complete(operation));
+                        }
                     }
+                    logger.debug(
+                            "{} operations processed and {} pollers completed (latency={}ns). ",
+                            operations.size(),
+                            completedFutures,
+                            System.nanoTime() - processStartTime);
                 }
-                logger.debug(
-                        "{} operations processed and {} pollers completed (latency={}ns). ",
-                        operations.size(),
-                        completedFutures,
-                        System.nanoTime() - processStartTime);
+            } catch (AwsServiceException e) {
+                throw DurableApiErrorClassifier.classifyException(e);
             }
         }
     }
diff --git a/sdk/src/main/java/software/amazon/lambda/durable/util/DurableApiErrorClassifier.java b/sdk/src/main/java/software/amazon/lambda/durable/util/DurableApiErrorClassifier.java
@@ -0,0 +1,109 @@
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+package software.amazon.lambda.durable.util;
+
+import java.util.Set;
+import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.services.lambda.model.ErrorObject;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
+
+/**
+ * Classifies AWS service exceptions from Durable Execution API calls as non-retryable or retryable.
+ *
+ * <p>Returns {@link UnrecoverableDurableExecutionException} with {@code retryable=false} for non-retryable customer
+ * errors (e.g., KMS key misconfiguration), or {@code retryable=true} for retryable errors (throttling, server errors,
+ * stale checkpoint tokens).
+ *
+ * <p>To add a new non-retryable error, add its error code to {@link #NON_RETRYABLE_ERROR_CODES}.
+ */
+public final class DurableApiErrorClassifier {
+
+    /**
+     * Error codes that represent non-retryable customer errors. When a Durable Execution API call fails with one of
+     * these error codes, the execution is terminated immediately.
+     *
+     * <p>These error codes are documented under the Lambda Invoke API but also apply to other Lambda APIs such as
+     * {@code CheckpointDurableExecution} and {@code GetDurableExecutionState}.
+     *
+     * @see <a href="https://docs.aws.amazon.com/lambda/latest/api/API_Invoke.html">Lambda Invoke API — Errors
+     *     (reference for KMS exception names)</a>
+     */
+    static final Set<String> NON_RETRYABLE_ERROR_CODES = Set.of(
+            "KMSAccessDeniedException", "KMSDisabledException", "KMSInvalidStateException", "KMSNotFoundException");
+
+    /** HTTP 429 (Too Many Requests) indicates throttling — a transient condition that resolves on retry. */
+    private static final int THROTTLING_STATUS_CODE = 429;
+
+    /**
+     * Error code for invalid checkpoint token errors. These occur when the SDK uses a stale checkpoint token that has
+     * been superseded by a newer invocation. Retrying with a fresh invocation resolves this.
+     *
+     * @see <a
+     *     href="https://github.com/aws/aws-durable-execution-sdk-js/blob/main/packages/aws-durable-execution-sdk-js/src/utils/checkpoint/checkpoint-manager.ts">JS
+     *     SDK classifyCheckpointError</a>
+     */
+    private static final String INVALID_CHECKPOINT_TOKEN_ERROR_CODE = "InvalidParameterValueException";
+
+    /**
+     * Message prefix that distinguishes invalid checkpoint token errors from other
+     * {@code InvalidParameterValueException} errors.
+     */
+    private static final String INVALID_CHECKPOINT_TOKEN_MESSAGE_PREFIX = "Invalid Checkpoint Token";
+
+    private DurableApiErrorClassifier() {}
+
+    /**
+     * Classifies the given exception and returns the appropriate exception to throw.
+     *
+     * <p>Returns {@link UnrecoverableDurableExecutionException} with {@code retryable=false} for non-retryable customer
+     * errors, or {@code retryable=true} for retryable errors.
+     *
+     * <p>Classification rules:
+     *
+     * <ul>
+     *   <li>Error code in {@link #NON_RETRYABLE_ERROR_CODES} → non-retryable ({@code retryable=false})
+     *   <li>4xx + "Invalid Checkpoint Token" → retryable ({@code retryable=true}, stale token resolves on retry)
+     *   <li>4xx (non-429) → non-retryable ({@code retryable=false}, customer error)
+     *   <li>429, 5xx, unknown → retryable ({@code retryable=true})
+     * </ul>
+     *
+     * @param e the AWS service exception from a Durable Execution API call
+     * @return an {@link UnrecoverableDurableExecutionException} for all cases, with the retryable flag set accordingly
+     */
+    public static UnrecoverableDurableExecutionException classifyException(AwsServiceException e) {
+        var errorCode = e.awsErrorDetails().errorCode();
+
+        // Non-retryable customer errors: execution is terminally broken (e.g., KMS key misconfiguration)
+        if (NON_RETRYABLE_ERROR_CODES.contains(errorCode)) {
+            return buildUnrecoverableDurableExecutionException(e, false);
+        }
+
+        var statusCode = e.awsErrorDetails().sdkHttpResponse().statusCode();
+        var message = e.getMessage();
+
+        // 4xx errors (excluding throttling) are non-retryable customer errors
+        if (statusCode >= 400 && statusCode < 500 && statusCode != THROTTLING_STATUS_CODE) {
+            // Stale checkpoint token: occurs when a newer invocation has superseded this one.
+            // Retrying with a fresh invocation resolves this, so treat as retryable.
+            if (INVALID_CHECKPOINT_TOKEN_ERROR_CODE.equals(errorCode)
+                    && message != null
+                    && message.startsWith(INVALID_CHECKPOINT_TOKEN_MESSAGE_PREFIX)) {
+                return buildUnrecoverableDurableExecutionException(e, true);
+            }
+            return buildUnrecoverableDurableExecutionException(e, false);
+        }
+
+        // 429 (throttling), 5xx (service errors), unknown — transient, retryable
+        return buildUnrecoverableDurableExecutionException(e, true);
+    }
+
+    private static UnrecoverableDurableExecutionException buildUnrecoverableDurableExecutionException(
+            AwsServiceException e, boolean retryable) {
+        return new UnrecoverableDurableExecutionException(
+                ErrorObject.builder()
+                        .errorType(e.awsErrorDetails().errorCode())
+                        .errorMessage(e.getMessage())
+                        .build(),
+                retryable);
+    }
+}
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/CheckpointManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/CheckpointManagerTest.java
@@ -15,6 +15,9 @@
 import java.util.concurrent.atomic.AtomicInteger;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
+import software.amazon.awssdk.awscore.exception.AwsErrorDetails;
+import software.amazon.awssdk.awscore.exception.AwsServiceException;
+import software.amazon.awssdk.http.SdkHttpResponse;
 import software.amazon.awssdk.services.lambda.model.CheckpointDurableExecutionResponse;
 import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState;
 import software.amazon.awssdk.services.lambda.model.GetDurableExecutionStateResponse;
@@ -25,6 +28,7 @@
 import software.amazon.awssdk.services.lambda.model.OperationUpdate;
 import software.amazon.lambda.durable.DurableConfig;
 import software.amazon.lambda.durable.client.DurableExecutionClient;
+import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException;
 import software.amazon.lambda.durable.retry.JitterStrategy;
 import software.amazon.lambda.durable.retry.PollingStrategies;
 
@@ -613,4 +617,54 @@ void pollForUpdate_withFixedDelay_intervalsAreConsistent() throws Exception {
                             + "ms, intervals=" + intervals);
         }
     }
+
+    // --- Checkpoint error classification wiring tests ---
+
+    @Test
+    void checkpointBatch_nonRetryableError_throwsUnrecoverable() {
+        when(client.checkpoint(anyString(), anyString(), anyList()))
+                .thenThrow(AwsServiceException.builder()
+                        .message("KMSAccessDeniedException: Lambda was unable to decrypt the environment variables")
+                        .awsErrorDetails(AwsErrorDetails.builder()
+                                .errorCode("KMSAccessDeniedException")
+                                .errorMessage("Lambda was unable to decrypt the environment variables")
+                                .sdkHttpResponse(SdkHttpResponse.builder()
+                                        .statusCode(502)
+                                        .build())
+                                .build())
+                        .statusCode(502)
+                        .build());
+
+        var future = batcher.checkpoint(OperationUpdate.builder()
+                .id("op-1")
+                .type(OperationType.STEP)
+                .action(OperationAction.START)
+                .build());
+
+        var ex = assertThrows(Exception.class, () -> future.get(200, TimeUnit.MILLISECONDS));
+        assertInstanceOf(UnrecoverableDurableExecutionException.class, ex.getCause());
+    }
+
+    @Test
+    void fetchAllPages_nonRetryableError_throwsUnrecoverable() {
+        when(client.getExecutionState(eq("arn:test"), eq("token-1"), eq("marker-1")))
+                .thenThrow(AwsServiceException.builder()
+                        .message("KMSAccessDeniedException: Lambda was unable to decrypt the environment variables")
+                        .awsErrorDetails(AwsErrorDetails.builder()
+                                .errorCode("KMSAccessDeniedException")
+                                .errorMessage("Lambda was unable to decrypt the environment variables")
+                                .sdkHttpResponse(SdkHttpResponse.builder()
+                                        .statusCode(502)
+                                        .build())
+                                .build())
+                        .statusCode(502)
+                        .build());
+
+        var state = CheckpointUpdatedExecutionState.builder()
+                .operations(List.of(Operation.builder().id("op-1").build()))
+                .nextMarker("marker-1")
+                .build();
+
+        assertThrows(UnrecoverableDurableExecutionException.class, () -> batcher.fetchAllPages(state));
+    }
 }
diff --git a/sdk/src/test/java/software/amazon/lambda/durable/util/DurableApiErrorClassifierTest.java b/sdk/src/test/java/software/amazon/lambda/durable/util/DurableApiErrorClassifierTest.java
