Skip to content

Commit c1455a8

Browse files
SilanHehsilan
andauthored
feat: withRetry helper (#343)
feat: withRetry helper (#343) --------- Co-authored-by: hsilan <hsilan@amazon.com>
1 parent 8e5d886 commit c1455a8

15 files changed

Lines changed: 2132 additions & 2 deletions

File tree

Lines changed: 50 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,50 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
package software.amazon.lambda.durable.examples.callback;
4+
5+
import java.time.Duration;
6+
import software.amazon.lambda.durable.DurableContext;
7+
import software.amazon.lambda.durable.DurableHandler;
8+
import software.amazon.lambda.durable.config.WithRetryConfig;
9+
import software.amazon.lambda.durable.examples.types.ApprovalRequest;
10+
import software.amazon.lambda.durable.retry.RetryDecision;
11+
12+
/**
13+
* Example demonstrating {@code context.withRetry} with {@code context.waitForCallback}.
14+
*
15+
* <p>Submits an approval request to an external system via a callback. If the callback fails (e.g., the external system
16+
* rejects the request), the helper retries the entire waitForCallback cycle — creating a fresh callback with a new ID
17+
* each time.
18+
*
19+
* <p>Each attempt uses a unique callback name ({@code "approval-1"}, {@code "approval-2"}, etc.) so the execution
20+
* history stays clean and replay-safe. A {@code null} name is used, so attempts are grouped under a default-named child
21+
* context.
22+
*/
23+
public class RetryWaitForCallbackExample extends DurableHandler<ApprovalRequest, String> {
24+
25+
private static final int MAX_ATTEMPTS = 3;
26+
27+
@Override
28+
public String handleRequest(ApprovalRequest input, DurableContext context) {
29+
// Step 1: Prepare the approval request
30+
var prepared = context.step(
31+
"prepare",
32+
String.class,
33+
stepCtx -> "Approval for: " + input.description() + " ($" + input.amount() + ")");
34+
35+
// Step 2: waitForCallback with retry — if the external system fails, try again with a fresh callback
36+
var approvalResult = context.withRetry(
37+
null,
38+
(attempt, ctx) -> ctx.waitForCallback(
39+
"approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
40+
.info("Attempt {}: sending callback {} to approval system", attempt, callbackId)),
41+
WithRetryConfig.builder()
42+
.retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS
43+
? RetryDecision.retry(Duration.ofSeconds(2))
44+
: RetryDecision.fail())
45+
.build());
46+
47+
// Step 3: Process the result
48+
return context.step("process-result", String.class, stepCtx -> prepared + " - Result: " + approvalResult);
49+
}
50+
}
Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
package software.amazon.lambda.durable.examples.invoke;
4+
5+
import java.time.Duration;
6+
import software.amazon.lambda.durable.DurableContext;
7+
import software.amazon.lambda.durable.DurableHandler;
8+
import software.amazon.lambda.durable.config.WithRetryConfig;
9+
import software.amazon.lambda.durable.examples.types.GreetingRequest;
10+
import software.amazon.lambda.durable.retry.RetryDecision;
11+
12+
/**
13+
* Example demonstrating {@code context.withRetry} with {@code context.invoke}.
14+
*
15+
* <p>Retries a chained Lambda invocation up to 3 times with a fixed 2-second backoff between attempts. Each attempt
16+
* uses a unique operation name ({@code "call-greeting-1"}, {@code "call-greeting-2"}, etc.) so the execution history
17+
* stays clean and replay-safe.
18+
*
19+
* <p>A {@code null} name is used, so attempts are grouped under a default-named child context.
20+
*/
21+
public class RetryInvokeExample extends DurableHandler<GreetingRequest, String> {
22+
23+
private static final int MAX_ATTEMPTS = 3;
24+
25+
@Override
26+
public String handleRequest(GreetingRequest input, DurableContext context) {
27+
return context.withRetry(
28+
null,
29+
(attempt, ctx) -> ctx.invoke(
30+
"call-greeting-" + attempt,
31+
"simple-step-example" + input.getName() + ":$LATEST",
32+
input,
33+
String.class),
34+
WithRetryConfig.builder()
35+
.retryStrategy((error, attempt) -> attempt < MAX_ATTEMPTS
36+
? RetryDecision.retry(Duration.ofSeconds(2))
37+
: RetryDecision.fail())
38+
.build());
39+
}
40+
}

examples/src/test/java/software/amazon/lambda/durable/examples/CloudBasedIntegrationTest.java

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -710,6 +710,89 @@ void testComplexFlatMapExample() {
710710
assertTrue(output.contains("reason=MIN_SUCCESSFUL_REACHED"));
711711
}
712712

713+
@Test
714+
void testRetryInvokeExample() {
715+
var runner = CloudDurableTestRunner.create(
716+
arn("retry-invoke-example"), GreetingRequest.class, String.class, lambdaClient);
717+
// The handler invokes "simple-step-example" + input.getName() + ":$LATEST",
718+
// so passing the functionNameSuffix as the name targets the deployed simple-step-example function
719+
var result = runner.run(new GreetingRequest(functionNameSuffix));
720+
721+
assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
722+
assertNotNull(result.getResult());
723+
}
724+
725+
@Test
726+
void testRetryWaitForCallbackExampleSucceeds() {
727+
var runner = CloudDurableTestRunner.create(
728+
arn("retry-wait-for-callback-example"), ApprovalRequest.class, String.class, lambdaClient);
729+
730+
var execution = runner.startAsync(new ApprovalRequest("Server upgrade", 5000.0));
731+
732+
// Wait for the first callback from attempt 1
733+
execution.pollUntil(exec -> exec.hasCallback("approval-1-callback"));
734+
var callbackId = execution.getCallbackId("approval-1-callback");
735+
assertNotNull(callbackId);
736+
737+
// Complete the callback on the first attempt
738+
execution.completeCallback(callbackId, "\"approved\"");
739+
740+
var result = execution.pollUntilComplete();
741+
assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
742+
743+
var finalResult = result.getResult();
744+
assertNotNull(finalResult);
745+
assertTrue(finalResult.contains("Approval for: Server upgrade"));
746+
assertTrue(finalResult.contains("5000"));
747+
assertTrue(finalResult.contains("approved"));
748+
749+
// Verify operations
750+
assertNotNull(execution.getOperation("prepare"));
751+
assertNotNull(execution.getOperation("process-result"));
752+
}
753+
754+
@Test
755+
void testRetryWaitForCallbackExampleRetriesAfterFailure() {
756+
var runner = CloudDurableTestRunner.create(
757+
arn("retry-wait-for-callback-example"), ApprovalRequest.class, String.class, lambdaClient);
758+
759+
var execution = runner.startAsync(new ApprovalRequest("Expensive item", 10000.0));
760+
761+
// Wait for the first callback from attempt 1
762+
execution.pollUntil(exec -> exec.hasCallback("approval-1-callback"));
763+
var callbackId1 = execution.getCallbackId("approval-1-callback");
764+
assertNotNull(callbackId1);
765+
766+
// Fail the first attempt
767+
execution.failCallback(
768+
callbackId1,
769+
ErrorObject.builder()
770+
.errorType("Rejected")
771+
.errorMessage("denied by reviewer")
772+
.build());
773+
774+
// Wait for the second callback from attempt 2 (after backoff)
775+
execution.pollUntil(exec -> exec.hasCallback("approval-2-callback"));
776+
var callbackId2 = execution.getCallbackId("approval-2-callback");
777+
assertNotNull(callbackId2);
778+
779+
// Complete the second attempt
780+
execution.completeCallback(callbackId2, "\"approved on retry\"");
781+
782+
var result = execution.pollUntilComplete();
783+
assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
784+
785+
var finalResult = result.getResult();
786+
assertNotNull(finalResult);
787+
assertTrue(finalResult.contains("Approval for: Expensive item"));
788+
assertTrue(finalResult.contains("10000"));
789+
assertTrue(finalResult.contains("approved on retry"));
790+
791+
// Verify operations
792+
assertNotNull(execution.getOperation("prepare"));
793+
assertNotNull(execution.getOperation("process-result"));
794+
}
795+
713796
@Test
714797
void testWaitForConditionExample() {
715798
var runner = CloudDurableTestRunner.create(
Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
2+
// SPDX-License-Identifier: Apache-2.0
3+
package software.amazon.lambda.durable.examples.callback;
4+
5+
import static org.junit.jupiter.api.Assertions.*;
6+
7+
import org.junit.jupiter.api.Test;
8+
import software.amazon.awssdk.services.lambda.model.ErrorObject;
9+
import software.amazon.lambda.durable.examples.types.ApprovalRequest;
10+
import software.amazon.lambda.durable.model.ExecutionStatus;
11+
import software.amazon.lambda.durable.testing.LocalDurableTestRunner;
12+
13+
class RetryWaitForCallbackExampleTest {
14+
15+
@Test
16+
void succeedsOnFirstAttempt() {
17+
var handler = new RetryWaitForCallbackExample();
18+
var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler);
19+
var input = new ApprovalRequest("New laptop", 1500.00);
20+
21+
// First run — prepares request, starts waitForCallback, suspends
22+
var result = runner.run(input);
23+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
24+
25+
// Complete the callback (waitForCallback names it "approval-1-callback" internally)
26+
var callbackId = runner.getCallbackId("approval-1-callback");
27+
assertNotNull(callbackId, "Callback 'approval-1-callback' should have been created");
28+
runner.completeCallback(callbackId, "\"Approved by manager\"");
29+
30+
// Run to completion
31+
result = runner.runUntilComplete(input);
32+
33+
assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
34+
assertEquals(
35+
"Approval for: New laptop ($1500.0) - Result: Approved by manager", result.getResult(String.class));
36+
}
37+
38+
@Test
39+
void retriesAfterFirstCallbackFails() {
40+
var handler = new RetryWaitForCallbackExample();
41+
var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler);
42+
var input = new ApprovalRequest("Server upgrade", 5000.00);
43+
44+
// First run — prepares, starts waitForCallback attempt 1, suspends
45+
var result = runner.run(input);
46+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
47+
48+
// Fail the first callback
49+
var callbackId1 = runner.getCallbackId("approval-1-callback");
50+
assertNotNull(callbackId1);
51+
runner.failCallback(
52+
callbackId1,
53+
ErrorObject.builder()
54+
.errorType("RejectedError")
55+
.errorMessage("Rejected by first reviewer")
56+
.build());
57+
58+
// Run — processes failure, hits backoff wait, suspends
59+
result = runner.run(input);
60+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
61+
62+
// Advance past the backoff wait
63+
runner.advanceTime();
64+
65+
// Run — starts waitForCallback attempt 2, suspends
66+
result = runner.run(input);
67+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
68+
69+
// Complete the second callback
70+
var callbackId2 = runner.getCallbackId("approval-2-callback");
71+
assertNotNull(callbackId2, "Callback 'approval-2-callback' should have been created after retry");
72+
runner.completeCallback(callbackId2, "\"Approved on second try\"");
73+
74+
// Run to completion
75+
result = runner.runUntilComplete(input);
76+
77+
assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus());
78+
assertEquals(
79+
"Approval for: Server upgrade ($5000.0) - Result: Approved on second try",
80+
result.getResult(String.class));
81+
}
82+
83+
@Test
84+
void failsAfterAllRetriesExhausted() {
85+
var handler = new RetryWaitForCallbackExample();
86+
var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler);
87+
var input = new ApprovalRequest("Expensive item", 10000.00);
88+
89+
// First run — starts waitForCallback attempt 1
90+
var result = runner.run(input);
91+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
92+
93+
// Fail callback attempt 1
94+
var callbackId1 = runner.getCallbackId("approval-1-callback");
95+
runner.failCallback(
96+
callbackId1,
97+
ErrorObject.builder()
98+
.errorType("Rejected")
99+
.errorMessage("fail 1")
100+
.build());
101+
result = runner.run(input);
102+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
103+
104+
// Advance past backoff 1, run to start attempt 2
105+
runner.advanceTime();
106+
result = runner.run(input);
107+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
108+
109+
// Fail callback attempt 2
110+
var callbackId2 = runner.getCallbackId("approval-2-callback");
111+
runner.failCallback(
112+
callbackId2,
113+
ErrorObject.builder()
114+
.errorType("Rejected")
115+
.errorMessage("fail 2")
116+
.build());
117+
result = runner.run(input);
118+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
119+
120+
// Advance past backoff 2, run to start attempt 3
121+
runner.advanceTime();
122+
result = runner.run(input);
123+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
124+
125+
// Fail callback attempt 3 — last attempt, retryStrategy returns fail()
126+
var callbackId3 = runner.getCallbackId("approval-3-callback");
127+
runner.failCallback(
128+
callbackId3,
129+
ErrorObject.builder()
130+
.errorType("Rejected")
131+
.errorMessage("fail 3")
132+
.build());
133+
result = runner.run(input);
134+
135+
assertEquals(ExecutionStatus.FAILED, result.getStatus());
136+
}
137+
138+
@Test
139+
void suspendsOnFirstRun() {
140+
var handler = new RetryWaitForCallbackExample();
141+
var runner = LocalDurableTestRunner.create(ApprovalRequest.class, handler);
142+
var input = new ApprovalRequest("Test item", 100.00);
143+
144+
var result = runner.run(input);
145+
146+
assertEquals(ExecutionStatus.PENDING, result.getStatus());
147+
}
148+
}

0 commit comments

Comments
 (0)