[Bug]: Remaining LMI cleanup bounds and worker recovery #726

// 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.*;
import static software.amazon.lambda.durable.model.ExecutionStatus.*;

import java.time.Duration;
import java.time.Instant;
import java.util.List;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import org.junit.jupiter.api.Test;
import software.amazon.awssdk.services.lambda.model.CheckpointUpdatedExecutionState;
import software.amazon.awssdk.services.lambda.model.ExecutionDetails;
import software.amazon.awssdk.services.lambda.model.Operation;
import software.amazon.awssdk.services.lambda.model.OperationStatus;
import software.amazon.awssdk.services.lambda.model.OperationType;
import software.amazon.lambda.durable.DurableConfig;
import software.amazon.lambda.durable.TestUtils;
import software.amazon.lambda.durable.TypeToken;
import software.amazon.lambda.durable.model.DurableExecutionInput;
import software.amazon.lambda.durable.plugin.DurableExecutionPlugin;
import software.amazon.lambda.durable.plugin.InvocationEndInfo;

/** Diagnostic assertions describe the observed defects, not the desired contract. */
class LmiInvocationLifecycleReproTest {
    private DurableExecutionInput input(String id) {
        var execution = Operation.builder()
                .id(id)
                .type(OperationType.EXECUTION)
                .status(OperationStatus.STARTED)
                .startTimestamp(Instant.now())
                .executionDetails(
                        ExecutionDetails.builder().inputPayload("\"input\"").build())
                .build();
        return new DurableExecutionInput(
                "arn:aws:lambda:us-east-1:123456789012:function:test/durable-execution/test/" + id,
                "token-" + id,
                CheckpointUpdatedExecutionState.builder()
                        .operations(List.of(execution))
                        .build());
    }

    private static void await(CountDownLatch latch) {
        try {
            assertTrue(latch.await(5, TimeUnit.SECONDS), "test synchronization timed out");
        } catch (InterruptedException e) {
            Thread.currentThread().interrupt();
            throw new AssertionError(e);
        }
    }

    @Test
    void pendingResponseCanPrecedeRootHandlerFinallyCompletion() throws Exception {
        var runtime = Executors.newSingleThreadExecutor();
        var users = Executors.newCachedThreadPool();
        var enteredFinally = new CountDownLatch(1);
        var releaseFinally = new CountDownLatch(1);
        var finishedFinally = new CountDownLatch(1);
        var config = DurableConfig.builder()
                .withDurableExecutionClient(TestUtils.createMockClient())
                .withExecutorService(users)
                .build();
        try {
            var response = runtime.submit(() -> DurableExecutor.execute(
                    input("pending"),
                    null,
                    TypeToken.get(String.class),
                    (value, context) -> {
                        try {
                            context.wait("pause", Duration.ofSeconds(30));
                            return "done";
                        } finally {
                            enteredFinally.countDown();
                            await(releaseFinally);
                            finishedFinally.countDown();
                        }
                    },
                    config));
            await(enteredFinally);
            assertEquals(PENDING, response.get(2, TimeUnit.SECONDS).status());
            assertEquals(1, finishedFinally.getCount(), "root cleanup still running after PENDING response");
        } finally {
            releaseFinally.countDown();
            users.shutdown();
            assertTrue(users.awaitTermination(5, TimeUnit.SECONDS));
            runtime.shutdownNow();
        }
    }

    @Test
    void rootHandlersExhaustSharedFixedPoolAndStarveTheirOwnSteps() throws Exception {
        var runtime = Executors.newFixedThreadPool(2);
        var users = (ThreadPoolExecutor) Executors.newFixedThreadPool(2);
        var bothHandlers = new CountDownLatch(2);
        var config = DurableConfig.builder()
                .withDurableExecutionClient(TestUtils.createMockClient())
                .withExecutorService(users)
                .build();
        try {
            var responses = List.of("A", "B").stream()
                    .map(id -> runtime.submit(() -> DurableExecutor.execute(
                            input(id),
                            null,
                            TypeToken.get(String.class),
                            (value, context) -> {
                                bothHandlers.countDown();
                                await(bothHandlers);
                                return context.step("step", String.class, stepContext -> "done");
                            },
                            config)))
                    .toList();
            await(bothHandlers);
            for (var response : responses) {
                assertThrows(TimeoutException.class, () -> response.get(250, TimeUnit.MILLISECONDS));
            }
            assertEquals(2, users.getActiveCount());
            assertEquals(2, users.getQueue().size(), "both step bodies are queued behind blocked handlers");
            users.setMaximumPoolSize(4);
            users.setCorePoolSize(4);
            for (var response : responses) {
                assertEquals(SUCCEEDED, response.get(3, TimeUnit.SECONDS).status());
            }
        } finally {
            users.shutdownNow();
            runtime.shutdownNow();
        }
    }

    @Test
    void invocationEndRunsBeforeUnfinishedStepIsDrainedAndCloseBlocks() throws Exception {
        var runtime = Executors.newSingleThreadExecutor();
        var users = Executors.newCachedThreadPool();
        var stepStarted = new CountDownLatch(1);
        var releaseStep = new CountDownLatch(1);
        var endCalled = new CountDownLatch(1);
        var config = DurableConfig.builder()
                .withDurableExecutionClient(TestUtils.createMockClient())
                .withExecutorService(users)
                .withPlugins(new DurableExecutionPlugin() {
                    @Override
                    public void onInvocationEnd(InvocationEndInfo info) {
                        endCalled.countDown();
                    }
                })
                .build();
        try {
            var response = runtime.submit(() -> DurableExecutor.execute(
                    input("unfinished"),
                    null,
                    TypeToken.get(String.class),
                    (value, context) -> {
                        context.stepAsync("slow-step", String.class, stepContext -> {
                            stepStarted.countDown();
                            await(releaseStep);
                            return "step-done";
                        });
                        await(stepStarted);
                        return "handler-done";
                    },
                    config));
            await(endCalled);
            assertThrows(TimeoutException.class, () -> response.get(250, TimeUnit.MILLISECONDS));
            releaseStep.countDown();
            assertEquals(SUCCEEDED, response.get(3, TimeUnit.SECONDS).status());
        } finally {
            releaseStep.countDown();
            users.shutdown();
            assertTrue(users.awaitTermination(5, TimeUnit.SECONDS));
            runtime.shutdownNow();
        }
    }
}