From f2934a343d245c137f251f47de64fa60e0028eeb Mon Sep 17 00:00:00 2001 From: Frank Chen Date: Wed, 2 Sep 2026 22:06:35 +0000 Subject: [PATCH] feat: add payload offloader core runtime --- .../amazon/lambda/durable/DurableConfig.java | 68 ++ .../lambda/durable/config/InvokeConfig.java | 48 +- .../lambda/durable/config/MapConfig.java | 16 + .../durable/config/ParallelBranchConfig.java | 17 +- .../lambda/durable/config/ParallelConfig.java | 18 +- .../config/RunInChildContextConfig.java | 17 +- .../lambda/durable/config/StepConfig.java | 20 +- .../config/WaitForConditionConfig.java | 16 + .../durable/context/DurableContextImpl.java | 40 +- .../exception/DurableOperationException.java | 19 + .../exception/PayloadOffloadException.java | 14 + .../RetryablePayloadOffloadException.java | 14 + .../durable/execution/DurableExecutor.java | 344 +++++++-- .../durable/execution/ExecutionManager.java | 16 + .../durable/execution/PayloadCodec.java | 513 +++++++++++++ .../durable/model/DurableExecutionInput.java | 19 +- .../durable/model/InvocationSource.java | 9 + .../durable/offload/OffloadedPayload.java | 263 +++++++ .../offload/PayloadOffloadContext.java | 101 +++ .../durable/offload/PayloadOffloader.java | 26 + .../durable/offload/PayloadOffloaders.java | 37 + .../durable/offload/PayloadStorageMode.java | 9 + .../durable/offload/SerDesPayloadKind.java | 23 + .../internal/ChainedInvokeOutputFrame.java | 39 + .../internal/ChainedInvokePayloadFrame.java | 33 + .../internal/PayloadOffloadTracking.java | 47 ++ .../operation/BaseDurableOperation.java | 23 +- .../durable/operation/CallbackOperation.java | 2 +- .../operation/ChildContextOperation.java | 197 +++-- .../operation/ConcurrencyOperation.java | 140 +++- .../durable/operation/InvokeOperation.java | 61 +- .../durable/operation/MapOperation.java | 16 +- .../durable/operation/ParallelOperation.java | 15 +- .../SerializableDurableOperation.java | 253 ++++++- .../durable/operation/StepOperation.java | 24 +- .../operation/WaitForConditionOperation.java | 27 +- .../lambda/durable/DurableConfigTest.java | 35 + .../execution/DurableExecutionTest.java | 618 +++++++++++++++ .../durable/execution/PayloadCodecTest.java | 704 ++++++++++++++++++ .../durable/offload/OffloadedPayloadTest.java | 47 ++ .../ChainedInvokeOutputFrameTest.java | 45 ++ .../ChainedInvokePayloadFrameTest.java | 38 + .../operation/CallbackOperationTest.java | 39 + .../operation/ChildContextOperationTest.java | 384 +++++++++- .../operation/ConcurrencyOperationTest.java | 187 +++++ .../operation/InvokeOperationTest.java | 380 ++++++++++ .../serde/DurableInputOutputSerDesTest.java | 21 + 47 files changed, 4868 insertions(+), 174 deletions(-) create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/exception/PayloadOffloadException.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/exception/RetryablePayloadOffloadException.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/execution/PayloadCodec.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/model/InvocationSource.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/offload/OffloadedPayload.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloadContext.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloader.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloaders.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadStorageMode.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/offload/SerDesPayloadKind.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/offload/internal/ChainedInvokeOutputFrame.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/offload/internal/ChainedInvokePayloadFrame.java create mode 100644 sdk/src/main/java/software/amazon/lambda/durable/offload/internal/PayloadOffloadTracking.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/execution/PayloadCodecTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/offload/OffloadedPayloadTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/offload/internal/ChainedInvokeOutputFrameTest.java create mode 100644 sdk/src/test/java/software/amazon/lambda/durable/offload/internal/ChainedInvokePayloadFrameTest.java diff --git a/sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java index 5101b9fda..88302e1cc 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/DurableConfig.java @@ -23,6 +23,7 @@ import software.amazon.lambda.durable.client.DurableExecutionClient; import software.amazon.lambda.durable.client.LambdaDurableFunctionsClient; import software.amazon.lambda.durable.logging.LoggerConfig; +import software.amazon.lambda.durable.offload.PayloadOffloader; import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; import software.amazon.lambda.durable.plugin.PluginRunner; import software.amazon.lambda.durable.retry.PollingStrategies; @@ -94,12 +95,15 @@ public final class DurableConfig { private final DurableExecutionClient durableExecutionClient; private final SerDes serDes; + private final PayloadOffloader payloadOffloader; private final ExecutorService executorService; + private final ExecutorService payloadOffloadExecutorService; private final LoggerConfig loggerConfig; private final PollingStrategy pollingStrategy; private final Duration checkpointDelay; private final boolean deserializeAfterSerialization; private final boolean checkpointEmptyMap; + private final boolean payloadOffloaderForChainedInvokePayloads; private final PluginRunner pluginRunner; private DurableConfig(Builder builder) { @@ -107,13 +111,16 @@ private DurableConfig(Builder builder) { this.durableExecutionClient = Objects.requireNonNullElseGet( builder.durableExecutionClient, DurableConfig::createDefaultDurableExecutionClient); this.serDes = Objects.requireNonNullElseGet(builder.serDes, JacksonSerDes::new); + this.payloadOffloader = builder.payloadOffloader; this.executorService = Objects.requireNonNullElseGet(builder.executorService, DurableConfig::createDefaultExecutor); + this.payloadOffloadExecutorService = builder.payloadOffloadExecutorService; this.loggerConfig = Objects.requireNonNullElseGet(builder.loggerConfig, LoggerConfig::defaults); this.pollingStrategy = Objects.requireNonNullElse(builder.pollingStrategy, PollingStrategies.Presets.DEFAULT); this.checkpointDelay = Objects.requireNonNullElseGet(builder.checkpointDelay, () -> Duration.ofSeconds(0)); this.deserializeAfterSerialization = builder.deserializeAfterSerialization; this.checkpointEmptyMap = builder.checkpointEmptyMap; + this.payloadOffloaderForChainedInvokePayloads = builder.payloadOffloaderForChainedInvokePayloads; this.pluginRunner = plugins.isEmpty() ? PluginRunner.noOp() : new PluginRunner(plugins); validateConfiguration(); @@ -155,6 +162,11 @@ public SerDes getSerDes() { return serDes; } + /** Gets the globally configured payload offloader, or null when payload offloading is disabled. */ + public PayloadOffloader getPayloadOffloader() { + return payloadOffloader; + } + /** * Gets the configured ExecutorService. * @@ -164,6 +176,11 @@ public ExecutorService getExecutorService() { return executorService; } + /** Gets the executor used for blocking payload offload and load operations, or null to execute inline. */ + public ExecutorService getPayloadOffloadExecutorService() { + return payloadOffloadExecutorService; + } + /** * Gets the configured LoggerConfig. * @@ -214,6 +231,11 @@ public boolean shouldCheckpointEmptyMap() { return checkpointEmptyMap; } + /** Returns whether framed chained-invoke inputs may use the configured payload offloader. */ + public boolean shouldUsePayloadOffloaderForChainedInvokePayloads() { + return payloadOffloaderForChainedInvokePayloads; + } + /** * Gets the plugin runner that dispatches lifecycle events to registered plugins. * @@ -235,6 +257,10 @@ public void validateConfiguration() { if (getExecutorService() == null) { throw new IllegalStateException("ExecutorService configuration failed"); } + if (getPayloadOffloadExecutorService() != null && getPayloadOffloadExecutorService() == getExecutorService()) { + throw new IllegalStateException( + "Payload offload ExecutorService must be different from the user operation ExecutorService"); + } } /** @@ -315,12 +341,15 @@ private static ExecutorService createDefaultExecutor() { public static final class Builder { private DurableExecutionClient durableExecutionClient; private SerDes serDes; + private PayloadOffloader payloadOffloader; private ExecutorService executorService; + private ExecutorService payloadOffloadExecutorService; private LoggerConfig loggerConfig; private PollingStrategy pollingStrategy; private Duration checkpointDelay; private boolean deserializeAfterSerialization = true; private boolean checkpointEmptyMap = false; + private boolean payloadOffloaderForChainedInvokePayloads; private List plugins = new ArrayList<>(); public Builder() {} @@ -381,6 +410,29 @@ public Builder withSerDes(SerDes serDes) { return this; } + /** + * Sets the global payload offloader applied after SerDes processing. + * + * @param payloadOffloader payload offloader + * @return this builder + */ + public Builder withPayloadOffloader(PayloadOffloader payloadOffloader) { + this.payloadOffloader = Objects.requireNonNull(payloadOffloader, "PayloadOffloader cannot be null"); + return this; + } + + /** + * Controls whether SDK-framed chained-invoke inputs and outputs may use this handler's configured payload + * offloader. + * + *

This is disabled by default. Enable it only for compatible durable callers that also opt in through + * {@link software.amazon.lambda.durable.config.InvokeConfig.Builder#usePayloadOffloaderForPayload(boolean)}. + */ + public Builder withPayloadOffloaderForChainedInvokePayloads(boolean enabled) { + payloadOffloaderForChainedInvokePayloads = enabled; + return this; + } + /** * Sets a custom ExecutorService for running user-defined operations. If not set, a default cached thread pool * will be created. @@ -396,6 +448,22 @@ public Builder withExecutorService(ExecutorService executorService) { return this; } + /** + * Sets the executor used for blocking payload storage operations. If not set, payload offloader calls execute + * inline on the calling thread. + * + *

This executor must be different from the user operation executor to prevent synchronous dispatch from + * deadlocking a saturated operation pool. + * + * @param executorService payload offload executor + * @return this builder + */ + public Builder withPayloadOffloadExecutorService(ExecutorService executorService) { + this.payloadOffloadExecutorService = + Objects.requireNonNull(executorService, "Payload offload ExecutorService cannot be null"); + return this; + } + /** * Sets a custom LoggerConfig. If not set, defaults to suppressing replay logs. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/InvokeConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/InvokeConfig.java index e9dc7af24..4afe0f6fe 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/InvokeConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/InvokeConfig.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.offload.PayloadOffloader; import software.amazon.lambda.durable.serde.SerDes; /** @@ -12,12 +13,16 @@ public class InvokeConfig { private final SerDes payloadSerDes; private final SerDes resultSerDes; + private final PayloadOffloader payloadOffloader; private final String tenantId; + private final boolean usePayloadOffloaderForPayload; public InvokeConfig(Builder builder) { this.payloadSerDes = builder.payloadSerDes; this.resultSerDes = builder.resultSerDes; + this.payloadOffloader = builder.payloadOffloader; this.tenantId = builder.tenantId; + this.usePayloadOffloaderForPayload = builder.usePayloadOffloaderForPayload; } public SerDes payloadSerDes() { @@ -28,28 +33,43 @@ public SerDes serDes() { return this.resultSerDes; } + /** Returns the offloader used for the invoke result, or null to inherit the global offloader. */ + public PayloadOffloader payloadOffloader() { + return payloadOffloader; + } + public String tenantId() { return tenantId; } + /** Returns whether this invoke should use the framed durable-target payload protocol. */ + public boolean usePayloadOffloaderForPayload() { + return usePayloadOffloaderForPayload; + } + public static Builder builder() { - return new Builder(null, null, null); + return new Builder(null, null, null, false); } public Builder toBuilder() { - return new Builder(payloadSerDes, resultSerDes, tenantId); + return new Builder(payloadSerDes, resultSerDes, tenantId, usePayloadOffloaderForPayload) + .payloadOffloader(payloadOffloader); } /** Builder for creating InvokeConfig instances. */ public static class Builder { private SerDes payloadSerDes; private SerDes resultSerDes; + private PayloadOffloader payloadOffloader; private String tenantId; + private boolean usePayloadOffloaderForPayload; - private Builder(SerDes payloadSerDes, SerDes resultSerDes, String tenantId) { + private Builder( + SerDes payloadSerDes, SerDes resultSerDes, String tenantId, boolean usePayloadOffloaderForPayload) { this.payloadSerDes = payloadSerDes; this.resultSerDes = resultSerDes; this.tenantId = tenantId; + this.usePayloadOffloaderForPayload = usePayloadOffloaderForPayload; } /** @@ -81,6 +101,19 @@ public Builder payloadSerDes(SerDes payloadSerDes) { return this; } + /** + * Selects whether a compatible durable target should use the framed payload protocol for this request and its + * result or error. + * + *

This enables payload offloading for the request and lets the caller distinguish SDK-owned result/error + * envelopes from ordinary Lambda data. It is disabled by default so standard Lambda functions and older SDK + * versions continue to exchange ordinary serialized values unchanged. + */ + public Builder usePayloadOffloaderForPayload(boolean enabled) { + usePayloadOffloaderForPayload = enabled; + return this; + } + /** * Sets a custom serializer for the invoke result. * @@ -96,6 +129,15 @@ public Builder serDes(SerDes resultSerDes) { return this; } + /** + * Sets the offloader for the invoke result and for the request when + * {@link #usePayloadOffloaderForPayload(boolean)} is enabled. + */ + public Builder payloadOffloader(PayloadOffloader payloadOffloader) { + this.payloadOffloader = payloadOffloader; + return this; + } + /** * Builds the InvokeConfig instance. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java index 78bdcc293..5894b1a9d 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/MapConfig.java @@ -4,6 +4,7 @@ import java.util.Objects; import java.util.function.BiFunction; +import software.amazon.lambda.durable.offload.PayloadOffloader; import software.amazon.lambda.durable.serde.SerDes; /** @@ -15,6 +16,7 @@ public class MapConfig { private final Integer maxConcurrency; private final CompletionConfig completionConfig; private final SerDes serDes; + private final PayloadOffloader payloadOffloader; private final NestingType nestingType; private final BiFunction itemNamer; @@ -23,6 +25,7 @@ private MapConfig(Builder builder) { this.completionConfig = Objects.requireNonNullElse(builder.completionConfig, CompletionConfig.allCompleted()); this.nestingType = Objects.requireNonNullElse(builder.nestingType, NestingType.NESTED); this.serDes = builder.serDes; + this.payloadOffloader = builder.payloadOffloader; this.itemNamer = builder.itemNamer; if (itemNamer != null && nestingType == NestingType.FLAT) { throw new IllegalArgumentException("itemNamer is not supported with FLAT map nesting"); @@ -44,6 +47,11 @@ public SerDes serDes() { return serDes; } + /** @return the map and iteration result offloader, or null to inherit the global offloader */ + public PayloadOffloader payloadOffloader() { + return payloadOffloader; + } + /** @return nesting type, defaults to {@link NestingType#NESTED} */ public NestingType nestingType() { return nestingType; @@ -70,6 +78,7 @@ public Builder toBuilder() { .maxConcurrency(maxConcurrency) .completionConfig(completionConfig) .serDes(serDes) + .payloadOffloader(payloadOffloader) .nestingType(nestingType) .itemNamer(itemNamer); } @@ -80,6 +89,7 @@ public static class Builder { private Integer maxConcurrency; private CompletionConfig completionConfig; private SerDes serDes; + private PayloadOffloader payloadOffloader; private BiFunction itemNamer; private Builder() {} @@ -114,6 +124,12 @@ public Builder serDes(SerDes serDes) { return this; } + /** Sets the payload offloader for map iteration and aggregate results. */ + public Builder payloadOffloader(PayloadOffloader payloadOffloader) { + this.payloadOffloader = payloadOffloader; + return this; + } + /** * Sets the nesting type for the map operation. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelBranchConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelBranchConfig.java index 689f9aa54..8e4316118 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelBranchConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelBranchConfig.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.offload.PayloadOffloader; import software.amazon.lambda.durable.serde.SerDes; /** @@ -11,9 +12,11 @@ */ public class ParallelBranchConfig { private final SerDes serDes; + private final PayloadOffloader payloadOffloader; private ParallelBranchConfig(Builder builder) { this.serDes = builder.serDes; + this.payloadOffloader = builder.payloadOffloader; } /** Returns the custom serializer for this step, or null if not specified (uses default SerDes). */ @@ -21,8 +24,13 @@ public SerDes serDes() { return serDes; } + /** Returns the branch result offloader, or null to inherit the global offloader. */ + public PayloadOffloader payloadOffloader() { + return payloadOffloader; + } + public Builder toBuilder() { - return new Builder(serDes); + return new Builder(serDes).payloadOffloader(payloadOffloader); } /** @@ -37,6 +45,7 @@ public static Builder builder() { /** Builder for creating StepConfig instances. */ public static class Builder { private SerDes serDes; + private PayloadOffloader payloadOffloader; public Builder(SerDes serDes) { this.serDes = serDes; @@ -57,6 +66,12 @@ public Builder serDes(SerDes serDes) { return this; } + /** Sets the payload offloader for the parallel branch result. */ + public Builder payloadOffloader(PayloadOffloader payloadOffloader) { + this.payloadOffloader = payloadOffloader; + return this; + } + /** * Builds the ParallelBranchConfig instance. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java index 863d36972..d504037bd 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/ParallelConfig.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable.config; import java.util.Objects; +import software.amazon.lambda.durable.offload.PayloadOffloader; /** * Configuration options for parallel operations in durable executions. @@ -14,11 +15,13 @@ public class ParallelConfig { private final int maxConcurrency; private final CompletionConfig completionConfig; private final NestingType nestingType; + private final PayloadOffloader payloadOffloader; private ParallelConfig(Builder builder) { this.maxConcurrency = Objects.requireNonNullElse(builder.maxConcurrency, Integer.MAX_VALUE); this.completionConfig = Objects.requireNonNullElseGet(builder.completionConfig, CompletionConfig::allCompleted); this.nestingType = Objects.requireNonNullElse(builder.nestingType, NestingType.NESTED); + this.payloadOffloader = builder.payloadOffloader; } /** @return the maximum number of branches running simultaneously, or -1 for unlimited */ @@ -36,6 +39,11 @@ public NestingType nestingType() { return nestingType; } + /** @return the aggregate result offloader, or null to inherit the global offloader */ + public PayloadOffloader payloadOffloader() { + return payloadOffloader; + } + /** * Creates a new builder for ParallelConfig. * @@ -49,7 +57,8 @@ public Builder toBuilder() { return new Builder() .maxConcurrency(maxConcurrency) .completionConfig(completionConfig) - .nestingType(nestingType); + .nestingType(nestingType) + .payloadOffloader(payloadOffloader); } /** Builder for creating ParallelConfig instances. */ @@ -57,6 +66,7 @@ public static class Builder { private Integer maxConcurrency; private CompletionConfig completionConfig; private NestingType nestingType; + private PayloadOffloader payloadOffloader; private Builder() {} @@ -101,6 +111,12 @@ public Builder nestingType(NestingType nestingType) { return this; } + /** Sets the payload offloader for the aggregate parallel result. */ + public Builder payloadOffloader(PayloadOffloader payloadOffloader) { + this.payloadOffloader = payloadOffloader; + return this; + } + /** * Builds the ParallelConfig instance. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/RunInChildContextConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/RunInChildContextConfig.java index 93fde31e8..092f21fa7 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/RunInChildContextConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/RunInChildContextConfig.java @@ -3,6 +3,7 @@ package software.amazon.lambda.durable.config; import java.util.Objects; +import software.amazon.lambda.durable.offload.PayloadOffloader; import software.amazon.lambda.durable.serde.SerDes; /** @@ -12,10 +13,12 @@ */ public class RunInChildContextConfig { private final SerDes serDes; + private final PayloadOffloader payloadOffloader; private final Boolean isVirtual; private RunInChildContextConfig(Builder builder) { this.serDes = builder.serDes; + this.payloadOffloader = builder.payloadOffloader; this.isVirtual = Objects.requireNonNullElse(builder.isVirtual, false); } @@ -27,13 +30,18 @@ public SerDes serDes() { return serDes; } + /** Returns the child context result offloader, or null to inherit the global offloader. */ + public PayloadOffloader payloadOffloader() { + return payloadOffloader; + } + /** Returns true if the context operation will not be checkpointed, false otherwise. */ public Boolean isVirtual() { return isVirtual; } public Builder toBuilder() { - return new Builder().serDes(serDes).isVirtual(isVirtual); + return new Builder().serDes(serDes).payloadOffloader(payloadOffloader).isVirtual(isVirtual); } /** @@ -48,6 +56,7 @@ public static Builder builder() { /** Builder for creating StepConfig instances. */ public static class Builder { private SerDes serDes; + private PayloadOffloader payloadOffloader; private Boolean isVirtual; private Builder() {} @@ -67,6 +76,12 @@ public Builder serDes(SerDes serDes) { return this; } + /** Sets the payload offloader for the child context result and exception. */ + public Builder payloadOffloader(PayloadOffloader payloadOffloader) { + this.payloadOffloader = payloadOffloader; + return this; + } + /** * Sets whether the context is virtual (not checkpointed) or not. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/StepConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/StepConfig.java index 92a90a643..49b53f771 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/StepConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/StepConfig.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.offload.PayloadOffloader; import software.amazon.lambda.durable.retry.RetryStrategies; import software.amazon.lambda.durable.retry.RetryStrategy; import software.amazon.lambda.durable.serde.SerDes; @@ -16,11 +17,13 @@ public class StepConfig { private final RetryStrategy retryStrategy; private final StepSemantics semanticsPerRetry; private final SerDes serDes; + private final PayloadOffloader payloadOffloader; private StepConfig(Builder builder) { this.retryStrategy = builder.retryStrategy; this.semanticsPerRetry = builder.semanticsPerRetry; this.serDes = builder.serDes; + this.payloadOffloader = builder.payloadOffloader; } /** Returns the retry strategy for this step, or the default strategy if not specified. */ @@ -38,8 +41,13 @@ public SerDes serDes() { return serDes; } + /** Returns the operation payload offloader, or null to inherit the handler configuration. */ + public PayloadOffloader payloadOffloader() { + return payloadOffloader; + } + public Builder toBuilder() { - return new Builder(retryStrategy, semanticsPerRetry, serDes); + return new Builder(retryStrategy, semanticsPerRetry, serDes).payloadOffloader(payloadOffloader); } /** @@ -56,6 +64,7 @@ public static class Builder { private RetryStrategy retryStrategy; private StepSemantics semanticsPerRetry; private SerDes serDes; + private PayloadOffloader payloadOffloader; public Builder(RetryStrategy retryStrategy, StepSemantics semanticsPerRetry, SerDes serDes) { this.retryStrategy = retryStrategy; @@ -100,6 +109,15 @@ public Builder serDes(SerDes serDes) { return this; } + /** + * Sets the payload offloader for this step. Use {@link PayloadOffloader#disabled()} to force inline storage + * when a global offloader is configured. + */ + public Builder payloadOffloader(PayloadOffloader payloadOffloader) { + this.payloadOffloader = payloadOffloader; + return this; + } + /** * Builds the StepConfig instance. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForConditionConfig.java b/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForConditionConfig.java index 1561199f9..18096fc14 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForConditionConfig.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/config/WaitForConditionConfig.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.config; +import software.amazon.lambda.durable.offload.PayloadOffloader; import software.amazon.lambda.durable.retry.WaitForConditionWaitStrategy; import software.amazon.lambda.durable.retry.WaitStrategies; import software.amazon.lambda.durable.serde.SerDes; @@ -16,11 +17,13 @@ public class WaitForConditionConfig { private final WaitForConditionWaitStrategy waitStrategy; private final SerDes serDes; + private final PayloadOffloader payloadOffloader; private final T initialState; private WaitForConditionConfig(Builder builder) { this.waitStrategy = builder.waitStrategy; this.serDes = builder.serDes; + this.payloadOffloader = builder.payloadOffloader; this.initialState = builder.initialState; } @@ -37,6 +40,11 @@ public SerDes serDes() { return serDes; } + /** Returns the state offloader, or null to inherit the global offloader. */ + public PayloadOffloader payloadOffloader() { + return payloadOffloader; + } + /** Returns the initial state object, or null if not specified. */ public T initialState() { return initialState; @@ -52,6 +60,7 @@ public Builder toBuilder() { var b = new Builder(); b.waitStrategy = this.waitStrategy; b.serDes = this.serDes; + b.payloadOffloader = this.payloadOffloader; b.initialState = this.initialState; return b; } @@ -69,6 +78,7 @@ public static Builder builder() { public static class Builder { private WaitForConditionWaitStrategy waitStrategy; private SerDes serDes; + private PayloadOffloader payloadOffloader; private T initialState; private Builder() {} @@ -100,6 +110,12 @@ public Builder serDes(SerDes serDes) { return this; } + /** Sets the payload offloader for checkpointed condition state. */ + public Builder payloadOffloader(PayloadOffloader payloadOffloader) { + this.payloadOffloader = payloadOffloader; + return this; + } + /** * Sets the initial state for the waitForCondition operation. The initial state will be null if it's not set. * diff --git a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java index 0c79165ec..e1b204760 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java @@ -26,6 +26,7 @@ import software.amazon.lambda.durable.config.WaitForCallbackConfig; import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.config.WithRetryConfig; +import software.amazon.lambda.durable.exception.PayloadOffloadException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.OperationIdGenerator; @@ -140,6 +141,11 @@ public DurableFuture stepAsync( if (config.serDes() == null) { config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); } + if (config.payloadOffloader() == null) { + config = config.toBuilder() + .payloadOffloader(getDurableConfig().getPayloadOffloader()) + .build(); + } var operationId = nextOperationId(); // Create and start step operation with TypeToken @@ -181,6 +187,11 @@ public DurableFuture invokeAsync( .payloadSerDes(getDurableConfig().getSerDes()) .build(); } + if (config.payloadOffloader() == null) { + config = config.toBuilder() + .payloadOffloader(getDurableConfig().getPayloadOffloader()) + .build(); + } var operationId = nextOperationId(); // Create and start invoke operation @@ -242,6 +253,11 @@ private DurableFuture runInChildContextAsync( if (config.serDes() == null) { config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); } + if (config.payloadOffloader() == null) { + config = config.toBuilder() + .payloadOffloader(getDurableConfig().getPayloadOffloader()) + .build(); + } var operationId = nextOperationId(); @@ -265,6 +281,11 @@ public DurableFuture> mapAsync( if (config.serDes() == null) { config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); } + if (config.payloadOffloader() == null) { + config = config.toBuilder() + .payloadOffloader(getDurableConfig().getPayloadOffloader()) + .build(); + } // Convert to List for deterministic index-based access var itemList = List.copyOf(items); @@ -286,6 +307,11 @@ public DurableFuture> mapAsync( @Override public ParallelDurableFuture parallel(String name, ParallelConfig config) { Objects.requireNonNull(config, "config cannot be null"); + if (config.payloadOffloader() == null) { + config = config.toBuilder() + .payloadOffloader(getDurableConfig().getPayloadOffloader()) + .build(); + } var operationId = nextOperationId(); var parallelOp = new ParallelOperation( @@ -357,6 +383,11 @@ public DurableFuture waitForConditionAsync( if (config.serDes() == null) { config = config.toBuilder().serDes(getDurableConfig().getSerDes()).build(); } + if (config.payloadOffloader() == null) { + config = config.toBuilder() + .payloadOffloader(getDurableConfig().getPayloadOffloader()) + .build(); + } var operationId = nextOperationId(); var operation = new WaitForConditionOperation<>( @@ -401,8 +432,9 @@ public DurableFuture withRetryAsync( * Core retry loop. Replay-safe because every side-effect is a durable operation: the user's operation calls durable * primitives, and backoff uses {@code context.wait()}. * - *

{@link SuspendExecutionException} and {@link UnrecoverableDurableExecutionException} are never retried — they - * are internal SDK control flow signals that must propagate immediately. + *

{@link SuspendExecutionException}, {@link UnrecoverableDurableExecutionException}, and + * {@link PayloadOffloadException} are never retried — they are SDK control flow or infrastructure failures that + * must propagate immediately. */ private static T executeRetryLoop( DurableContext context, @@ -413,8 +445,8 @@ private static T executeRetryLoop( while (true) { try { return operation.apply(attempt, context); - } catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) { - // Internal SDK control flow — never retry, always propagate + } catch (SuspendExecutionException | UnrecoverableDurableExecutionException | PayloadOffloadException e) { + // SDK control flow and payload infrastructure failures — never retry, always propagate throw e; } catch (Exception e) { RetryDecision decision = config.retryStrategy().makeRetryDecision(e, attempt); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/exception/DurableOperationException.java b/sdk/src/main/java/software/amazon/lambda/durable/exception/DurableOperationException.java index 73078ea1d..7bc430d16 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/exception/DurableOperationException.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/exception/DurableOperationException.java @@ -5,12 +5,16 @@ import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.lambda.durable.offload.PayloadOffloadContext; +import software.amazon.lambda.durable.offload.PayloadOffloader; import software.amazon.lambda.durable.util.ExceptionHelper; /** Exception associated with a specific durable operation, carrying the operation and error details. */ public class DurableOperationException extends DurableExecutionException { private final Operation operation; private final ErrorObject errorObject; + private PayloadOffloader payloadOffloader; + private PayloadOffloadContext payloadOffloadContext; public DurableOperationException(Operation operation, ErrorObject errorObject) { this(operation, errorObject, errorObject != null ? errorObject.errorMessage() : null); @@ -60,4 +64,19 @@ public OperationStatus getOperationStatus() { public String getOperationId() { return operation.id(); } + + /** Retains the operation-level payload policy that produced this exception's serialized error data. */ + public DurableOperationException withPayloadSource(PayloadOffloader offloader, PayloadOffloadContext context) { + payloadOffloader = offloader; + payloadOffloadContext = context; + return this; + } + + public PayloadOffloader getPayloadOffloader() { + return payloadOffloader; + } + + public PayloadOffloadContext getPayloadOffloadContext() { + return payloadOffloadContext; + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/exception/PayloadOffloadException.java b/sdk/src/main/java/software/amazon/lambda/durable/exception/PayloadOffloadException.java new file mode 100644 index 000000000..1a94ddfd6 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/exception/PayloadOffloadException.java @@ -0,0 +1,14 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.exception; + +/** Thrown when a serialized payload cannot be stored in or loaded from external storage. */ +public class PayloadOffloadException extends DurableExecutionException { + public PayloadOffloadException(String message, Throwable cause) { + super(message, cause); + } + + public PayloadOffloadException(String message) { + super(message); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/exception/RetryablePayloadOffloadException.java b/sdk/src/main/java/software/amazon/lambda/durable/exception/RetryablePayloadOffloadException.java new file mode 100644 index 000000000..428352efd --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/exception/RetryablePayloadOffloadException.java @@ -0,0 +1,14 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.exception; + +/** Signals a transient payload storage failure that may be retried by a configured offloader wrapper. */ +public class RetryablePayloadOffloadException extends PayloadOffloadException { + public RetryablePayloadOffloadException(String message, Throwable cause) { + super(message, cause); + } + + public RetryablePayloadOffloadException(String message) { + super(message); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java index d8db91326..106b9e6fb 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/DurableExecutor.java @@ -12,7 +12,6 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; import software.amazon.awssdk.services.lambda.model.ErrorObject; -import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationAction; import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.awssdk.services.lambda.model.OperationUpdate; @@ -22,16 +21,24 @@ import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.DurableOperationException; import software.amazon.lambda.durable.exception.IllegalDurableOperationException; +import software.amazon.lambda.durable.exception.PayloadOffloadException; +import software.amazon.lambda.durable.exception.RetryablePayloadOffloadException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.logging.DurableLogger; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.DurableExecutionOutput; +import software.amazon.lambda.durable.model.InvocationSource; +import software.amazon.lambda.durable.offload.PayloadOffloadContext; +import software.amazon.lambda.durable.offload.PayloadOffloader; +import software.amazon.lambda.durable.offload.PayloadOffloaders; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; +import software.amazon.lambda.durable.offload.internal.ChainedInvokeOutputFrame; +import software.amazon.lambda.durable.offload.internal.ChainedInvokePayloadFrame; import software.amazon.lambda.durable.plugin.InvocationEndInfo; import software.amazon.lambda.durable.plugin.InvocationInfo; import software.amazon.lambda.durable.plugin.InvocationStatus; import software.amazon.lambda.durable.plugin.PluginInfoConverter; import software.amazon.lambda.durable.plugin.PluginRunner; -import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ExceptionHelper; /** @@ -60,6 +67,11 @@ public static DurableExecutionOutput execute( var isFirstInvocation = !executionManager.isReplaying(); var requestId = lambdaContext != null ? lambdaContext.getAwsRequestId() : null; var executionArn = input.durableExecutionArn(); + var frameChainedInvokeOutput = shouldFrameChainedInvokeOutput(executionManager, config); + var outputOffloader = + input.invocationSource() == InvocationSource.CHAINED_INVOKE && !frameChainedInvokeOutput + ? PayloadOffloaders.disabled() + : config.getPayloadOffloader(); executionManager.registerActiveThread(null); // Captured for onInvocationEnd, which runs outside the handler thread below. @@ -77,8 +89,7 @@ public static DurableExecutionOutput execute( I userInput = null; Throwable inputFailure = null; try { - userInput = extractUserInput( - executionManager.getExecutionOperation(), config.getSerDes(), inputType); + userInput = extractUserInput(executionManager, config, inputType); } catch (Throwable t) { inputFailure = t; } @@ -122,62 +133,53 @@ public static DurableExecutionOutput execute( .runUntilCompleteOrSuspend(handlerFuture) .handle((result, ex) -> { if (ex != null) { - // an exception thrown from handlerFuture or suspension/termination occurred - Throwable cause = ExceptionHelper.unwrapCompletableFuture(ex); - - // return PENDING if it's SuspendExecutionException - if (cause instanceof SuspendExecutionException) { - fireOnInvocationEnd( - pluginRunner, - executionManager, - requestId, - executionArn, - isFirstInvocation, - InvocationStatus.PENDING, - null, - pluginExecutionInput.get(), - null); - return DurableExecutionOutput.pending(); - } - - // let the backend retry the invocation if the exception is retryable - if (cause - instanceof - UnrecoverableDurableExecutionException - unrecoverableDurableExecutionException - && unrecoverableDurableExecutionException.isRetryable()) { - fireOnInvocationEnd( - pluginRunner, - executionManager, - requestId, - executionArn, - isFirstInvocation, - InvocationStatus.RETRYING, - cause, - pluginExecutionInput.get(), - null); - throw unrecoverableDurableExecutionException; - } - - // fail the execution otherwise - logger.debug("Execution failed: {}", cause.getMessage()); - fireOnInvocationEnd( + return handleExecutionFailure( + ExceptionHelper.unwrapCompletableFuture(ex), + pluginRunner, + executionManager, + config, + requestId, + executionArn, + isFirstInvocation, + frameChainedInvokeOutput, + outputOffloader, + pluginExecutionInput.get()); + } + final String outputPayload; + try { + outputPayload = executionManager + .getPayloadCodec() + .serialize( + result, + config.getSerDes(), + outputOffloader, + executionContext(executionManager, SerDesPayloadKind.OUTPUT)); + } catch (PayloadOffloadException outputFailure) { + return handleExecutionFailure( + outputFailure, pluginRunner, executionManager, + config, requestId, executionArn, isFirstInvocation, - InvocationStatus.FAILED, - cause, - pluginExecutionInput.get(), - null); - return DurableExecutionOutput.failure(buildErrorObject(cause, config.getSerDes())); + frameChainedInvokeOutput, + outputOffloader, + pluginExecutionInput.get()); } - // user handler complete successfully + + // User handler and output serialization completed successfully. Infrastructure failures + // while + // publishing a large root result must escape for invocation retry rather than being + // converted + // into a terminal user failure. logger.debug("Execution completed"); - var outputPayload = config.getSerDes().serialize(result); - var output = - DurableExecutionOutput.success(handleLargePayload(executionManager, outputPayload)); + var responsePayload = frameChainedInvokeOutput + ? ChainedInvokeOutputFrame.encode( + outputPayload, PayloadCodec.isOffloadEnvelope(outputPayload)) + : outputPayload; + var output = DurableExecutionOutput.success( + handleLargePayload(executionManager, responsePayload)); fireOnInvocationEnd( pluginRunner, executionManager, @@ -199,6 +201,140 @@ public static DurableExecutionOutput execute( } } + private static DurableExecutionOutput handleExecutionFailure( + Throwable cause, + PluginRunner pluginRunner, + ExecutionManager executionManager, + DurableConfig config, + String requestId, + String executionArn, + boolean isFirstInvocation, + boolean frameChainedInvokeOutput, + PayloadOffloader outputOffloader, + Object executionInput) { + if (cause instanceof SuspendExecutionException) { + fireOnInvocationEnd( + pluginRunner, + executionManager, + requestId, + executionArn, + isFirstInvocation, + InvocationStatus.PENDING, + null, + executionInput, + null); + return DurableExecutionOutput.pending(); + } + + if (cause instanceof RetryablePayloadOffloadException retryablePayloadOffloadException) { + fireOnInvocationEnd( + pluginRunner, + executionManager, + requestId, + executionArn, + isFirstInvocation, + InvocationStatus.RETRYING, + cause, + executionInput, + null); + throw retryablePayloadOffloadException; + } + + if (cause instanceof UnrecoverableDurableExecutionException unrecoverable && unrecoverable.isRetryable()) { + fireOnInvocationEnd( + pluginRunner, + executionManager, + requestId, + executionArn, + isFirstInvocation, + InvocationStatus.RETRYING, + cause, + executionInput, + null); + throw unrecoverable; + } + + logger.debug("Execution failed: {}", cause.getMessage()); + Throwable reportedCause = cause; + EncodedError encodedError; + try { + encodedError = buildErrorObject(cause, executionManager, config, outputOffloader); + } catch (PayloadOffloadException payloadFailure) { + if (payloadFailure instanceof RetryablePayloadOffloadException retryablePayloadFailure) { + fireOnInvocationEnd( + pluginRunner, + executionManager, + requestId, + executionArn, + isFirstInvocation, + InvocationStatus.RETRYING, + retryablePayloadFailure, + executionInput, + null); + throw retryablePayloadFailure; + } + reportedCause = payloadFailure; + try { + encodedError = buildErrorObject(payloadFailure, executionManager, config, outputOffloader); + } catch (Throwable encodingFailure) { + return rethrowEncodingFailureAfterInvocationEnd( + encodingFailure, + cause, + pluginRunner, + executionManager, + requestId, + executionArn, + isFirstInvocation, + executionInput); + } + } catch (Throwable encodingFailure) { + return rethrowEncodingFailureAfterInvocationEnd( + encodingFailure, + cause, + pluginRunner, + executionManager, + requestId, + executionArn, + isFirstInvocation, + executionInput); + } + fireOnInvocationEnd( + pluginRunner, + executionManager, + requestId, + executionArn, + isFirstInvocation, + InvocationStatus.FAILED, + reportedCause, + executionInput, + null); + var errorObject = frameChainedInvokeOutput ? frameChainedInvokeError(encodedError) : encodedError.error(); + return DurableExecutionOutput.failure(errorObject); + } + + private static T rethrowEncodingFailureAfterInvocationEnd( + Throwable encodingFailure, + Throwable originalCause, + PluginRunner pluginRunner, + ExecutionManager executionManager, + String requestId, + String executionArn, + boolean isFirstInvocation, + Object executionInput) { + fireOnInvocationEnd( + pluginRunner, + executionManager, + requestId, + executionArn, + isFirstInvocation, + InvocationStatus.FAILED, + originalCause, + executionInput, + null); + ExceptionHelper.sneakyThrow(encodingFailure); + return null; + } + private static void fireOnInvocationEnd( PluginRunner pluginRunner, ExecutionManager executionManager, @@ -254,25 +390,115 @@ private static String handleLargePayload(ExecutionManager executionManager, Stri return outputPayload; } - private static ErrorObject buildErrorObject(Throwable e, SerDes serDes) { + private static EncodedError buildErrorObject( + Throwable e, ExecutionManager executionManager, DurableConfig config, PayloadOffloader outputOffloader) { // exceptions thrown from operations, e.g. Step if (e instanceof DurableOperationException durableOperationException) { - return durableOperationException.getErrorObject(); + var error = durableOperationException.getErrorObject(); + if (error == null || error.errorData() == null) { + return new EncodedError(error, false); + } + var targetContext = executionContext(executionManager, SerDesPayloadKind.EXCEPTION); + final String errorData; + final boolean usesPayloadCodec; + if (durableOperationException.getPayloadOffloadContext() == null) { + errorData = executionManager + .getPayloadCodec() + .offloadSerializedPayload(error.errorData(), outputOffloader, targetContext); + usesPayloadCodec = hasActivePayloadOffloader(outputOffloader); + } else { + errorData = executionManager + .getPayloadCodec() + .rebindSerializedPayload( + error.errorData(), + durableOperationException.getPayloadOffloader(), + durableOperationException.getPayloadOffloadContext(), + outputOffloader, + targetContext); + usesPayloadCodec = PayloadCodec.isOffloadEnvelope(errorData); + } + return new EncodedError(error.toBuilder().errorData(errorData).build(), usesPayloadCodec); } if (e instanceof UnrecoverableDurableExecutionException unrecoverableDurableExecutionException) { - return unrecoverableDurableExecutionException.getErrorObject(); + return new EncodedError(unrecoverableDurableExecutionException.getErrorObject(), false); } // exceptions thrown from non-operation code - return ExceptionHelper.buildErrorObject(e, serDes); + final String errorData; + final boolean usesPayloadCodec; + if (e instanceof PayloadOffloadException) { + errorData = config.getSerDes().serialize(e); + usesPayloadCodec = false; + } else { + errorData = executionManager + .getPayloadCodec() + .serialize( + e, + config.getSerDes(), + outputOffloader, + executionContext(executionManager, SerDesPayloadKind.EXCEPTION)); + usesPayloadCodec = PayloadCodec.isOffloadEnvelope(errorData); + } + return new EncodedError( + ErrorObject.builder() + .errorType(e.getClass().getName()) + .errorMessage(e.getMessage()) + .errorData(errorData) + .stackTrace(ExceptionHelper.serializeStackTrace(e.getStackTrace())) + .build(), + usesPayloadCodec); + } + + private static ErrorObject frameChainedInvokeError(EncodedError encodedError) { + var error = encodedError.error(); + if (error == null || error.errorData() == null) { + return error; + } + return error.toBuilder() + .errorData(ChainedInvokeOutputFrame.encode(error.errorData(), encodedError.usesPayloadCodec())) + .build(); + } + + private static boolean shouldFrameChainedInvokeOutput(ExecutionManager executionManager, DurableConfig config) { + var details = executionManager.getExecutionOperation().executionDetails(); + return config.shouldUsePayloadOffloaderForChainedInvokePayloads() + && details != null + && ChainedInvokePayloadFrame.isFramed(details.inputPayload()); + } + + private static boolean hasActivePayloadOffloader(PayloadOffloader offloader) { + return offloader != null && !PayloadOffloaders.isDisabled(offloader); } - private static I extractUserInput(Operation executionOp, SerDes serDes, TypeToken inputType) { + private record EncodedError(ErrorObject error, boolean usesPayloadCodec) {} + + private static I extractUserInput( + ExecutionManager executionManager, DurableConfig config, TypeToken inputType) { + var executionOp = executionManager.getExecutionOperation(); if (executionOp.executionDetails() == null) { throw new IllegalDurableOperationException("EXECUTION operation missing executionDetails"); } var inputPayload = executionOp.executionDetails().inputPayload(); - return serDes.deserialize(inputPayload, inputType); + if (!config.shouldUsePayloadOffloaderForChainedInvokePayloads() + || !ChainedInvokePayloadFrame.isFramed(inputPayload)) { + return config.getSerDes().deserialize(inputPayload, inputType); + } + inputPayload = ChainedInvokePayloadFrame.decode(inputPayload); + return executionManager + .getPayloadCodec() + .deserialize( + inputPayload, + inputType, + config.getSerDes(), + config.getPayloadOffloader(), + executionContext(executionManager, SerDesPayloadKind.INPUT)); + } + + private static PayloadOffloadContext executionContext( + ExecutionManager executionManager, SerDesPayloadKind payloadKind) { + var operation = executionManager.getExecutionOperation(); + return PayloadOffloadContext.forExecution( + executionManager.getDurableExecutionArn(), operation.id(), operation.name(), payloadKind); } /** diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java index 0e9d8426e..557db7612 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutionManager.java @@ -63,6 +63,7 @@ public class ExecutionManager implements SafeCloseable { private final Context lambdaContext; private final AtomicReference executionMode; private final DurableConfig durableConfig; + private final PayloadCodec payloadCodec; private final Set updatedOperationIdsSinceLastInvocation; private final Set initialOperationIds; @@ -79,6 +80,8 @@ public class ExecutionManager implements SafeCloseable { public ExecutionManager(DurableExecutionInput input, DurableConfig config, Context lambdaContext) { durableConfig = config; + payloadCodec = + new PayloadCodec(config.getPayloadOffloadExecutorService(), () -> currentThreadContext.get() != null); this.durableExecutionArn = input.durableExecutionArn(); this.lambdaContext = lambdaContext; @@ -267,6 +270,11 @@ public Operation getExecutionOperation() { return executionOp; } + /** Returns the invocation-scoped payload serialization and offload pipeline. */ + public PayloadCodec getPayloadCodec() { + return payloadCodec; + } + /** * Checks whether there are any cached operations for the given parent context ID. Used to initialize per-context * replay state — a context starts in replay mode if the ExecutionManager has cached operations belonging to it. @@ -412,6 +420,7 @@ public void close() { validateRunningThreads(); checkpointManager.shutdown(); + payloadCodec.clear(); } private void validateRunningThreads() { @@ -464,6 +473,13 @@ public void terminateExecution(UnrecoverableDurableExecutionException exception) throw exception; } + /** Fails the current invocation without converting an infrastructure failure into an operation outcome. */ + public void failInvocation(RuntimeException exception) { + stopAllOperations(exception); + executionExceptionFuture.completeExceptionally(exception); + throw exception; + } + /** Suspends the execution by completing the execution exception future with a {@link SuspendExecutionException}. */ public void suspendExecution() { throw signalSuspension(); diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/PayloadCodec.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/PayloadCodec.java new file mode 100644 index 000000000..0b57708a8 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/PayloadCodec.java @@ -0,0 +1,513 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.databind.DeserializationFeature; +import com.fasterxml.jackson.databind.MapperFeature; +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.json.JsonMapper; +import java.lang.ref.WeakReference; +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.Collections; +import java.util.HexFormat; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Objects; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; +import java.util.function.BooleanSupplier; +import java.util.function.Supplier; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.PayloadOffloadException; +import software.amazon.lambda.durable.exception.SerDesException; +import software.amazon.lambda.durable.offload.OffloadedPayload; +import software.amazon.lambda.durable.offload.PayloadOffloadContext; +import software.amazon.lambda.durable.offload.PayloadOffloader; +import software.amazon.lambda.durable.offload.PayloadOffloaders; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; +import software.amazon.lambda.durable.offload.internal.PayloadOffloadTracking; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ExceptionHelper; + +/** + * Invocation-scoped pipeline that composes object serialization, payload offloading, loading, and bounded caching. + * + *

Legacy raw serialized strings remain readable. New offloaded values use a reserved, versioned prefix so arbitrary + * JSON payloads cannot be mistaken for SDK envelopes. + */ +public final class PayloadCodec { + static final int MAX_COMPLETED_CACHE_ENTRIES = 256; + private static final String ENVELOPE_MARKER = "@aws-durable-payload:"; + private static final String ENVELOPE_PREFIX = "@aws-durable-payload:v1:"; + private static final Object CACHE_MISS = new Object(); + private static final Object NULL_VALUE = new Object(); + private static final TypeToken OFFLOADED_PAYLOAD_TYPE = TypeToken.get(OffloadedPayload.class); + private static final ObjectMapper ENVELOPE_OBJECT_MAPPER = JsonMapper.builder() + .enable(DeserializationFeature.FAIL_ON_TRAILING_TOKENS) + .enable(DeserializationFeature.FAIL_ON_NUMBERS_FOR_ENUMS) + .disable(MapperFeature.ALLOW_COERCION_OF_SCALARS) + .disable(DeserializationFeature.ACCEPT_FLOAT_AS_INT) + .disable(DeserializationFeature.FAIL_ON_UNKNOWN_PROPERTIES) + .build(); + private static final SerDes ENVELOPE_SER_DES = new JacksonSerDes(ENVELOPE_OBJECT_MAPPER); + + private final ExecutorService executorService; + private final BooleanSupplier runInlineOnCurrentThread; + private final Map> inFlightLoads = new ConcurrentHashMap<>(); + private final Map> inFlightDeserializations = + new ConcurrentHashMap<>(); + private final BoundedWeakCache loadedPayloadCache = new BoundedWeakCache<>(); + private final BoundedWeakCache deserializedPayloadCache = new BoundedWeakCache<>(); + + /** + * Creates an invocation-scoped codec. + * + * @param executorService executor for blocking offloader calls, or null to execute inline + */ + public PayloadCodec(ExecutorService executorService) { + this(executorService, () -> false); + } + + PayloadCodec(ExecutorService executorService, BooleanSupplier runInlineOnCurrentThread) { + this.executorService = executorService; + this.runInlineOnCurrentThread = + Objects.requireNonNull(runInlineOnCurrentThread, "runInlineOnCurrentThread cannot be null"); + } + + /** Returns whether a checkpoint value uses the SDK payload offload envelope. */ + public static boolean isOffloadEnvelope(String checkpointPayload) { + return checkpointPayload != null && checkpointPayload.startsWith(ENVELOPE_MARKER); + } + + /** Returns the UTF-8 byte size of an SDK payload envelope. */ + public static int envelopeSizeBytes(OffloadedPayload payload) { + return encodeEnvelope(payload).getBytes(StandardCharsets.UTF_8).length; + } + + /** Serializes and optionally offloads a value. */ + public String serialize(Object value, SerDes serDes, PayloadOffloader offloader, PayloadOffloadContext context) { + var serialized = serDes.serialize(value); + return offloadSerialized(serialized, value, offloader, context, true); + } + + /** Rebinds already serialized payload data from a producer operation to a forwarding operation's policy. */ + public String rebindSerializedPayload( + String checkpointPayload, + PayloadOffloader sourceOffloader, + PayloadOffloadContext sourceContext, + PayloadOffloader targetOffloader, + PayloadOffloadContext targetContext) { + var serialized = resolve(checkpointPayload, sourceOffloader, sourceContext); + return offloadSerialized(serialized, null, targetOffloader, targetContext, true); + } + + /** Encodes already serialized SDK payload data, escaping the reserved envelope marker when necessary. */ + public String serializePreEncodedPayload( + String serializedPayload, PayloadOffloader offloader, PayloadOffloadContext context) { + return offloadSerialized(serializedPayload, null, offloader, context, true); + } + + /** Offloads already serialized data without interpreting it as an existing SDK payload envelope. */ + public String offloadSerializedPayload( + String serializedPayload, PayloadOffloader offloader, PayloadOffloadContext context) { + return offloadSerialized(serializedPayload, null, offloader, context, false); + } + + private String offloadSerialized( + String serialized, + Object originalValue, + PayloadOffloader offloader, + PayloadOffloadContext context, + boolean escapeReservedMarker) { + var effectiveOffloader = effectiveOffloader(offloader); + if (serialized == null) { + return serialized; + } + if (effectiveOffloader == null) { + if (!escapeReservedMarker || !serialized.startsWith(ENVELOPE_MARKER)) { + return serialized; + } + if (offloader != null) { + PayloadOffloadTracking.record(context, offloader); + } + return encodeEnvelope(OffloadedPayload.inline(serialized).bindProducer(context, hash(serialized), false)); + } + + PayloadOffloadTracking.record(context, effectiveOffloader); + var serializationContext = context.withOriginalValue(originalValue); + var payload = runOffloadTask( + () -> effectiveOffloader.offload(serialized, serializationContext), "store", serializationContext); + if (payload == null) { + throw new PayloadOffloadException("Payload offloader returned null for " + describe(context)); + } + try { + payload = payload.bindProducer(context, hash(serialized), true); + } catch (IllegalArgumentException e) { + throw new PayloadOffloadException( + "Payload offloader returned inconsistent metadata for " + describe(context), e); + } + validateOwner(payload, context); + verifyDigest(payload, serialized, context); + return encodeEnvelope(payload); + } + + /** Deserializes a raw legacy payload or an SDK offload envelope. */ + @SuppressWarnings("unchecked") + public T deserialize( + String checkpointPayload, + TypeToken typeToken, + SerDes serDes, + PayloadOffloader offloader, + PayloadOffloadContext context) { + Objects.requireNonNull(typeToken, "typeToken cannot be null"); + Objects.requireNonNull(serDes, "serDes cannot be null"); + var serialized = resolve(checkpointPayload, offloader, context); + var key = new DeserializedCacheKey( + serDes, + context.durableExecutionArn(), + context.entityId(), + context.payloadKind(), + context.attempt(), + typeToken.getType().getTypeName(), + hash(serialized)); + var cached = deserializedPayloadCache.get(key); + if (cached != CACHE_MISS) { + return (T) unmaskNull(cached); + } + + return (T) unmaskNull(loadOnce( + key, + inFlightDeserializations, + deserializedPayloadCache, + () -> maskNull(serDes.deserialize(serialized, typeToken)))); + } + + /** Resolves an SDK envelope to the serialized text produced by SerDes without deserializing the object. */ + public String resolveSerializedPayload( + String checkpointPayload, PayloadOffloader offloader, PayloadOffloadContext context) { + return resolve(checkpointPayload, offloader, context); + } + + /** Validates that a value is a supported, well-formed SDK envelope owned by the supplied payload context. */ + public void validateEnvelope(String checkpointPayload, PayloadOffloadContext context) { + if (checkpointPayload == null || !checkpointPayload.startsWith(ENVELOPE_MARKER)) { + throw new PayloadOffloadException("Expected payload offload envelope for " + describe(context)); + } + if (!checkpointPayload.startsWith(ENVELOPE_PREFIX)) { + throw new PayloadOffloadException( + "Unsupported or malformed payload offload envelope version for " + describe(context)); + } + validateOwner(decodeEnvelope(checkpointPayload, context), context); + } + + /** Resolves an envelope using its embedded producer context when no consuming execution context is available. */ + public String resolveSerializedPayloadUsingProducerContext(String checkpointPayload, PayloadOffloader offloader) { + if (checkpointPayload == null || !checkpointPayload.startsWith(ENVELOPE_MARKER)) { + return checkpointPayload; + } + if (!checkpointPayload.startsWith(ENVELOPE_PREFIX)) { + throw new PayloadOffloadException("Unsupported or malformed payload offload envelope version"); + } + var payload = decodeEnvelope(checkpointPayload, null); + if (payload.producerContext() == null) { + throw new PayloadOffloadException("Payload envelope is missing producer context"); + } + return resolve(checkpointPayload, offloader, payload.producerContext()); + } + + /** Clears invocation-scoped payload caches. */ + public void clear() { + inFlightLoads.clear(); + inFlightDeserializations.clear(); + loadedPayloadCache.clear(); + deserializedPayloadCache.clear(); + } + + private String resolve(String checkpointPayload, PayloadOffloader offloader, PayloadOffloadContext context) { + if (checkpointPayload == null || !checkpointPayload.startsWith(ENVELOPE_MARKER)) { + return checkpointPayload; + } + if (!checkpointPayload.startsWith(ENVELOPE_PREFIX)) { + throw new PayloadOffloadException( + "Unsupported or malformed payload offload envelope version for " + describe(context)); + } + + var payload = decodeEnvelope(checkpointPayload, context); + validateOwner(payload, context); + var effectiveOffloader = effectiveOffloader(offloader); + var loadOffloader = payload.requiresLoad() ? effectiveOffloader : null; + var key = payloadCacheKey(context, checkpointPayload, loadOffloader); + var cached = loadedPayloadCache.get(key); + if (cached != CACHE_MISS) { + return (String) cached; + } + + return (String) loadOnce(key, inFlightLoads, loadedPayloadCache, () -> { + final String serialized; + if (payload.requiresLoad()) { + if (loadOffloader == null) { + throw new PayloadOffloadException( + "Payload requires its producing offloader but no payload offloader is configured for " + + describe(context)); + } + var loadContext = payload.producerContext() != null ? payload.producerContext() : context; + serialized = runOffloadTask(() -> loadOffloader.load(payload, loadContext), "load", context); + } else { + serialized = payload.data(); + } + if (serialized == null) { + throw new PayloadOffloadException("Payload offloader returned null while loading " + describe(context)); + } + verifyDigest(payload, serialized, context); + return serialized; + }); + } + + private static OffloadedPayload decodeEnvelope(String checkpointPayload, PayloadOffloadContext context) { + try { + var envelopeJson = checkpointPayload.substring(ENVELOPE_PREFIX.length()); + var envelopeNode = ENVELOPE_OBJECT_MAPPER.readTree(envelopeJson); + if (envelopeNode == null + || !envelopeNode.isObject() + || !envelopeNode.path("mode").isTextual() + || !envelopeNode.path("requiresLoad").isBoolean()) { + throw new PayloadOffloadException( + "Payload envelope has invalid scalar field types for " + describe(context)); + } + var payload = ENVELOPE_SER_DES.deserialize(envelopeJson, OFFLOADED_PAYLOAD_TYPE); + if (payload == null) { + throw new PayloadOffloadException("Payload envelope decoded to null for " + describe(context)); + } + return payload; + } catch (JsonProcessingException | SerDesException e) { + throw new PayloadOffloadException("Invalid payload offload envelope for " + describe(context), e); + } + } + + private static void validateOwner(OffloadedPayload payload, PayloadOffloadContext context) { + if (!payload.hasIntegrityMetadata() || payload.producerContext() == null) { + throw new PayloadOffloadException( + "Payload envelope is missing producer ownership or integrity metadata for " + describe(context)); + } + var sameOwner = payload.ownerDurableExecutionArn().equals(context.durableExecutionArn()) + && payload.ownerEntityId().equals(context.entityId()); + if (!sameOwner + && context.payloadKind() != SerDesPayloadKind.INPUT + && context.operationType() != OperationType.CHAINED_INVOKE) { + throw new PayloadOffloadException("Payload belongs to a different durable entity"); + } + } + + private static void verifyDigest(OffloadedPayload payload, String serialized, PayloadOffloadContext context) { + if (payload.hasIntegrityMetadata() && !hash(serialized).equals(payload.payloadDigest())) { + throw new PayloadOffloadException("Payload digest does not match stored content for " + describe(context)); + } + } + + private T runOffloadTask(Supplier task, String action, PayloadOffloadContext context) { + try { + if (executorService == null || runInlineOnCurrentThread.getAsBoolean()) { + return task.get(); + } + return CompletableFuture.supplyAsync(task, executorService).join(); + } catch (Throwable throwable) { + var cause = ExceptionHelper.unwrapCompletableFuture(throwable); + if (cause instanceof Error error) { + throw error; + } + if (cause instanceof PayloadOffloadException payloadOffloadException) { + throw payloadOffloadException; + } + throw new PayloadOffloadException("Failed to " + action + " payload for " + describe(context), cause); + } + } + + private static Object loadOnce( + K key, Map> inFlight, BoundedWeakCache completed, Supplier loader) { + var cached = completed.get(key); + if (cached != CACHE_MISS) { + return cached; + } + + var pending = new CompletableFuture(); + var existing = inFlight.putIfAbsent(key, pending); + if (existing != null) { + return join(existing); + } + + try { + cached = completed.get(key); + if (cached != CACHE_MISS) { + pending.complete(cached); + return cached; + } + var value = loader.get(); + completed.put(key, value); + pending.complete(value); + return value; + } catch (Throwable failure) { + pending.completeExceptionally(failure); + ExceptionHelper.sneakyThrow(failure); + return null; + } finally { + inFlight.remove(key, pending); + } + } + + private static Object join(CompletableFuture future) { + try { + return future.join(); + } catch (Throwable failure) { + ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(failure)); + return null; + } + } + + private static Object maskNull(Object value) { + return value == null ? NULL_VALUE : value; + } + + private static Object unmaskNull(Object value) { + return value == NULL_VALUE ? null : value; + } + + private static PayloadOffloader effectiveOffloader(PayloadOffloader offloader) { + return offloader == null || PayloadOffloaders.isDisabled(offloader) ? null : offloader; + } + + private static String encodeEnvelope(OffloadedPayload payload) { + Objects.requireNonNull(payload, "payload cannot be null"); + try { + return ENVELOPE_PREFIX + ENVELOPE_SER_DES.serialize(payload); + } catch (SerDesException e) { + throw new PayloadOffloadException("Failed to encode payload offload envelope", e); + } + } + + private static PayloadCacheKey payloadCacheKey( + PayloadOffloadContext context, String checkpointPayload, PayloadOffloader offloader) { + return new PayloadCacheKey( + offloader, + context.durableExecutionArn(), + context.entityId(), + context.payloadKind(), + context.attempt(), + hash(checkpointPayload)); + } + + private static String describe(PayloadOffloadContext context) { + return context == null ? "payload" : context.payloadKind() + " payload '" + context.entityId() + "'"; + } + + private static String hash(String value) { + if (value == null) { + return "null"; + } + try { + var digest = MessageDigest.getInstance("SHA-256"); + return HexFormat.of().formatHex(digest.digest(value.getBytes(StandardCharsets.UTF_8))); + } catch (NoSuchAlgorithmException e) { + throw new IllegalStateException("SHA-256 is unavailable", e); + } + } + + private static final class BoundedWeakCache { + private final Map> values = + Collections.synchronizedMap(new LinkedHashMap<>(16, 0.75f, true) { + @Override + protected boolean removeEldestEntry(Map.Entry> eldest) { + return size() > MAX_COMPLETED_CACHE_ENTRIES; + } + }); + + private Object get(K key) { + synchronized (values) { + var reference = values.get(key); + if (reference == null) { + return CACHE_MISS; + } + var value = reference.get(); + if (value == null) { + values.remove(key); + return CACHE_MISS; + } + return value; + } + } + + private void put(K key, Object value) { + values.put(key, new WeakReference<>(value)); + } + + private void clear() { + values.clear(); + } + } + + private record PayloadCacheKey( + PayloadOffloader offloader, + String durableExecutionArn, + String entityId, + SerDesPayloadKind payloadKind, + Integer attempt, + String checkpointPayloadHash) { + @Override + public boolean equals(Object other) { + return other instanceof PayloadCacheKey that + && offloader == that.offloader + && Objects.equals(durableExecutionArn, that.durableExecutionArn) + && Objects.equals(entityId, that.entityId) + && payloadKind == that.payloadKind + && Objects.equals(attempt, that.attempt) + && Objects.equals(checkpointPayloadHash, that.checkpointPayloadHash); + } + + @Override + public int hashCode() { + int result = System.identityHashCode(offloader); + result = 31 * result + Objects.hashCode(durableExecutionArn); + result = 31 * result + Objects.hashCode(entityId); + result = 31 * result + Objects.hashCode(payloadKind); + result = 31 * result + Objects.hashCode(attempt); + return 31 * result + Objects.hashCode(checkpointPayloadHash); + } + } + + private record DeserializedCacheKey( + SerDes serDes, + String durableExecutionArn, + String entityId, + SerDesPayloadKind payloadKind, + Integer attempt, + String targetType, + String serializedPayloadHash) { + @Override + public boolean equals(Object other) { + return other instanceof DeserializedCacheKey that + && serDes == that.serDes + && Objects.equals(durableExecutionArn, that.durableExecutionArn) + && Objects.equals(entityId, that.entityId) + && payloadKind == that.payloadKind + && Objects.equals(attempt, that.attempt) + && Objects.equals(targetType, that.targetType) + && Objects.equals(serializedPayloadHash, that.serializedPayloadHash); + } + + @Override + public int hashCode() { + int result = System.identityHashCode(serDes); + result = 31 * result + Objects.hashCode(durableExecutionArn); + result = 31 * result + Objects.hashCode(entityId); + result = 31 * result + Objects.hashCode(payloadKind); + result = 31 * result + Objects.hashCode(attempt); + result = 31 * result + Objects.hashCode(targetType); + return 31 * result + Objects.hashCode(serializedPayloadHash); + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/model/DurableExecutionInput.java b/sdk/src/main/java/software/amazon/lambda/durable/model/DurableExecutionInput.java index f6eaea9e6..ab6d49b04 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/model/DurableExecutionInput.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/model/DurableExecutionInput.java @@ -13,12 +13,27 @@ * @param initialExecutionState snapshot of operations already completed in previous invocations * @param updatedOperationIds IDs of operations that changed since the previous successful invocation; empty list if * nothing changed + * @param invocationSource whether this execution was invoked directly or by a chained invoke */ public record DurableExecutionInput( String durableExecutionArn, String checkpointToken, CheckpointUpdatedExecutionState initialExecutionState, - List updatedOperationIds) { + List updatedOperationIds, + InvocationSource invocationSource) { + + public DurableExecutionInput { + invocationSource = invocationSource == null ? InvocationSource.DIRECT : invocationSource; + } + + /** Constructor that defaults invocation source to direct execution. */ + public DurableExecutionInput( + String durableExecutionArn, + String checkpointToken, + CheckpointUpdatedExecutionState initialExecutionState, + List updatedOperationIds) { + this(durableExecutionArn, checkpointToken, initialExecutionState, updatedOperationIds, InvocationSource.DIRECT); + } /** * Constructor that defaults updatedOperationIds to empty list. Used by tests that don't need to supply updated @@ -26,6 +41,6 @@ public record DurableExecutionInput( */ public DurableExecutionInput( String durableExecutionArn, String checkpointToken, CheckpointUpdatedExecutionState initialExecutionState) { - this(durableExecutionArn, checkpointToken, initialExecutionState, List.of()); + this(durableExecutionArn, checkpointToken, initialExecutionState, List.of(), InvocationSource.DIRECT); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/model/InvocationSource.java b/sdk/src/main/java/software/amazon/lambda/durable/model/InvocationSource.java new file mode 100644 index 000000000..4030020b4 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/model/InvocationSource.java @@ -0,0 +1,9 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.model; + +/** Identifies how a durable execution invocation was started. */ +public enum InvocationSource { + DIRECT, + CHAINED_INVOKE +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/offload/OffloadedPayload.java b/sdk/src/main/java/software/amazon/lambda/durable/offload/OffloadedPayload.java new file mode 100644 index 000000000..571c71a2d --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/offload/OffloadedPayload.java @@ -0,0 +1,263 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload; + +import java.lang.reflect.Array; +import java.math.BigDecimal; +import java.math.BigInteger; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.IdentityHashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; + +/** + * SDK-owned representation of inline or externally stored serialized payload data. + * + * @param mode payload storage mode + * @param data inline serialized data, when {@code mode} is {@link PayloadStorageMode#INLINE} + * @param reference external storage reference, when {@code mode} is {@link PayloadStorageMode#REFERENCE} + * @param preview optional inline preview metadata + * @param ownerDurableExecutionArn producing durable execution ARN, when integrity metadata is present + * @param ownerEntityId producing payload entity, when integrity metadata is present + * @param payloadDigest lowercase SHA-256 digest of the serialized payload, when integrity metadata is present + * @param producerContext exact producing payload context, when the SDK bound the envelope + * @param requiresLoad whether the producing offloader must restore the serialized payload + */ +public record OffloadedPayload( + PayloadStorageMode mode, + String data, + String reference, + Map preview, + String ownerDurableExecutionArn, + String ownerEntityId, + String payloadDigest, + PayloadOffloadContext producerContext, + Boolean requiresLoad) { + + public OffloadedPayload { + Objects.requireNonNull(mode, "mode cannot be null"); + Objects.requireNonNull(requiresLoad, "requiresLoad cannot be null"); + preview = preview == null ? null : immutablePreview(preview); + if (mode == PayloadStorageMode.INLINE) { + Objects.requireNonNull(data, "data cannot be null for an inline payload"); + if (reference != null) { + throw new IllegalArgumentException("reference must be null for an inline payload"); + } + } else { + if (reference == null || reference.isBlank()) { + throw new IllegalArgumentException("reference cannot be blank for an externally stored payload"); + } + if (data != null) { + throw new IllegalArgumentException("data must be null for an externally stored payload"); + } + if (!requiresLoad) { + throw new IllegalArgumentException("externally stored payloads must require load"); + } + } + var metadataCount = (ownerDurableExecutionArn == null ? 0 : 1) + + (ownerEntityId == null ? 0 : 1) + + (payloadDigest == null ? 0 : 1); + if (metadataCount != 0 && metadataCount != 3) { + throw new IllegalArgumentException("payload owner and digest metadata must be provided together"); + } + if (metadataCount == 3) { + if (ownerDurableExecutionArn.isBlank()) { + throw new IllegalArgumentException("ownerDurableExecutionArn cannot be blank"); + } + if (ownerEntityId.isBlank()) { + throw new IllegalArgumentException("ownerEntityId cannot be blank"); + } + if (!payloadDigest.matches("[0-9a-f]{64}")) { + throw new IllegalArgumentException("payloadDigest must be lowercase SHA-256 hex"); + } + } + if (producerContext != null) { + if (producerContext.originalValue() != null) { + throw new IllegalArgumentException("producerContext.originalValue must be null"); + } + if (metadataCount != 3 + || !producerContext.durableExecutionArn().equals(ownerDurableExecutionArn) + || !producerContext.entityId().equals(ownerEntityId)) { + throw new IllegalArgumentException("producerContext must match payload owner metadata"); + } + } + } + + /** Creates a payload with the legacy constructor shape. */ + public OffloadedPayload( + PayloadStorageMode mode, + String data, + String reference, + Map preview, + String ownerDurableExecutionArn, + String ownerEntityId, + String payloadDigest, + PayloadOffloadContext producerContext) { + this( + mode, + data, + reference, + preview, + ownerDurableExecutionArn, + ownerEntityId, + payloadDigest, + producerContext, + true); + } + + /** Creates an inline payload. */ + public static OffloadedPayload inline(String data) { + return new OffloadedPayload(PayloadStorageMode.INLINE, data, null, null, null, null, null, null, true); + } + + /** Creates an integrity-bound inline payload. */ + public static OffloadedPayload inline( + String data, String ownerDurableExecutionArn, String ownerEntityId, String payloadDigest) { + return new OffloadedPayload( + PayloadStorageMode.INLINE, + data, + null, + null, + ownerDurableExecutionArn, + ownerEntityId, + payloadDigest, + null, + true); + } + + /** Creates an externally stored payload. */ + public static OffloadedPayload reference(String reference, Map preview) { + return new OffloadedPayload( + PayloadStorageMode.REFERENCE, null, reference, preview, null, null, null, null, true); + } + + /** Creates an integrity-bound externally stored payload. */ + public static OffloadedPayload reference( + String reference, + Map preview, + String ownerDurableExecutionArn, + String ownerEntityId, + String payloadDigest) { + return new OffloadedPayload( + PayloadStorageMode.REFERENCE, + null, + reference, + preview, + ownerDurableExecutionArn, + ownerEntityId, + payloadDigest, + null, + true); + } + + /** Returns whether producer ownership and content integrity metadata is present. */ + public boolean hasIntegrityMetadata() { + return payloadDigest != null; + } + + /** Returns this payload bound to the exact SDK context that produced the serialized content. */ + public OffloadedPayload bindProducer(PayloadOffloadContext context, String digest) { + return bindProducer(context, digest, requiresLoad); + } + + /** Returns this payload bound to the exact SDK context and load semantics that produced the serialized content. */ + public OffloadedPayload bindProducer(PayloadOffloadContext context, String digest, boolean requiresLoad) { + Objects.requireNonNull(context, "context cannot be null"); + Objects.requireNonNull(digest, "digest cannot be null"); + var cleanContext = context.withOriginalValue(null); + if (hasIntegrityMetadata() + && (!ownerDurableExecutionArn.equals(cleanContext.durableExecutionArn()) + || !ownerEntityId.equals(cleanContext.entityId()) + || !payloadDigest.equals(digest) + || this.requiresLoad != requiresLoad)) { + throw new IllegalArgumentException( + "payload owner, digest, or load semantics do not match the producing context"); + } + if (producerContext != null && !producerContext.equals(cleanContext)) { + throw new IllegalArgumentException("producerContext does not match the producing context"); + } + return new OffloadedPayload( + mode, + data, + reference, + preview, + cleanContext.durableExecutionArn(), + cleanContext.entityId(), + digest, + cleanContext, + requiresLoad); + } + + @SuppressWarnings("unchecked") + private static Map immutablePreview(Map preview) { + return (Map) immutableValue(preview, new IdentityHashMap<>()); + } + + private static Object immutableValue(Object value, IdentityHashMap copies) { + if (value == null) { + return null; + } + var existing = copies.get(value); + if (existing != null) { + return existing; + } + if (value instanceof Map map) { + var target = new LinkedHashMap(); + var immutable = Collections.unmodifiableMap(target); + copies.put(value, immutable); + for (var entry : map.entrySet()) { + var key = Objects.requireNonNull(entry.getKey(), "preview map keys cannot be null"); + if (!(key instanceof String stringKey)) { + throw new IllegalArgumentException("preview map keys must be strings"); + } + target.put(stringKey, immutableValue(entry.getValue(), copies)); + } + return immutable; + } + if (value instanceof Collection collection) { + var target = new ArrayList<>(); + List immutable = Collections.unmodifiableList(target); + copies.put(value, immutable); + for (var item : collection) { + target.add(immutableValue(item, copies)); + } + return immutable; + } + if (value.getClass().isArray()) { + var target = new ArrayList<>(); + List immutable = Collections.unmodifiableList(target); + copies.put(value, immutable); + for (int index = 0; index < Array.getLength(value); index++) { + target.add(immutableValue(Array.get(value, index), copies)); + } + return immutable; + } + if (value instanceof CharSequence sequence) { + return sequence.toString(); + } + if (value instanceof Character character) { + return character.toString(); + } + if (value instanceof Boolean + || value instanceof Byte + || value instanceof Short + || value instanceof Integer + || value instanceof Long + || value instanceof Float + || value instanceof Double + || value instanceof BigInteger + || value instanceof BigDecimal) { + return value; + } + if (value instanceof Enum enumValue) { + return enumValue.name(); + } + throw new IllegalArgumentException( + "preview values must be JSON-compatible maps, collections, arrays, or scalars: " + + value.getClass().getName()); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloadContext.java b/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloadContext.java new file mode 100644 index 000000000..cc7711913 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloadContext.java @@ -0,0 +1,101 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload; + +import java.util.Objects; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.model.OperationSubType; + +/** + * Stable identity and operation metadata for a payload being stored or loaded. + * + * @param durableExecutionArn durable execution ARN + * @param entityId stable identifier for this payload within the durable execution + * @param payloadKind role of the payload + * @param operationId operation identifier, or the execution operation identifier for root payloads + * @param operationName operation name, when present + * @param parentId parent context identifier, when present + * @param operationType durable operation type + * @param operationSubType durable operation subtype, or null for root execution payloads + * @param attempt current attempt for retryable operations, when available + * @param originalValue original object supplied to SerDes during serialization, or null during loading + */ +public record PayloadOffloadContext( + String durableExecutionArn, + String entityId, + SerDesPayloadKind payloadKind, + String operationId, + String operationName, + String parentId, + OperationType operationType, + OperationSubType operationSubType, + Integer attempt, + Object originalValue) { + + public PayloadOffloadContext { + Objects.requireNonNull(durableExecutionArn, "durableExecutionArn cannot be null"); + Objects.requireNonNull(entityId, "entityId cannot be null"); + Objects.requireNonNull(payloadKind, "payloadKind cannot be null"); + Objects.requireNonNull(operationId, "operationId cannot be null"); + Objects.requireNonNull(operationType, "operationType cannot be null"); + } + + /** Creates context for a root execution input, output, or exception payload. */ + public static PayloadOffloadContext forExecution( + String durableExecutionArn, String executionOperationId, String executionName, SerDesPayloadKind kind) { + return new PayloadOffloadContext( + durableExecutionArn, + "execution/" + executionOperationId + "/" + kind.entitySuffix(), + kind, + executionOperationId, + executionName, + null, + OperationType.EXECUTION, + null, + null, + null); + } + + /** Creates context for a durable operation payload. */ + public static PayloadOffloadContext forOperation( + String durableExecutionArn, + OperationIdentifier operation, + String parentId, + SerDesPayloadKind kind, + Integer attempt) { + var entityId = "operation/" + operation.operationId() + "/" + kind.entitySuffix(); + if (attempt != null) { + entityId += "/attempt-" + attempt; + } + return new PayloadOffloadContext( + durableExecutionArn, + entityId, + kind, + operation.operationId(), + operation.name(), + parentId, + operation.operationType(), + operation.subType(), + attempt, + null); + } + + /** Returns this context with the original value available to preview generators. */ + public PayloadOffloadContext withOriginalValue(Object originalValue) { + if (this.originalValue == originalValue) { + return this; + } + return new PayloadOffloadContext( + durableExecutionArn, + entityId, + payloadKind, + operationId, + operationName, + parentId, + operationType, + operationSubType, + attempt, + originalValue); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloader.java b/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloader.java new file mode 100644 index 000000000..67e0476fd --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloader.java @@ -0,0 +1,26 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload; + +/** + * Stores and loads serialized durable execution payloads. + * + *

Implementations receive serialized text after {@code SerDes} processing. They may keep it inline or replace it + * with a reference to external storage. A returned reference must keep the same meaning for the lifetime of every + * checkpoint that contains it. Implementations that overwrite storage should include + * {@link PayloadOffloadContext#attempt()} or another immutable version in the storage key when a payload can be + * updated. + */ +public interface PayloadOffloader { + + /** Stores serialized payload data or returns it inline. */ + OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context); + + /** Restores serialized payload data from either an inline value or an external reference. */ + String load(OffloadedPayload payload, PayloadOffloadContext context); + + /** Returns a sentinel that disables a globally configured offloader for a specific operation. */ + static PayloadOffloader disabled() { + return PayloadOffloaders.disabled(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloaders.java b/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloaders.java new file mode 100644 index 000000000..b51db0596 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadOffloaders.java @@ -0,0 +1,37 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload; + +/** Factory methods for built-in payload offloader policies. */ +public final class PayloadOffloaders { + private static final PayloadOffloader DISABLED = new DisabledPayloadOffloader(); + + private PayloadOffloaders() {} + + /** Returns a sentinel that forces payloads to remain in the normal inline checkpoint format. */ + public static PayloadOffloader disabled() { + return DISABLED; + } + + /** Returns whether the supplied offloader is the disabled sentinel. */ + public static boolean isDisabled(PayloadOffloader offloader) { + return offloader == DISABLED; + } + + private static final class DisabledPayloadOffloader implements PayloadOffloader { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.inline(serializedPayload); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return payload.data(); + } + + @Override + public String toString() { + return "PayloadOffloader.disabled()"; + } + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadStorageMode.java b/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadStorageMode.java new file mode 100644 index 000000000..82f798a94 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/offload/PayloadStorageMode.java @@ -0,0 +1,9 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload; + +/** Identifies whether a serialized payload is stored inline or in external storage. */ +public enum PayloadStorageMode { + INLINE, + REFERENCE +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/offload/SerDesPayloadKind.java b/sdk/src/main/java/software/amazon/lambda/durable/offload/SerDesPayloadKind.java new file mode 100644 index 000000000..aa9b56e5c --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/offload/SerDesPayloadKind.java @@ -0,0 +1,23 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload; + +/** Identifies the role of a serialized user payload in a durable execution. */ +public enum SerDesPayloadKind { + INPUT("input"), + OUTPUT("output"), + RESULT("result"), + INVOKE_PAYLOAD("invoke-payload"), + STATE("state"), + EXCEPTION("exception"); + + private final String entitySuffix; + + SerDesPayloadKind(String entitySuffix) { + this.entitySuffix = entitySuffix; + } + + String entitySuffix() { + return entitySuffix; + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/offload/internal/ChainedInvokeOutputFrame.java b/sdk/src/main/java/software/amazon/lambda/durable/offload/internal/ChainedInvokeOutputFrame.java new file mode 100644 index 000000000..8d0269c8d --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/offload/internal/ChainedInvokeOutputFrame.java @@ -0,0 +1,39 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload.internal; + +import software.amazon.lambda.durable.exception.PayloadOffloadException; + +/** Versioned frame identifying results and errors returned by a compatible durable chained-invoke target. */ +public final class ChainedInvokeOutputFrame { + private static final String FRAME_MARKER = "__durable_execution_chained_invoke_output:"; + private static final String FRAME_PREFIX = FRAME_MARKER + "1:"; + private static final String CODEC_PREFIX = FRAME_PREFIX + "codec:"; + private static final String RAW_PREFIX = FRAME_PREFIX + "raw:"; + + private ChainedInvokeOutputFrame() {} + + public static String encode(String payload, boolean usesPayloadCodec) { + if (payload == null) { + return null; + } + return (usesPayloadCodec ? CODEC_PREFIX : RAW_PREFIX) + payload; + } + + public static boolean isFramed(String payload) { + return payload != null && payload.startsWith(FRAME_MARKER); + } + + public static Decoded decode(String payload) { + if (payload != null && payload.startsWith(CODEC_PREFIX)) { + return new Decoded(payload.substring(CODEC_PREFIX.length()), true); + } + if (payload != null && payload.startsWith(RAW_PREFIX)) { + return new Decoded(payload.substring(RAW_PREFIX.length()), false); + } + throw new PayloadOffloadException("Unsupported or malformed chained-invoke output frame"); + } + + /** Decoded chained-invoke output and whether its payload uses the SDK payload codec. */ + public record Decoded(String payload, boolean usesPayloadCodec) {} +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/offload/internal/ChainedInvokePayloadFrame.java b/sdk/src/main/java/software/amazon/lambda/durable/offload/internal/ChainedInvokePayloadFrame.java new file mode 100644 index 000000000..8bd286b9b --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/offload/internal/ChainedInvokePayloadFrame.java @@ -0,0 +1,33 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload.internal; + +import software.amazon.lambda.durable.exception.PayloadOffloadException; + +/** Versioned source frame for chained-invoke payloads produced through the persisted payload offloader. */ +public final class ChainedInvokePayloadFrame { + private static final String FRAME_MARKER = "__durable_execution_chained_invoke_payload:"; + private static final String FRAME_PREFIX = FRAME_MARKER + "1:"; + private static final String NULL_FRAME = FRAME_PREFIX + "null"; + private static final String VALUE_PREFIX = FRAME_PREFIX + "value:"; + + private ChainedInvokePayloadFrame() {} + + public static String encode(String payload) { + return payload == null ? NULL_FRAME : VALUE_PREFIX + payload; + } + + public static boolean isFramed(String payload) { + return payload != null && payload.startsWith(FRAME_MARKER); + } + + public static String decode(String payload) { + if (NULL_FRAME.equals(payload)) { + return null; + } + if (payload != null && payload.startsWith(VALUE_PREFIX)) { + return payload.substring(VALUE_PREFIX.length()); + } + throw new PayloadOffloadException("Unsupported or malformed chained-invoke payload frame"); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/offload/internal/PayloadOffloadTracking.java b/sdk/src/main/java/software/amazon/lambda/durable/offload/internal/PayloadOffloadTracking.java new file mode 100644 index 000000000..cc4045604 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/offload/internal/PayloadOffloadTracking.java @@ -0,0 +1,47 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload.internal; + +import java.util.Objects; +import java.util.concurrent.ConcurrentHashMap; +import java.util.function.BiConsumer; +import software.amazon.lambda.durable.offload.PayloadOffloadContext; +import software.amazon.lambda.durable.offload.PayloadOffloader; + +/** + * Internal hook used by local testing utilities to retain the offloader selected for each persisted payload. + * + *

Applications should not use this class. + */ +public final class PayloadOffloadTracking { + private static final ConcurrentHashMap> OBSERVERS = + new ConcurrentHashMap<>(); + + private PayloadOffloadTracking() {} + + /** Registers an observer for one local durable execution until the returned registration is closed. */ + public static Registration observe( + String durableExecutionArn, BiConsumer observer) { + Objects.requireNonNull(durableExecutionArn, "durableExecutionArn cannot be null"); + Objects.requireNonNull(observer, "observer cannot be null"); + if (OBSERVERS.putIfAbsent(durableExecutionArn, observer) != null) { + throw new IllegalStateException("Payload offload tracking is already active for " + durableExecutionArn); + } + return () -> OBSERVERS.remove(durableExecutionArn, observer); + } + + /** Records the selected offloader when an observer is active for the payload's execution. */ + public static void record(PayloadOffloadContext context, PayloadOffloader offloader) { + var observer = OBSERVERS.get(context.durableExecutionArn()); + if (observer != null) { + observer.accept(context, offloader); + } + } + + /** Registration handle for an active observer. */ + @FunctionalInterface + public interface Registration extends AutoCloseable { + @Override + void close(); + } +} diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java index 5cd40820e..55d346b6e 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/BaseDurableOperation.java @@ -20,6 +20,7 @@ import software.amazon.lambda.durable.exception.DurableOperationException; import software.amazon.lambda.durable.exception.IllegalDurableOperationException; import software.amazon.lambda.durable.exception.NonDeterministicExecutionException; +import software.amazon.lambda.durable.exception.PayloadOffloadException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.SuspendExecutionException; @@ -113,6 +114,11 @@ public String getName() { return operationIdentifier.name(); } + /** Gets the complete operation identifier. */ + protected OperationIdentifier getOperationIdentifier() { + return operationIdentifier; + } + /** Gets the parent context. */ protected DurableContextImpl getContext() { return durableContext; @@ -216,6 +222,11 @@ protected boolean isOperationCompleted() { return completionFuture.isDone(); } + /** Returns whether this operation has an in-memory completion outcome that does not require a stored operation. */ + protected boolean hasInMemoryCompletion() { + return false; + } + /** * Waits for the operation to complete. Deregisters the current thread to allow Lambda suspension if the operation * is still in progress, then re-registers when the operation completes. @@ -254,8 +265,8 @@ protected Operation waitForOperationCompletion() { ExceptionHelper.sneakyThrow(ExceptionHelper.unwrapCompletableFuture(throwable)); } - if (isVirtual) { - // We don't store virtual operations so they don't corresponding Operation in storage + if (isVirtual || hasInMemoryCompletion()) { + // Virtual and explicitly in-memory outcomes do not require a corresponding stored operation. return null; } else { // Get result based on status @@ -287,6 +298,9 @@ protected void runUserHandler(Runnable runnable, ThreadType threadType) { try { runnable.run(); } catch (Throwable throwable) { + if (throwable instanceof PayloadOffloadException payloadOffloadException) { + executionManager.failInvocation(payloadOffloadException); + } // Operations wrap the user function and handle all outcomes except for SuspendExecutionException. // Anything else reaching here is unexpected and terminates the execution. if (!executionManager.isExecutionCompletedExceptionally() @@ -525,6 +539,11 @@ protected CompletableFuture sendOperationUpdateAsync(OperationUpdate.Build } } + /** Returns whether an update sent by this operation would be persisted. */ + protected boolean shouldPersistUpdate() { + return !replayCompletedOperation.get(); + } + /** Validates that current operation matches checkpointed operation during replay. */ protected void validateReplay(Operation checkpointed) { if (checkpointed == null || checkpointed.type() == null) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java index 9d9481fb9..f2226e5d0 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/CallbackOperation.java @@ -76,7 +76,7 @@ public T get() { var op = waitForOperationCompletion(); return switch (op.status()) { - case SUCCEEDED -> deserializeResult(op.callbackDetails().result()); + case SUCCEEDED -> deserializeExternalResult(op.callbackDetails().result()); case FAILED -> throw new CallbackFailedException(op); case TIMED_OUT -> throw new CallbackTimeoutException(op); default -> diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java index 8c299cfa4..834fa69a6 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ChildContextOperation.java @@ -27,6 +27,7 @@ import software.amazon.lambda.durable.exception.DurableOperationException; import software.amazon.lambda.durable.exception.MapIterationFailedException; import software.amazon.lambda.durable.exception.ParallelBranchFailedException; +import software.amazon.lambda.durable.exception.PayloadOffloadException; import software.amazon.lambda.durable.exception.StepFailedException; import software.amazon.lambda.durable.exception.StepInterruptedException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; @@ -35,6 +36,10 @@ import software.amazon.lambda.durable.logging.DurableLogger; import software.amazon.lambda.durable.model.DeserializedOperationResult; import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.offload.PayloadOffloader; +import software.amazon.lambda.durable.offload.PayloadOffloaders; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; import software.amazon.lambda.durable.util.ExceptionHelper; /** @@ -48,10 +53,16 @@ * parent operation has already succeeded. */ public class ChildContextOperation extends SerializableDurableOperation { + private enum PersistenceOutcome { + PERSISTED, + NOT_APPLICABLE, + REJECTED_BY_PARENT + } private static final int LARGE_RESULT_THRESHOLD = 256 * 1024; private final Function function; + private final ConcurrencyOperation concurrencyParent; private final AtomicBoolean replayChildren = new AtomicBoolean(false); private final AtomicReference> cachedOperationResult = new AtomicReference<>(null); @@ -77,10 +88,12 @@ public ChildContextOperation( operationIdentifier, resultTypeToken, config.serDes(), + config.payloadOffloader(), durableContext, parentOperation, config.isVirtual()); this.function = function; + this.concurrencyParent = parentOperation; } /** Starts the operation. */ @@ -140,6 +153,13 @@ private void executeChildContext() { T result = runUserFunction(null, () -> function.apply(childContext)); handleChildContextSuccess(result); + } catch (PayloadOffloadException e) { + if (concurrencyParent == null || concurrencyParent.claimChildPayloadFailure(e)) { + throw e; + } + cachedOperationResult.set(DeserializedOperationResult.failed(e)); + fireOnOperationEnd(null, e, false); + markAlreadyCompleted(); } catch (Throwable e) { handleChildContextFailure(e); } @@ -151,23 +171,23 @@ private void executeChildContext() { } private void handleChildContextSuccess(T result) { - var serializedResult = serializeAndDeserializeResult(result); - - if (replayChildren.get() || isVirtual || parentOperation != null && parentOperation.isOperationCompleted()) { - // Skip checkpointing if - // - parent ConcurrencyOperation has already completed, preventing race conditions where a child finishes - // after the parent has already completed. - // - replaying a SUCCEEDED child with replayChildren=true — skip checkpointing. - // - nestingType is FLAT - // Mark the completableFuture completed so get() doesn't block waiting for a checkpoint response. - cachedOperationResult.set(DeserializedOperationResult.succeeded(serializedResult.deserialized())); - if (isVirtual) { - fireOnOperationEnd(null, null, false); - } - markAlreadyCompleted(); - } else { + var persistenceOutcome = persistCompletionIfAllowed(() -> { + var serializedResult = serializeAndDeserializeResult(result); checkpointSuccess(serializedResult.deserialized(), serializedResult.serialized()); + }); + if (persistenceOutcome == PersistenceOutcome.PERSISTED) { + return; + } + + // The result is not persisted when this is virtual, replaying children, or the parent has begun early + // completion. Normalize without offloading so the caller still observes SerDes semantics without an orphaned + // storage side effect. + var normalizedResult = normalizeResult(result); + cachedOperationResult.set(DeserializedOperationResult.succeeded(normalizedResult.deserialized())); + if (isVirtual || persistenceOutcome == PersistenceOutcome.REJECTED_BY_PARENT) { + fireOnOperationEnd(null, null, false); } + markAlreadyCompleted(); } private void checkpointSuccess(T result, String serialized) { @@ -187,40 +207,60 @@ private void checkpointSuccess(T result, String serialized) { } private void handleChildContextFailure(Throwable exception) { - exception = ExceptionHelper.unwrapCompletableFuture(exception); - if (exception instanceof SuspendExecutionException suspendExecutionException) { + var unwrappedException = ExceptionHelper.unwrapCompletableFuture(exception); + if (unwrappedException instanceof SuspendExecutionException suspendExecutionException) { // Rethrow Error immediately — do not checkpoint throw suspendExecutionException; } - if (exception instanceof UnrecoverableDurableExecutionException unrecoverableDurableExecutionException) { + if (unwrappedException + instanceof UnrecoverableDurableExecutionException unrecoverableDurableExecutionException) { // terminate the execution and throw the exception if it's not recoverable throw terminateExecution(unrecoverableDurableExecutionException); } - final ErrorObject errorObject; - if (exception instanceof DurableOperationException opEx) { - errorObject = opEx.getErrorObject(); - } else { - errorObject = serializeException(exception); + var persistenceOutcome = persistCompletionIfAllowed(() -> { + var errorObject = unwrappedException instanceof DurableOperationException opEx + ? rebindForwardedError(opEx) + : serializeException(unwrappedException); + var op = createVirtualOperation(errorObject); + cachedOperationResult.set(DeserializedOperationResult.failed(translateException(op, errorObject))); + sendOperationUpdate( + OperationUpdate.builder().action(OperationAction.FAIL).error(errorObject)); + }); + if (persistenceOutcome == PersistenceOutcome.PERSISTED) { + return; } - var op = createVirtualOperation(errorObject); - cachedOperationResult.set(DeserializedOperationResult.failed(translateException(op, errorObject))); + if (unwrappedException instanceof DurableOperationException operationException) { + cachedOperationResult.set(DeserializedOperationResult.failed(operationException)); + } else { + var errorObject = serializeExceptionWithoutOffloading(unwrappedException); + var op = createVirtualOperation(errorObject); + cachedOperationResult.set(DeserializedOperationResult.failed( + translateException(op, errorObject, PayloadOffloaders.disabled()))); + } + if (isVirtual || persistenceOutcome == PersistenceOutcome.REJECTED_BY_PARENT) { + fireOnOperationEnd(null, unwrappedException, false); + } + markAlreadyCompleted(); + } - // Skip checkpointing if - // - parent ConcurrencyOperation has already completed, preventing race conditions where a child finishes after - // the parent has already succeeded. - // - this child is not a direct child of a parent context (i.e. nestingType == FLAT), such as a parallel branch. - if ((parentOperation != null && parentOperation.isOperationCompleted()) || isVirtual) { - if (isVirtual) { - fireOnOperationEnd(null, exception, false); - } - markAlreadyCompleted(); - return; + private PersistenceOutcome persistCompletionIfAllowed(Runnable persistence) { + if (!shouldPersistUpdate() || replayChildren.get() || isVirtual) { + return PersistenceOutcome.NOT_APPLICABLE; } + if (concurrencyParent == null) { + persistence.run(); + return PersistenceOutcome.PERSISTED; + } + return concurrencyParent.persistChildCompletion(persistence) + ? PersistenceOutcome.PERSISTED + : PersistenceOutcome.REJECTED_BY_PARENT; + } - sendOperationUpdate( - OperationUpdate.builder().action(OperationAction.FAIL).error(errorObject)); + @Override + protected boolean hasInMemoryCompletion() { + return cachedOperationResult.get() != null; } @Override @@ -245,8 +285,22 @@ public T get() { } private Throwable translateException(Operation op, ErrorObject errorObject) { + return translateException(op, errorObject, null); + } + + private Throwable translateException( + Operation op, ErrorObject errorObject, PayloadOffloader sourceOffloaderOverride) { + if (getSubType() == OperationSubType.WAIT_FOR_CALLBACK) { + var callbackFailure = externalCallbackFailure(); + if (callbackFailure != null) { + return callbackFailure; + } + } + // Attempt to reconstruct and throw the original exception - Throwable original = deserializeException(errorObject); + Throwable original = sourceOffloaderOverride == null + ? deserializeException(errorObject) + : deserializeException(errorObject, null, sourceOffloaderOverride); if (original != null) { return original; } @@ -254,9 +308,12 @@ private Throwable translateException(Operation op, ErrorObject errorObject) { // throw a general failed exception if a user exception is not reconstructed return switch (getSubType()) { case WAIT_FOR_CALLBACK -> handleWaitForCallbackFailure(); - case MAP_ITERATION -> new MapIterationFailedException(op); - case PARALLEL_BRANCH -> new ParallelBranchFailedException(op); - case RUN_IN_CHILD_CONTEXT, WITH_RETRY -> new ChildContextFailedException(op); + case MAP_ITERATION -> + attachTranslatedPayloadSource(new MapIterationFailedException(op), sourceOffloaderOverride); + case PARALLEL_BRANCH -> + attachTranslatedPayloadSource(new ParallelBranchFailedException(op), sourceOffloaderOverride); + case RUN_IN_CHILD_CONTEXT, WITH_RETRY -> + attachTranslatedPayloadSource(new ChildContextFailedException(op), sourceOffloaderOverride); // the following subtypes should not be able to reach here case PARALLEL, MAP, WAIT_FOR_CONDITION, STEP, WAIT, CALLBACK, CHAINED_INVOKE -> @@ -264,6 +321,13 @@ private Throwable translateException(Operation op, ErrorObject errorObject) { }; } + private E attachTranslatedPayloadSource( + E exception, PayloadOffloader sourceOffloaderOverride) { + return sourceOffloaderOverride == null + ? attachPayloadSource(exception, SerDesPayloadKind.EXCEPTION, null) + : attachPayloadSource(exception, SerDesPayloadKind.EXCEPTION, null, sourceOffloaderOverride); + } + private Operation createVirtualOperation(ErrorObject errorObject) { return Operation.builder() .id(getOperationId()) @@ -275,28 +339,21 @@ private Operation createVirtualOperation(ErrorObject errorObject) { } private Throwable handleWaitForCallbackFailure() { + var callbackFailure = externalCallbackFailure(); + if (callbackFailure != null) { + return callbackFailure; + } + var childrenOps = getChildOperations(); - var callbackOp = childrenOps.stream() - .filter(o -> o.type() == OperationType.CALLBACK) - .findFirst() - .orElse(null); var submitterOp = childrenOps.stream() .filter(o -> o.type() == OperationType.STEP) .findFirst() .orElse(null); + var callbackOp = childrenOps.stream() + .filter(o -> o.type() == OperationType.CALLBACK) + .findFirst() + .orElse(null); if (callbackOp != null) { - // if callback failed - if (isTerminalStatus(callbackOp.status())) { - switch (callbackOp.status()) { - case FAILED -> { - return new CallbackFailedException(callbackOp); - } - case TIMED_OUT -> { - return new CallbackTimeoutException(callbackOp); - } - } - } - // if submitter failed if (submitterOp != null && isTerminalStatus(submitterOp.status()) @@ -312,4 +369,30 @@ && isTerminalStatus(submitterOp.status()) return new IllegalStateException("Unknown waitForCallback status"); } + + private Throwable externalCallbackFailure() { + var childrenOps = getChildOperations(); + var submitterOp = childrenOps.stream() + .filter(operation -> operation.type() == OperationType.STEP) + .findFirst() + .orElse(null); + if (submitterOp != null + && isTerminalStatus(submitterOp.status()) + && submitterOp.status() != OperationStatus.SUCCEEDED) { + return null; + } + + var callbackOp = childrenOps.stream() + .filter(operation -> operation.type() == OperationType.CALLBACK) + .findFirst() + .orElse(null); + if (callbackOp == null) { + return null; + } + return switch (callbackOp.status()) { + case FAILED -> new CallbackFailedException(callbackOp); + case TIMED_OUT -> new CallbackTimeoutException(callbackOp); + default -> null; + }; + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java index 646b0685a..7f9760ee6 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ConcurrencyOperation.java @@ -22,12 +22,14 @@ import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.PayloadOffloadException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.OperationIdGenerator; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.offload.PayloadOffloader; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ExceptionHelper; @@ -94,6 +96,11 @@ private static final class CoordinatorState { private boolean stateChangedQueued; private boolean coordinatorWaiting; + // coordinates parent completion with child result persistence + private final Object childPersistenceLock = new Object(); + private boolean completionInitiated; + private PayloadOffloadException claimedChildPayloadFailure; + // set by context thread and used by consumer thread protected final AtomicBoolean isJoined = new AtomicBoolean(false); @@ -105,7 +112,27 @@ protected ConcurrencyOperation( int maxConcurrency, Function shouldComplete, NestingType nestingType) { - super(operationIdentifier, resultTypeToken, resultSerDes, durableContext); + this( + operationIdentifier, + resultTypeToken, + resultSerDes, + null, + durableContext, + maxConcurrency, + shouldComplete, + nestingType); + } + + protected ConcurrencyOperation( + OperationIdentifier operationIdentifier, + TypeToken resultTypeToken, + SerDes resultSerDes, + PayloadOffloader payloadOffloader, + DurableContextImpl durableContext, + int maxConcurrency, + Function shouldComplete, + NestingType nestingType) { + super(operationIdentifier, resultTypeToken, resultSerDes, payloadOffloader, durableContext); this.maxConcurrency = maxConcurrency; this.shouldComplete = Objects.requireNonNull(shouldComplete, "shouldComplete cannot be null"); this.operationIdGenerator = new OperationIdGenerator(getOperationId()); @@ -151,6 +178,30 @@ protected ChildContextOperation createItem( this); } + protected ChildContextOperation createItem( + String operationId, + String name, + Function function, + TypeToken resultType, + SerDes serDes, + PayloadOffloader payloadOffloader, + OperationSubType branchSubType) { + if (payloadOffloader == null) { + return createItem(operationId, name, function, resultType, serDes, branchSubType); + } + return new ChildContextOperation<>( + OperationIdentifier.of(operationId, name, branchSubType), + function, + resultType, + RunInChildContextConfig.builder() + .serDes(serDes) + .payloadOffloader(payloadOffloader) + .isVirtual(nestingType == NestingType.FLAT) + .build(), + rootContext, + this); + } + /** Called when the concurrency operation completes. Subclasses define checkpointing behavior. */ protected abstract void handleCompletion(CompletionConfig.CompletionDecision completionDecision); @@ -168,8 +219,19 @@ protected ChildContextOperation enqueueItem( SerDes serDes, OperationSubType branchSubType, boolean skipped) { + return enqueueItem(name, function, resultType, serDes, null, branchSubType, skipped); + } + + protected ChildContextOperation enqueueItem( + String name, + Function function, + TypeToken resultType, + SerDes serDes, + PayloadOffloader payloadOffloader, + OperationSubType branchSubType, + boolean skipped) { var operationId = this.operationIdGenerator.nextOperationId(); - var childOp = createItem(operationId, name, function, resultType, serDes, branchSubType); + var childOp = createItem(operationId, name, function, resultType, serDes, payloadOffloader, branchSubType); branches.add(childOp); if (!skipped) { logger.debug("Item enqueued {}", name); @@ -222,7 +284,7 @@ private void runCoordinator(ExpectedCompletionStatus expectedCompletionStatus) { while (!isOperationCompleted()) { var completionDecision = canComplete(state.succeededCount, state.failedCount, expectedCompletionStatus); if (completionDecision != null) { - handleCompletion(completionDecision); + initiateCompletion(completionDecision); return; } startPendingItems(state); @@ -291,11 +353,77 @@ private void handleException(Throwable ex) { if (throwable instanceof UnrecoverableDurableExecutionException unrecoverableDurableExecutionException) { throw terminateExecution(unrecoverableDurableExecutionException); } + if (throwable instanceof PayloadOffloadException payloadOffloadException) { + throw payloadOffloadException; + } throw terminateExecutionWithIllegalDurableOperationException( String.format("Unexpected exception in concurrency operation: %s", throwable)); } + void initiateCompletion(CompletionConfig.CompletionDecision completionDecision) { + synchronized (childPersistenceLock) { + if (completionInitiated || isOperationCompleted()) { + return; + } + completionInitiated = true; + handleCompletion(completionDecision); + } + } + + /** + * Persists one child outcome only when parent completion has not started. + * + *

The lock remains held through offloading and checkpoint publication so parent completion cannot win between + * those two actions and orphan an external payload. A payload failure claims parent completion before the lock is + * released. + */ + boolean persistChildCompletion(Runnable persistence) { + synchronized (childPersistenceLock) { + if (completionInitiated || isOperationCompleted()) { + return false; + } + try { + persistence.run(); + } catch (Throwable failure) { + var unwrapped = ExceptionHelper.unwrapCompletableFuture(failure); + if (unwrapped instanceof PayloadOffloadException payloadFailure) { + completionInitiated = true; + claimedChildPayloadFailure = payloadFailure; + ExceptionHelper.sneakyThrow(payloadFailure); + } + if (unwrapped instanceof UnrecoverableDurableExecutionException) { + completionInitiated = true; + ExceptionHelper.sneakyThrow(unwrapped); + } + throw failure; + } + return true; + } + } + + /** + * Atomically claims parent completion for a child payload failure. + * + *

This also recognizes a failure already claimed by {@link #persistChildCompletion(Runnable)}. A false result + * means successful early completion already won, so the late child must be treated as skipped. + * + * @param failure the payload failure being claimed + */ + boolean claimChildPayloadFailure(PayloadOffloadException failure) { + synchronized (childPersistenceLock) { + if (claimedChildPayloadFailure == failure) { + claimedChildPayloadFailure = null; + return true; + } + if (completionInitiated || isOperationCompleted()) { + return false; + } + completionInitiated = true; + return true; + } + } + /** * Called by a ChildContextOperation BEFORE it closes its child context. Updates counters, checks completion * criteria, and either triggers the next queued item or completes the operation. @@ -310,7 +438,11 @@ private void onItemComplete(CoordinatorState state, ChildContextOperation chi logger.debug("Result succeeded - {}", child.getName()); state.succeededCount++; } catch (Throwable e) { - logger.debug("Child operation {} failed: {}", child.getOperationId(), e.getMessage()); + var failure = ExceptionHelper.unwrapCompletableFuture(e); + if (failure instanceof PayloadOffloadException payloadOffloadException) { + throw payloadOffloadException; + } + logger.debug("Child operation {} failed: {}", child.getOperationId(), failure.getMessage()); state.failedCount++; } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java index 9e2c54ace..0b6a150a2 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/InvokeOperation.java @@ -2,6 +2,7 @@ // SPDX-License-Identifier: Apache-2.0 package software.amazon.lambda.durable.operation; +import java.util.function.Function; import software.amazon.awssdk.services.lambda.model.ChainedInvokeOptions; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationAction; @@ -9,11 +10,15 @@ import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.InvokeConfig; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.DurableOperationException; import software.amazon.lambda.durable.exception.InvokeException; import software.amazon.lambda.durable.exception.InvokeFailedException; import software.amazon.lambda.durable.exception.InvokeStoppedException; import software.amazon.lambda.durable.exception.InvokeTimedOutException; import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; +import software.amazon.lambda.durable.offload.internal.ChainedInvokeOutputFrame; +import software.amazon.lambda.durable.offload.internal.ChainedInvokePayloadFrame; import software.amazon.lambda.durable.serde.SerDes; /** @@ -35,7 +40,7 @@ public InvokeOperation( TypeToken resultTypeToken, InvokeConfig config, DurableContextImpl durableContext) { - super(operationIdentifier, resultTypeToken, config.serDes(), durableContext); + super(operationIdentifier, resultTypeToken, config.serDes(), config.payloadOffloader(), durableContext); this.functionName = functionName; this.payload = payload; @@ -64,13 +69,19 @@ protected void replay(Operation existing) { } private void startInvocation() { + var serializedPayload = invokeConfig.usePayloadOffloaderForPayload() + ? serializePayload(payload, payloadSerDes, SerDesPayloadKind.INVOKE_PAYLOAD, null) + : payloadSerDes.serialize(payload); var update = OperationUpdate.builder() .action(OperationAction.START) .chainedInvokeOptions(ChainedInvokeOptions.builder() .functionName(functionName) .tenantId(invokeConfig.tenantId()) .build()) - .payload(payloadSerDes.serialize(this.payload)); + .payload( + invokeConfig.usePayloadOffloaderForPayload() + ? ChainedInvokePayloadFrame.encode(serializedPayload) + : serializedPayload); sendOperationUpdate(update); } @@ -86,12 +97,48 @@ public T get() { var invokeDetails = op.chainedInvokeDetails(); var result = invokeDetails != null ? invokeDetails.result() : null; return switch (op.status()) { - case SUCCEEDED -> deserializeResult(result); - case FAILED -> throw new InvokeFailedException(op); - case TIMED_OUT -> throw new InvokeTimedOutException(op); - case STOPPED -> throw new InvokeStoppedException(op); + case SUCCEEDED -> deserializeInvokeResult(result); + case FAILED -> throw createInvokeFailure(op, InvokeFailedException::new); + case TIMED_OUT -> throw createInvokeFailure(op, InvokeTimedOutException::new); + case STOPPED -> throw createInvokeFailure(op, InvokeStoppedException::new); // Unexpected status which should not happen. This is added for forward-compatibility. - default -> throw new InvokeException(op); + default -> throw createInvokeFailure(op, InvokeException::new); }; } + + private T deserializeInvokeResult(String result) { + if (!invokeConfig.usePayloadOffloaderForPayload() || !ChainedInvokeOutputFrame.isFramed(result)) { + return deserializeExternalResult(result); + } + var decoded = ChainedInvokeOutputFrame.decode(result); + if (decoded.usesPayloadCodec()) { + validatePayloadEnvelope(decoded.payload(), SerDesPayloadKind.RESULT, null); + } + return decoded.usesPayloadCodec() + ? deserializeResult(decoded.payload()) + : deserializeExternalResult(decoded.payload()); + } + + private E createInvokeFailure( + Operation operation, Function exceptionFactory) { + var details = operation.chainedInvokeDetails(); + var error = details != null ? details.error() : null; + if (!invokeConfig.usePayloadOffloaderForPayload() + || error == null + || !ChainedInvokeOutputFrame.isFramed(error.errorData())) { + return exceptionFactory.apply(operation); + } + + var decoded = ChainedInvokeOutputFrame.decode(error.errorData()); + var errorData = decoded.payload(); + if (decoded.usesPayloadCodec()) { + validatePayloadEnvelope(errorData, SerDesPayloadKind.EXCEPTION, null); + errorData = resolveSerializedPayload(errorData, SerDesPayloadKind.EXCEPTION, null); + } + var decodedError = error.toBuilder().errorData(errorData).build(); + var decodedOperation = operation.toBuilder() + .chainedInvokeDetails(details.toBuilder().error(decodedError).build()) + .build(); + return exceptionFactory.apply(decodedOperation); + } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/MapOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/MapOperation.java index 2f665547a..df88c0039 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/MapOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/MapOperation.java @@ -20,12 +20,14 @@ import software.amazon.lambda.durable.config.MapConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.NonDeterministicExecutionException; +import software.amazon.lambda.durable.exception.PayloadOffloadException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; import software.amazon.lambda.durable.model.MapResult; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.offload.PayloadOffloader; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ExceptionHelper; import software.amazon.lambda.durable.util.ParameterValidator; @@ -50,6 +52,7 @@ public class MapOperation extends ConcurrencyOperation> { private final DurableContext.MapFunction function; private final TypeToken itemResultType; private final SerDes serDes; + private final PayloadOffloader payloadOffloader; private final List iterationNames; private volatile MapResult cachedResult; @@ -82,6 +85,7 @@ public MapOperation( operationIdentifier, new TypeToken<>() {}, config.serDes(), + config.payloadOffloader(), durableContext, config.maxConcurrency(), config.completionConfig().completionDecisionFunction(), @@ -96,6 +100,7 @@ public MapOperation( this.function = function; this.itemResultType = itemResultType; this.serDes = config.serDes(); + this.payloadOffloader = config.payloadOffloader(); this.iterationNames = Collections.unmodifiableList(new ArrayList<>(iterationNames)); if (this.iterationNames.size() != this.items.size()) { throw new IllegalArgumentException("iterationNames must have one entry per item"); @@ -152,6 +157,7 @@ private void addUnskippedItems(List resultItems) childCtx -> function.apply(item, index, childCtx), itemResultType, serDes, + payloadOffloader, OperationSubType.MAP_ITERATION, skip); } @@ -251,7 +257,8 @@ protected void replay(Operation existing) { @Override protected void handleCompletion(CompletionConfig.CompletionDecision completionDecision) { this.cachedResult = constructMapResult(completionDecision); - var serializedResult = serializeAndDeserializeResult(cachedResult); + var serializedResult = + shouldPersistUpdate() ? serializeAndDeserializeResult(cachedResult) : normalizeResult(cachedResult); this.cachedResult = serializedResult.deserialized(); var serializedBytes = serializedResult.serialized().getBytes(StandardCharsets.UTF_8); @@ -262,7 +269,9 @@ protected void handleCompletion(CompletionConfig.CompletionDecision completionDe .payload(serializedResult.serialized())); } else { // Large result: checkpoint with stripped payload + replayChildren flag - var strippedResult = serializeAndDeserializeResult(stripMapResult(cachedResult)); + var strippedResult = shouldPersistUpdate() + ? serializeAndDeserializeResult(stripMapResult(cachedResult)) + : normalizeResult(stripMapResult(cachedResult)); sendOperationUpdate(OperationUpdate.builder() .action(OperationAction.SUCCEED) .subType(getSubType().getValue()) @@ -303,6 +312,9 @@ private MapResult constructMapResult(CompletionConfig.CompletionDecision comp // terminate the execution and throw the exception if it's not recoverable throw terminateExecution(unrecoverableDurableExecutionException); } + if (throwable instanceof PayloadOffloadException payloadOffloadException) { + throw payloadOffloadException; + } resultItems.set(i, MapResult.MapResultItem.failed(MapResult.MapError.of(throwable))); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java index 41a687e6b..b0a92bcd2 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/ParallelOperation.java @@ -17,11 +17,13 @@ import software.amazon.lambda.durable.config.ParallelBranchConfig; import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.PayloadOffloadException; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.model.ParallelResult; import software.amazon.lambda.durable.serde.SerDes; +import software.amazon.lambda.durable.util.ExceptionHelper; /** * Manages parallel execution of multiple branches as child context operations. @@ -59,6 +61,7 @@ public ParallelOperation( operationIdentifier, TypeToken.get(ParallelResult.class), resultSerDes, + config.payloadOffloader(), durableContext, config.maxConcurrency(), config.completionConfig().completionDecisionFunction(), @@ -83,7 +86,8 @@ protected void handleCompletion(CompletionConfig.CompletionDecision completionDe skippedCount, completionDecision.completionStatus(), statuses); - var serializedResult = serializeAndDeserializeResult(cachedResult); + var serializedResult = + shouldPersistUpdate() ? serializeAndDeserializeResult(cachedResult) : normalizeResult(cachedResult); cachedResult = serializedResult.deserialized(); // Branches added after checkpoint will not exist in the checkpointed result, but they'll be in the returned @@ -103,6 +107,10 @@ private ParallelResult.Status getParallelItemStatus(ChildContextOperation chi childContextOperation.get(); return ParallelResult.Status.SUCCEEDED; } catch (Throwable t) { + var failure = ExceptionHelper.unwrapCompletableFuture(t); + if (failure instanceof PayloadOffloadException payloadOffloadException) { + throw payloadOffloadException; + } return ParallelResult.Status.FAILED; } } @@ -187,6 +195,9 @@ public DurableFuture branch( && (partialResult.statuses().size() <= nextBranchIndex || partialResult.statuses().get(nextBranchIndex) == ParallelResult.Status.SKIPPED); var serDes = config.serDes() == null ? getContext().getDurableConfig().getSerDes() : config.serDes(); - return enqueueItem(name, func, resultType, serDes, OperationSubType.PARALLEL_BRANCH, skip); + var offloader = config.payloadOffloader() == null + ? getContext().getDurableConfig().getPayloadOffloader() + : config.payloadOffloader(); + return enqueueItem(name, func, resultType, serDes, offloader, OperationSubType.PARALLEL_BRANCH, skip); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java index 6457c996d..db7c87eaa 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/SerializableDurableOperation.java @@ -8,8 +8,14 @@ import software.amazon.lambda.durable.DurableFuture; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.DurableOperationException; import software.amazon.lambda.durable.exception.SerDesException; +import software.amazon.lambda.durable.execution.PayloadCodec; import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.offload.PayloadOffloadContext; +import software.amazon.lambda.durable.offload.PayloadOffloader; +import software.amazon.lambda.durable.offload.PayloadOffloaders; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; import software.amazon.lambda.durable.serde.SerDes; import software.amazon.lambda.durable.util.ExceptionHelper; @@ -38,6 +44,7 @@ protected record SerializedResult(String serialized, T deserialized) {} private final TypeToken resultTypeToken; private final SerDes resultSerDes; + private final PayloadOffloader payloadOffloader; /** * Constructs a new durable operation. @@ -52,7 +59,16 @@ protected SerializableDurableOperation( TypeToken resultTypeToken, SerDes resultSerDes, DurableContextImpl durableContext) { - this(operationIdentifier, resultTypeToken, resultSerDes, durableContext, null, false); + this(operationIdentifier, resultTypeToken, resultSerDes, null, durableContext, null, false); + } + + protected SerializableDurableOperation( + OperationIdentifier operationIdentifier, + TypeToken resultTypeToken, + SerDes resultSerDes, + PayloadOffloader payloadOffloader, + DurableContextImpl durableContext) { + this(operationIdentifier, resultTypeToken, resultSerDes, payloadOffloader, durableContext, null, false); } /** @@ -72,9 +88,21 @@ protected SerializableDurableOperation( DurableContextImpl durableContext, BaseDurableOperation parentOperation, boolean isVirtual) { + this(operationIdentifier, resultTypeToken, resultSerDes, null, durableContext, parentOperation, isVirtual); + } + + protected SerializableDurableOperation( + OperationIdentifier operationIdentifier, + TypeToken resultTypeToken, + SerDes resultSerDes, + PayloadOffloader payloadOffloader, + DurableContextImpl durableContext, + BaseDurableOperation parentOperation, + boolean isVirtual) { super(operationIdentifier, durableContext, parentOperation, isVirtual); this.resultTypeToken = resultTypeToken; this.resultSerDes = resultSerDes; + this.payloadOffloader = payloadOffloader; } /** @@ -85,8 +113,28 @@ protected SerializableDurableOperation( * @throws SerDesException if deserialization fails */ protected T deserializeResult(String result) { + return deserializeResult(result, SerDesPayloadKind.RESULT, null); + } + + /** Deserializes externally supplied data without interpreting SDK payload envelopes. */ + protected T deserializeExternalResult(String result) { + return resultSerDes.deserialize(result, resultTypeToken); + } + + /** Deserializes a result with explicit payload kind and attempt metadata. */ + protected T deserializeResult(String result, SerDesPayloadKind payloadKind, Integer attempt) { try { - return resultSerDes.deserialize(result, resultTypeToken); + if (!usesPayloadCodec(result)) { + return resultSerDes.deserialize(result, resultTypeToken); + } + return executionManager + .getPayloadCodec() + .deserialize( + result, + resultTypeToken, + resultSerDes, + payloadOffloader, + payloadContext(payloadKind, attempt)); } catch (SerDesException e) { logger.warn( "Failed to deserialize {} result for operation name '{}'. Ensure the result is properly encoded.", @@ -106,11 +154,47 @@ protected T deserializeResult(String result) { * @return the serialized string and the deserialized result */ protected SerializedResult serializeAndDeserializeResult(T result) { + return serializeAndDeserializeResult(result, SerDesPayloadKind.RESULT, null); + } + + /** Serializes a result with explicit payload kind and attempt metadata. */ + protected SerializedResult serializeAndDeserializeResult( + T result, SerDesPayloadKind payloadKind, Integer attempt) { + var serialized = serializePayload(result, resultSerDes, payloadKind, attempt); + var deserialized = + shouldDeserializeAfterSerialization() ? deserializeResult(serialized, payloadKind, attempt) : result; + return new SerializedResult<>(serialized, deserialized); + } + + /** + * Normalizes a result through SerDes without offloading it. + * + *

Use this when an operation needs the same first-execution value normalization as replay but will not persist + * an update. + */ + protected SerializedResult normalizeResult(T result) { var serialized = resultSerDes.serialize(result); - var deserialized = shouldDeserializeAfterSerialization() ? deserializeResult(serialized) : result; + var deserialized = + shouldDeserializeAfterSerialization() ? resultSerDes.deserialize(serialized, resultTypeToken) : result; return new SerializedResult<>(serialized, deserialized); } + /** Serializes an operation-owned payload with the operation's offloader policy. */ + protected String serializePayload(Object value, SerDes serDes, SerDesPayloadKind payloadKind, Integer attempt) { + if (!hasActivePayloadOffloader()) { + var serialized = serDes.serialize(value); + return PayloadCodec.isOffloadEnvelope(serialized) + ? executionManager + .getPayloadCodec() + .serializePreEncodedPayload( + serialized, payloadOffloader, payloadContext(payloadKind, attempt)) + : serialized; + } + return executionManager + .getPayloadCodec() + .serialize(value, serDes, payloadOffloader, payloadContext(payloadKind, attempt)); + } + /** * Serializes a throwable into an {@link ErrorObject} for checkpointing. * @@ -119,9 +203,121 @@ protected SerializedResult serializeAndDeserializeResult(T result) { */ @SuppressWarnings("ThrowableNotThrown") protected ErrorObject serializeException(Throwable throwable) { - var error = ExceptionHelper.buildErrorObject(throwable, resultSerDes); + return serializeException(throwable, null); + } + + /** Serializes a throwable with attempt metadata. */ + protected ErrorObject serializeException(Throwable throwable, Integer attempt) { + return serializeException(throwable, attempt, true); + } + + /** Serializes an exception for local result normalization without offloading it. */ + protected ErrorObject serializeExceptionWithoutOffloading(Throwable throwable) { + return serializeException(throwable, null, false); + } + + /** Rebinds a forwarded operation error to this operation's payload policy before checkpointing it. */ + protected ErrorObject rebindForwardedError(DurableOperationException exception) { + return rebindForwardedError(exception, null); + } + + /** Rebinds a forwarded operation error with attempt metadata before checkpointing it. */ + protected ErrorObject rebindForwardedError(DurableOperationException exception, Integer attempt) { + var error = exception.getErrorObject(); + if (error == null || error.errorData() == null) { + return error; + } + var targetContext = payloadContext(SerDesPayloadKind.EXCEPTION, attempt); + var errorData = exception.getPayloadOffloadContext() == null + ? executionManager + .getPayloadCodec() + .serializePreEncodedPayload(error.errorData(), payloadOffloader, targetContext) + : executionManager + .getPayloadCodec() + .rebindSerializedPayload( + error.errorData(), + exception.getPayloadOffloader(), + exception.getPayloadOffloadContext(), + payloadOffloader, + targetContext); + return error.toBuilder().errorData(errorData).build(); + } + + /** Resolves forwarded error data for local use without creating a target-owned external payload. */ + protected ErrorObject resolveForwardedErrorWithoutOffloading(DurableOperationException exception) { + var error = exception.getErrorObject(); + if (error == null || error.errorData() == null) { + return error; + } + var serialized = exception.getPayloadOffloadContext() == null + ? error.errorData() + : executionManager + .getPayloadCodec() + .resolveSerializedPayload( + error.errorData(), + exception.getPayloadOffloader(), + exception.getPayloadOffloadContext()); + var errorData = executionManager + .getPayloadCodec() + .serializePreEncodedPayload( + serialized, PayloadOffloaders.disabled(), payloadContext(SerDesPayloadKind.EXCEPTION, null)); + return error.toBuilder().errorData(errorData).build(); + } + + protected E attachPayloadSource( + E exception, SerDesPayloadKind kind, Integer attempt) { + return attachPayloadSource(exception, kind, attempt, payloadOffloader); + } + + protected E attachPayloadSource( + E exception, SerDesPayloadKind kind, Integer attempt, PayloadOffloader sourceOffloader) { + var error = exception.getErrorObject(); + if (error == null || !PayloadCodec.isOffloadEnvelope(error.errorData())) { + return exception; + } + exception.withPayloadSource(sourceOffloader, payloadContext(kind, attempt)); + return exception; + } + + /** Validates an SDK payload envelope against this operation's payload identity. */ + protected void validatePayloadEnvelope(String checkpointPayload, SerDesPayloadKind kind, Integer attempt) { + executionManager.getPayloadCodec().validateEnvelope(checkpointPayload, payloadContext(kind, attempt)); + } + + /** Resolves operation-owned payload data through the configured offloader and verifies its envelope. */ + protected String resolveSerializedPayload(String checkpointPayload, SerDesPayloadKind kind, Integer attempt) { + return executionManager + .getPayloadCodec() + .resolveSerializedPayload(checkpointPayload, payloadOffloader, payloadContext(kind, attempt)); + } + + private ErrorObject serializeException(Throwable throwable, Integer attempt, boolean allowOffloading) { + final String errorData; + if (allowOffloading) { + errorData = serializePayload(throwable, resultSerDes, SerDesPayloadKind.EXCEPTION, attempt); + } else { + var serialized = resultSerDes.serialize(throwable); + errorData = PayloadCodec.isOffloadEnvelope(serialized) + ? executionManager + .getPayloadCodec() + .serializePreEncodedPayload( + serialized, + PayloadOffloaders.disabled(), + payloadContext(SerDesPayloadKind.EXCEPTION, attempt)) + : serialized; + } + var error = ErrorObject.builder() + .errorType(throwable.getClass().getName()) + .errorMessage(throwable.getMessage()) + .errorData(errorData) + .stackTrace(ExceptionHelper.serializeStackTrace(throwable.getStackTrace())) + .build(); if (shouldDeserializeAfterSerialization()) { - deserializeException(error); + if (allowOffloading) { + deserializeException(error, attempt); + } else { + deserializeException(error, attempt, PayloadOffloaders.disabled()); + } } return error; } @@ -139,6 +335,16 @@ private boolean shouldDeserializeAfterSerialization() { * @return the reconstructed throwable, or null if reconstruction is not possible */ protected Throwable deserializeException(ErrorObject errorObject) { + return deserializeException(errorObject, null); + } + + /** Deserializes a throwable with attempt metadata. */ + protected Throwable deserializeException(ErrorObject errorObject, Integer attempt) { + return deserializeException(errorObject, attempt, payloadOffloader); + } + + protected Throwable deserializeException( + ErrorObject errorObject, Integer attempt, PayloadOffloader sourceOffloader) { Throwable original = null; if (errorObject == null) { return original; @@ -153,8 +359,19 @@ protected Throwable deserializeException(ErrorObject errorObject) { Class exceptionClass = Class.forName(errorType); if (Throwable.class.isAssignableFrom(exceptionClass)) { - original = - resultSerDes.deserialize(errorData, TypeToken.get(exceptionClass.asSubclass(Throwable.class))); + var exceptionType = TypeToken.get(exceptionClass.asSubclass(Throwable.class)); + if (usesPayloadCodec(errorData, sourceOffloader)) { + original = executionManager + .getPayloadCodec() + .deserialize( + errorData, + exceptionType, + resultSerDes, + sourceOffloader, + payloadContext(SerDesPayloadKind.EXCEPTION, attempt)); + } else { + original = resultSerDes.deserialize(errorData, exceptionType); + } if (original != null) { original.setStackTrace(ExceptionHelper.deserializeStackTrace(errorObject.stackTrace())); @@ -168,5 +385,27 @@ protected Throwable deserializeException(ErrorObject errorObject) { return original; } + private boolean usesPayloadCodec(String checkpointPayload) { + return usesPayloadCodec(checkpointPayload, payloadOffloader); + } + + private boolean usesPayloadCodec(String checkpointPayload, PayloadOffloader sourceOffloader) { + return sourceOffloader != null && !PayloadOffloaders.isDisabled(sourceOffloader) + || PayloadCodec.isOffloadEnvelope(checkpointPayload); + } + + private boolean hasActivePayloadOffloader() { + return payloadOffloader != null && !PayloadOffloaders.isDisabled(payloadOffloader); + } + + private PayloadOffloadContext payloadContext(SerDesPayloadKind kind, Integer attempt) { + return PayloadOffloadContext.forOperation( + executionManager.getDurableExecutionArn(), + getOperationIdentifier(), + getContext().getParentId(), + kind, + attempt); + } + public abstract T get(); } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java index 467a87b94..7a54c47e3 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/StepOperation.java @@ -18,6 +18,7 @@ import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.DurableOperationException; +import software.amazon.lambda.durable.exception.PayloadOffloadException; import software.amazon.lambda.durable.exception.StepFailedException; import software.amazon.lambda.durable.exception.StepInterruptedException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; @@ -25,6 +26,7 @@ import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.logging.DurableLogger; import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; import software.amazon.lambda.durable.util.ExceptionHelper; /** @@ -47,7 +49,7 @@ public StepOperation( TypeToken resultTypeToken, StepConfig config, DurableContextImpl durableContext) { - super(operationIdentifier, resultTypeToken, config.serDes(), durableContext); + super(operationIdentifier, resultTypeToken, config.serDes(), config.payloadOffloader(), durableContext); this.function = function; this.config = config; @@ -117,7 +119,9 @@ private void executeStepLogic(int attempt) { // through onUserFunctionEnd; retry/checkpoint handling stays outside the boundary. T result = runUserFunction(attempt, () -> function.apply(stepContext)); - handleStepSucceeded(result); + handleStepSucceeded(result, attempt); + } catch (PayloadOffloadException e) { + throw e; } catch (Throwable e) { handleStepFailure(e, attempt); } @@ -144,8 +148,8 @@ private void checkpointStarted() { } } - private void handleStepSucceeded(T result) { - var serializedResult = serializeAndDeserializeResult(result); + private void handleStepSucceeded(T result, int attempt) { + var serializedResult = serializeAndDeserializeResult(result, SerDesPayloadKind.RESULT, attempt); // Send SUCCEED var successUpdate = @@ -168,9 +172,9 @@ private void handleStepFailure(Throwable exception, int attempt) { final ErrorObject errorObject; if (exception instanceof DurableOperationException durableOperationException) { - errorObject = durableOperationException.getErrorObject(); + errorObject = rebindForwardedError(durableOperationException, attempt); } else { - errorObject = serializeException(exception); + errorObject = serializeException(exception, attempt); } var retryDecision = config.retryStrategy().makeRetryDecision(exception, attempt); @@ -205,8 +209,9 @@ public T get() { if (op.status() == OperationStatus.SUCCEEDED) { var stepDetails = op.stepDetails(); var result = (stepDetails != null) ? stepDetails.result() : null; + var attempt = stepDetails != null ? stepDetails.attempt() : null; - return deserializeResult(result); + return deserializeResult(result, SerDesPayloadKind.RESULT, attempt); } else { var errorObject = op.stepDetails().error(); @@ -216,12 +221,13 @@ public T get() { } // Attempt to reconstruct and throw the original exception - Throwable original = deserializeException(errorObject); + var attempt = op.stepDetails() != null ? op.stepDetails().attempt() : null; + Throwable original = deserializeException(errorObject, attempt); if (original != null) { ExceptionHelper.sneakyThrow(original); } // Fallback: wrap in StepFailedException - throw new StepFailedException(op); + throw attachPayloadSource(new StepFailedException(op), SerDesPayloadKind.EXCEPTION, attempt); } } diff --git a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java index 72eb65b8e..dbf70a408 100644 --- a/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java +++ b/sdk/src/main/java/software/amazon/lambda/durable/operation/WaitForConditionOperation.java @@ -16,6 +16,7 @@ import software.amazon.lambda.durable.context.BaseContextImpl; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.DurableOperationException; +import software.amazon.lambda.durable.exception.PayloadOffloadException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.exception.WaitForConditionFailedException; import software.amazon.lambda.durable.execution.SuspendExecutionException; @@ -23,6 +24,7 @@ import software.amazon.lambda.durable.logging.DurableLogger; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.WaitForConditionResult; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; import software.amazon.lambda.durable.util.ExceptionHelper; /** @@ -46,7 +48,7 @@ public WaitForConditionOperation( TypeToken resultTypeToken, WaitForConditionConfig config, DurableContextImpl durableContext) { - super(operationIdentifier, resultTypeToken, config.serDes(), durableContext); + super(operationIdentifier, resultTypeToken, config.serDes(), config.payloadOffloader(), durableContext); this.checkFunc = checkFunc; this.config = config; @@ -79,17 +81,19 @@ public T get() { if (op.status() == OperationStatus.SUCCEEDED) { var stepDetails = op.stepDetails(); var result = (stepDetails != null) ? stepDetails.result() : null; - return deserializeResult(result); + var attempt = stepDetails != null ? stepDetails.attempt() : null; + return deserializeResult(result, SerDesPayloadKind.STATE, attempt); } else { var errorObject = op.stepDetails().error(); // Attempt to reconstruct and throw the original exception - Throwable original = deserializeException(errorObject); + var attempt = op.stepDetails() != null ? op.stepDetails().attempt() : null; + Throwable original = deserializeException(errorObject, attempt); if (original != null) { ExceptionHelper.sneakyThrow(original); } // Fallback: wrap in WaitForConditionFailedException - throw new WaitForConditionFailedException(op); + throw attachPayloadSource(new WaitForConditionFailedException(op), SerDesPayloadKind.EXCEPTION, attempt); } } @@ -100,7 +104,7 @@ private void resumeCheckLoop(Operation existing) { var checkpointData = stepDetails != null ? stepDetails.result() : null; T currentState; // Get current state if (checkpointData != null) { - currentState = deserializeResult(checkpointData); + currentState = deserializeResult(checkpointData, SerDesPayloadKind.STATE, attempt - 1); } else { currentState = config.initialState(); } @@ -134,7 +138,8 @@ private void executeCheckLogic(T currentState, int attempt) { runUserFunction(attempt, () -> checkFunc.apply(currentState, stepContext)); // Normalize the value through SerDes so first execution matches replay. - var serializedState = serializeAndDeserializeResult(result.value()); + var serializedState = + serializeAndDeserializeResult(result.value(), SerDesPayloadKind.STATE, attempt); T deserializedValue = serializedState.deserialized(); if (result.isDone()) { @@ -163,8 +168,10 @@ private void executeCheckLogic(T currentState, int attempt) { : pollForOperationUpdates()) .thenRun(() -> executeCheckLogic(deserializedValue, attempt + 1)); } + } catch (PayloadOffloadException e) { + throw e; } catch (Throwable e) { - handleCheckFailure(e); + handleCheckFailure(e, attempt); } } }; @@ -172,7 +179,7 @@ private void executeCheckLogic(T currentState, int attempt) { runUserHandler(userHandler, ThreadType.STEP); } - private void handleCheckFailure(Throwable exception) { + private void handleCheckFailure(Throwable exception, int attempt) { exception = ExceptionHelper.unwrapCompletableFuture(exception); if (exception instanceof SuspendExecutionException suspendExecutionException) { throw suspendExecutionException; @@ -182,8 +189,8 @@ private void handleCheckFailure(Throwable exception) { } final var errorObject = (exception instanceof DurableOperationException durableOpEx) - ? durableOpEx.getErrorObject() - : serializeException(exception); + ? rebindForwardedError(durableOpEx, attempt) + : serializeException(exception, attempt); // Checkpoint FAIL var failUpdate = OperationUpdate.builder().action(OperationAction.FAIL).error(errorObject); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/DurableConfigTest.java b/sdk/src/test/java/software/amazon/lambda/durable/DurableConfigTest.java index 266e43a6e..e936345e6 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/DurableConfigTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/DurableConfigTest.java @@ -87,6 +87,41 @@ void testBuilder_WithCustomExecutorService() { assertNotNull(config.getSerDes()); } + @Test + void payloadOffloadExecutorDefaultsToInline() { + var config = + DurableConfig.builder().withDurableExecutionClient(mockClient).build(); + + assertEquals(null, config.getPayloadOffloadExecutorService()); + } + + @Test + void payloadOffloadExecutorCannotAliasUserExecutor() { + var builder = DurableConfig.builder() + .withDurableExecutionClient(mockClient) + .withExecutorService(mockExecutor) + .withPayloadOffloadExecutorService(mockExecutor); + + var error = assertThrows(IllegalStateException.class, builder::build); + + assertEquals( + "Payload offload ExecutorService must be different from the user operation ExecutorService", + error.getMessage()); + } + + @Test + void chainedInvokePayloadOffloaderAcceptanceIsOptIn() { + var defaults = + DurableConfig.builder().withDurableExecutionClient(mockClient).build(); + var enabled = DurableConfig.builder() + .withDurableExecutionClient(mockClient) + .withPayloadOffloaderForChainedInvokePayloads(true) + .build(); + + assertFalse(defaults.shouldUsePayloadOffloaderForChainedInvokePayloads()); + assertTrue(enabled.shouldUsePayloadOffloaderForChainedInvokePayloads()); + } + @Test void testBuilder_DeserializeAfterSerializationDefaultsToTrue() { var config = diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java index 23f677e05..b6d352e95 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutionTest.java @@ -8,14 +8,20 @@ import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.when; import static software.amazon.lambda.durable.TypeToken.get; +import java.nio.charset.StandardCharsets; import java.time.Instant; +import java.util.Base64; import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutorService; import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; @@ -28,12 +34,28 @@ import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.TestUtils; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.DurableOperationException; +import software.amazon.lambda.durable.exception.PayloadOffloadException; +import software.amazon.lambda.durable.exception.RetryablePayloadOffloadException; +import software.amazon.lambda.durable.exception.SerDesException; import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.model.InvocationSource; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.offload.OffloadedPayload; +import software.amazon.lambda.durable.offload.PayloadOffloadContext; +import software.amazon.lambda.durable.offload.PayloadOffloader; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; +import software.amazon.lambda.durable.offload.internal.ChainedInvokeOutputFrame; +import software.amazon.lambda.durable.offload.internal.ChainedInvokePayloadFrame; import software.amazon.lambda.durable.operation.BaseDurableOperation; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.InvocationEndInfo; +import software.amazon.lambda.durable.plugin.InvocationStatus; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; class DurableExecutionTest { @@ -81,6 +103,555 @@ void testExecuteSuccess() { assertTrue(output.result().contains("Hello test-input")); } + @Test + void framedChainedInvokeInputUsesPayloadOffloaderWhenEnabled() { + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.inline(serializedPayload); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return payload.data(); + } + }; + var callerContext = PayloadOffloadContext.forOperation( + "arn:aws:lambda:us-east-1:123456789012:function:caller:$LATEST/durable-execution/name/id", + OperationIdentifier.of("invoke", "invoke", OperationSubType.CHAINED_INVOKE), + null, + SerDesPayloadKind.INVOKE_PAYLOAD, + null); + var encoded = new PayloadCodec(null).serialize("test-input", new JacksonSerDes(), offloader, callerContext); + var executionOp = executionOp().toBuilder() + .executionDetails(ExecutionDetails.builder() + .inputPayload(ChainedInvokePayloadFrame.encode(encoded)) + .build()) + .build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp)) + .build()); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPayloadOffloader(offloader) + .withPayloadOffloaderForChainedInvokePayloads(true) + .build(); + + var output = DurableExecutor.execute(input, null, get(String.class), (userInput, ctx) -> userInput, config); + var framedResult = ChainedInvokeOutputFrame.decode(output.result()); + var result = new PayloadCodec(null) + .deserialize( + framedResult.payload(), + get(String.class), + new JacksonSerDes(), + offloader, + PayloadOffloadContext.forExecution( + EXECUTION_ARN, EXECUTION_OP_ID, EXECUTION_NAME, SerDesPayloadKind.OUTPUT)); + + assertEquals(ExecutionStatus.SUCCEEDED, output.status()); + assertTrue(framedResult.usesPayloadCodec()); + assertEquals("test-input", result); + } + + @Test + void unframedInputDoesNotEnterPayloadOffloaderPipeline() { + var loadCount = new int[1]; + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.inline(serializedPayload); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + loadCount[0]++; + return payload.data(); + } + }; + var rawInput = new PayloadCodec(null) + .serialize( + "domain-input", + new JacksonSerDes(), + offloader, + PayloadOffloadContext.forExecution( + "arn:aws:lambda:us-east-1:123456789012:function:caller:$LATEST/durable-execution/name/id", + "id", + "caller", + SerDesPayloadKind.INPUT)); + SerDes passThroughSerDes = new SerDes() { + @Override + public String serialize(Object value) { + return (String) value; + } + + @Override + @SuppressWarnings("unchecked") + public T deserialize(String data, software.amazon.lambda.durable.TypeToken typeToken) { + return (T) data; + } + }; + var executionOp = executionOp().toBuilder() + .executionDetails( + ExecutionDetails.builder().inputPayload(rawInput).build()) + .build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp)) + .build()); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withSerDes(passThroughSerDes) + .withPayloadOffloader(offloader) + .withPayloadOffloaderForChainedInvokePayloads(true) + .build(); + var observedInput = new AtomicReference(); + + var output = DurableExecutor.execute( + input, + null, + get(String.class), + (userInput, ctx) -> { + observedInput.set(userInput); + return "result"; + }, + config); + + assertEquals(ExecutionStatus.SUCCEEDED, output.status()); + assertEquals(rawInput, observedInput.get()); + assertEquals(0, loadCount[0]); + } + + @Test + void unframedChainedInvocationKeepsOutputOnOrdinaryWire() { + var offloadCount = new AtomicInteger(); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPayloadOffloader(countingOffloader(offloadCount)) + .build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp())) + .build(), + List.of(), + InvocationSource.CHAINED_INVOKE); + + var output = DurableExecutor.execute(input, null, get(String.class), (userInput, ctx) -> "result", config); + + assertEquals(ExecutionStatus.SUCCEEDED, output.status()); + assertEquals("\"result\"", output.result()); + assertEquals(0, offloadCount.get()); + } + + @Test + void unframedChainedInvocationKeepsErrorOnOrdinaryWire() { + var offloadCount = new AtomicInteger(); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPayloadOffloader(countingOffloader(offloadCount)) + .build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp())) + .build(), + List.of(), + InvocationSource.CHAINED_INVOKE); + + var output = DurableExecutor.execute( + input, + null, + get(String.class), + (userInput, ctx) -> { + throw new IllegalStateException("failed"); + }, + config); + + assertEquals(ExecutionStatus.FAILED, output.status()); + assertFalse(PayloadCodec.isOffloadEnvelope(output.error().errorData())); + assertEquals(0, offloadCount.get()); + } + + @Test + void framedInputRejectsUnsupportedPayloadEnvelopeVersion() { + var executionOp = executionOp().toBuilder() + .executionDetails(ExecutionDetails.builder() + .inputPayload(ChainedInvokePayloadFrame.encode("@aws-durable-payload:v2:{}")) + .build()) + .build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp)) + .build()); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPayloadOffloaderForChainedInvokePayloads(true) + .build(); + + var output = DurableExecutor.execute(input, null, get(String.class), (userInput, ctx) -> userInput, config); + + assertEquals(ExecutionStatus.FAILED, output.status()); + assertTrue(output.error().errorMessage().contains("Unsupported or malformed")); + } + + @Test + void framedChainedInvokeFailureMarksExternalErrorDataAsRaw() { + var marker = "@aws-durable-payload:v2:{}"; + var executionOp = executionOp().toBuilder() + .executionDetails(ExecutionDetails.builder() + .inputPayload(ChainedInvokePayloadFrame.encode("\"test-input\"")) + .build()) + .build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp)) + .build(), + List.of(), + InvocationSource.CHAINED_INVOKE); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPayloadOffloaderForChainedInvokePayloads(true) + .build(); + var remoteError = ErrorObject.builder() + .errorType("RemoteError") + .errorMessage("remote failure") + .errorData(marker) + .build(); + var operation = Operation.builder() + .id("invoke") + .type(OperationType.CHAINED_INVOKE) + .status(OperationStatus.FAILED) + .build(); + + var output = DurableExecutor.execute( + input, + null, + get(String.class), + (userInput, ctx) -> { + throw new DurableOperationException(operation, remoteError); + }, + config); + var decodedError = ChainedInvokeOutputFrame.decode(output.error().errorData()); + + assertEquals(ExecutionStatus.FAILED, output.status()); + assertEquals(marker, decodedError.payload()); + assertFalse(decodedError.usesPayloadCodec()); + } + + @Test + void permanentRootOutputOffloadFailureProducesFailedOutputAndPluginEnd() { + var statuses = new CopyOnWriteArrayList(); + var config = configWithFailingPayloadKind(SerDesPayloadKind.OUTPUT, false, statuses); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp())) + .build()); + + var output = DurableExecutor.execute(input, null, get(String.class), (userInput, ctx) -> "result", config); + + assertEquals(ExecutionStatus.FAILED, output.status()); + assertEquals(PayloadOffloadException.class.getName(), output.error().errorType()); + assertEquals(List.of(InvocationStatus.FAILED), statuses); + } + + @Test + void retryableRootOutputOffloadFailureRetriesInvocationAndPluginEnd() { + var statuses = new CopyOnWriteArrayList(); + var config = configWithFailingPayloadKind(SerDesPayloadKind.OUTPUT, true, statuses); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp())) + .build()); + + assertThrows( + RetryablePayloadOffloadException.class, + () -> DurableExecutor.execute(input, null, get(String.class), (userInput, ctx) -> "result", config)); + assertEquals(List.of(InvocationStatus.RETRYING), statuses); + } + + @Test + void rootOutputSerDesFailureEscapesInvocation() { + var serDes = new JacksonSerDes() { + @Override + public String serialize(Object value) { + if ("result".equals(value)) { + throw new SerDesException("cannot serialize root output"); + } + return super.serialize(value); + } + }; + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withSerDes(serDes) + .build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp())) + .build()); + + var failure = assertThrows( + SerDesException.class, + () -> DurableExecutor.execute(input, null, get(String.class), (userInput, ctx) -> "result", config)); + + assertEquals("cannot serialize root output", failure.getMessage()); + } + + @Test + void errorDataOffloadFailureIsClassifiedBeforePluginEnd() { + var statuses = new CopyOnWriteArrayList(); + var config = configWithFailingPayloadKind(SerDesPayloadKind.EXCEPTION, false, statuses); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp())) + .build()); + + var output = DurableExecutor.execute( + input, + null, + get(String.class), + (userInput, ctx) -> { + throw new IllegalStateException("user failure"); + }, + config); + + assertEquals(ExecutionStatus.FAILED, output.status()); + assertEquals(PayloadOffloadException.class.getName(), output.error().errorType()); + assertEquals(List.of(InvocationStatus.FAILED), statuses); + } + + @Test + void errorSerDesFailureStillFiresInvocationEnd() { + var statuses = new CopyOnWriteArrayList(); + var reportedError = new AtomicReference(); + var original = new IllegalStateException("user failure"); + var serDes = new JacksonSerDes() { + @Override + public String serialize(Object value) { + if (value == original) { + throw new SerDesException("cannot serialize handler failure"); + } + return super.serialize(value); + } + }; + var plugin = new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + statuses.add(info.invocationStatus()); + reportedError.set(info.executionError()); + } + }; + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withSerDes(serDes) + .withPlugins(plugin) + .build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp())) + .build()); + + var failure = assertThrows( + SerDesException.class, + () -> DurableExecutor.execute( + input, + null, + get(String.class), + (userInput, ctx) -> { + throw original; + }, + config)); + + assertEquals("cannot serialize handler failure", failure.getMessage()); + assertEquals(List.of(InvocationStatus.FAILED), statuses); + assertEquals(original, reportedError.get()); + } + + @Test + void operationErrorIsReboundFromOverrideOffloaderToGlobalOffloader() { + var sourceLoads = new AtomicInteger(); + var targetOffloads = new AtomicInteger(); + var sourceOffloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.inline( + Base64.getEncoder().encodeToString(serializedPayload.getBytes(StandardCharsets.UTF_8))); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + sourceLoads.incrementAndGet(); + return new String(Base64.getDecoder().decode(payload.data()), StandardCharsets.UTF_8); + } + }; + var targetOffloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + targetOffloads.incrementAndGet(); + return OffloadedPayload.inline(serializedPayload); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return payload.data(); + } + }; + var producerContext = PayloadOffloadContext.forOperation( + EXECUTION_ARN, + OperationIdentifier.of("inner-step", "inner", OperationSubType.STEP), + null, + SerDesPayloadKind.EXCEPTION, + 1); + var sourcePayload = new PayloadCodec(null) + .serialize( + new IllegalStateException("nested failure"), + new JacksonSerDes(), + sourceOffloader, + producerContext); + var error = ErrorObject.builder() + .errorType(IllegalStateException.class.getName()) + .errorMessage("nested failure") + .errorData(sourcePayload) + .build(); + var operation = Operation.builder() + .id("inner-step") + .type(OperationType.STEP) + .status(OperationStatus.FAILED) + .build(); + var operationFailure = + new DurableOperationException(operation, error).withPayloadSource(sourceOffloader, producerContext); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPayloadOffloader(targetOffloader) + .build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp())) + .build()); + + var output = DurableExecutor.execute( + input, + null, + get(String.class), + (userInput, ctx) -> { + throw operationFailure; + }, + config); + var restored = new PayloadCodec(null) + .deserialize( + output.error().errorData(), + get(IllegalStateException.class), + new JacksonSerDes(), + targetOffloader, + PayloadOffloadContext.forExecution( + EXECUTION_ARN, EXECUTION_OP_ID, EXECUTION_NAME, SerDesPayloadKind.EXCEPTION)); + + assertEquals(ExecutionStatus.FAILED, output.status()); + assertEquals("nested failure", restored.getMessage()); + assertEquals(1, sourceLoads.get()); + assertEquals(1, targetOffloads.get()); + } + + @Test + void rawOperationErrorIsOffloadedAtExecutionBoundary() { + var targetOffloads = new AtomicInteger(); + var storedPayload = new AtomicReference(); + var targetOffloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + targetOffloads.incrementAndGet(); + storedPayload.set(serializedPayload); + return OffloadedPayload.inline(serializedPayload); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return payload.data(); + } + }; + var error = ErrorObject.builder() + .errorType("RemoteError") + .errorMessage("remote failure") + .errorData("raw-error-data") + .build(); + var operation = Operation.builder() + .id("invoke") + .type(OperationType.CHAINED_INVOKE) + .status(OperationStatus.FAILED) + .build(); + var operationFailure = new DurableOperationException(operation, error); + var config = DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPayloadOffloader(targetOffloader) + .build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp())) + .build()); + + var output = DurableExecutor.execute( + input, + null, + get(String.class), + (userInput, ctx) -> { + throw operationFailure; + }, + config); + var context = PayloadOffloadContext.forExecution( + EXECUTION_ARN, EXECUTION_OP_ID, EXECUTION_NAME, SerDesPayloadKind.EXCEPTION); + var restored = + new PayloadCodec(null).resolveSerializedPayload(output.error().errorData(), targetOffloader, context); + + assertEquals(ExecutionStatus.FAILED, output.status()); + assertEquals("raw-error-data", restored); + assertEquals("raw-error-data", storedPayload.get()); + assertEquals(1, targetOffloads.get()); + } + + @Test + void largeRootOutputCheckpointTransportFailureEscapesForInvocationRetry() { + var client = TestUtils.createMockClient(); + when(client.checkpoint(any(), any(), any())).thenThrow(new IllegalStateException("transport unavailable")); + var config = DurableConfig.builder().withDurableExecutionClient(client).build(); + var input = new DurableExecutionInput( + EXECUTION_ARN, + "token1", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(executionOp())) + .build()); + var largeResult = "x".repeat(6 * 1024 * 1024); + + var failure = assertThrows( + IllegalStateException.class, + () -> DurableExecutor.execute(input, null, get(String.class), (userInput, ctx) -> largeResult, config)); + + assertEquals("transport unavailable", failure.getMessage()); + } + @Test void testExecutePending() { var executionOp = Operation.builder() @@ -450,4 +1021,51 @@ private Operation executionOp() { .build()) .build(); } + + private static PayloadOffloader countingOffloader(AtomicInteger offloadCount) { + return new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + offloadCount.incrementAndGet(); + return OffloadedPayload.inline(serializedPayload); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return payload.data(); + } + }; + } + + private DurableConfig configWithFailingPayloadKind( + SerDesPayloadKind failingKind, boolean retryable, List statuses) { + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + if (context.payloadKind() == failingKind) { + if (retryable) { + throw new RetryablePayloadOffloadException("storage unavailable"); + } + throw new PayloadOffloadException("storage unavailable"); + } + return OffloadedPayload.inline(serializedPayload); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return payload.data(); + } + }; + var plugin = new DurableExecutionPlugin() { + @Override + public void onInvocationEnd(InvocationEndInfo info) { + statuses.add(info.invocationStatus()); + } + }; + return DurableConfig.builder() + .withDurableExecutionClient(TestUtils.createMockClient()) + .withPayloadOffloader(offloader) + .withPlugins(plugin) + .build(); + } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/PayloadCodecTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/PayloadCodecTest.java new file mode 100644 index 000000000..8117d49fb --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/PayloadCodecTest.java @@ -0,0 +1,704 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.fasterxml.jackson.databind.node.ObjectNode; +import java.util.ArrayList; +import java.util.Base64; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Consumer; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.exception.PayloadOffloadException; +import software.amazon.lambda.durable.model.OperationIdentifier; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.offload.OffloadedPayload; +import software.amazon.lambda.durable.offload.PayloadOffloadContext; +import software.amazon.lambda.durable.offload.PayloadOffloader; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; +import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; + +class PayloadCodecTest { + private static final String ENVELOPE_PREFIX = "@aws-durable-payload:v1:"; + private static final ObjectMapper TEST_MAPPER = new ObjectMapper(); + + private ExecutorService executor; + + @AfterEach + void shutdownExecutor() { + if (executor != null) { + executor.shutdownNow(); + } + } + + @Test + void legacyPayloadRemainsReadable() { + var codec = codec(); + var value = codec.deserialize( + "{\"value\":\"legacy\"}", TypeToken.get(TestValue.class), new JacksonSerDes(), null, context()); + + assertEquals("legacy", value.value()); + } + + @Test + void referencePayloadIsLoadedOnceAndDeserializedObjectIsCached() { + var offloader = new InMemoryOffloader(); + var writer = codec(); + var serDes = new JacksonSerDes(); + var payload = writer.serialize(new TestValue("stored"), serDes, offloader, context()); + writer.clear(); + + var reader = codec(); + var first = reader.deserialize(payload, TypeToken.get(TestValue.class), serDes, offloader, context()); + var second = reader.deserialize(payload, TypeToken.get(TestValue.class), serDes, offloader, context()); + + assertEquals("stored", first.value()); + assertSame(first, second); + assertEquals(1, offloader.loadCount.get()); + } + + @Test + void referenceIsLoadedAndVerifiedBeforeFirstDeserialization() { + var offloader = new InMemoryOffloader(); + var codec = codec(); + var serDes = new JacksonSerDes(); + var payload = codec.serialize(new TestValue("stored"), serDes, offloader, context()); + + var value = codec.deserialize(payload, TypeToken.get(TestValue.class), serDes, offloader, context()); + + assertEquals("stored", value.value()); + assertEquals(1, offloader.loadCount.get()); + } + + @Test + void inlinePayloadUsesLoadPathOnInitialAndFreshCodecDeserialization() { + var loadCount = new AtomicInteger(); + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.inline(Base64.getEncoder() + .encodeToString(serializedPayload.getBytes(java.nio.charset.StandardCharsets.UTF_8))); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + loadCount.incrementAndGet(); + return new String(Base64.getDecoder().decode(payload.data()), java.nio.charset.StandardCharsets.UTF_8); + } + }; + var serDes = new JacksonSerDes(); + var writer = codec(); + var payload = writer.serialize(new TestValue("stored"), serDes, offloader, context()); + + var initial = writer.deserialize(payload, TypeToken.get(TestValue.class), serDes, offloader, context()); + var replay = new PayloadCodec(null) + .deserialize(payload, TypeToken.get(TestValue.class), serDes, offloader, context()); + + assertEquals("stored", initial.value()); + assertEquals("stored", replay.value()); + assertEquals(2, loadCount.get()); + } + + @Test + void concurrentDeserializationSharesOneLoadAndOneObject() { + var offloader = new InMemoryOffloader() { + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + try { + Thread.sleep(100); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + return super.load(payload, context); + } + }; + var serDes = new JacksonSerDes(); + var writer = codec(); + var payload = writer.serialize(new TestValue("stored"), serDes, offloader, context()); + writer.clear(); + var reader = codec(); + + var first = CompletableFuture.supplyAsync( + () -> reader.deserialize(payload, TypeToken.get(TestValue.class), serDes, offloader, context())); + var second = CompletableFuture.supplyAsync( + () -> reader.deserialize(payload, TypeToken.get(TestValue.class), serDes, offloader, context())); + + assertSame(first.join(), second.join()); + assertEquals(1, offloader.loadCount.get()); + } + + @Test + void differentSerDesInstancesHaveIndependentObjectCaches() { + var codec = codec(); + SerDes firstSerDes = new FixedSerDes("one"); + SerDes secondSerDes = new FixedSerDes("two"); + + var first = codec.deserialize("payload", TypeToken.get(TestValue.class), firstSerDes, null, context()); + var second = codec.deserialize("payload", TypeToken.get(TestValue.class), secondSerDes, null, context()); + + assertEquals("one", first.value()); + assertEquals("two", second.value()); + } + + @Test + void changedSerializedDataInvalidatesObjectCache() { + var codec = codec(); + var serDes = new JacksonSerDes(); + + var first = codec.deserialize("{\"value\":\"one\"}", TypeToken.get(TestValue.class), serDes, null, context()); + var second = codec.deserialize("{\"value\":\"two\"}", TypeToken.get(TestValue.class), serDes, null, context()); + + assertEquals("one", first.value()); + assertEquals("two", second.value()); + } + + @Test + void attemptsHaveIndependentCachesWhenReferenceIsReused() { + var codec = codec(); + var serDes = new JacksonSerDes(); + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.reference("memory://shared", null); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return context.attempt() == 1 ? "{\"value\":\"one\"}" : "{\"value\":\"two\"}"; + } + }; + + var firstPayload = codec.serialize(new TestValue("one"), serDes, offloader, context(1)); + var secondPayload = codec.serialize(new TestValue("two"), serDes, offloader, context(2)); + + assertTrue(!firstPayload.equals(secondPayload)); + assertEquals( + "one", + codec.deserialize(firstPayload, TypeToken.get(TestValue.class), serDes, offloader, context(1)) + .value()); + assertEquals( + "two", + codec.deserialize(secondPayload, TypeToken.get(TestValue.class), serDes, offloader, context(2)) + .value()); + } + + @Test + void disabledOffloaderKeepsLegacyInlineFormat() { + var codec = codec(); + var payload = + codec.serialize(new TestValue("inline"), new JacksonSerDes(), PayloadOffloader.disabled(), context()); + + assertEquals("{\"value\":\"inline\"}", payload); + } + + @Test + void disabledOffloaderEscapesReservedMarkerAndRoundTripsCustomSerDes() { + var marker = "@aws-durable-payload:v2:{}"; + SerDes serDes = new PassThroughSerDes(); + var writer = codec(); + + var payload = writer.serialize(marker, serDes, PayloadOffloader.disabled(), context()); + var initial = writer.deserialize( + payload, TypeToken.get(String.class), serDes, PayloadOffloader.disabled(), context()); + var replay = new PayloadCodec(null) + .deserialize(payload, TypeToken.get(String.class), serDes, PayloadOffloader.disabled(), context()); + + assertTrue(payload.startsWith("@aws-durable-payload:v1:")); + assertEquals(marker, initial); + assertEquals(marker, replay); + } + + @Test + void escapedInlineEnvelopeBypassesConsumerOffloader() { + var marker = "@aws-durable-payload:v2:{}"; + var loadCount = new AtomicInteger(); + var consumerOffloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + throw new AssertionError("offload should not be called"); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + loadCount.incrementAndGet(); + return "transformed"; + } + }; + var payload = codec().serialize(marker, new PassThroughSerDes(), PayloadOffloader.disabled(), context()); + + var restored = codec().deserialize( + payload, TypeToken.get(String.class), new PassThroughSerDes(), consumerOffloader, context()); + + assertEquals(marker, restored); + assertEquals(0, loadCount.get()); + } + + @Test + void activeInlineOffloaderIsRequiredDuringReplay() { + var producerOffloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.inline(Base64.getEncoder() + .encodeToString(serializedPayload.getBytes(java.nio.charset.StandardCharsets.UTF_8))); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return new String(Base64.getDecoder().decode(payload.data()), java.nio.charset.StandardCharsets.UTF_8); + } + }; + var payload = codec().serialize(new TestValue("stored"), new JacksonSerDes(), producerOffloader, context()); + + var error = assertThrows(PayloadOffloadException.class, () -> codec().deserialize( + payload, TypeToken.get(TestValue.class), new JacksonSerDes(), null, context())); + + assertTrue(error.getMessage().contains("requires its producing offloader")); + } + + @Test + void versionedEnvelopeWithoutLoadSemanticsFailsClosed() { + var marker = "@aws-durable-payload:v2:{}"; + var payload = codec().serialize(marker, new PassThroughSerDes(), PayloadOffloader.disabled(), context()); + var missingLoadSemantics = payload.replace(",\"requiresLoad\":false", ""); + assertTrue(!payload.equals(missingLoadSemantics)); + + assertThrows(PayloadOffloadException.class, () -> codec().deserialize( + missingLoadSemantics, + TypeToken.get(String.class), + new PassThroughSerDes(), + PayloadOffloader.disabled(), + context())); + } + + @Test + void numericEnvelopeModeFailsClosed() { + var payload = codec().serialize( + "@aws-durable-payload:v2:{}", new PassThroughSerDes(), PayloadOffloader.disabled(), context()); + var numericMode = payload.replace("\"mode\":\"INLINE\"", "\"mode\":0"); + assertTrue(!payload.equals(numericMode)); + + assertThrows(PayloadOffloadException.class, () -> codec().deserialize( + numericMode, + TypeToken.get(String.class), + new PassThroughSerDes(), + PayloadOffloader.disabled(), + context())); + } + + @Test + void stringEnvelopeLoadFlagFailsClosed() { + var payload = codec().serialize( + "@aws-durable-payload:v2:{}", new PassThroughSerDes(), PayloadOffloader.disabled(), context()); + var stringLoadFlag = payload.replace("\"requiresLoad\":false", "\"requiresLoad\":\"false\""); + assertTrue(!payload.equals(stringLoadFlag)); + + assertThrows(PayloadOffloadException.class, () -> codec().deserialize( + stringLoadFlag, + TypeToken.get(String.class), + new PassThroughSerDes(), + PayloadOffloader.disabled(), + context())); + } + + @Test + void coercedEnvelopeScalarTypesFailClosed() { + var disabled = PayloadOffloader.disabled(); + var inlinePayload = + codec().serialize("@aws-durable-payload:v2:{}", new PassThroughSerDes(), disabled, context()); + var referenceOffloader = new InMemoryOffloader(); + var referencePayload = + codec().serialize(new TestValue("stored"), new JacksonSerDes(), referenceOffloader, context()); + + assertInvalidEnvelope(mutateEnvelope(inlinePayload, envelope -> envelope.put("data", 123)), disabled); + assertInvalidEnvelope( + mutateEnvelope(referencePayload, envelope -> envelope.put("reference", 123)), referenceOffloader); + assertInvalidEnvelope( + mutateEnvelope(inlinePayload, envelope -> envelope.put("ownerDurableExecutionArn", 123)), disabled); + assertInvalidEnvelope( + mutateEnvelope( + inlinePayload, envelope -> producerContext(envelope).put("attempt", "1")), + disabled); + assertInvalidEnvelope( + mutateEnvelope( + inlinePayload, envelope -> producerContext(envelope).put("attempt", 1.5)), + disabled); + assertInvalidEnvelope( + mutateEnvelope( + inlinePayload, envelope -> producerContext(envelope).put("operationType", 1)), + disabled); + assertInvalidEnvelope( + mutateEnvelope( + inlinePayload, envelope -> producerContext(envelope).put("operationSubType", 1)), + disabled); + } + + @Test + void opaqueSerializedPayloadPreservesReservedMarkerWithoutOffloader() { + var marker = "@aws-durable-payload:v2:{}"; + + var payload = codec().offloadSerializedPayload(marker, PayloadOffloader.disabled(), context()); + + assertEquals(marker, payload); + } + + @Test + void nullPayloadSkipsOffloaderAndRemainsNull() { + var codec = codec(); + var offloadCount = new AtomicInteger(); + var offloader = new InMemoryOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + offloadCount.incrementAndGet(); + return super.offload(serializedPayload, context); + } + }; + + var payload = codec.serialize(null, new JacksonSerDes(), offloader, context()); + + assertNull(payload); + assertNull(codec.deserialize(payload, TypeToken.get(String.class), new JacksonSerDes(), offloader, context())); + assertEquals(0, offloadCount.get()); + } + + @Test + void externalEnvelopeRequiresConfiguredOffloader() { + var offloader = new InMemoryOffloader(); + var writer = codec(); + var payload = writer.serialize(new TestValue("stored"), new JacksonSerDes(), offloader, context()); + writer.clear(); + + var reader = codec(); + assertThrows( + PayloadOffloadException.class, + () -> reader.deserialize( + payload, TypeToken.get(TestValue.class), new JacksonSerDes(), null, context())); + } + + @Test + void offloadRunsOnConfiguredExecutor() { + var threadName = new AtomicReference(); + executor = Executors.newSingleThreadExecutor(r -> new Thread(r, "payload-io-test")); + var codec = new PayloadCodec(executor); + var offloader = new InMemoryOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + threadName.set(Thread.currentThread().getName()); + return super.offload(serializedPayload, context); + } + }; + + codec.serialize(new TestValue("stored"), new JacksonSerDes(), offloader, context()); + + assertTrue(threadName.get().startsWith("payload-io-test")); + } + + @Test + void managedUserThreadRunsInlineWhenPayloadExecutorWrapsSameBackingPool() throws Exception { + executor = Executors.newSingleThreadExecutor(r -> new Thread(r, "shared-user-payload-test")); + var wrappedPayloadExecutor = Executors.unconfigurableExecutorService(executor); + var codec = new PayloadCodec(wrappedPayloadExecutor, () -> true); + var offloadThread = new AtomicReference(); + var offloader = new InMemoryOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + offloadThread.set(Thread.currentThread().getName()); + return super.offload(serializedPayload, context); + } + }; + + var payload = executor.submit( + () -> codec.serialize(new TestValue("stored"), new JacksonSerDes(), offloader, context())) + .get(5, TimeUnit.SECONDS); + + assertTrue(payload.startsWith(ENVELOPE_PREFIX)); + assertEquals("shared-user-payload-test", offloadThread.get()); + } + + @Test + void malformedEnvelopeReportsPayloadIdentity() { + var codec = codec(); + var error = assertThrows( + PayloadOffloadException.class, + () -> codec.deserialize( + "@aws-durable-payload:v1:not-json", + TypeToken.get(TestValue.class), + new JacksonSerDes(), + null, + context())); + + assertTrue(error.getMessage().contains("operation/op-1/result")); + } + + @Test + void envelopeWithTrailingJsonTokenFailsClosed() { + var payload = + codec().serialize(new TestValue("stored"), new JacksonSerDes(), new InMemoryOffloader(), context()); + + assertThrows(PayloadOffloadException.class, () -> codec().deserialize( + payload + "{}", + TypeToken.get(TestValue.class), + new JacksonSerDes(), + new InMemoryOffloader(), + context())); + } + + @Test + void unsupportedEnvelopeVersionFailsClosed() { + var codec = codec(); + + var error = assertThrows( + PayloadOffloadException.class, + () -> codec.deserialize( + "@aws-durable-payload:v2:{}", + TypeToken.get(TestValue.class), + new JacksonSerDes(), + null, + context())); + + assertTrue(error.getMessage().contains("Unsupported or malformed")); + } + + @Test + void versionedEnvelopeWithoutProducerMetadataFailsClosed() { + var codec = codec(); + var envelope = "@aws-durable-payload:v1:" + + new JacksonSerDes().serialize(OffloadedPayload.inline("{\"value\":\"stored\"}")); + + var error = assertThrows( + PayloadOffloadException.class, + () -> codec.deserialize( + envelope, TypeToken.get(TestValue.class), new JacksonSerDes(), null, context())); + + assertTrue(error.getMessage().contains("missing producer ownership or integrity metadata")); + } + + @Test + void mismatchedExceptionEnvelopeOwnerIsRejected() { + var codec = codec(); + var offloader = new InMemoryOffloader(); + var producer = PayloadOffloadContext.forOperation( + context().durableExecutionArn(), + OperationIdentifier.of("inner-step", "inner", OperationSubType.STEP), + "child-context", + SerDesPayloadKind.EXCEPTION, + 1); + var consumer = PayloadOffloadContext.forOperation( + context().durableExecutionArn(), + OperationIdentifier.of("child-context", "child", OperationSubType.RUN_IN_CHILD_CONTEXT), + null, + SerDesPayloadKind.EXCEPTION, + null); + var payload = + codec.serialize(new IllegalStateException("nested failure"), new JacksonSerDes(), offloader, producer); + + var error = assertThrows( + PayloadOffloadException.class, + () -> codec.deserialize( + payload, TypeToken.get(IllegalStateException.class), new JacksonSerDes(), offloader, consumer)); + + assertTrue(error.getMessage().contains("different durable entity")); + } + + @Test + void forwardedExceptionPayloadIsReboundToTargetOffloader() { + var codec = codec(); + var sourceOffloader = new InMemoryOffloader(); + var targetOffloader = new InMemoryOffloader(); + var producer = PayloadOffloadContext.forOperation( + context().durableExecutionArn(), + OperationIdentifier.of("inner-step", "inner", OperationSubType.STEP), + "child-context", + SerDesPayloadKind.EXCEPTION, + 1); + var consumer = PayloadOffloadContext.forOperation( + context().durableExecutionArn(), + OperationIdentifier.of("child-context", "child", OperationSubType.RUN_IN_CHILD_CONTEXT), + null, + SerDesPayloadKind.EXCEPTION, + null); + var sourcePayload = codec.serialize( + new IllegalStateException("nested failure"), new JacksonSerDes(), sourceOffloader, producer); + + var rebound = + codec.rebindSerializedPayload(sourcePayload, sourceOffloader, producer, targetOffloader, consumer); + var restored = codec.deserialize( + rebound, TypeToken.get(IllegalStateException.class), new JacksonSerDes(), targetOffloader, consumer); + + assertEquals("nested failure", restored.getMessage()); + assertEquals(1, sourceOffloader.loadCount.get()); + assertEquals(1, targetOffloader.loadCount.get()); + } + + @Test + void envelopeEncodingFailureIsClassifiedAsPayloadFailure() { + var recursivePreview = new java.util.HashMap(); + recursivePreview.put("self", recursivePreview); + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.reference("memory://payload", recursivePreview); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + throw new AssertionError("load should not be called"); + } + }; + + var error = assertThrows(PayloadOffloadException.class, () -> codec().serialize( + new TestValue("stored"), new JacksonSerDes(), offloader, context())); + + assertTrue(error.getMessage().contains("Failed to encode payload offload envelope")); + } + + @Test + @SuppressWarnings("unchecked") + void nullablePreviewMetadataRoundTripsThroughEnvelope() { + var preview = new LinkedHashMap(); + preview.put("nullable", null); + var nested = new LinkedHashMap(); + nested.put("values", new ArrayList<>(java.util.Arrays.asList("present", null))); + preview.put("nested", nested); + var storedPayload = new AtomicReference(); + var loadedPreview = new AtomicReference>(); + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + storedPayload.set(serializedPayload); + return OffloadedPayload.reference("memory://preview", preview); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + loadedPreview.set(payload.preview()); + return storedPayload.get(); + } + }; + var writer = codec(); + var payload = writer.serialize(new TestValue("stored"), new JacksonSerDes(), offloader, context()); + writer.clear(); + + var restored = + codec().deserialize(payload, TypeToken.get(TestValue.class), new JacksonSerDes(), offloader, context()); + + assertEquals("stored", restored.value()); + assertTrue(loadedPreview.get().containsKey("nullable")); + assertNull(loadedPreview.get().get("nullable")); + var restoredNested = (Map) loadedPreview.get().get("nested"); + var restoredValues = (List) restoredNested.get("values"); + assertEquals("present", restoredValues.get(0)); + assertNull(restoredValues.get(1)); + assertThrows( + UnsupportedOperationException.class, () -> loadedPreview.get().put("later", "value")); + assertThrows(UnsupportedOperationException.class, () -> ((List) restoredValues).add("later")); + } + + private PayloadCodec codec() { + executor = Executors.newCachedThreadPool(); + return new PayloadCodec(executor); + } + + private static void assertInvalidEnvelope(String payload, PayloadOffloader offloader) { + assertThrows(PayloadOffloadException.class, () -> new PayloadCodec(null) + .deserialize(payload, TypeToken.get(String.class), new PassThroughSerDes(), offloader, context())); + } + + private static String mutateEnvelope(String payload, Consumer mutation) { + try { + var envelope = (ObjectNode) TEST_MAPPER.readTree(payload.substring(ENVELOPE_PREFIX.length())); + mutation.accept(envelope); + return ENVELOPE_PREFIX + TEST_MAPPER.writeValueAsString(envelope); + } catch (Exception e) { + throw new AssertionError(e); + } + } + + private static ObjectNode producerContext(ObjectNode envelope) { + return (ObjectNode) envelope.get("producerContext"); + } + + private static PayloadOffloadContext context() { + return context(1); + } + + private static PayloadOffloadContext context(int attempt) { + return PayloadOffloadContext.forOperation( + "arn:aws:lambda:us-east-1:123456789012:function:test:$LATEST/durable-execution/name/invocation", + OperationIdentifier.of("op-1", "step", OperationSubType.STEP), + null, + SerDesPayloadKind.RESULT, + attempt); + } + + record TestValue(String value) {} + + private static final class FixedSerDes implements SerDes { + private final String value; + + private FixedSerDes(String value) { + this.value = value; + } + + @Override + public String serialize(Object value) { + return "payload"; + } + + @Override + @SuppressWarnings("unchecked") + public T deserialize(String data, TypeToken typeToken) { + return (T) new TestValue(value); + } + } + + private static final class PassThroughSerDes implements SerDes { + @Override + public String serialize(Object value) { + return (String) value; + } + + @Override + @SuppressWarnings("unchecked") + public T deserialize(String data, TypeToken typeToken) { + return (T) data; + } + } + + private static class InMemoryOffloader implements PayloadOffloader { + private final Map values = new ConcurrentHashMap<>(); + private final AtomicInteger sequence = new AtomicInteger(); + final AtomicInteger loadCount = new AtomicInteger(); + + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + var reference = "memory://" + sequence.incrementAndGet(); + values.put(reference, serializedPayload); + return OffloadedPayload.reference(reference, null); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + loadCount.incrementAndGet(); + return values.get(payload.reference()); + } + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/offload/OffloadedPayloadTest.java b/sdk/src/test/java/software/amazon/lambda/durable/offload/OffloadedPayloadTest.java new file mode 100644 index 000000000..180e306b0 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/offload/OffloadedPayloadTest.java @@ -0,0 +1,47 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; + +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import org.junit.jupiter.api.Test; + +class OffloadedPayloadTest { + @Test + @SuppressWarnings("unchecked") + void previewArraysAreDetachedAndNormalizedToImmutableLists() { + var tags = new String[] {"original", "stable"}; + var counts = new int[] {1, 2}; + var preview = new LinkedHashMap(); + preview.put("tags", tags); + preview.put("nested", Map.of("counts", counts)); + + var payload = OffloadedPayload.reference("memory://payload", preview); + tags[0] = "changed"; + counts[0] = 99; + + var storedTags = (List) payload.preview().get("tags"); + var nested = (Map) payload.preview().get("nested"); + var storedCounts = (List) nested.get("counts"); + assertEquals(List.of("original", "stable"), storedTags); + assertEquals(List.of(1, 2), storedCounts); + assertThrows(UnsupportedOperationException.class, () -> storedTags.add("later")); + assertThrows(UnsupportedOperationException.class, () -> storedCounts.set(0, 99)); + } + + @Test + void mutableScalarLikeValuesAreSnapshottedOrRejected() { + var text = new StringBuilder("original"); + var payload = OffloadedPayload.reference("memory://payload", Map.of("text", text)); + text.append("-changed"); + + assertEquals("original", payload.preview().get("text")); + assertThrows( + IllegalArgumentException.class, + () -> OffloadedPayload.reference("memory://payload", Map.of("unsupported", new java.util.Date()))); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/offload/internal/ChainedInvokeOutputFrameTest.java b/sdk/src/test/java/software/amazon/lambda/durable/offload/internal/ChainedInvokeOutputFrameTest.java new file mode 100644 index 000000000..e0520e0f7 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/offload/internal/ChainedInvokeOutputFrameTest.java @@ -0,0 +1,45 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload.internal; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.exception.PayloadOffloadException; + +class ChainedInvokeOutputFrameTest { + @Test + void codecAndRawPayloadsRoundTrip() { + var codec = ChainedInvokeOutputFrame.encode("@aws-durable-payload:v1:value", true); + var raw = ChainedInvokeOutputFrame.encode("@aws-durable-payload:v2:external", false); + + assertEquals("__durable_execution_chained_invoke_output:1:codec:@aws-durable-payload:v1:value", codec); + assertEquals("__durable_execution_chained_invoke_output:1:raw:@aws-durable-payload:v2:external", raw); + assertTrue(ChainedInvokeOutputFrame.isFramed(codec)); + assertTrue(ChainedInvokeOutputFrame.decode(codec).usesPayloadCodec()); + assertEquals( + "@aws-durable-payload:v1:value", + ChainedInvokeOutputFrame.decode(codec).payload()); + assertFalse(ChainedInvokeOutputFrame.decode(raw).usesPayloadCodec()); + assertEquals( + "@aws-durable-payload:v2:external", + ChainedInvokeOutputFrame.decode(raw).payload()); + } + + @Test + void nullRemainsUnframed() { + assertNull(ChainedInvokeOutputFrame.encode(null, true)); + assertFalse(ChainedInvokeOutputFrame.isFramed(null)); + } + + @Test + void unsupportedVersionFailsClosed() { + assertThrows( + PayloadOffloadException.class, + () -> ChainedInvokeOutputFrame.decode("__durable_execution_chained_invoke_output:2:raw:value")); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/offload/internal/ChainedInvokePayloadFrameTest.java b/sdk/src/test/java/software/amazon/lambda/durable/offload/internal/ChainedInvokePayloadFrameTest.java new file mode 100644 index 000000000..8a3a97147 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/offload/internal/ChainedInvokePayloadFrameTest.java @@ -0,0 +1,38 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.offload.internal; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import org.junit.jupiter.api.Test; +import software.amazon.lambda.durable.exception.PayloadOffloadException; + +class ChainedInvokePayloadFrameTest { + @Test + void frameRoundTripsOpaquePayload() { + var framed = ChainedInvokePayloadFrame.encode("value:with\nseparators"); + + assertEquals("__durable_execution_chained_invoke_payload:1:value:value:with\nseparators", framed); + assertTrue(ChainedInvokePayloadFrame.isFramed(framed)); + assertEquals("value:with\nseparators", ChainedInvokePayloadFrame.decode(framed)); + } + + @Test + void nullUsesExplicitFrame() { + var framed = ChainedInvokePayloadFrame.encode(null); + + assertEquals("__durable_execution_chained_invoke_payload:1:null", framed); + assertTrue(ChainedInvokePayloadFrame.isFramed(framed)); + assertNull(ChainedInvokePayloadFrame.decode(framed)); + } + + @Test + void unsupportedFrameFailsClosed() { + assertThrows( + PayloadOffloadException.class, + () -> ChainedInvokePayloadFrame.decode("__durable_execution_chained_invoke_payload:2:value")); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/CallbackOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/CallbackOperationTest.java index c8cea7934..fb39fa385 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/CallbackOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/CallbackOperationTest.java @@ -191,6 +191,45 @@ void getReturnsDeserializedResultWhenSucceeded() { assertEquals("approved", result); } + @Test + void callbackResultDoesNotInterpretPayloadOffloadEnvelopeMarkers() { + var markerPayload = "@aws-durable-payload:v2:{}"; + var existingCallback = Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .type(OperationType.CALLBACK) + .subType(OperationSubType.CALLBACK.getValue()) + .status(OperationStatus.SUCCEEDED) + .callbackDetails(CallbackDetails.builder() + .callbackId("callback-id") + .result(markerPayload) + .build()) + .build(); + var executionManager = createExecutionManager(List.of(existingCallback)); + when(durableContext.getExecutionManager()).thenReturn(executionManager); + var passThroughSerDes = new SerDes() { + @Override + public String serialize(Object value) { + return (String) value; + } + + @Override + @SuppressWarnings("unchecked") + public T deserialize(String data, TypeToken typeToken) { + return (T) data; + } + }; + var operation = new CallbackOperation<>( + OPERATION_IDENTIFIER, + TypeToken.get(String.class), + CallbackConfig.builder().serDes(passThroughSerDes).build(), + durableContext); + + operation.execute(); + + assertEquals(markerPayload, operation.get()); + } + @Test void getThrowsCallbackExceptionWhenFailed() { var existingCallback = Operation.builder() diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java index 99d994538..0d9d79f53 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/ChildContextOperationTest.java @@ -6,8 +6,13 @@ import static org.mockito.Mockito.*; import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import java.util.function.Function; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -20,16 +25,28 @@ import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; import software.amazon.lambda.durable.TypeToken; +import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.DurableContextImpl; import software.amazon.lambda.durable.exception.ChildContextFailedException; +import software.amazon.lambda.durable.exception.DurableOperationException; import software.amazon.lambda.durable.exception.NonDeterministicExecutionException; +import software.amazon.lambda.durable.exception.PayloadOffloadException; import software.amazon.lambda.durable.exception.SerDesException; import software.amazon.lambda.durable.execution.ExecutionManager; +import software.amazon.lambda.durable.execution.PayloadCodec; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.offload.OffloadedPayload; +import software.amazon.lambda.durable.offload.PayloadOffloadContext; +import software.amazon.lambda.durable.offload.PayloadOffloader; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; +import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; +import software.amazon.lambda.durable.plugin.OperationEndInfo; import software.amazon.lambda.durable.serde.JacksonSerDes; import software.amazon.lambda.durable.serde.SerDes; @@ -117,11 +134,28 @@ private ChildContextOperation createVirtualOperation(Function createOperationWithParent( Function func, ConcurrencyOperation parent) { + return createOperationWithParent(func, parent, false); + } + + private ChildContextOperation createOperationWithParent( + Function func, ConcurrencyOperation parent, boolean isVirtual) { + return createOperationWithParent(func, parent, isVirtual, null); + } + + private ChildContextOperation createOperationWithParent( + Function func, + ConcurrencyOperation parent, + boolean isVirtual, + PayloadOffloader payloadOffloader) { return new ChildContextOperation<>( OPERATION_IDENTIFIER, func, TypeToken.get(String.class), - RunInChildContextConfig.builder().serDes(SERDES).build(), + RunInChildContextConfig.builder() + .serDes(SERDES) + .payloadOffloader(payloadOffloader) + .isVirtual(isVirtual) + .build(), durableContext, parent); } @@ -414,4 +448,352 @@ void childSkipsFailureCheckpointWhenParentAlreadyCompleted() throws Exception { verify(executionManager, never()) .sendOperationUpdate(argThat(update -> update.action() == OperationAction.FAIL)); } + + @Test + void parentRejectedNonVirtualSuccessFiresOperationEnd() throws Exception { + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(null); + var operationEnd = new AtomicReference(); + when(durableContext.getDurableConfig()) + .thenReturn(DurableConfig.builder() + .withExecutorService(Executors.newCachedThreadPool()) + .withPlugins(new DurableExecutionPlugin() { + @Override + public void onOperationEnd(OperationEndInfo info) { + operationEnd.set(info); + } + }) + .build()); + var parent = new RecordingCompletionParent(durableContext); + parent.beginCompletion(); + var operation = createOperationWithParent(ctx -> "result", parent); + + operation.execute(); + operation.getRunningUserHandler().get(5, TimeUnit.SECONDS); + + assertNotNull(operationEnd.get()); + assertNull(operationEnd.get().error()); + } + + @Test + void parentRejectedNonVirtualFailureFiresOperationEnd() throws Exception { + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(null); + var operationEnd = new AtomicReference(); + when(durableContext.getDurableConfig()) + .thenReturn(DurableConfig.builder() + .withExecutorService(Executors.newCachedThreadPool()) + .withPlugins(new DurableExecutionPlugin() { + @Override + public void onOperationEnd(OperationEndInfo info) { + operationEnd.set(info); + } + }) + .build()); + var parent = new RecordingCompletionParent(durableContext); + parent.beginCompletion(); + var branchFailure = new IllegalStateException("branch failed"); + var operation = createOperationWithParent( + ctx -> { + throw branchFailure; + }, + parent); + + operation.execute(); + operation.getRunningUserHandler().get(5, TimeUnit.SECONDS); + + assertNotNull(operationEnd.get()); + assertSame(branchFailure, operationEnd.get().error()); + } + + @Test + void latePayloadFailureIsSkippedAfterParentCompletionStarts() throws Exception { + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(null); + var parent = new BlockingCompletionParent(durableContext); + var parentCompletion = CompletableFuture.runAsync(parent::beginCompletion); + assertTrue(parent.awaitCompletionStarted()); + + var lateFailure = new PayloadOffloadException("late payload failure"); + try { + var operation = createOperationWithParent( + ctx -> { + throw lateFailure; + }, + parent); + operation.execute(); + + Thread.sleep(100); + assertFalse(operation.getCompletionFuture().isDone()); + parent.releaseCompletion(); + parentCompletion.get(5, TimeUnit.SECONDS); + operation.getCompletionFuture().get(5, TimeUnit.SECONDS); + + assertTrue(operation.isOperationCompleted()); + assertSame(lateFailure, assertThrows(PayloadOffloadException.class, operation::get)); + verify(executionManager, never()).failInvocation(any()); + } finally { + parent.releaseCompletion(); + parentCompletion.get(5, TimeUnit.SECONDS); + } + } + + @Test + void lateVirtualPayloadFailureFiresOperationEndAndRemainsObservable() throws Exception { + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(null); + var operationEnd = new AtomicReference(); + var plugin = new DurableExecutionPlugin() { + @Override + public void onOperationEnd(OperationEndInfo info) { + if ("1".equals(info.id())) { + operationEnd.set(info); + } + } + }; + when(durableContext.getDurableConfig()) + .thenReturn(DurableConfig.builder() + .withExecutorService(Executors.newCachedThreadPool()) + .withPlugins(plugin) + .build()); + var parent = new BlockingCompletionParent(durableContext); + var parentCompletion = CompletableFuture.runAsync(parent::beginCompletion); + assertTrue(parent.awaitCompletionStarted()); + var lateFailure = new PayloadOffloadException("late virtual payload failure"); + + try { + var operation = createOperationWithParent( + ctx -> { + throw lateFailure; + }, + parent, + true); + operation.execute(); + + Thread.sleep(100); + assertFalse(operation.getCompletionFuture().isDone()); + parent.releaseCompletion(); + parentCompletion.get(5, TimeUnit.SECONDS); + assertSame(lateFailure, assertThrows(PayloadOffloadException.class, operation::get)); + assertNotNull(operationEnd.get()); + assertSame(lateFailure, operationEnd.get().error()); + verify(executionManager, never()).failInvocation(any()); + } finally { + parent.releaseCompletion(); + parentCompletion.get(5, TimeUnit.SECONDS); + } + } + + @Test + void lateForwardedFailureDoesNotLoadUnavailableSourcePayload() throws Exception { + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(null); + var parent = new BlockingCompletionParent(durableContext); + var parentCompletion = CompletableFuture.runAsync(parent::beginCompletion); + assertTrue(parent.awaitCompletionStarted()); + var loadCount = new AtomicInteger(); + var sourceOffloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.reference("memory://source-error", null); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + loadCount.incrementAndGet(); + throw new PayloadOffloadException("source unavailable"); + } + }; + var sourceContext = PayloadOffloadContext.forOperation( + "arn:aws:lambda:us-east-1:123456789012:function:test:$LATEST/durable-execution/name/id", + OperationIdentifier.of("source", "source", OperationSubType.STEP), + null, + SerDesPayloadKind.EXCEPTION, + 1); + var sourcePayload = new PayloadCodec(null) + .serialize(new IllegalStateException("source"), SERDES, sourceOffloader, sourceContext); + var sourceError = ErrorObject.builder() + .errorType(IllegalStateException.class.getName()) + .errorMessage("source") + .errorData(sourcePayload) + .build(); + var sourceOperation = Operation.builder() + .id("source") + .type(OperationType.STEP) + .status(OperationStatus.FAILED) + .build(); + var sourceFailure = new DurableOperationException(sourceOperation, sourceError) + .withPayloadSource(sourceOffloader, sourceContext); + + try { + var operation = createOperationWithParent( + ctx -> { + throw sourceFailure; + }, + parent); + operation.execute(); + + Thread.sleep(100); + assertFalse(operation.getCompletionFuture().isDone()); + parent.releaseCompletion(); + parentCompletion.get(5, TimeUnit.SECONDS); + + assertSame(sourceFailure, assertThrows(DurableOperationException.class, operation::get)); + assertEquals(0, loadCount.get()); + verify(executionManager, never()).failInvocation(any()); + } finally { + parent.releaseCompletion(); + parentCompletion.get(5, TimeUnit.SECONDS); + } + } + + @Test + void exceptionOffloadFailureClaimsParentBeforeEarlyCompletion() throws Exception { + when(executionManager.getOperationAndUpdateReplayState("1")).thenReturn(null); + when(executionManager.getDurableExecutionArn()) + .thenReturn("arn:aws:lambda:us-east-1:123456789012:function:test:$LATEST/durable-execution/name/id"); + when(executionManager.getPayloadCodec()).thenReturn(new PayloadCodec(null)); + var invocationFailed = new AtomicBoolean(); + doAnswer(invocation -> { + invocationFailed.set(true); + return null; + }) + .when(executionManager) + .failInvocation(any()); + when(executionManager.isExecutionCompletedExceptionally()).thenAnswer(invocation -> invocationFailed.get()); + + var offloadStarted = new CountDownLatch(1); + var releaseOffload = new CountDownLatch(1); + var payloadFailure = new PayloadOffloadException("exception offload failed"); + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + offloadStarted.countDown(); + try { + assertTrue(releaseOffload.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + throw payloadFailure; + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + throw new AssertionError("load should not be called"); + } + }; + var parent = new RecordingCompletionParent(durableContext); + var operation = createOperationWithParent( + ctx -> { + throw new IllegalStateException("branch failed"); + }, + parent, + false, + offloader); + + try { + operation.execute(); + assertTrue(offloadStarted.await(5, TimeUnit.SECONDS)); + + var parentCompletion = CompletableFuture.runAsync(parent::beginCompletion); + Thread.sleep(100); + assertFalse(parentCompletion.isDone(), "Early completion must wait for child exception persistence"); + + releaseOffload.countDown(); + operation.getRunningUserHandler().get(5, TimeUnit.SECONDS); + parentCompletion.get(5, TimeUnit.SECONDS); + + verify(executionManager).failInvocation(same(payloadFailure)); + assertFalse(parent.isCompletionHandled()); + } finally { + releaseOffload.countDown(); + } + } + + private static final class BlockingCompletionParent extends ConcurrencyOperation { + private final CountDownLatch completionStarted = new CountDownLatch(1); + private final CountDownLatch releaseCompletion = new CountDownLatch(1); + + private BlockingCompletionParent(DurableContextImpl durableContext) { + super( + OperationIdentifier.of("parent", "parent", OperationSubType.PARALLEL), + TypeToken.get(Void.class), + SERDES, + durableContext, + 1, + CompletionConfig.allSuccessful().completionDecisionFunction(), + NestingType.NESTED); + } + + private void beginCompletion() { + initiateCompletion( + CompletionConfig.CompletionDecision.complete(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED)); + } + + private boolean awaitCompletionStarted() throws InterruptedException { + return completionStarted.await(5, TimeUnit.SECONDS); + } + + private void releaseCompletion() { + releaseCompletion.countDown(); + } + + @Override + protected void handleCompletion(CompletionConfig.CompletionDecision completionDecision) { + completionStarted.countDown(); + try { + assertTrue(releaseCompletion.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + + @Override + protected void start() {} + + @Override + protected void replay(Operation existing) {} + + @Override + public Void get() { + return null; + } + } + + private static final class RecordingCompletionParent extends ConcurrencyOperation { + private final AtomicBoolean completionHandled = new AtomicBoolean(); + + private RecordingCompletionParent(DurableContextImpl durableContext) { + super( + OperationIdentifier.of("parent", "parent", OperationSubType.PARALLEL), + TypeToken.get(Void.class), + SERDES, + durableContext, + 1, + CompletionConfig.allSuccessful().completionDecisionFunction(), + NestingType.NESTED); + } + + private void beginCompletion() { + initiateCompletion( + CompletionConfig.CompletionDecision.complete(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED)); + } + + private boolean isCompletionHandled() { + return completionHandled.get(); + } + + @Override + protected void handleCompletion(CompletionConfig.CompletionDecision completionDecision) { + completionHandled.set(true); + } + + @Override + protected void start() {} + + @Override + protected void replay(Operation existing) {} + + @Override + public Void get() { + return null; + } + } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java index b6488139f..04f4ea5a7 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/ConcurrencyOperationTest.java @@ -12,6 +12,7 @@ import java.util.ArrayList; import java.util.List; import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CompletionException; import java.util.concurrent.CopyOnWriteArrayList; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; @@ -22,6 +23,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.ContextDetails; +import software.amazon.awssdk.services.lambda.model.ErrorObject; import software.amazon.awssdk.services.lambda.model.Operation; import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.OperationType; @@ -33,11 +35,14 @@ import software.amazon.lambda.durable.config.NestingType; import software.amazon.lambda.durable.config.RunInChildContextConfig; import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.exception.PayloadOffloadException; +import software.amazon.lambda.durable.exception.UnrecoverableDurableExecutionException; import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.OperationIdGenerator; import software.amazon.lambda.durable.execution.SuspendExecutionException; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; +import software.amazon.lambda.durable.model.ConcurrencyCompletionStatus; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; import software.amazon.lambda.durable.serde.JacksonSerDes; @@ -351,6 +356,153 @@ void exceptionalCompletionWakesWaitingCoordinator() throws Exception { assertTrue(op.awaitCoordinatorStopped()); } + @Test + void parentCompletionWaitsForReservedChildPersistence() throws Exception { + var operation = createOperation(CompletionConfig.firstSuccessful()); + var persistenceStarted = new CountDownLatch(1); + var releasePersistence = new CountDownLatch(1); + var persistenceCompleted = new AtomicBoolean(); + + var childPersistence = CompletableFuture.runAsync(() -> assertTrue(operation.persistChildCompletion(() -> { + persistenceStarted.countDown(); + try { + assertTrue(releasePersistence.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + persistenceCompleted.set(true); + }))); + assertTrue(persistenceStarted.await(5, TimeUnit.SECONDS)); + + var parentCompletion = CompletableFuture.runAsync(() -> operation.initiateCompletion( + CompletionConfig.CompletionDecision.complete(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED))); + Thread.sleep(50); + assertFalse(parentCompletion.isDone(), "Parent completion must wait until child persistence finishes"); + + releasePersistence.countDown(); + childPersistence.get(5, TimeUnit.SECONDS); + parentCompletion.get(5, TimeUnit.SECONDS); + + assertTrue(persistenceCompleted.get()); + assertTrue(operation.isSuccessHandled()); + assertFalse( + operation.persistChildCompletion(() -> fail("Persistence must be rejected after completion starts"))); + } + + @Test + void parentCompletionHoldsLockThroughAggregateCheckpointing() throws Exception { + var operation = new BlockingAggregateConcurrencyOperation(durableContext); + var parentCompletion = CompletableFuture.runAsync(() -> operation.initiateCompletion( + CompletionConfig.CompletionDecision.complete(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED))); + assertTrue(operation.awaitCompletionStarted()); + var lateChildRan = new AtomicBoolean(); + var lateChild = + CompletableFuture.supplyAsync(() -> operation.persistChildCompletion(() -> lateChildRan.set(true))); + + Thread.sleep(100); + assertFalse(lateChild.isDone()); + + operation.releaseCompletion(); + parentCompletion.get(5, TimeUnit.SECONDS); + + assertFalse(lateChild.get(5, TimeUnit.SECONDS)); + assertFalse(lateChildRan.get()); + } + + @Test + void childPayloadFailureClaimPreventsParentCompletion() throws Exception { + var operation = createOperation(CompletionConfig.firstSuccessful()); + var failure = new PayloadOffloadException("payload failure"); + + assertTrue(operation.claimChildPayloadFailure(failure)); + operation.initiateCompletion( + CompletionConfig.CompletionDecision.complete(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED)); + + assertFalse(operation.isSuccessHandled()); + assertFalse(operation.claimChildPayloadFailure(new PayloadOffloadException("later payload failure"))); + } + + @Test + void payloadFailureDuringChildPersistenceBeatsWaitingEarlyCompletion() throws Exception { + var operation = createOperation(CompletionConfig.firstSuccessful()); + var persistenceStarted = new CountDownLatch(1); + var releasePersistence = new CountDownLatch(1); + var payloadFailure = new PayloadOffloadException("exception offload failed"); + var observedFailure = new java.util.concurrent.atomic.AtomicReference(); + + var childPersistence = CompletableFuture.runAsync(() -> { + try { + operation.persistChildCompletion(() -> { + persistenceStarted.countDown(); + try { + assertTrue(releasePersistence.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + throw payloadFailure; + }); + } catch (Throwable failure) { + observedFailure.set(failure); + } + }); + assertTrue(persistenceStarted.await(5, TimeUnit.SECONDS)); + + var parentCompletion = CompletableFuture.runAsync(() -> operation.initiateCompletion( + CompletionConfig.CompletionDecision.complete(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED))); + Thread.sleep(100); + assertFalse(parentCompletion.isDone(), "Early completion must wait while exception persistence holds the lock"); + + releasePersistence.countDown(); + childPersistence.get(5, TimeUnit.SECONDS); + parentCompletion.get(5, TimeUnit.SECONDS); + + assertSame(payloadFailure, observedFailure.get()); + assertFalse(operation.isSuccessHandled()); + assertTrue(operation.claimChildPayloadFailure(payloadFailure)); + } + + @Test + void checkpointFailureDuringChildPersistenceBeatsWaitingEarlyCompletion() throws Exception { + var operation = createOperation(CompletionConfig.firstSuccessful()); + var persistenceStarted = new CountDownLatch(1); + var releasePersistence = new CountDownLatch(1); + var checkpointFailure = new UnrecoverableDurableExecutionException( + ErrorObject.builder().errorMessage("checkpoint failed").build()); + var observedFailure = new java.util.concurrent.atomic.AtomicReference(); + + var childPersistence = CompletableFuture.runAsync(() -> { + try { + operation.persistChildCompletion(() -> { + persistenceStarted.countDown(); + try { + assertTrue(releasePersistence.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + throw new CompletionException(checkpointFailure); + }); + } catch (Throwable failure) { + observedFailure.set(failure); + } + }); + assertTrue(persistenceStarted.await(5, TimeUnit.SECONDS)); + + var parentCompletion = CompletableFuture.runAsync(() -> operation.initiateCompletion( + CompletionConfig.CompletionDecision.complete(ConcurrencyCompletionStatus.MIN_SUCCESSFUL_REACHED))); + Thread.sleep(100); + assertFalse(parentCompletion.isDone()); + + releasePersistence.countDown(); + childPersistence.get(5, TimeUnit.SECONDS); + parentCompletion.get(5, TimeUnit.SECONDS); + + assertSame(checkpointFailure, observedFailure.get()); + assertFalse(operation.isSuccessHandled()); + } + // ===== Test subclass ===== static class TestConcurrencyOperation extends ConcurrencyOperation { @@ -521,6 +673,41 @@ boolean awaitCoordinatorStopped() throws Exception { } } + static class BlockingAggregateConcurrencyOperation extends TestConcurrencyOperation { + private final CountDownLatch completionStarted = new CountDownLatch(1); + private final CountDownLatch releaseCompletion = new CountDownLatch(1); + + BlockingAggregateConcurrencyOperation(DurableContextImpl durableContext) { + super( + OperationIdentifier.of(OPERATION_ID, "test-concurrency", OperationSubType.PARALLEL), + RESULT_TYPE, + SER_DES, + durableContext, + Integer.MAX_VALUE, + CompletionConfig.firstSuccessful()); + } + + @Override + protected void handleCompletion(CompletionConfig.CompletionDecision completionDecision) { + completionStarted.countDown(); + try { + assertTrue(releaseCompletion.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + super.handleCompletion(completionDecision); + } + + boolean awaitCompletionStarted() throws InterruptedException { + return completionStarted.await(5, TimeUnit.SECONDS); + } + + void releaseCompletion() { + releaseCompletion.countDown(); + } + } + static class ControlledChildOperation extends ChildContextOperation { private final AtomicInteger activeCount; diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/InvokeOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/InvokeOperationTest.java index 2c1d76c74..26b20b368 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/InvokeOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/InvokeOperationTest.java @@ -3,10 +3,17 @@ package software.amazon.lambda.durable.operation; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.argThat; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import software.amazon.awssdk.services.lambda.model.ChainedInvokeDetails; @@ -20,12 +27,21 @@ import software.amazon.lambda.durable.exception.InvokeFailedException; import software.amazon.lambda.durable.exception.InvokeStoppedException; import software.amazon.lambda.durable.exception.InvokeTimedOutException; +import software.amazon.lambda.durable.exception.PayloadOffloadException; import software.amazon.lambda.durable.execution.ExecutionManager; +import software.amazon.lambda.durable.execution.PayloadCodec; import software.amazon.lambda.durable.execution.ThreadContext; import software.amazon.lambda.durable.execution.ThreadType; import software.amazon.lambda.durable.model.OperationIdentifier; import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.offload.OffloadedPayload; +import software.amazon.lambda.durable.offload.PayloadOffloadContext; +import software.amazon.lambda.durable.offload.PayloadOffloader; +import software.amazon.lambda.durable.offload.SerDesPayloadKind; +import software.amazon.lambda.durable.offload.internal.ChainedInvokeOutputFrame; +import software.amazon.lambda.durable.offload.internal.ChainedInvokePayloadFrame; import software.amazon.lambda.durable.serde.JacksonSerDes; +import software.amazon.lambda.durable.serde.SerDes; class InvokeOperationTest { private static final String OPERATION_ID = "2"; @@ -100,6 +116,220 @@ void getInvokeFailedExceptionWhenInvocationFailed() { assertEquals("errorMessage", ex.getMessage()); } + @Test + void standardLambdaMarkerResultRemainsExternalData() { + var marker = "@aws-durable-payload:v2:{}"; + var loadCount = new AtomicInteger(); + var offloader = countingOffloader(loadCount); + var op = Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .status(OperationStatus.SUCCEEDED) + .chainedInvokeDetails( + ChainedInvokeDetails.builder().result(marker).build()) + .build(); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); + var operation = new InvokeOperation<>( + OPERATION_IDENTIFIER, + "standard-function", + "{}", + TypeToken.get(String.class), + InvokeConfig.builder() + .serDes(new PassThroughSerDes()) + .payloadOffloader(offloader) + .build(), + durableContext); + operation.onCheckpointComplete(op); + + assertEquals(marker, operation.get()); + assertEquals(0, loadCount.get()); + } + + @Test + void standardLambdaMarkerErrorDoesNotAttachPayloadSource() { + var marker = "@aws-durable-payload:v2:{}"; + var op = Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .status(OperationStatus.FAILED) + .chainedInvokeDetails(ChainedInvokeDetails.builder() + .error(ErrorObject.builder() + .errorType("RemoteError") + .errorMessage("remote failure") + .errorData(marker) + .build()) + .build()) + .build(); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); + var operation = new InvokeOperation<>( + OPERATION_IDENTIFIER, + "standard-function", + "{}", + TypeToken.get(String.class), + InvokeConfig.builder() + .serDes(new PassThroughSerDes()) + .payloadOffloader(countingOffloader(new AtomicInteger())) + .build(), + durableContext); + operation.onCheckpointComplete(op); + + var failure = assertThrows(InvokeFailedException.class, operation::get); + + assertEquals(marker, failure.getErrorObject().errorData()); + assertNull(failure.getPayloadOffloadContext()); + } + + @Test + void durableTargetRawMarkerErrorRemainsExternalData() { + var marker = "@aws-durable-payload:v2:{}"; + var framedError = ChainedInvokeOutputFrame.encode(marker, false); + var op = Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .status(OperationStatus.FAILED) + .chainedInvokeDetails(ChainedInvokeDetails.builder() + .error(ErrorObject.builder() + .errorType("RemoteError") + .errorMessage("remote failure") + .errorData(framedError) + .build()) + .build()) + .build(); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); + var operation = new InvokeOperation<>( + OPERATION_IDENTIFIER, + "durable-function", + "{}", + TypeToken.get(String.class), + InvokeConfig.builder() + .serDes(new PassThroughSerDes()) + .payloadOffloader(countingOffloader(new AtomicInteger())) + .usePayloadOffloaderForPayload(true) + .build(), + durableContext); + operation.onCheckpointComplete(op); + + var failure = assertThrows(InvokeFailedException.class, operation::get); + + assertEquals(marker, failure.getErrorObject().errorData()); + assertNull(failure.getPayloadOffloadContext()); + } + + @Test + void malformedCodecResultFrameFailsClosed() { + var op = Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .status(OperationStatus.SUCCEEDED) + .chainedInvokeDetails(ChainedInvokeDetails.builder() + .result(ChainedInvokeOutputFrame.encode("\"ordinary-json\"", true)) + .build()) + .build(); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); + configurePayloadCodec(); + var operation = new InvokeOperation<>( + OPERATION_IDENTIFIER, + "durable-function", + "{}", + TypeToken.get(String.class), + InvokeConfig.builder() + .serDes(new JacksonSerDes()) + .usePayloadOffloaderForPayload(true) + .build(), + durableContext); + operation.onCheckpointComplete(op); + + assertThrows(PayloadOffloadException.class, operation::get); + } + + @Test + void malformedCodecErrorFrameFailsClosed() { + var op = Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .status(OperationStatus.FAILED) + .chainedInvokeDetails(ChainedInvokeDetails.builder() + .error(ErrorObject.builder() + .errorType("RemoteError") + .errorMessage("remote failure") + .errorData(ChainedInvokeOutputFrame.encode("ordinary-error", true)) + .build()) + .build()) + .build(); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); + configurePayloadCodec(); + var operation = new InvokeOperation<>( + OPERATION_IDENTIFIER, + "durable-function", + "{}", + TypeToken.get(String.class), + InvokeConfig.builder() + .serDes(new PassThroughSerDes()) + .usePayloadOffloaderForPayload(true) + .build(), + durableContext); + operation.onCheckpointComplete(op); + + assertThrows(PayloadOffloadException.class, operation::get); + } + + @Test + void codecFramedErrorLoadsReferenceBeforeThrowing() { + configurePayloadCodec(); + var stored = new AtomicReference(); + var loadCount = new AtomicInteger(); + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + stored.set(serializedPayload); + return OffloadedPayload.reference("memory://invoke-error", null); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + loadCount.incrementAndGet(); + return stored.get(); + } + }; + var codec = executionManager.getPayloadCodec(); + var payload = codec.serializePreEncodedPayload("serialized-error", offloader, invokeErrorContext()); + var op = failedInvoke(ChainedInvokeOutputFrame.encode(payload, true)); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); + var operation = invokeOperation(offloader); + operation.onCheckpointComplete(op); + + var failure = assertThrows(InvokeFailedException.class, operation::get); + + assertEquals("serialized-error", failure.getErrorObject().errorData()); + assertEquals(1, loadCount.get()); + assertNull(failure.getPayloadOffloadContext()); + } + + @Test + void codecFramedErrorRejectsTamperedReferenceContent() { + configurePayloadCodec(); + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.reference("memory://invoke-error", null); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return "tampered"; + } + }; + var payload = executionManager + .getPayloadCodec() + .serializePreEncodedPayload("serialized-error", offloader, invokeErrorContext()); + var op = failedInvoke(ChainedInvokeOutputFrame.encode(payload, true)); + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(op); + var operation = invokeOperation(offloader); + operation.onCheckpointComplete(op); + + assertThrows(PayloadOffloadException.class, operation::get); + } + @Test void getInvokeTimedOutExceptionWhenInvocationTimedOut() { var op = Operation.builder() @@ -189,4 +419,154 @@ void getInvokeFailedExceptionWhenInvocationEndedUnexpectedly() { assertThrows(InvokeException.class, () -> operation.get()); } + + @Test + void invokeRequestPayloadIsNotOffloaded() { + var offloadCount = new AtomicInteger(); + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + offloadCount.incrementAndGet(); + return OffloadedPayload.inline(serializedPayload); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return payload.data(); + } + }; + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(null); + when(executionManager.sendOperationUpdate(any())).thenReturn(CompletableFuture.completedFuture(null)); + + var operation = new InvokeOperation<>( + OPERATION_IDENTIFIER, + "test-function", + new InvokePayload("request"), + TypeToken.get(String.class), + InvokeConfig.builder() + .serDes(new JacksonSerDes()) + .payloadOffloader(offloader) + .build(), + durableContext); + + operation.execute(); + + verify(executionManager) + .sendOperationUpdate(argThat(update -> "{\"value\":\"request\"}".equals(update.payload()))); + assertEquals(0, offloadCount.get()); + } + + @Test + void invokeRequestCanExplicitlyUsePayloadOffloader() { + var offloader = new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.inline(serializedPayload); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + return payload.data(); + } + }; + when(executionManager.getOperationAndUpdateReplayState(OPERATION_ID)).thenReturn(null); + when(executionManager.sendOperationUpdate(any())).thenReturn(CompletableFuture.completedFuture(null)); + when(executionManager.getPayloadCodec()).thenReturn(new PayloadCodec(null)); + when(executionManager.getDurableExecutionArn()) + .thenReturn("arn:aws:lambda:us-east-1:123456789012:function:test:$LATEST/durable-execution/name/id"); + + var operation = new InvokeOperation<>( + OPERATION_IDENTIFIER, + "test-function", + new InvokePayload("request"), + TypeToken.get(String.class), + InvokeConfig.builder() + .serDes(new JacksonSerDes()) + .payloadOffloader(offloader) + .usePayloadOffloaderForPayload(true) + .build(), + durableContext); + + operation.execute(); + + verify(executionManager).sendOperationUpdate(argThat(update -> { + var payload = update.payload(); + return ChainedInvokePayloadFrame.isFramed(payload) + && ChainedInvokePayloadFrame.decode(payload).startsWith("@aws-durable-payload:v1:"); + })); + } + + private static PayloadOffloader countingOffloader(AtomicInteger loadCount) { + return new PayloadOffloader() { + @Override + public OffloadedPayload offload(String serializedPayload, PayloadOffloadContext context) { + return OffloadedPayload.inline(serializedPayload); + } + + @Override + public String load(OffloadedPayload payload, PayloadOffloadContext context) { + loadCount.incrementAndGet(); + return payload.data(); + } + }; + } + + private void configurePayloadCodec() { + when(executionManager.getPayloadCodec()).thenReturn(new PayloadCodec(null)); + when(executionManager.getDurableExecutionArn()) + .thenReturn("arn:aws:lambda:us-east-1:123456789012:function:test:$LATEST/durable-execution/name/id"); + } + + private InvokeOperation invokeOperation(PayloadOffloader offloader) { + return new InvokeOperation<>( + OPERATION_IDENTIFIER, + "durable-function", + "{}", + TypeToken.get(String.class), + InvokeConfig.builder() + .serDes(new PassThroughSerDes()) + .payloadOffloader(offloader) + .usePayloadOffloaderForPayload(true) + .build(), + durableContext); + } + + private static Operation failedInvoke(String errorData) { + return Operation.builder() + .id(OPERATION_ID) + .name(OPERATION_NAME) + .status(OperationStatus.FAILED) + .chainedInvokeDetails(ChainedInvokeDetails.builder() + .error(ErrorObject.builder() + .errorType("RemoteError") + .errorMessage("remote failure") + .errorData(errorData) + .build()) + .build()) + .build(); + } + + private PayloadOffloadContext invokeErrorContext() { + return PayloadOffloadContext.forOperation( + executionManager.getDurableExecutionArn(), + OPERATION_IDENTIFIER, + null, + SerDesPayloadKind.EXCEPTION, + null); + } + + private static final class PassThroughSerDes implements SerDes { + @Override + public String serialize(Object value) { + return (String) value; + } + + @Override + @SuppressWarnings("unchecked") + public T deserialize(String data, TypeToken typeToken) { + return (T) data; + } + } + + private record InvokePayload(String value) {} } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/serde/DurableInputOutputSerDesTest.java b/sdk/src/test/java/software/amazon/lambda/durable/serde/DurableInputOutputSerDesTest.java index d64a88665..922df87a5 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/serde/DurableInputOutputSerDesTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/serde/DurableInputOutputSerDesTest.java @@ -12,6 +12,7 @@ import software.amazon.lambda.durable.model.DurableExecutionInput; import software.amazon.lambda.durable.model.DurableExecutionOutput; import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.model.InvocationSource; class DurableInputOutputSerDesTest { @@ -50,6 +51,26 @@ void testObjectMapperDeserializesDurableExecutionInput() { assertEquals("arn:aws:lambda:us-east-1:123456789012:function:my-function", input.durableExecutionArn()); assertEquals("token-123", input.checkpointToken()); assertNotNull(input.initialExecutionState()); + assertEquals(InvocationSource.DIRECT, input.invocationSource()); + } + + @Test + void testObjectMapperDeserializesChainedInvokeSource() { + var json = """ + { + "DurableExecutionArn": "arn:aws:lambda:us-east-1:123456789012:function:my-function", + "CheckpointToken": "token-123", + "InitialExecutionState": { + "Operations": [], + "NextMarker": null + }, + "InvocationSource": "CHAINED_INVOKE" + } + """; + + var input = serDes.deserialize(json, TypeToken.get(DurableExecutionInput.class)); + + assertEquals(InvocationSource.CHAINED_INVOKE, input.invocationSource()); } @Test