diff --git a/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/InvocationDrainingIntegrationTest.java b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/InvocationDrainingIntegrationTest.java new file mode 100644 index 000000000..90ccada1f --- /dev/null +++ b/sdk-integration-tests/src/test/java/software/amazon/lambda/durable/execution/InvocationDrainingIntegrationTest.java @@ -0,0 +1,213 @@ +// 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.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +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.BiFunction; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.CsvSource; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.config.CompletionConfig; +import software.amazon.lambda.durable.config.MapConfig; +import software.amazon.lambda.durable.config.NestingType; +import software.amazon.lambda.durable.config.ParallelConfig; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.model.ExecutionStatus; +import software.amazon.lambda.durable.testing.LocalDurableTestRunner; + +class InvocationDrainingIntegrationTest { + @ParameterizedTest + @CsvSource({ + "false, FLAT, false", "false, NESTED, false", "true, FLAT, false", "true, NESTED, false", + "false, FLAT, true", "false, NESTED, true", "true, FLAT, true", "true, NESTED, true" + }) + void handlerReturnDrainsQueuedBranches(boolean parallel, NestingType nestingType, boolean nestedSteps) + throws Exception { + var users = Executors.newCachedThreadPool(); + var runtime = Executors.newSingleThreadExecutor(); + var firstStarted = new CountDownLatch(1); + var releaseFirst = new CountDownLatch(1); + var manager = new AtomicReference(); + var completed = new CopyOnWriteArrayList(); + var config = DurableConfig.builder().withExecutorService(users).build(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + manager.set(((DurableContextImpl) context).getExecutionManager()); + startQueuedWork(context, parallel, nestingType, (item, child) -> { + if (item.equals("a")) { + firstStarted.countDown(); + await(releaseFirst); + } + if (nestedSteps) { + return child.step("nested-step", String.class, step -> { + completed.add(item); + return item; + }); + } + completed.add(item); + return item; + }); + await(firstStarted); + return "done"; + }, + config); + + try { + var invocation = runtime.submit(() -> runner.runUntilComplete("input")); + await(firstStarted); + awaitDraining(manager.get()); + releaseFirst.countDown(); + var result = invocation.get(10, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals("done", result.getResult(String.class)); + assertEquals(List.of("a", "b", "c"), completed); + assertEquals( + OperationStatus.SUCCEEDED, + result.getOperation("queued-work").getStatus()); + if (nestedSteps) { + var steps = result.getOperations().stream() + .filter(op -> "nested-step".equals(op.getName())) + .toList(); + assertEquals(3, steps.size()); + assertTrue(steps.stream().allMatch(op -> op.getStatus() == OperationStatus.SUCCEEDED)); + var replay = runner.runUntilComplete("input"); + assertEquals(ExecutionStatus.SUCCEEDED, replay.getStatus()); + assertEquals(List.of("a", "b", "c"), completed, "Replay must not repeat nested step bodies"); + assertEquals( + steps.stream().map(op -> op.getId()).toList(), + replay.getOperations().stream() + .filter(op -> "nested-step".equals(op.getName())) + .map(op -> op.getId()) + .toList()); + } + } finally { + releaseFirst.countDown(); + stop(runtime); + stop(users); + } + } + + @Test + void drainsAsyncDescendantsCreatedAfterTheRootReturns() throws Exception { + var users = Executors.newCachedThreadPool(); + var runtime = Executors.newSingleThreadExecutor(); + var childEntered = new CountDownLatch(1); + var releaseChild = new CountDownLatch(1); + var leafEntered = new CountDownLatch(1); + var releaseLeaf = new CountDownLatch(1); + var effects = new AtomicInteger(); + var manager = new AtomicReference(); + var runner = LocalDurableTestRunner.create( + String.class, + (input, context) -> { + manager.set(((DurableContextImpl) context).getExecutionManager()); + context.runInChildContextAsync("child", String.class, child -> { + childEntered.countDown(); + await(releaseChild); + child.runInChildContextAsync("grandchild", String.class, grandchild -> { + grandchild.stepAsync("leaf", String.class, step -> { + leafEntered.countDown(); + await(releaseLeaf); + effects.incrementAndGet(); + return "leaf-result"; + }); + return "grandchild-result"; + }); + return "child-result"; + }); + await(childEntered); + return "done"; + }, + DurableConfig.builder().withExecutorService(users).build()); + try { + var invocation = runtime.submit(() -> runner.runUntilComplete("input")); + await(childEntered); + awaitDraining(manager.get()); + releaseChild.countDown(); + await(leafEntered); + assertFalse(invocation.isDone(), "Shutdown must wait for the newly registered async leaf"); + releaseLeaf.countDown(); + var result = invocation.get(10, TimeUnit.SECONDS); + assertEquals(ExecutionStatus.SUCCEEDED, result.getStatus()); + assertEquals(OperationStatus.SUCCEEDED, result.getOperation("leaf").getStatus()); + assertEquals("leaf-result", result.getOperation("leaf").getStepResult(String.class)); + var replay = runner.runUntilComplete("input"); + assertEquals(ExecutionStatus.SUCCEEDED, replay.getStatus()); + assertEquals(1, effects.get(), "The checkpointed descendant must not repeat on replay"); + } finally { + releaseChild.countDown(); + releaseLeaf.countDown(); + stop(runtime); + stop(users); + } + } + + private static void startQueuedWork( + DurableContext context, + boolean parallel, + NestingType nestingType, + BiFunction action) { + // Default allCompleted requires a caller to join. Use an explicit threshold to complete without joining. + var completion = CompletionConfig.minSuccessful(3); + if (parallel) { + var operation = context.parallel( + "queued-work", + ParallelConfig.builder() + .maxConcurrency(1) + .nestingType(nestingType) + .completionConfig(completion) + .build()); + for (var item : List.of("a", "b", "c")) { + operation.branch(item, String.class, child -> action.apply(item, child)); + } + } else { + context.mapAsync( + "queued-work", + List.of("a", "b", "c"), + String.class, + (item, index, child) -> action.apply(item, child), + MapConfig.builder() + .maxConcurrency(1) + .nestingType(nestingType) + .completionConfig(completion) + .build()); + } + } + + private static void awaitDraining(ExecutionManager manager) { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (manager.getLifecycleState() != ExecutionManager.LifecycleState.DRAINING) { + assertTrue(System.nanoTime() < deadline, "Invocation did not begin draining"); + Thread.yield(); + } + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(10, TimeUnit.SECONDS), "Test coordination timed out"); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + + private static void stop(ExecutorService executor) throws Exception { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(10, TimeUnit.SECONDS)); + } +} 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 24a8a8958..9e65467ab 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 @@ -143,8 +143,8 @@ private CompletableFuture execute() { return finishInvocation(selected.value(), selected.failure()); }; // Cleanup is awaited even without plugins. Keep that path free of plugin MDC handling. - if (plugins.isEmpty()) return CompletableFuture.supplyAsync(task, config.getExecutorService()); - return supplyHandler(task, this::finishFailure, config.getExecutorService()); + if (plugins.isEmpty()) return manager.submitRootTask(task); + return supplyHandler(task, this::finishFailure, manager::executeRootTask); } private O invokeHandler() { @@ -489,7 +489,7 @@ private static String handleLargePayload(ExecutionManager executionManager, Stri // Checkpoint the large result and wait for it to complete executionManager - .sendOperationUpdate(OperationUpdate.builder() + .sendExecutionResultUpdate(OperationUpdate.builder() .type(OperationType.EXECUTION) .id(executionManager.getExecutionOperation().id()) .action(OperationAction.SUCCEED) 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 1743aefaa..91cbc48a5 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 @@ -3,22 +3,32 @@ package software.amazon.lambda.durable.execution; import com.amazonaws.services.lambda.runtime.Context; +import java.time.Duration; import java.time.Instant; import java.util.ArrayList; import java.util.Collection; import java.util.Collections; +import java.util.Comparator; import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Objects; +import java.util.Optional; import java.util.Set; import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ExecutionException; import java.util.concurrent.Executor; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; import java.util.stream.Collectors; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -62,6 +72,24 @@ public class ExecutionManager implements SafeCloseable { private static final Logger logger = LoggerFactory.getLogger(ExecutionManager.class); + private static final AtomicLong LOCAL_INVOCATION_SEQUENCE = new AtomicLong(); + + enum LifecycleState { + /** The invocation accepts new durable operations, executor tasks, checkpoints, and polls. */ + OPEN, + + /** + * Invocation shutdown has started. Admitted child/coordinator tasks may create descendants, and registered + * operations may submit tasks, checkpoint, or poll as they finish. New root and unrelated work is rejected. + */ + DRAINING, + + /** + * Work admission has ended; cleanup finishes enqueueing and flushing already-admitted checkpoints and polls. + */ + CLOSED + } + // ===== Execution State ===== private final Map operationStorage; private final Operation executionOp; @@ -72,6 +100,18 @@ public class ExecutionManager implements SafeCloseable { private final Set updatedOperationIdsSinceLastInvocation; private final Set initialOperationIds; + // ===== Invocation Lifecycle ===== + private final Object admissionLock = new Object(); + private final String invocationId; + private final Long deadlineNanos; + private final AtomicLong taskSequence = new AtomicLong(); + private final Map> activeExecutorTasks = new ConcurrentHashMap<>(); + // The current task is shared across managers on a thread; admission verifies the exact handle's owner. + private static final ThreadLocal> currentExecutorTask = new ThreadLocal<>(); + private LifecycleState lifecycleState = LifecycleState.OPEN; + // Guarded by admissionLock. Counts enqueue calls, not the lifetime of their returned futures. + private int checkpointAdmissionsInFlight; + // ===== Thread Coordination ===== private final Map registeredOperations = new ConcurrentHashMap<>(); private final Map operationCompletionLocks = new ConcurrentHashMap<>(); @@ -106,6 +146,8 @@ public ExecutionManager(DurableExecutionInput input, DurableConfig config, Conte durableConfig = config; this.durableExecutionArn = input.durableExecutionArn(); this.lambdaContext = lambdaContext; + this.invocationId = resolveInvocationId(lambdaContext, durableExecutionArn); + this.deadlineNanos = resolveDeadlineNanos(lambdaContext); // Store the set of operation IDs updated since the last successful invocation this.updatedOperationIdsSinceLastInvocation = @@ -148,7 +190,25 @@ public ExecutionManager(DurableExecutionInput input, DurableConfig config, Conte logger.debug("EXECUTION operation found: {}", executionOp.id()); } - // ===== State Management ===== + private static String resolveInvocationId(Context lambdaContext, String durableExecutionArn) { + if (lambdaContext != null) { + var requestId = lambdaContext.getAwsRequestId(); + if (requestId != null && !requestId.isBlank()) { + return requestId; + } + } + return durableExecutionArn + "#local-" + LOCAL_INVOCATION_SEQUENCE.incrementAndGet(); + } + + private static Long resolveDeadlineNanos(Context lambdaContext) { + if (lambdaContext == null) { + return null; + } + var remainingMillis = Math.max(0L, lambdaContext.getRemainingTimeInMillis()); + var remainingNanos = TimeUnit.MILLISECONDS.toNanos(remainingMillis); + var now = System.nanoTime(); + return now > Long.MAX_VALUE - remainingNanos ? Long.MAX_VALUE : now + remainingNanos; + } // ===== State Management ===== /** Returns the ARN of the durable execution being managed. */ public String getDurableExecutionArn() { @@ -225,12 +285,20 @@ public void registerOperation(BaseDurableOperation operation) { * @param completionLockParent parent operation whose completion lock should also serialize this operation, or null */ public void registerOperation(BaseDurableOperation operation, BaseDurableOperation completionLockParent) { + // Identity/parent lock lookup can be implemented by an operation subclass. Keep it outside admissionLock, + // just like checkpoint enqueue calls, so an in-flight token withdrawal cannot be blocked by registration. + synchronized (admissionLock) { + requireOperationAdmission(); + } var operationId = operation.getOperationId(); var completionLock = completionLockParent == null ? operationCompletionLocks.computeIfAbsent(operationId, ignored -> new Object()) : completionLockFor(completionLockParent); - operationCompletionLocks.put(operationId, completionLock); - registeredOperations.put(operationId, operation); + synchronized (admissionLock) { + requireOperationAdmission(); + operationCompletionLocks.put(operationId, completionLock); + registeredOperations.put(operationId, operation); + } // Revocation is published before existing operations are suspended. Check after insertion so an // operation registered too late for that notification still releases its waiters. if (isCheckpointTokenRevoked()) { @@ -238,6 +306,131 @@ public void registerOperation(BaseDurableOperation operation, BaseDurableOperati } } + /** Submits the invocation's root handler and records its executor task separately from its logical result. */ + CompletableFuture submitRootTask(Supplier action) { + return submitExecutorTask(ExecutorTaskHandle.Role.ROOT, null, durableConfig.getExecutorService(), action); + } + + /** Submits an operation handler and records its executor task separately from its logical result. */ + public CompletableFuture submitOperationTask(BaseDurableOperation operation, Runnable action) { + var role = operation.getType() == OperationType.STEP + ? ExecutorTaskHandle.Role.STEP + : switch (operation.getSubType()) { + case MAP, PARALLEL -> ExecutorTaskHandle.Role.COORDINATOR; + default -> ExecutorTaskHandle.Role.CHILD_CONTEXT; + }; + return submitExecutorTask(role, operation, durableConfig.getExecutorService(), () -> { + action.run(); + return null; + }); + } + + private CompletableFuture submitExecutorTask( + ExecutorTaskHandle.Role role, + BaseDurableOperation operation, + ExecutorService executor, + Supplier action) { + return submitExecutorTask(role, operation, executor, action, false); + } + + /** Tracks the existing plugin worker command without absorbing its intentional worker failures. */ + void executeRootTask(Runnable command) { + submitExecutorTask( + ExecutorTaskHandle.Role.ROOT, + null, + durableConfig.getExecutorService(), + () -> { + command.run(); + return null; + }, + true); + } + + private CompletableFuture submitExecutorTask( + ExecutorTaskHandle.Role role, + BaseDurableOperation operation, + ExecutorService executor, + Supplier action, + boolean preserveWorkerFailure) { + ExecutorTaskHandle task; + synchronized (admissionLock) { + requireTaskAdmission(operation); + var taskId = taskSequence.incrementAndGet(); + var operationId = operation == null ? null : operation.getOperationId(); + task = new ExecutorTaskHandle<>(taskId, role, operationId, () -> activeExecutorTasks.remove(taskId)); + activeExecutorTasks.put(task.id(), task); + } + + try { + dispatchExecutorTask(task, action, executor, preserveWorkerFailure); + return task.completion(); + } catch (RuntimeException | Error failure) { + task.submissionFailed(failure); + throw failure; + } + } + + private void dispatchExecutorTask( + ExecutorTaskHandle task, Supplier action, ExecutorService executor, boolean preserveWorkerFailure) { + var escapedFailure = new AtomicReference(); + var execution = new FutureTask( + () -> runExecutorTask(task, () -> { + try { + return action.get(); + } catch (Throwable failure) { + if (preserveWorkerFailure) escapedFailure.set(failure); + throw failure; + } + }), + null); + // Binding precedes execute because CallerRunsPolicy may execute synchronously. + task.bindExecution(execution); + executor.execute(() -> { + execution.run(); + // FutureTask observes failure for cancellation/exit bookkeeping, but a lifecycle worker's deliberate + // escape must still reach the executor so a dirty or fatally failed worker is not silently reused. + var failure = escapedFailure.get(); + if (failure != null) ExceptionHelper.sneakyThrow(failure); + }); + } + + private void runExecutorTask(ExecutorTaskHandle task, Supplier action) { + var previous = currentExecutorTask.get(); + currentExecutorTask.set(task); + try { + task.run(action); + } finally { + if (previous == null) { + currentExecutorTask.remove(); + } else { + currentExecutorTask.set(previous); + } + } + } + + String getInvocationId() { + return invocationId; + } + + Optional getRemainingInvocationTime() { + if (deadlineNanos == null) { + return Optional.empty(); + } + return Optional.of(Duration.ofNanos(Math.max(0, deadlineNanos - System.nanoTime()))); + } + + LifecycleState getLifecycleState() { + synchronized (admissionLock) { + return lifecycleState; + } + } + + List> getActiveExecutorTasks() { + return activeExecutorTasks.values().stream() + .sorted(Comparator.comparingLong(ExecutorTaskHandle::id)) + .toList(); + } + // ===== Checkpoint Completion Handler ===== /** Called by CheckpointManager when a checkpoint completes. Updates operationStorage and notify operations . */ void onCheckpointComplete(List newOperations) { @@ -714,7 +907,7 @@ private void preSuspendCheck() { // This method will checkpoint the operation updates to the durable backend and return a future which completes // when the checkpoint completes. public CompletableFuture sendOperationUpdate(OperationUpdate update) { - return checkpointManager.checkpoint(update); + return admitCheckpoint(() -> checkpointManager.checkpoint(update)); } // ===== Polling ===== @@ -725,7 +918,9 @@ public CompletableFuture sendOperationUpdate(OperationUpdate update) { // wait while another thread is still running, and we therefore are not // re-invoked because we never suspended. public CompletableFuture pollForOperationUpdates(String operationId) { - return checkpointManager.pollForUpdate(operationId); + return admitCheckpoint( + () -> checkpointManager.pollForUpdate(operationId), + () -> CompletableFuture.failedFuture(new SuspendExecutionException())); } /** @@ -736,7 +931,14 @@ public CompletableFuture pollForOperationUpdates(String operationId) * @return a completable future that completes with the operation update */ public CompletableFuture pollForOperationUpdates(String operationId, Instant at) { - return checkpointManager.pollForUpdate(operationId, at); + return admitCheckpoint( + () -> checkpointManager.pollForUpdate(operationId, at), + () -> CompletableFuture.failedFuture(new SuspendExecutionException())); + } + + /** Reserved result delivery after operation admission has closed; the caller observes this checkpoint. */ + CompletableFuture sendExecutionResultUpdate(OperationUpdate update) { + return checkpointManager.checkpoint(update); } // ===== Utilities ===== @@ -754,8 +956,13 @@ public void close() { void drainOperations() { // Closing admission can skip queued continuations. Stop their owners before joining user handlers, // or a child waiting on a skipped continuation would prevent the drain from completing. + beginDraining(); stopCheckpointContinuations(); waitForRunningUserHandlers(); + drainExecutorTasks(); + // An interrupted wait leaves admitted work draining. The caller's subsequent close must still allow its + // completion checkpoints; closing here on failure would discard results and repeat completed side effects. + closeLifecycle(); validateExecutorPool(); checkpointManager.shutdown(); } @@ -803,6 +1010,125 @@ private void waitForRunningUserHandlers() { } } + void beginDraining() { + synchronized (admissionLock) { + if (lifecycleState == LifecycleState.OPEN) { + lifecycleState = LifecycleState.DRAINING; + } + } + } + + private void closeLifecycle() { + var interrupted = false; + synchronized (admissionLock) { + lifecycleState = LifecycleState.CLOSED; + // Admission must stop before the final flush, including calls admitted before CLOSED but not yet enqueued. + while (checkpointAdmissionsInFlight > 0) { + try { + admissionLock.wait(); + } catch (InterruptedException e) { + interrupted = true; + } + } + } + if (interrupted) { + Thread.currentThread().interrupt(); + } + } + + private T admitCheckpoint(Supplier request) { + return admitCheckpoint(request, () -> { + throw rejected("checkpoint request"); + }); + } + + private T admitCheckpoint(Supplier request, Supplier whenClosed) { + synchronized (admissionLock) { + if (lifecycleState == LifecycleState.CLOSED) { + return whenClosed.get(); + } + checkpointAdmissionsInFlight++; + } + try { + // CheckpointManager completes continuations under its polling monitor. Those continuations may submit + // tasks or poll again, so never call it while holding admissionLock. + return request.get(); + } finally { + synchronized (admissionLock) { + if (--checkpointAdmissionsInFlight == 0) { + admissionLock.notifyAll(); + } + } + } + } + + private void requireOperationAdmission() { + var task = currentExecutorTask.get(); + if (lifecycleState == LifecycleState.DRAINING + && task != null + && task.state() == ExecutorTaskHandle.State.RUNNING + && activeExecutorTasks.get(task.id()) == task + && (task.role() == ExecutorTaskHandle.Role.CHILD_CONTEXT + || task.role() == ExecutorTaskHandle.Role.COORDINATOR)) { + // Descendants are part of the admitted task's work, including branches that start during draining. + return; + } + requireOpen("durable operation"); + } + + private void requireTaskAdmission(BaseDurableOperation operation) { + if (lifecycleState == LifecycleState.DRAINING + && operation != null + && registeredOperations.get(operation.getOperationId()) == operation) { + // Map/parallel branches are registered before their coordinator submits them to the executor. + return; + } + requireOpen("task"); + } + + private void requireOpen(String workType) { + if (lifecycleState != LifecycleState.OPEN) { + throw rejected(workType); + } + } + + private RejectedExecutionException rejected(String workType) { + return new RejectedExecutionException( + "Invocation " + invocationId + " is " + lifecycleState + "; cannot admit new " + workType); + } + + private void drainExecutorTasks() { + while (true) { + var tasks = operationTasksToDrain(); + if (tasks.isEmpty()) return; + for (var task : tasks) awaitTaskExit(task); + } + } + + private List> operationTasksToDrain() { + synchronized (admissionLock) { + var tasks = activeExecutorTasks.values().stream() + .filter(task -> task.role() != ExecutorTaskHandle.Role.ROOT) + .toList(); + if (tasks.isEmpty()) { + // Close under the submission lock so no task can enter after the final empty snapshot. + lifecycleState = LifecycleState.CLOSED; + } + return tasks; + } + } + + private void awaitTaskExit(ExecutorTaskHandle task) { + try { + task.exit().get(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException("Interrupted while draining operation task: " + task.operationId(), e); + } catch (ExecutionException e) { + throw new IllegalStateException("Could not drain operation task: " + task.operationId(), e); + } + } + private void validateExecutorPool() { // double check if the thread pool is empty if (durableConfig.getExecutorService() instanceof ThreadPoolExecutor threadPoolExecutor) { diff --git a/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutorTaskHandle.java b/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutorTaskHandle.java new file mode 100644 index 000000000..61c4453e9 --- /dev/null +++ b/sdk/src/main/java/software/amazon/lambda/durable/execution/ExecutorTaskHandle.java @@ -0,0 +1,200 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 +package software.amazon.lambda.durable.execution; + +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.Future; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; + +/** + * Tracks one executor task owned by a single Lambda invocation. + * + *

