Skip to content

Commit eae41bd

Browse files
hsilanhsilan
committed
feat: refactor WithRetry to be implemented by consumers of DurableContext
1 parent 344f828 commit eae41bd

18 files changed

Lines changed: 627 additions & 1326 deletions

File tree

‎examples/src/main/java/software/amazon/lambda/durable/examples/callback/RetryWaitForCallbackExample.java‎

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,17 +8,17 @@
88
import software.amazon.lambda.durable.config.WithRetryConfig;
99
import software.amazon.lambda.durable.examples.types.ApprovalRequest;
1010
import software.amazon.lambda.durable.retry.RetryDecision;
11-
import software.amazon.lambda.durable.util.WithRetryHelper;
1211

1312
/**
14-
* Example demonstrating {@link WithRetryHelper} with {@code context.waitForCallback}.
13+
* Example demonstrating {@code context.withRetry} with {@code context.waitForCallback}.
1514
*
1615
* <p>Submits an approval request to an external system via a callback. If the callback fails (e.g., the external system
1716
* rejects the request), the helper retries the entire waitForCallback cycle — creating a fresh callback with a new ID
1817
* each time.
1918
*
2019
* <p>Each attempt uses a unique callback name ({@code "approval-1"}, {@code "approval-2"}, etc.) so the execution
21-
* history stays clean and replay-safe. The anonymous form is used, so attempts run directly in the caller's context.
20+
* history stays clean and replay-safe. The anonymous form is used, so attempts are grouped under a default-named child
21+
* context.
2222
*/
2323
public class RetryWaitForCallbackExample extends DurableHandler<ApprovalRequest, String> {
2424

@@ -33,8 +33,7 @@ public String handleRequest(ApprovalRequest input, DurableContext context) {
3333
stepCtx -> "Approval for: " + input.description() + " ($" + input.amount() + ")");
3434

3535
// Step 2: waitForCallback with retry — if the external system fails, try again with a fresh callback
36-
var approvalResult = WithRetryHelper.withRetry(
37-
context,
36+
var approvalResult = context.withRetry(
3837
(ctx, attempt) -> ctx.waitForCallback(
3938
"approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
4039
.info("Attempt {}: sending callback {} to approval system", attempt, callbackId)),

‎examples/src/main/java/software/amazon/lambda/durable/examples/invoke/RetryInvokeExample.java‎

Lines changed: 3 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -8,25 +8,23 @@
88
import software.amazon.lambda.durable.config.WithRetryConfig;
99
import software.amazon.lambda.durable.examples.types.GreetingRequest;
1010
import software.amazon.lambda.durable.retry.RetryDecision;
11-
import software.amazon.lambda.durable.util.WithRetryHelper;
1211

1312
/**
14-
* Example demonstrating {@link WithRetryHelper} with {@code context.invoke}.
13+
* Example demonstrating {@code context.withRetry} with {@code context.invoke}.
1514
*
1615
* <p>Retries a chained Lambda invocation up to 3 times with a fixed 2-second backoff between attempts. Each attempt
1716
* uses a unique operation name ({@code "call-greeting-1"}, {@code "call-greeting-2"}, etc.) so the execution history
1817
* stays clean and replay-safe.
1918
*
20-
* <p>The anonymous form is used, so attempts run directly in the caller's context without child-context wrapping.
19+
* <p>The anonymous form is used, so attempts are grouped under a default-named child context.
2120
*/
2221
public class RetryInvokeExample extends DurableHandler<GreetingRequest, String> {
2322

2423
private static final int MAX_ATTEMPTS = 3;
2524

2625
@Override
2726
public String handleRequest(GreetingRequest input, DurableContext context) {
28-
return WithRetryHelper.withRetry(
29-
context,
27+
return context.withRetry(
3028
(ctx, attempt) -> ctx.invoke(
3129
"call-greeting-" + attempt,
3230
"simple-step-example" + input.getName() + ":$LATEST",

‎sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryInvokeIntegrationTest.java‎

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -13,16 +13,14 @@
1313
import software.amazon.lambda.durable.retry.RetryDecision;
1414
import software.amazon.lambda.durable.retry.RetryStrategies;
1515
import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
16-
import software.amazon.lambda.durable.util.WithRetryHelper;
1716

1817
class RetryInvokeIntegrationTest {
1918

2019
@Test
2120
void invokeSucceedsOnFirstAttempt() {
2221
var runner = LocalDurableTestRunner.create(
2322
String.class,
24-
(input, context) -> WithRetryHelper.withRetry(
25-
context,
23+
(input, context) -> context.withRetry(
2624
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
2725
WithRetryConfig.builder()
2826
.retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
@@ -42,8 +40,7 @@ void invokeSucceedsOnFirstAttempt() {
4240
void invokeRetriesAfterFailure() {
4341
var runner = LocalDurableTestRunner.create(
4442
String.class,
45-
(input, context) -> WithRetryHelper.withRetry(
46-
context,
43+
(input, context) -> context.withRetry(
4744
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
4845
WithRetryConfig.builder()
4946
.retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
@@ -80,8 +77,7 @@ void invokeRetriesAfterFailure() {
8077
void invokeFailsAfterAllRetriesExhausted() {
8178
var runner = LocalDurableTestRunner.create(
8279
String.class,
83-
(input, context) -> WithRetryHelper.withRetry(
84-
context,
80+
(input, context) -> context.withRetry(
8581
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
8682
WithRetryConfig.builder()
8783
.retryStrategy((error, attempt) ->
@@ -114,8 +110,7 @@ void invokeFailsAfterAllRetriesExhausted() {
114110
void invokeRetryWithCustomBackoffDelay() {
115111
var runner = LocalDurableTestRunner.create(
116112
String.class,
117-
(input, context) -> WithRetryHelper.withRetry(
118-
context,
113+
(input, context) -> context.withRetry(
119114
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
120115
WithRetryConfig.builder()
121116
.retryStrategy((error, attempt) -> attempt < 3
@@ -148,8 +143,7 @@ void invokeRetryWithStepsBeforeAndAfter() {
148143
var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
149144
var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
150145

151-
var invokeResult = WithRetryHelper.withRetry(
152-
context,
146+
var invokeResult = context.withRetry(
153147
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
154148
WithRetryConfig.builder()
155149
.retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
@@ -174,8 +168,7 @@ void invokeRetryWithStepsBeforeAndAfter() {
174168
void invokeRetryPreservesOriginalExceptionType() {
175169
var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
176170
try {
177-
return WithRetryHelper.withRetry(
178-
context,
171+
return context.withRetry(
179172
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
180173
WithRetryConfig.builder()
181174
.retryStrategy(RetryStrategies.Presets.NO_RETRY)

‎sdk-integration-tests/src/test/java/software/amazon/lambda/durable/RetryWaitForCallbackIntegrationTest.java‎

Lines changed: 6 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -12,16 +12,14 @@
1212
import software.amazon.lambda.durable.retry.RetryDecision;
1313
import software.amazon.lambda.durable.retry.RetryStrategies;
1414
import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
15-
import software.amazon.lambda.durable.util.WithRetryHelper;
1615

1716
class RetryWaitForCallbackIntegrationTest {
1817

1918
@Test
2019
void waitForCallbackSucceedsOnFirstAttempt() {
2120
var runner = LocalDurableTestRunner.create(
2221
String.class,
23-
(input, context) -> WithRetryHelper.withRetry(
24-
context,
22+
(input, context) -> context.withRetry(
2523
(ctx, attempt) -> ctx.waitForCallback(
2624
"approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
2725
.info("Submitting callback {}", callbackId)),
@@ -47,8 +45,7 @@ void waitForCallbackSucceedsOnFirstAttempt() {
4745
void waitForCallbackRetriesAfterFailure() {
4846
var runner = LocalDurableTestRunner.create(
4947
String.class,
50-
(input, context) -> WithRetryHelper.withRetry(
51-
context,
48+
(input, context) -> context.withRetry(
5249
(ctx, attempt) -> ctx.waitForCallback(
5350
"approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
5451
.info("Attempt {} callback {}", attempt, callbackId)),
@@ -94,8 +91,7 @@ void waitForCallbackRetriesAfterFailure() {
9491
void waitForCallbackFailsAfterAllRetriesExhausted() {
9592
var runner = LocalDurableTestRunner.create(
9693
String.class,
97-
(input, context) -> WithRetryHelper.withRetry(
98-
context,
94+
(input, context) -> context.withRetry(
9995
(ctx, attempt) ->
10096
ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
10197
WithRetryConfig.builder()
@@ -132,8 +128,7 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() {
132128
var runner = LocalDurableTestRunner.create(String.class, (input, context) -> {
133129
var prefix = context.step("prepare", String.class, stepCtx -> "prepared");
134130

135-
var callbackResult = WithRetryHelper.withRetry(
136-
context,
131+
var callbackResult = context.withRetry(
137132
(ctx, attempt) ->
138133
ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
139134
WithRetryConfig.builder()
@@ -161,8 +156,7 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() {
161156
void waitForCallbackRetryMultipleFailuresThenSuccess() {
162157
var runner = LocalDurableTestRunner.create(
163158
String.class,
164-
(input, context) -> WithRetryHelper.withRetry(
165-
context,
159+
(input, context) -> context.withRetry(
166160
(ctx, attempt) ->
167161
ctx.waitForCallback("cb-" + attempt, String.class, (callbackId, stepCtx) -> {}),
168162
WithRetryConfig.builder()
@@ -207,8 +201,7 @@ void waitForCallbackRetryWithSubmitterLogic() {
207201
// Verify the submitter runs on each retry attempt
208202
var runner = LocalDurableTestRunner.create(
209203
String.class,
210-
(input, context) -> WithRetryHelper.withRetry(
211-
context,
204+
(input, context) -> context.withRetry(
212205
(ctx, attempt) ->
213206
ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {
214207
// Submitter runs each attempt — in a real scenario this would

‎sdk/src/main/java/software/amazon/lambda/durable/DurableContext.java‎

Lines changed: 66 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,11 @@
1616
import software.amazon.lambda.durable.config.StepConfig;
1717
import software.amazon.lambda.durable.config.WaitForCallbackConfig;
1818
import software.amazon.lambda.durable.config.WaitForConditionConfig;
19+
import software.amazon.lambda.durable.config.WithRetryConfig;
1920
import software.amazon.lambda.durable.context.BaseContext;
2021
import software.amazon.lambda.durable.model.MapResult;
2122
import software.amazon.lambda.durable.model.WaitForConditionResult;
23+
import software.amazon.lambda.durable.model.WithRetry;
2224

2325
public interface DurableContext extends BaseContext {
2426
/**
@@ -727,6 +729,70 @@ <T> DurableFuture<T> waitForConditionAsync(
727729
BiFunction<T, StepContext, WaitForConditionResult<T>> checkFunc,
728730
WaitForConditionConfig<T> config);
729731

732+
// =============== withRetry ================
733+
734+
/**
735+
* Replay-safe retry loop for any durable operation (named form, sync).
736+
*
737+
* <p>Provides the same retry-with-backoff pattern that {@code step()} has built in, but for operations that cannot
738+
* live inside a step ({@code waitForCallback}, {@code invoke}, {@code waitForCondition}, etc.).
739+
*
740+
* <p>Every side-effect in the loop is a durable operation, so the loop is replay-safe by construction. On replay,
741+
* completed operations return cached results instantly and the loop fast-forwards to the current attempt.
742+
*
743+
* <p>By default, the retry loop runs directly on the caller's context. If
744+
* {@link WithRetryConfig#wrapInChildContext()} is enabled, the loop is wrapped in a child context so all attempts
745+
* are grouped under a single named operation in execution history.
746+
*
747+
* @param <T> the result type
748+
* @param name operation name (used for backoff wait names, and as the child context name when wrapping)
749+
* @param operation the retryable operation — receives the context and 1-based attempt number
750+
* @param config retry configuration including the retry strategy and child context wrapping
751+
* @return the operation result
752+
*/
753+
<T> T withRetry(String name, WithRetry<T> operation, WithRetryConfig config);
754+
755+
/**
756+
* Replay-safe retry loop for any durable operation (anonymous form, sync).
757+
*
758+
* <p>By default, the retry loop runs directly on the caller's context. If
759+
* {@link WithRetryConfig#wrapInChildContext()} is enabled, the loop is wrapped in a child context with a default
760+
* name so all attempts are grouped under a single operation in execution history.
761+
*
762+
* @param <T> the result type
763+
* @param operation the retryable operation — receives the context and 1-based attempt number
764+
* @param config retry configuration including the retry strategy and child context wrapping
765+
* @return the operation result
766+
*/
767+
<T> T withRetry(WithRetry<T> operation, WithRetryConfig config);
768+
769+
/**
770+
* Replay-safe retry loop for any durable operation (named form, async).
771+
*
772+
* <p>Wraps the retry loop in {@code runInChildContextAsync} so all attempts are grouped under a single named
773+
* operation in execution history, and returns a {@link DurableFuture} that can be composed or blocked on.
774+
*
775+
* @param <T> the result type
776+
* @param name operation name (used for child context and backoff wait names)
777+
* @param operation the retryable operation — receives the context and 1-based attempt number
778+
* @param config retry configuration including the retry strategy
779+
* @return a future representing the operation result
780+
*/
781+
<T> DurableFuture<T> withRetryAsync(String name, WithRetry<T> operation, WithRetryConfig config);
782+
783+
/**
784+
* Replay-safe retry loop for any durable operation (anonymous form, async).
785+
*
786+
* <p>Wraps the retry loop in {@code runInChildContextAsync} with a default name so all attempts are grouped under a
787+
* single operation in execution history, and returns a {@link DurableFuture} that can be composed or blocked on.
788+
*
789+
* @param <T> the result type
790+
* @param operation the retryable operation — receives the context and 1-based attempt number
791+
* @param config retry configuration including the retry strategy
792+
* @return a future representing the operation result
793+
*/
794+
<T> DurableFuture<T> withRetryAsync(WithRetry<T> operation, WithRetryConfig config);
795+
730796
/**
731797
* Function applied to each item in a map operation.
732798
*

‎sdk/src/main/java/software/amazon/lambda/durable/config/RetryOperationConfig.java‎

Lines changed: 0 additions & 103 deletions
This file was deleted.

0 commit comments

Comments
 (0)