Skip to content

Commit 3b98eec

Browse files
committed
chore: use BiFunction<Integer,DurableContext, T> instead of WithRetry<DurableContext,Integer>
1 parent fbcd329 commit 3b98eec

9 files changed

Lines changed: 75 additions & 147 deletions

File tree

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ public String handleRequest(ApprovalRequest input, DurableContext context) {
3535
// Step 2: waitForCallback with retry — if the external system fails, try again with a fresh callback
3636
var approvalResult = context.withRetry(
3737
null,
38-
(ctx, attempt) -> ctx.waitForCallback(
38+
(attempt, ctx) -> ctx.waitForCallback(
3939
"approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
4040
.info("Attempt {}: sending callback {} to approval system", attempt, callbackId)),
4141
WithRetryConfig.builder()

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ public class RetryInvokeExample extends DurableHandler<GreetingRequest, String>
2626
public String handleRequest(GreetingRequest input, DurableContext context) {
2727
return context.withRetry(
2828
null,
29-
(ctx, attempt) -> ctx.invoke(
29+
(attempt, ctx) -> ctx.invoke(
3030
"call-greeting-" + attempt,
3131
"simple-step-example" + input.getName() + ":$LATEST",
3232
input,

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ void invokeSucceedsOnFirstAttempt() {
2222
String.class,
2323
(input, context) -> context.withRetry(
2424
null,
25-
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
25+
(attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
2626
WithRetryConfig.builder()
2727
.retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
2828
.build()));
@@ -43,7 +43,7 @@ void invokeRetriesAfterFailure() {
4343
String.class,
4444
(input, context) -> context.withRetry(
4545
null,
46-
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
46+
(attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
4747
WithRetryConfig.builder()
4848
.retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(2)))
4949
.build()));
@@ -81,7 +81,7 @@ void invokeFailsAfterAllRetriesExhausted() {
8181
String.class,
8282
(input, context) -> context.withRetry(
8383
null,
84-
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
84+
(attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
8585
WithRetryConfig.builder()
8686
.retryStrategy((error, attempt) ->
8787
attempt < 2 ? RetryDecision.retry(Duration.ofSeconds(1)) : RetryDecision.fail())
@@ -115,7 +115,7 @@ void invokeRetryWithCustomBackoffDelay() {
115115
String.class,
116116
(input, context) -> context.withRetry(
117117
null,
118-
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
118+
(attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
119119
WithRetryConfig.builder()
120120
.retryStrategy((error, attempt) -> attempt < 3
121121
? RetryDecision.retry(Duration.ofSeconds(attempt * 5L))
@@ -149,7 +149,7 @@ void invokeRetryWithStepsBeforeAndAfter() {
149149

150150
var invokeResult = context.withRetry(
151151
null,
152-
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
152+
(attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
153153
WithRetryConfig.builder()
154154
.retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
155155
.build());
@@ -175,7 +175,7 @@ void invokeRetryPreservesOriginalExceptionType() {
175175
try {
176176
return context.withRetry(
177177
null,
178-
(ctx, attempt) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
178+
(attempt, ctx) -> ctx.invoke("invoke-" + attempt, "target-fn", "{}", String.class),
179179
WithRetryConfig.builder()
180180
.retryStrategy(RetryStrategies.Presets.NO_RETRY)
181181
.build());

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

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ void waitForCallbackSucceedsOnFirstAttempt() {
2121
String.class,
2222
(input, context) -> context.withRetry(
2323
null,
24-
(ctx, attempt) -> ctx.waitForCallback(
24+
(attempt, ctx) -> ctx.waitForCallback(
2525
"approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
2626
.info("Submitting callback {}", callbackId)),
2727
WithRetryConfig.builder()
@@ -48,7 +48,7 @@ void waitForCallbackRetriesAfterFailure() {
4848
String.class,
4949
(input, context) -> context.withRetry(
5050
null,
51-
(ctx, attempt) -> ctx.waitForCallback(
51+
(attempt, ctx) -> ctx.waitForCallback(
5252
"approval-" + attempt, String.class, (callbackId, stepCtx) -> stepCtx.getLogger()
5353
.info("Attempt {} callback {}", attempt, callbackId)),
5454
WithRetryConfig.builder()
@@ -95,7 +95,7 @@ void waitForCallbackFailsAfterAllRetriesExhausted() {
9595
String.class,
9696
(input, context) -> context.withRetry(
9797
null,
98-
(ctx, attempt) ->
98+
(attempt, ctx) ->
9999
ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
100100
WithRetryConfig.builder()
101101
.retryStrategy((error, attempt) ->
@@ -133,7 +133,7 @@ void waitForCallbackRetryWithStepsBeforeAndAfter() {
133133

134134
var callbackResult = context.withRetry(
135135
null,
136-
(ctx, attempt) ->
136+
(attempt, ctx) ->
137137
ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {}),
138138
WithRetryConfig.builder()
139139
.retryStrategy(RetryStrategies.fixedDelay(3, Duration.ofSeconds(1)))
@@ -162,7 +162,7 @@ void waitForCallbackRetryMultipleFailuresThenSuccess() {
162162
String.class,
163163
(input, context) -> context.withRetry(
164164
null,
165-
(ctx, attempt) ->
165+
(attempt, ctx) ->
166166
ctx.waitForCallback("cb-" + attempt, String.class, (callbackId, stepCtx) -> {}),
167167
WithRetryConfig.builder()
168168
.retryStrategy(RetryStrategies.fixedDelay(4, Duration.ofSeconds(1)))
@@ -208,7 +208,7 @@ void waitForCallbackRetryWithSubmitterLogic() {
208208
String.class,
209209
(input, context) -> context.withRetry(
210210
null,
211-
(ctx, attempt) ->
211+
(attempt, ctx) ->
212212
ctx.waitForCallback("approval-" + attempt, String.class, (callbackId, stepCtx) -> {
213213
// Submitter runs each attempt — in a real scenario this would
214214
// send the callbackId to an external system

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

Lines changed: 11 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,6 @@
2020
import software.amazon.lambda.durable.context.BaseContext;
2121
import software.amazon.lambda.durable.model.MapResult;
2222
import software.amazon.lambda.durable.model.WaitForConditionResult;
23-
import software.amazon.lambda.durable.model.WithRetry;
2423

2524
public interface DurableContext extends BaseContext {
2625
/**
@@ -741,11 +740,11 @@ <T> DurableFuture<T> waitForConditionAsync(
741740
* @param <T> the result type
742741
* @param name operation name (used for backoff wait names, and as the child context name when wrapping); pass
743742
* {@code null} for an anonymous retry whose backoff waits use default names
744-
* @param operation the retryable operation — receives the context and 1-based attempt number
743+
* @param operation the retryable operation — receives the 1-based attempt number and the durable context
745744
* @return the operation result
746-
* @see #withRetry(String, WithRetry, WithRetryConfig)
745+
* @see #withRetry(String, BiFunction, WithRetryConfig)
747746
*/
748-
default <T> T withRetry(String name, WithRetry<T> operation) {
747+
default <T> T withRetry(String name, BiFunction<Integer, DurableContext, T> operation) {
749748
return withRetry(name, operation, WithRetryConfig.builder().build());
750749
}
751750

@@ -766,11 +765,11 @@ default <T> T withRetry(String name, WithRetry<T> operation) {
766765
* @param <T> the result type
767766
* @param name operation name (used for backoff wait names, and as the child context name when wrapping); pass
768767
* {@code null} for an anonymous retry whose backoff waits use default names
769-
* @param operation the retryable operation — receives the context and 1-based attempt number
768+
* @param operation the retryable operation — receives the 1-based attempt number and the durable context
770769
* @param config retry configuration including the retry strategy and child context wrapping
771770
* @return the operation result
772771
*/
773-
default <T> T withRetry(String name, WithRetry<T> operation, WithRetryConfig config) {
772+
default <T> T withRetry(String name, BiFunction<Integer, DurableContext, T> operation, WithRetryConfig config) {
774773
return withRetryAsync(name, operation, config).get();
775774
}
776775

@@ -784,11 +783,11 @@ default <T> T withRetry(String name, WithRetry<T> operation, WithRetryConfig con
784783
* @param <T> the result type
785784
* @param name operation name (used for child context and backoff wait names); pass {@code null} for an anonymous
786785
* retry whose backoff waits use default names
787-
* @param operation the retryable operation — receives the context and 1-based attempt number
786+
* @param operation the retryable operation — receives the 1-based attempt number and the durable context
788787
* @return a future representing the operation result
789-
* @see #withRetryAsync(String, WithRetry, WithRetryConfig)
788+
* @see #withRetryAsync(String, BiFunction, WithRetryConfig)
790789
*/
791-
default <T> DurableFuture<T> withRetryAsync(String name, WithRetry<T> operation) {
790+
default <T> DurableFuture<T> withRetryAsync(String name, BiFunction<Integer, DurableContext, T> operation) {
792791
return withRetryAsync(name, operation, WithRetryConfig.builder().build());
793792
}
794793

@@ -803,11 +802,12 @@ default <T> DurableFuture<T> withRetryAsync(String name, WithRetry<T> operation)
803802
* @param <T> the result type
804803
* @param name operation name (used for child context and backoff wait names); pass {@code null} for an anonymous
805804
* retry whose backoff waits use default names
806-
* @param operation the retryable operation — receives the context and 1-based attempt number
805+
* @param operation the retryable operation — receives the 1-based attempt number and the durable context
807806
* @param config retry configuration including the retry strategy
808807
* @return a future representing the operation result
809808
*/
810-
<T> DurableFuture<T> withRetryAsync(String name, WithRetry<T> operation, WithRetryConfig config);
809+
<T> DurableFuture<T> withRetryAsync(
810+
String name, BiFunction<Integer, DurableContext, T> operation, WithRetryConfig config);
811811

812812
/**
813813
* Function applied to each item in a map operation.

sdk/src/main/java/software/amazon/lambda/durable/context/DurableContextImpl.java

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,6 @@
3838
import software.amazon.lambda.durable.model.OperationIdentifier;
3939
import software.amazon.lambda.durable.model.OperationSubType;
4040
import software.amazon.lambda.durable.model.WaitForConditionResult;
41-
import software.amazon.lambda.durable.model.WithRetry;
4241
import software.amazon.lambda.durable.operation.CallbackOperation;
4342
import software.amazon.lambda.durable.operation.ChildContextOperation;
4443
import software.amazon.lambda.durable.operation.InvokeOperation;
@@ -380,7 +379,8 @@ public <T> DurableFuture<T> waitForConditionAsync(
380379

381380
@Override
382381
@SuppressWarnings("unchecked")
383-
public <T> DurableFuture<T> withRetryAsync(String name, WithRetry<T> operation, WithRetryConfig config) {
382+
public <T> DurableFuture<T> withRetryAsync(
383+
String name, BiFunction<Integer, DurableContext, T> operation, WithRetryConfig config) {
384384
Objects.requireNonNull(operation, "operation cannot be null");
385385
Objects.requireNonNull(config, "config cannot be null");
386386

@@ -404,11 +404,14 @@ public <T> DurableFuture<T> withRetryAsync(String name, WithRetry<T> operation,
404404
* are internal SDK control flow signals that must propagate immediately.
405405
*/
406406
private static <T> T executeRetryLoop(
407-
DurableContext context, String name, WithRetry<T> operation, WithRetryConfig config) {
407+
DurableContext context,
408+
String name,
409+
BiFunction<Integer, DurableContext, T> operation,
410+
WithRetryConfig config) {
408411
var attempt = 1;
409412
while (true) {
410413
try {
411-
return operation.execute(context, attempt);
414+
return operation.apply(attempt, context);
412415
} catch (SuspendExecutionException | UnrecoverableDurableExecutionException e) {
413416
// Internal SDK control flow — never retry, always propagate
414417
throw e;

sdk/src/main/java/software/amazon/lambda/durable/model/WithRetry.java

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

0 commit comments

Comments
 (0)