This is runtime lifecycle metadata, not a checkpointed durable operation and not a request to invoke another + * Lambda function. Tasks include the root durable handler, user code for steps and child contexts, and SDK-owned + * map/parallel coordination. A durable operation may have no active task while it is replayed or waiting, and may + * create multiple tasks over its lifetime. + * + *

The logical {@link #completion()} future reports the action's result to SDK code. The separate {@link #exit()} + * future proves that the executor wrapper has actually stopped running. {@link #execution()} retains the underlying + * executor handle so a later invocation-cleanup phase can request interruption without confusing logical completion + * with actual task exit. + */ +final class ExecutorTaskHandle { + + /** + * Runtime role used for invocation-local cleanup, diagnostics, and executor routing. This is deliberately separate + * from a persisted durable operation type: the root has no operation, and map/parallel coordinators share the + * CONTEXT operation type with child user code. + */ + enum Role { + /** The top-level durable handler submitted by {@link DurableExecutor}. */ + ROOT, + + /** User code for an operation whose type is STEP, including wait-for-condition checks. */ + STEP, + + /** User code running in a child context, map iteration, parallel branch, or retry context. */ + CHILD_CONTEXT, + + /** The SDK-owned scheduling loop for a map or parallel operation. */ + COORDINATOR + } + + enum State { + /** Admitted to the invocation scope but not yet started by the executor. */ + REGISTERED, + + /** The executor has entered the task wrapper and the action may still be running. */ + RUNNING, + + /** The wrapper has returned, or the executor accepted cancellation before the wrapper started. */ + EXITED + } + + private enum Publication { + PENDING, + RESULT, + CANCELLED + } + + /** Scope-local sequence number; this is not a durable operation ID. */ + private final long id; + + private final Role role; + + /** Associated durable operation ID, or {@code null} for the root handler. */ + private final String operationId; + + /** Removes this task from the invocation scope's active-task registry. */ + private final Runnable onExit; + + /** Logical action result consumed by the SDK; cancellation can complete it before the action exits. */ + private final CompletableFuture completion = new CompletableFuture<>(); + + /** Completes only when no executor thread can still be running this task. */ + private final CompletableFuture exit = new CompletableFuture<>(); + + /** Interruptible handle bound before dispatch, including when an executor runs the task inline. */ + private final AtomicReference> execution = new AtomicReference<>(); + + private final AtomicReference state = new AtomicReference<>(State.REGISTERED); + private final AtomicReference publication = new AtomicReference<>(Publication.PENDING); + private final AtomicBoolean interruptRequested = new AtomicBoolean(); + + ExecutorTaskHandle(long id, Role role, String operationId, Runnable onExit) { + this.id = id; + this.role = role; + this.operationId = operationId; + this.onExit = onExit; + } + + /** Executor wrapper that captures the logical result and always records actual exit. */ + void run(Supplier action) { + if (!state.compareAndSet(State.REGISTERED, State.RUNNING)) { + return; + } + try { + var value = action.get(); + // Select the publisher atomically: cancellation must not win between a flag check and completion. + // A cancel-winning worker leaves publication to the cancelling thread after physical interruption. + if (publication.compareAndSet(Publication.PENDING, Publication.RESULT)) completion.complete(value); + } catch (Throwable throwable) { + if (publication.compareAndSet(Publication.PENDING, Publication.RESULT)) + completion.completeExceptionally(throwable); + } finally { + markExited(); + } + } + + /** + * Binds the executor handle before dispatch. A cancellation racing with registration is remembered and applied as + * soon as this handle becomes available. + */ + void bindExecution(Future future) { + if (!execution.compareAndSet(null, future)) { + throw new IllegalStateException("Invocation task already has an execution future"); + } + if (publication.get() == Publication.CANCELLED) { + cancelExecution(future); + } + } + + /** Records a synchronous executor rejection and releases the scope registration. */ + void submissionFailed(Throwable failure) { + if (publication.compareAndSet(Publication.PENDING, Publication.RESULT)) + completion.completeExceptionally(failure); + markExited(); + } + + /** + * Requests cooperative cancellation through the executor handle. Logical completion is cancelled immediately; + * actual exit remains pending until running code returns. Once the worker has selected its result, cancellation + * returns false even if its completion observers have not run yet. Observers of an ordinary result still run on the + * publishing worker and must not wait for that same worker to exit. User code that ignores interruption cannot be + * forcibly stopped. + */ + boolean cancel(boolean mayInterruptIfRunning) { + while (true) { + var selected = publication.get(); + if (selected == Publication.RESULT) return false; + if (selected == Publication.CANCELLED + || publication.compareAndSet(Publication.PENDING, Publication.CANCELLED)) break; + } + if (mayInterruptIfRunning) { + interruptRequested.set(true); + } + var future = execution.get(); + // CompletableFuture cancellation runs observers inline. Physical interruption must already be requested + // before an observer is allowed to wait for this task's actual exit. + var cancelled = future == null || cancelExecution(future); + completion.cancel(false); + return cancelled; + } + + private boolean cancelExecution(Future future) { + if (!future.cancel(interruptRequested.get())) { + return false; + } + if (state.compareAndSet(State.REGISTERED, State.EXITED)) { + onExit.run(); + exit.complete(null); + } + return true; + } + + private void markExited() { + state.set(State.EXITED); + onExit.run(); + exit.complete(null); + } + + long id() { + return id; + } + + Role role() { + return role; + } + + String operationId() { + return operationId; + } + + State state() { + return state.get(); + } + + CompletableFuture completion() { + return completion; + } + + CompletableFuture exit() { + return exit; + } + + Future execution() { + return execution.get(); + } +} 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 ecefc5736..3418375c7 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 @@ -325,8 +325,7 @@ protected void runUserHandler(Runnable runnable, ThreadType threadType) { registerActiveThread(operationId); try { - runningUserHandler.set(CompletableFuture.runAsync( - wrapped, getContext().getDurableConfig().getExecutorService())); + runningUserHandler.set(executionManager.submitOperationTask(this, wrapped)); } catch (RuntimeException | Error dispatchFailure) { // Admission failed, so no worker will execute its deregistration finally block. The submitting // coordinator need not have a user ThreadContext; undo the registration directly on the manager. diff --git a/sdk/src/test/java/software/amazon/lambda/durable/TestUtils.java b/sdk/src/test/java/software/amazon/lambda/durable/TestUtils.java index 44f0078a9..f64dcf783 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/TestUtils.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/TestUtils.java @@ -8,8 +8,11 @@ import java.util.LinkedHashMap; import java.util.List; import java.util.UUID; +import java.util.concurrent.CompletableFuture; import software.amazon.awssdk.services.lambda.model.*; import software.amazon.lambda.durable.client.DurableExecutionClient; +import software.amazon.lambda.durable.context.DurableContextImpl; +import software.amazon.lambda.durable.execution.ExecutionManager; import software.amazon.lambda.durable.execution.OperationIdGenerator; public class TestUtils { @@ -70,4 +73,11 @@ public static DurableExecutionClient createMockClient() { public static String hashOperationId(String rawId) { return OperationIdGenerator.hashOperationId(rawId); } + + /** Makes a mocked ExecutionManager execute operation tasks like the real manager. */ + public static void executeOperationTasks(ExecutionManager manager, DurableContextImpl context) { + when(manager.submitOperationTask(any(), any())) + .thenAnswer(invocation -> CompletableFuture.runAsync( + invocation.getArgument(1), context.getDurableConfig().getExecutorService())); + } } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleRegressionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleRegressionTest.java index ea50f1f2a..851e780c0 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleRegressionTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/DurableExecutorLifecycleRegressionTest.java @@ -13,6 +13,9 @@ import static software.amazon.lambda.durable.TypeToken.get; import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; @@ -36,6 +39,7 @@ import software.amazon.lambda.durable.client.DurableExecutionClient; 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.operation.BaseDurableOperation; import software.amazon.lambda.durable.plugin.DurableExecutionPlugin; import software.amazon.lambda.durable.plugin.InvocationEndInfo; @@ -158,6 +162,66 @@ public String serialize(Object value) { }; } + @Test + void interruptedDrainPreservesLateStepResultForReplay() throws Exception { + var recorded = new ConcurrentHashMap(); + var delegate = TestUtils.createMockClient(); + var client = mock(DurableExecutionClient.class); + when(client.checkpoint(any(), any(), any())).thenAnswer(call -> { + var response = delegate.checkpoint(call.getArgument(0), call.getArgument(1), call.getArgument(2)); + response.newExecutionState().operations().forEach(operation -> recorded.put(operation.id(), operation)); + return response; + }); + var users = Executors.newFixedThreadPool(2); + var caller = Executors.newSingleThreadExecutor(); + var stepEntered = new CountDownLatch(1); + var handlerReturned = new CountDownLatch(1); + var release = new CountDownLatch(1); + var root = new AtomicReference(); + var bodies = new AtomicInteger(); + var config = DurableConfig.builder() + .withExecutorService(users) + .withDurableExecutionClient(client) + .build(); + BiFunction handler = (value, context) -> { + context.stepAsync("late-result", String.class, step -> { + bodies.incrementAndGet(); + stepEntered.countDown(); + await(release); + return "stored"; + }); + root.set(Thread.currentThread()); + handlerReturned.countDown(); + return "root-result"; + }; + try { + var first = caller.submit(() -> execute(config, handler)); + assertTrue(stepEntered.await(5, TimeUnit.SECONDS)); + assertTrue(handlerReturned.await(5, TimeUnit.SECONDS)); + root.get().interrupt(); + assertThrows(TimeoutException.class, () -> first.get(100, TimeUnit.MILLISECONDS)); + release.countDown(); + assertInstanceOf( + IllegalStateException.class, + assertThrows(ExecutionException.class, () -> first.get(5, TimeUnit.SECONDS)) + .getCause()); + var step = recorded.values().stream() + .filter(operation -> "late-result".equals(operation.name())) + .findFirst() + .orElseThrow(); + assertEquals(OperationStatus.SUCCEEDED, step.status()); + var replay = execute(config, handler, List.copyOf(recorded.values())); + assertEquals(ExecutionStatus.SUCCEEDED, replay.status()); + assertEquals(1, bodies.get(), "The acknowledged late result must replay without repeating its side effect"); + } finally { + release.countDown(); + caller.shutdownNow(); + users.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + assertTrue(users.awaitTermination(5, TimeUnit.SECONDS)); + } + } + private static DurableExecutionPlugin recordingPlugin(AtomicInteger ended) { return new DurableExecutionPlugin() { @Override @@ -169,6 +233,13 @@ public void onInvocationEnd(InvocationEndInfo info) { private static DurableExecutionOutput execute( DurableConfig config, BiFunction handler) { + return execute(config, handler, List.of()); + } + + private static DurableExecutionOutput execute( + DurableConfig config, + BiFunction handler, + List replayOperations) { var execution = Operation.builder() .id("exec") .type(OperationType.EXECUTION) @@ -177,10 +248,12 @@ private static DurableExecutionOutput execute( .executionDetails( ExecutionDetails.builder().inputPayload("\"input\"").build()) .build(); + var operations = new ArrayList<>(replayOperations); + operations.add(execution); var input = new DurableExecutionInput( "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/lifecycle/exec", "token", - CheckpointUpdatedExecutionState.builder().operations(execution).build()); + CheckpointUpdatedExecutionState.builder().operations(operations).build()); return DurableExecutor.execute(input, null, get(String.class), handler, config); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java new file mode 100644 index 000000000..76a584591 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerAdmissionTest.java @@ -0,0 +1,355 @@ +// 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.assertFalse; +import static org.junit.jupiter.api.Assertions.assertInstanceOf; +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.doAnswer; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; + +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.concurrent.CompletionException; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.awssdk.services.lambda.model.CheckpointDurableExecutionResponse; +import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationAction; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.awssdk.services.lambda.model.OperationUpdate; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.client.DurableExecutionClient; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.model.OperationSubType; +import software.amazon.lambda.durable.operation.BaseDurableOperation; + +class ExecutionManagerAdmissionTest { + private final ExecutorService users = Executors.newCachedThreadPool(); + private final ExecutorService callers = Executors.newCachedThreadPool(); + private final DurableExecutionClient client = mock(DurableExecutionClient.class); + private DurableConfig config; + private ExecutionManager manager; + + @BeforeEach + void setUp() { + config = spy(DurableConfig.builder() + .withDurableExecutionClient(client) + .withExecutorService(users) + .withCheckpointDelay(Duration.ofHours(1)) + .withPollingStrategy(attempt -> Duration.ofHours(1)) + .build()); + manager = createManager(); + manager.registerActiveThread(null); + when(client.checkpoint(any(), any(), any())).thenReturn(response()); + } + + @AfterEach + void tearDown() throws Exception { + manager.close(); + stop(callers); + stop(users); + } + + @Test + void pollAdmissionDoesNotBlockRetryTaskSubmission() throws Exception { + var checkpointEntered = new CountDownLatch(1); + var releaseCheckpoint = new CountDownLatch(1); + var pollEntered = new CountDownLatch(1); + var step = step("retry"); + manager.registerOperation(step); + when(client.checkpoint(any(), any(), any())).thenAnswer(invocation -> { + checkpointEntered.countDown(); + await(releaseCheckpoint); + return response(Operation.builder() + .id("retry") + .type(OperationType.STEP) + .status(OperationStatus.READY) + .build()); + }); + doAnswer(invocation -> { + pollEntered.countDown(); + return invocation.callRealMethod(); + }) + .when(config) + .getPollingStrategy(); + + try { + var retryPoll = manager.pollForOperationUpdates("retry", Instant.now()); + await(checkpointEntered); + var concurrentPoll = callers.submit(() -> manager.pollForOperationUpdates("other")); + await(pollEntered); + // The other poll is waiting for CheckpointManager's monitor. It must not retain admissionLock. + callers.submit(() -> manager.submitOperationTask(step, () -> {})).get(5, TimeUnit.SECONDS); + var retries = new AtomicInteger(); + var retry = retryPoll.thenCompose(op -> manager.submitOperationTask(step, retries::incrementAndGet)); + releaseCheckpoint.countDown(); + retry.get(5, TimeUnit.SECONDS); + concurrentPoll.get(5, TimeUnit.SECONDS); + assertEquals(1, retries.get()); + } finally { + releaseCheckpoint.countDown(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void shutdownWaitsForAdmittedPollAndPreservesInterruption(boolean interrupted) throws Exception { + var admissionEntered = new CountDownLatch(1); + var releaseAdmission = new CountDownLatch(1); + doAnswer(invocation -> { + admissionEntered.countDown(); + await(releaseAdmission); + return invocation.callRealMethod(); + }) + .when(config) + .getPollingStrategy(); + + try { + var admission = callers.submit(() -> manager.pollForOperationUpdates("pending")); + await(admissionEntered); + var shutdown = callers.submit(() -> { + if (interrupted) { + Thread.currentThread().interrupt(); + } + manager.close(); + return Thread.interrupted(); + }); + callers.submit(this::awaitClosedAdmission).get(5, TimeUnit.SECONDS); + assertFalse(shutdown.isDone()); + assertThrows(RejectedExecutionException.class, () -> manager.sendOperationUpdate(update())); + releaseAdmission.countDown(); + var poll = admission.get(5, TimeUnit.SECONDS); + assertEquals(interrupted, shutdown.get(5, TimeUnit.SECONDS)); + assertTrue(poll.isCompletedExceptionally()); + } finally { + releaseAdmission.countDown(); + } + } + + @Test + void failedAdmissionDoesNotPreventShutdown() throws Exception { + doThrow(new IllegalArgumentException("polling configuration failed")) + .when(config) + .getPollingStrategy(); + assertThrows(IllegalArgumentException.class, () -> manager.pollForOperationUpdates("pending")); + callers.submit(manager::close).get(5, TimeUnit.SECONDS); + assertEquals(ExecutionManager.LifecycleState.CLOSED, manager.getLifecycleState()); + } + + @Test + void finalCheckpointFlushCannotAdmitMoreCheckpointsOrPolls() throws Exception { + when(client.checkpoint(any(), any(), any())).thenAnswer(invocation -> { + assertThrows(RejectedExecutionException.class, () -> manager.sendOperationUpdate(update())); + for (var poll : List.of( + manager.pollForOperationUpdates("late"), manager.pollForOperationUpdates("late", Instant.now()))) { + assertTrue(poll.isCompletedExceptionally()); + assertInstanceOf( + SuspendExecutionException.class, + assertThrows(CompletionException.class, poll::join).getCause()); + } + return response(); + }); + var checkpoint = manager.sendOperationUpdate(update()); + manager.close(); + checkpoint.get(5, TimeUnit.SECONDS); + } + + @Test + void drainingAllowsOnlyTasksOfPreviouslyRegisteredOperations() throws Exception { + var registered = step("registered"); + manager.registerOperation(registered); + manager.beginDraining(); + + var executions = new AtomicInteger(); + manager.submitOperationTask(registered, executions::incrementAndGet).get(5, TimeUnit.SECONDS); + assertEquals(1, executions.get()); + assertThrows(RejectedExecutionException.class, () -> manager.submitOperationTask(step("registered"), () -> {})); + assertThrows(RejectedExecutionException.class, () -> manager.submitOperationTask(step("new"), () -> {})); + assertThrows(RejectedExecutionException.class, () -> manager.submitRootTask(() -> null)); + manager.close(); + assertThrows(RejectedExecutionException.class, () -> manager.submitOperationTask(registered, () -> {})); + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void admittedContextTasksCanRegisterDescendantsWithoutLeakingAdmission(boolean coordinator) throws Exception { + var executor = Executors.newSingleThreadExecutor(); + when(config.getExecutorService()).thenReturn(executor); + var parent = + contextOperation("parent", coordinator ? OperationSubType.MAP : OperationSubType.RUN_IN_CHILD_CONTEXT); + var descendant = step("descendant"); + manager.registerOperation(parent); + manager.beginDraining(); + try { + manager.submitOperationTask(parent, () -> manager.registerOperation(descendant)) + .get(5, TimeUnit.SECONDS); + manager.submitOperationTask(descendant, () -> {}).get(5, TimeUnit.SECONDS); + executor.submit(() -> assertThrows( + RejectedExecutionException.class, () -> manager.registerOperation(step("unrelated")))) + .get(5, TimeUnit.SECONDS); + } finally { + stop(executor); + } + } + + @Test + void admittedTaskCannotRegisterOperationsInAnotherDrainingInvocation() throws Exception { + var other = createManager(); + var parent = contextOperation("parent", OperationSubType.RUN_IN_CHILD_CONTEXT); + manager.registerOperation(parent); + manager.beginDraining(); + other.beginDraining(); + try { + manager.submitOperationTask(parent, () -> { + assertThrows(RejectedExecutionException.class, () -> other.registerOperation(step("foreign"))); + manager.registerOperation(step("owned")); + }) + .get(5, TimeUnit.SECONDS); + } finally { + other.close(); + } + } + + @Test + void inlineNestedInvocationCannotUseItsCallersDrainingPermission() throws Exception { + var inline = mock(ExecutorService.class); + doAnswer(invocation -> { + invocation.getArgument(0, Runnable.class).run(); + return null; + }) + .when(inline) + .execute(any(Runnable.class)); + when(config.getExecutorService()).thenReturn(inline); + var other = createManager(); + var parent = contextOperation("parent", OperationSubType.RUN_IN_CHILD_CONTEXT); + var nested = contextOperation("nested", OperationSubType.RUN_IN_CHILD_CONTEXT); + manager.registerOperation(parent); + other.registerOperation(nested); + manager.beginDraining(); + other.beginDraining(); + try { + manager.submitOperationTask(parent, () -> { + other.submitOperationTask( + nested, + () -> assertThrows( + RejectedExecutionException.class, + () -> manager.registerOperation(step("foreign")))) + .join(); + manager.registerOperation(step("restored-owner")); + }) + .get(5, TimeUnit.SECONDS); + } finally { + other.close(); + } + } + + @Test + void alreadyRunningRootCannotAdmitOperationsAfterDrainingStarts() throws Exception { + var entered = new CountDownLatch(1); + var release = new CountDownLatch(1); + var root = manager.submitRootTask(() -> { + entered.countDown(); + await(release); + assertThrows(RejectedExecutionException.class, () -> manager.registerOperation(step("late-root"))); + return null; + }); + try { + await(entered); + manager.beginDraining(); + release.countDown(); + root.get(5, TimeUnit.SECONDS); + } finally { + release.countDown(); + } + } + + private ExecutionManager createManager() { + var execution = Operation.builder() + .id("execution") + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .build(); + return new ExecutionManager( + new DurableExecutionInput( + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/test/execution", + "token", + CheckpointUpdatedExecutionState.builder() + .operations(execution) + .build()), + config, + null); + } + + private static BaseDurableOperation contextOperation(String id, OperationSubType subType) { + var operation = mock(BaseDurableOperation.class); + when(operation.getOperationId()).thenReturn(id); + when(operation.getType()).thenReturn(OperationType.CONTEXT); + when(operation.getSubType()).thenReturn(subType); + return operation; + } + + private void awaitClosedAdmission() { + var deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (manager.getLifecycleState() != ExecutionManager.LifecycleState.CLOSED) { + assertTrue(System.nanoTime() < deadline, "Checkpoint admission did not close"); + Thread.yield(); + } + } + + private static BaseDurableOperation step(String id) { + var operation = mock(BaseDurableOperation.class); + when(operation.getOperationId()).thenReturn(id); + when(operation.getType()).thenReturn(OperationType.STEP); + return operation; + } + + private static OperationUpdate update() { + return OperationUpdate.builder() + .id("step") + .type(OperationType.STEP) + .action(OperationAction.START) + .build(); + } + + private static CheckpointDurableExecutionResponse response(Operation... operations) { + return CheckpointDurableExecutionResponse.builder() + .checkpointToken("next-token") + .newExecutionState(CheckpointUpdatedExecutionState.builder() + .operations(List.of(operations)) + .build()) + .build(); + } + + private static void await(CountDownLatch latch) { + try { + assertTrue(latch.await(15, TimeUnit.SECONDS), "Test coordination timed out"); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + + private static void stop(ExecutorService executor) throws Exception { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java new file mode 100644 index 000000000..edd1cd741 --- /dev/null +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerLifecycleTest.java @@ -0,0 +1,399 @@ +// 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.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +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 static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.amazonaws.services.lambda.runtime.Context; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.FutureTask; +import java.util.concurrent.RejectedExecutionException; +import java.util.concurrent.SynchronousQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState; +import software.amazon.awssdk.services.lambda.model.Operation; +import software.amazon.awssdk.services.lambda.model.OperationAction; +import software.amazon.awssdk.services.lambda.model.OperationStatus; +import software.amazon.awssdk.services.lambda.model.OperationType; +import software.amazon.awssdk.services.lambda.model.OperationUpdate; +import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; +import software.amazon.lambda.durable.model.DurableExecutionInput; +import software.amazon.lambda.durable.operation.BaseDurableOperation; + +class ExecutionManagerLifecycleTest { + private static final String EXECUTION_ID = "execution"; + private static final String EXECUTION_ARN = + "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/test/" + EXECUTION_ID; + + @Test + void capturesRequestIdAndInvocationDeadline() { + var context = mock(Context.class); + when(context.getAwsRequestId()).thenReturn("request-id"); + when(context.getRemainingTimeInMillis()).thenReturn(2_000); + var manager = createManager(context, null); + + try { + assertEquals("request-id", manager.getInvocationId()); + var remaining = manager.getRemainingInvocationTime().orElseThrow(); + assertTrue(remaining.compareTo(Duration.ZERO) > 0); + assertTrue(remaining.compareTo(Duration.ofSeconds(2)) <= 0); + } finally { + manager.close(); + } + } + + @Test + void createsDistinctLocalIdsWithoutInventingADeadline() { + var first = createManager(null, null); + var second = createManager(null, null); + + try { + assertNotEquals(first.getInvocationId(), second.getInvocationId()); + assertTrue(first.getRemainingInvocationTime().isEmpty()); + } finally { + first.close(); + second.close(); + } + } + + @Test + void registersTaskBeforeQueuedExecutionAndTracksItsExit() throws Exception { + var executor = Executors.newSingleThreadExecutor(); + var blockerEntered = new CountDownLatch(1); + var releaseBlocker = new CountDownLatch(1); + executor.submit(() -> { + blockerEntered.countDown(); + await(releaseBlocker); + }); + assertTrue(blockerEntered.await(5, TimeUnit.SECONDS)); + var manager = createManager(null, executor); + + try { + var completion = manager.submitRootTask(() -> "result"); + var task = manager.getActiveExecutorTasks().get(0); + + assertEquals(ExecutorTaskHandle.State.REGISTERED, task.state()); + assertNotNull(task.execution()); + assertFalse(completion.isDone()); + assertFalse(task.exit().isDone()); + + releaseBlocker.countDown(); + assertEquals("result", completion.get(5, TimeUnit.SECONDS)); + task.exit().get(5, TimeUnit.SECONDS); + assertEquals(ExecutorTaskHandle.State.EXITED, task.state()); + assertTrue(manager.getActiveExecutorTasks().isEmpty()); + } finally { + releaseBlocker.countDown(); + manager.close(); + stop(executor); + } + } + + @Test + void taskCancellationSeparatesLogicalCompletionFromActualExit() throws Exception { + var executor = Executors.newSingleThreadExecutor(); + var entered = new CountDownLatch(1); + var interrupted = new CountDownLatch(1); + var release = new CountDownLatch(1); + var manager = createManager(null, executor); + try { + var completion = manager.submitRootTask(() -> { + entered.countDown(); + while (release.getCount() > 0) { + try { + release.await(); + } catch (InterruptedException expected) { + interrupted.countDown(); + } + } + return "ignored"; + }); + assertTrue(entered.await(5, TimeUnit.SECONDS)); + var task = manager.getActiveExecutorTasks().get(0); + + assertTrue(task.cancel(true)); + assertTrue(interrupted.await(5, TimeUnit.SECONDS)); + assertTrue(completion.isCancelled()); + assertFalse(task.exit().isDone()); + + release.countDown(); + task.exit().get(5, TimeUnit.SECONDS); + assertEquals(ExecutorTaskHandle.State.EXITED, task.state()); + } finally { + release.countDown(); + manager.close(); + stop(executor); + } + } + + @Test + void cancellationRequestedBeforeExecutorHandleIsBoundIsNotLost() throws Exception { + var exits = new AtomicInteger(); + var task = new ExecutorTaskHandle(1, ExecutorTaskHandle.Role.ROOT, null, exits::incrementAndGet); + var execution = new FutureTask(() -> null); + + assertTrue(task.cancel(true)); + assertTrue(task.completion().isCancelled()); + assertFalse(task.exit().isDone()); + + task.bindExecution(execution); + + assertTrue(execution.isCancelled()); + task.exit().get(5, TimeUnit.SECONDS); + assertEquals(1, exits.get()); + } + + @Test + void cancellationInterruptsCallerRunsBeforeSubmissionReturns() throws Exception { + var executor = new ThreadPoolExecutor( + 1, 1, 0, TimeUnit.MILLISECONDS, new SynchronousQueue<>(), new ThreadPoolExecutor.CallerRunsPolicy()); + var caller = Executors.newSingleThreadExecutor(); + var workerEntered = new CountDownLatch(1); + var releaseWorker = new CountDownLatch(1); + var bodyEntered = new CountDownLatch(1); + var interrupted = new CountDownLatch(1); + var releaseBody = new CountDownLatch(1); + var manager = createManager(null, executor); + executor.execute(() -> { + workerEntered.countDown(); + await(releaseWorker); + }); + assertTrue(workerEntered.await(5, TimeUnit.SECONDS)); + try { + var submitted = caller.submit(() -> manager.submitRootTask(() -> { + bodyEntered.countDown(); + while (releaseBody.getCount() > 0) { + try { + releaseBody.await(); + } catch (InterruptedException expected) { + interrupted.countDown(); + } + } + return "ignored"; + })); + assertTrue(bodyEntered.await(5, TimeUnit.SECONDS)); + assertFalse(submitted.isDone(), "CallerRunsPolicy has not returned from task submission"); + var task = manager.getActiveExecutorTasks().get(0); + assertNotNull(task.execution(), "The interruptible handle must exist before inline execution starts"); + assertTrue(task.cancel(true)); + assertTrue(interrupted.await(5, TimeUnit.SECONDS)); + assertTrue(task.completion().isCancelled()); + assertFalse(task.exit().isDone(), "Cancellation is not proof that user code has exited"); + releaseBody.countDown(); + task.exit().get(5, TimeUnit.SECONDS); + assertTrue(submitted.get(5, TimeUnit.SECONDS).isCancelled()); + assertTrue(manager.getActiveExecutorTasks().isEmpty()); + } finally { + releaseBody.countDown(); + releaseWorker.countDown(); + manager.close(); + stop(caller); + stop(executor); + } + } + + @Test + void failedSubmissionIsRemovedFromManager() { + var executor = Executors.newSingleThreadExecutor(); + executor.shutdown(); + var manager = createManager(null, executor); + + try { + assertThrows(RejectedExecutionException.class, () -> manager.submitRootTask(() -> "result")); + assertTrue(manager.getActiveExecutorTasks().isEmpty()); + } finally { + manager.close(); + } + } + + @Test + void cancellationInterruptsBeforeLogicalObserversWaitForExit() throws Exception { + var executor = Executors.newSingleThreadExecutor(); + var cancelling = Executors.newSingleThreadExecutor(); + var manager = createManager(null, executor); + var entered = new CountDownLatch(1); + var interrupted = new CountDownLatch(1); + var release = new CountDownLatch(1); + try { + manager.submitRootTask(() -> { + entered.countDown(); + try { + release.await(); + } catch (InterruptedException expected) { + interrupted.countDown(); + } + return "ignored"; + }); + assertTrue(entered.await(5, TimeUnit.SECONDS)); + var task = manager.getActiveExecutorTasks().get(0); + var observer = task.completion() + .whenComplete((result, failure) -> task.exit().join()); + var cancellation = cancelling.submit(() -> task.cancel(true)); + assertTrue(interrupted.await(2, TimeUnit.SECONDS), "A completion observer must not delay interruption"); + assertTrue(cancellation.get(5, TimeUnit.SECONDS)); + assertTrue(task.completion().isCancelled()); + assertTrue(observer.isDone()); + assertTrue(manager.getActiveExecutorTasks().isEmpty()); + } finally { + release.countDown(); + stop(executor); + stop(cancelling); + manager.close(); + } + } + + @ParameterizedTest + @ValueSource(booleans = {false, true}) + void selectedResultCannotBeOvertakenByCancellationBeforePublication(boolean failAction) throws Exception { + var executor = Executors.newSingleThreadExecutor(); + var publishing = new CountDownLatch(1); + var releasePublication = new CountDownLatch(1); + var expectedFailure = new IllegalStateException("selected failure"); + var completion = new CompletableFuture() { + private void waitBeforePublication() { + publishing.countDown(); + boolean interrupted = false; + while (releasePublication.getCount() > 0) { + try { + releasePublication.await(); + } catch (InterruptedException ignored) { + interrupted = true; + } + } + if (interrupted) Thread.currentThread().interrupt(); + } + + @Override + public boolean complete(String value) { + waitBeforePublication(); + return super.complete(value); + } + + @Override + public boolean completeExceptionally(Throwable failure) { + waitBeforePublication(); + return super.completeExceptionally(failure); + } + }; + var task = new ExecutorTaskHandle(1, ExecutorTaskHandle.Role.ROOT, null, () -> {}); + // Hold publication after the worker selected its outcome, before CompletableFuture runs observers. + var completionField = ExecutorTaskHandle.class.getDeclaredField("completion"); + completionField.setAccessible(true); + completionField.set(task, completion); + var execution = new FutureTask( + () -> task.run(() -> { + if (failAction) throw expectedFailure; + return "selected value"; + }), + null); + task.bindExecution(execution); + try { + executor.execute(execution); + assertTrue(publishing.await(5, TimeUnit.SECONDS)); + assertFalse(task.cancel(true), "A selected result must win before its observers are published"); + assertFalse(execution.isCancelled()); + assertFalse(completion.isDone()); + assertFalse(task.exit().isDone(), "Outcome selection is not actual wrapper exit"); + releasePublication.countDown(); + if (failAction) { + var failure = assertThrows(ExecutionException.class, () -> completion.get(5, TimeUnit.SECONDS)); + assertSame(expectedFailure, failure.getCause()); + } else { + assertEquals("selected value", completion.get(5, TimeUnit.SECONDS)); + } + task.exit().get(5, TimeUnit.SECONDS); + } finally { + releasePublication.countDown(); + stop(executor); + } + } + + @Test + void drainingRejectsNewWorkButAllowsExistingCheckpointCleanup() throws Exception { + var executor = Executors.newSingleThreadExecutor(); + var manager = createManager(null, executor); + manager.beginDraining(); + + try { + assertEquals(ExecutionManager.LifecycleState.DRAINING, manager.getLifecycleState()); + assertThrows( + RejectedExecutionException.class, + () -> manager.registerOperation(mock(BaseDurableOperation.class))); + assertThrows(RejectedExecutionException.class, () -> manager.submitRootTask(() -> "result")); + + var checkpoint = manager.sendOperationUpdate(OperationUpdate.builder() + .id("step") + .name("step") + .type(OperationType.STEP) + .subType("Step") + .action(OperationAction.START) + .build()); + manager.close(); + checkpoint.get(5, TimeUnit.SECONDS); + + assertEquals(ExecutionManager.LifecycleState.CLOSED, manager.getLifecycleState()); + assertThrows( + RejectedExecutionException.class, + () -> manager.sendOperationUpdate(OperationUpdate.builder().build())); + } finally { + stop(executor); + } + } + + private ExecutionManager createManager(Context context, ExecutorService executor) { + var config = DurableConfig.builder().withDurableExecutionClient(TestUtils.createMockClient()); + if (executor != null) { + config.withExecutorService(executor); + } + var execution = Operation.builder() + .id(EXECUTION_ID) + .type(OperationType.EXECUTION) + .status(OperationStatus.STARTED) + .build(); + return new ExecutionManager( + new DurableExecutionInput( + EXECUTION_ARN, + "token", + CheckpointUpdatedExecutionState.builder() + .operations(List.of(execution)) + .build()), + config.build(), + context); + } + + private static void await(CountDownLatch latch) { + try { + if (!latch.await(5, TimeUnit.SECONDS)) { + throw new IllegalStateException("Test synchronization timed out"); + } + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(interrupted); + } + } + + private static void stop(ExecutorService executor) throws InterruptedException { + executor.shutdownNow(); + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } +} diff --git a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java index 52e93353f..c2c86c04a 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/execution/ExecutionManagerTest.java @@ -1002,6 +1002,72 @@ protected void replay(Operation existing) { assertTrue(operation.getCompletionFuture().isDone()); } + @Test + void tracksRootStepChildAndCoordinatorTasks() throws Exception { + var manager = createManager(List.of(executionOp())); + var entered = new CountDownLatch(5); + var release = new CountDownLatch(1); + try { + var root = manager.submitRootTask(() -> { + waitForRelease(entered, release); + return "root-result"; + }); + var step = manager.submitOperationTask( + operation("step", OperationSubType.STEP), () -> waitForRelease(entered, release)); + var waitForCondition = manager.submitOperationTask( + operation("wait", OperationSubType.WAIT_FOR_CONDITION), () -> waitForRelease(entered, release)); + var child = manager.submitOperationTask( + operation("child", OperationSubType.RUN_IN_CHILD_CONTEXT), () -> waitForRelease(entered, release)); + var coordinator = manager.submitOperationTask( + operation("map", OperationSubType.MAP), () -> waitForRelease(entered, release)); + + assertTrue(entered.await(5, TimeUnit.SECONDS)); + + var tasks = manager.getActiveExecutorTasks(); + assertEquals( + List.of( + ExecutorTaskHandle.Role.ROOT, + ExecutorTaskHandle.Role.STEP, + ExecutorTaskHandle.Role.STEP, + ExecutorTaskHandle.Role.CHILD_CONTEXT, + ExecutorTaskHandle.Role.COORDINATOR), + tasks.stream().map(ExecutorTaskHandle::role).toList()); + assertTrue(tasks.stream().allMatch(task -> task.state() == ExecutorTaskHandle.State.RUNNING)); + + release.countDown(); + assertEquals("root-result", root.get(5, TimeUnit.SECONDS)); + step.get(5, TimeUnit.SECONDS); + waitForCondition.get(5, TimeUnit.SECONDS); + child.get(5, TimeUnit.SECONDS); + coordinator.get(5, TimeUnit.SECONDS); + for (var task : tasks) { + task.exit().get(5, TimeUnit.SECONDS); + } + assertTrue(manager.getActiveExecutorTasks().isEmpty()); + } finally { + release.countDown(); + manager.close(); + } + } + + private void waitForRelease(CountDownLatch entered, CountDownLatch release) { + entered.countDown(); + try { + assertTrue(release.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException interrupted) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(interrupted); + } + } + + private BaseDurableOperation operation(String id, OperationSubType subType) { + var operation = mock(BaseDurableOperation.class); + when(operation.getOperationId()).thenReturn(id); + when(operation.getSubType()).thenReturn(subType); + when(operation.getType()).thenReturn(subType.getOperationType()); + return operation; + } + @Test void childRegistrationNeverPublishesATemporaryCompletionLock() throws Exception { var manager = new LockPublicationExecutionManager( 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 4ccc4805d..de97bb286 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 @@ -20,6 +20,7 @@ import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.lambda.durable.DurableConfig; import software.amazon.lambda.durable.DurableContext; +import software.amazon.lambda.durable.TestUtils; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.ParallelConfig; import software.amazon.lambda.durable.config.RunInChildContextConfig; @@ -75,6 +76,7 @@ void setUp() { when(durableContext.getExecutionManager()).thenReturn(executionManager); when(executionManager.getCurrentThreadContext()).thenReturn(new ThreadContext("Root", ThreadType.CONTEXT)); when(durableContext.getDurableConfig()).thenReturn(createConfig()); + TestUtils.executeOperationTasks(executionManager, durableContext); MockExecutionManagerSupport.stubWaitForOperationCompletion(executionManager); } 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 4a7b4c5ba..9a7395a67 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 @@ -90,6 +90,7 @@ void setUp() { .status(OperationStatus.SUCCEEDED) .build()); when(executionManager.sendOperationUpdate(any())).thenReturn(CompletableFuture.completedFuture(null)); + TestUtils.executeOperationTasks(executionManager, durableContext); MockExecutionManagerSupport.stubWaitForOperationCompletion(executionManager); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java index 1a2903251..2ec5c7675 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/ParallelOperationTest.java @@ -86,6 +86,7 @@ void setUp() { .thenReturn(childContext); when(childContext.createChildContext(anyString(), anyString(), anyBoolean())) .thenReturn(childContext); + TestUtils.executeOperationTasks(executionManager, durableContext); // Capture registered operations so we can drive onCheckpointComplete callbacks. var registeredOps = new ConcurrentHashMap(); diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/StepOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/StepOperationTest.java index 72db03028..cfb1e40aa 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/StepOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/StepOperationTest.java @@ -14,6 +14,7 @@ import software.amazon.awssdk.services.lambda.model.OperationStatus; import software.amazon.awssdk.services.lambda.model.StepDetails; import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.StepConfig; import software.amazon.lambda.durable.context.DurableContextImpl; @@ -46,6 +47,7 @@ void setUp() { .thenReturn(DurableConfig.builder() .withExecutorService(Executors.newCachedThreadPool()) .build()); + TestUtils.executeOperationTasks(executionManager, durableContext); MockExecutionManagerSupport.completeOperationsOnManager(executionManager); } diff --git a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java b/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java index 39ea05e8d..d86c05b1a 100644 --- a/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java +++ b/sdk/src/test/java/software/amazon/lambda/durable/operation/WaitForConditionOperationTest.java @@ -19,6 +19,7 @@ import software.amazon.awssdk.services.lambda.model.OperationType; import software.amazon.awssdk.services.lambda.model.StepDetails; import software.amazon.lambda.durable.DurableConfig; +import software.amazon.lambda.durable.TestUtils; import software.amazon.lambda.durable.TypeToken; import software.amazon.lambda.durable.config.WaitForConditionConfig; import software.amazon.lambda.durable.context.DurableContextImpl; @@ -54,6 +55,7 @@ void setUp() { .thenReturn(DurableConfig.builder() .withExecutorService(Executors.newCachedThreadPool()) .build()); + TestUtils.executeOperationTasks(executionManager, durableContext); MockExecutionManagerSupport.completeOperationsOnManager(executionManager); }