diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerCompletionPaginationTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerCompletionPaginationTest.java new file mode 100644 index 0000000000..dccd44fb2e --- /dev/null +++ b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerCompletionPaginationTest.java @@ -0,0 +1,389 @@ +package io.temporal.internal.worker; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import com.google.common.util.concurrent.Futures; +import com.google.protobuf.ByteString; +import com.uber.m3.tally.NoopScope; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.temporal.api.command.v1.Command; +import io.temporal.api.command.v1.RecordMarkerCommandAttributes; +import io.temporal.api.common.v1.Payload; +import io.temporal.api.common.v1.Payloads; +import io.temporal.api.common.v1.WorkflowExecution; +import io.temporal.api.common.v1.WorkflowType; +import io.temporal.api.enums.v1.CommandType; +import io.temporal.api.errordetails.v1.WorkflowTaskCompletionBufferLostFailure; +import io.temporal.api.namespace.v1.NamespaceInfo; +import io.temporal.api.protocol.v1.Message; +import io.temporal.api.workflowservice.v1.GetSystemInfoResponse; +import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueRequest; +import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse; +import io.temporal.api.workflowservice.v1.RespondWorkflowTaskCompletedRequest; +import io.temporal.api.workflowservice.v1.RespondWorkflowTaskCompletedResponse; +import io.temporal.api.workflowservice.v1.ShutdownWorkerRequest; +import io.temporal.api.workflowservice.v1.ShutdownWorkerResponse; +import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc; +import io.temporal.serviceclient.StatusUtils; +import io.temporal.serviceclient.WorkflowServiceStubs; +import io.temporal.worker.tuning.FixedSizeSlotSupplier; +import io.temporal.worker.tuning.PollerBehaviorSimpleMaximum; +import io.temporal.worker.tuning.SlotMarkUsedContext; +import io.temporal.worker.tuning.SlotPermit; +import io.temporal.worker.tuning.SlotReleaseContext; +import io.temporal.worker.tuning.SlotReserveContext; +import io.temporal.worker.tuning.SlotSupplier; +import io.temporal.worker.tuning.SlotSupplierFuture; +import io.temporal.worker.tuning.WorkflowSlotInfo; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Optional; +import java.util.Set; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import org.junit.Test; + +public class WorkflowWorkerCompletionPaginationTest { + + private static final ByteString TASK_TOKEN = ByteString.copyFrom("token", UTF_8); + private static final String FINAL_PAGE = "final"; + + // Four 1 MiB commands exceed MAX_PAGE_BYTES, so the completion is split into several pages. + private static final int COMMAND_COUNT = 4; + + @Test + public void aPaginatedCompletionSendsEveryIntermediatePageBeforeTheFinalPage() throws Exception { + RespondWorkflowTaskCompletedRequest completion = + oversizedCompletion().toBuilder() + .addMessages(Message.newBuilder().setId("message").build()) + .build(); + + List sent = + completeTask(true, completion, (page, worker) -> success()); + + assertEquals(attempts(1), labels(sent)); + int commandCount = 0; + for (RespondWorkflowTaskCompletedRequest page : sent) { + assertEquals(TASK_TOKEN, page.getTaskToken()); + if (page.getIntermediatePage()) { + assertEquals(0, page.getMessagesCount()); + commandCount += page.getCommandsCount(); + } else { + assertEquals(0, page.getCommandsCount()); + assertEquals(1, page.getMessagesCount()); + } + } + assertEquals(COMMAND_COUNT, commandCount); + } + + @Test + public void aCompletionIsSentAsOneRequestWhenTheNamespaceDoesNotSupportPagination() + throws Exception { + List sent = + completeTask(false, oversizedCompletion(), (page, worker) -> success()); + + assertEquals(1, sent.size()); + assertFalse(sent.get(0).getIntermediatePage()); + assertEquals(0, sent.get(0).getPageNumber()); + assertEquals(COMMAND_COUNT, sent.get(0).getCommandsCount()); + } + + @Test + public void bufferLossOnTheFinalPageResendsEveryPageFromTheFirst() throws Exception { + AtomicInteger finalAttempts = new AtomicInteger(); + List sent = + completeTask( + true, + oversizedCompletion(), + (page, worker) -> { + if (!page.getIntermediatePage() && finalAttempts.incrementAndGet() <= 2) { + throw bufferLost(); + } + return success(); + }); + + // Each attempt sends the final page exactly once, so the gRPC retryer must not have retried it. + assertEquals(attempts(3), labels(sent)); + } + + @Test + public void bufferLossOnAnIntermediatePageResendsEveryPageFromTheFirst() throws Exception { + AtomicInteger lastPageAttempts = new AtomicInteger(); + int lastIntermediatePage = intermediatePageCount() - 1; + List sent = + completeTask( + true, + oversizedCompletion(), + (page, worker) -> { + if (page.getIntermediatePage() + && page.getPageNumber() == lastIntermediatePage + && lastPageAttempts.incrementAndGet() == 1) { + throw bufferLost(); + } + return success(); + }); + + List expected = new ArrayList<>(attempts(1)); + expected.remove(FINAL_PAGE); + expected.addAll(attempts(1)); + assertEquals(expected, labels(sent)); + } + + @Test + public void otherErrorsAreNotResent() throws Exception { + AtomicInteger finalAttempts = new AtomicInteger(); + List sent = + completeTask( + true, + oversizedCompletion(), + (page, worker) -> { + // Accepting a resend lets a regression fail the assertion instead of looping. + if (!page.getIntermediatePage() && finalAttempts.incrementAndGet() == 1) { + // What the server returns once the task has timed out and the token is stale. + throw Status.NOT_FOUND + .withDescription("workflow task not found") + .asRuntimeException(); + } + return success(); + }); + + assertEquals(attempts(1), labels(sent)); + } + + @Test + public void shutdownStopsBufferLossResends() throws Exception { + AtomicInteger finalAttempts = new AtomicInteger(); + List sent = + completeTask( + true, + oversizedCompletion(), + (page, worker) -> { + // Accepting a resend lets a regression fail the assertion instead of looping. + if (!page.getIntermediatePage() && finalAttempts.incrementAndGet() == 1) { + // Shut down while the first attempt is in flight, so the resend loop observes it. + worker.shutdown(new ShutdownManager(), false); + throw bufferLost(); + } + return success(); + }); + + assertEquals(attempts(1), labels(sent)); + } + + @FunctionalInterface + private interface CompletionAnswer { + RespondWorkflowTaskCompletedResponse answer( + RespondWorkflowTaskCompletedRequest page, WorkflowWorker worker); + } + + /** + * Runs one workflow task that produces {@code completion} and returns every completion request + * the worker sent. The worker is shut down only after the task finishes, because the resend loop + * stops early on shutdown. + */ + private static List completeTask( + boolean paginationEnabled, + RespondWorkflowTaskCompletedRequest completion, + CompletionAnswer answer) + throws Exception { + NamespaceCapabilities namespaceCapabilities = new NamespaceCapabilities(); + namespaceCapabilities.setFromCapabilities( + NamespaceInfo.Capabilities.newBuilder() + .setWorkflowTaskCompletionPagination(paginationEnabled) + .build()); + + WorkflowServiceStubs client = mock(WorkflowServiceStubs.class); + when(client.getServerCapabilities()) + .thenReturn(() -> GetSystemInfoResponse.Capabilities.newBuilder().build()); + WorkflowRunLockManager runLockManager = new WorkflowRunLockManager(); + WorkflowExecutorCache cache = new WorkflowExecutorCache(10, runLockManager, new NoopScope()); + WorkflowTaskHandler taskHandler = mock(WorkflowTaskHandler.class); + when(taskHandler.isAnyTypeSupported()).thenReturn(true); + when(taskHandler.handleWorkflowTask(any(PollWorkflowTaskQueueResponse.class), any())) + .thenReturn( + new WorkflowTaskHandler.Result( + "workflow-type", completion, null, null, null, false, null, null)); + TaskFinishedSlotSupplier slotSupplier = new TaskFinishedSlotSupplier(); + + WorkflowWorker worker = + new WorkflowWorker( + client, + "default", + "task_queue", + "sticky_task_queue", + SingleWorkerOptions.newBuilder() + .setIdentity("test_identity") + .setBuildId(UUID.randomUUID().toString()) + .setWorkerInstanceKey(UUID.randomUUID().toString()) + .setPollerOptions( + PollerOptions.newBuilder() + .setPollerBehavior(new PollerBehaviorSimpleMaximum(1)) + .build()) + .build(), + runLockManager, + cache, + taskHandler, + mock(EagerActivityDispatcher.class), + 3, + slotSupplier, + namespaceCapabilities); + + WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub = + mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class); + when(futureStub.shutdownWorker(any(ShutdownWorkerRequest.class))) + .thenReturn(Futures.immediateFuture(ShutdownWorkerResponse.newBuilder().build())); + WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub = + mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class); + when(client.blockingStub()).thenReturn(blockingStub); + when(client.futureStub()).thenReturn(futureStub); + when(blockingStub.withOption(any(), any())).thenReturn(blockingStub); + + List sent = new CopyOnWriteArrayList<>(); + when(blockingStub.respondWorkflowTaskCompleted(any(RespondWorkflowTaskCompletedRequest.class))) + .thenAnswer( + invocation -> { + RespondWorkflowTaskCompletedRequest page = invocation.getArgument(0); + sent.add(page); + return answer.answer(page, worker); + }); + + PollWorkflowTaskQueueResponse pollResponse = + PollWorkflowTaskQueueResponse.newBuilder() + .setTaskToken(TASK_TOKEN) + .setWorkflowExecution( + WorkflowExecution.newBuilder().setWorkflowId("workflow-id").setRunId("run-id")) + .setWorkflowType(WorkflowType.newBuilder().setName("workflow-type")) + .build(); + CountDownLatch blockPolls = new CountDownLatch(1); + when(blockingStub.pollWorkflowTaskQueue(any(PollWorkflowTaskQueueRequest.class))) + .thenReturn(pollResponse) + .thenAnswer( + invocation -> { + blockPolls.await(); + return null; + }); + + assertTrue(worker.start()); + assertTrue( + "workflow task did not finish", slotSupplier.taskFinished.await(30, TimeUnit.SECONDS)); + worker.shutdown(new ShutdownManager(), false).get(); + return sent; + } + + private static RespondWorkflowTaskCompletedRequest oversizedCompletion() { + RespondWorkflowTaskCompletedRequest.Builder completion = + RespondWorkflowTaskCompletedRequest.newBuilder(); + for (int i = 0; i < COMMAND_COUNT; i++) { + completion.addCommands( + Command.newBuilder() + .setCommandType(CommandType.COMMAND_TYPE_RECORD_MARKER) + .setRecordMarkerCommandAttributes( + RecordMarkerCommandAttributes.newBuilder() + .setMarkerName("marker") + .putDetails( + "data", + Payloads.newBuilder() + .addPayloads( + Payload.newBuilder() + .setData(ByteString.copyFrom(new byte[1024 * 1024]))) + .build()))); + } + return completion.build(); + } + + private static int intermediatePageCount() { + int count = + WorkflowTaskCompletionPaginator.paginate( + oversizedCompletion(), WorkflowTaskCompletionPaginator.MAX_PAGE_BYTES) + .intermediatePages + .size(); + assertTrue("expected at least two intermediate pages, got " + count, count >= 2); + return count; + } + + /** The page labels of {@code count} complete send attempts, each ending with the final page. */ + private static List attempts(int count) { + List labels = new ArrayList<>(); + int pages = intermediatePageCount(); + for (int attempt = 0; attempt < count; attempt++) { + for (int page = 0; page < pages; page++) { + labels.add(Integer.toString(page)); + } + labels.add(FINAL_PAGE); + } + return labels; + } + + private static List labels(List sent) { + List labels = new ArrayList<>(); + for (RespondWorkflowTaskCompletedRequest page : sent) { + if (page.getIntermediatePage()) { + labels.add(Integer.toString(page.getPageNumber())); + } else { + assertEquals(intermediatePageCount(), page.getPageNumber()); + labels.add(FINAL_PAGE); + } + } + return Collections.unmodifiableList(labels); + } + + private static RespondWorkflowTaskCompletedResponse success() { + return RespondWorkflowTaskCompletedResponse.getDefaultInstance(); + } + + private static StatusRuntimeException bufferLost() { + return StatusUtils.newException( + Status.ABORTED.withDescription("buffer lost"), + WorkflowTaskCompletionBufferLostFailure.getDefaultInstance(), + WorkflowTaskCompletionBufferLostFailure.getDescriptor()); + } + + /** Signals once a slot that ran a task is released, after the task's completion is reported. */ + private static final class TaskFinishedSlotSupplier implements SlotSupplier { + private final SlotSupplier inner = new FixedSizeSlotSupplier<>(10); + private final Set usedPermits = ConcurrentHashMap.newKeySet(); + final CountDownLatch taskFinished = new CountDownLatch(1); + + @Override + public SlotSupplierFuture reserveSlot(SlotReserveContext ctx) + throws Exception { + return inner.reserveSlot(ctx); + } + + @Override + public Optional tryReserveSlot(SlotReserveContext ctx) { + return inner.tryReserveSlot(ctx); + } + + @Override + public void markSlotUsed(SlotMarkUsedContext ctx) { + usedPermits.add(ctx.getSlotPermit()); + inner.markSlotUsed(ctx); + } + + @Override + public void releaseSlot(SlotReleaseContext ctx) { + inner.releaseSlot(ctx); + if (usedPermits.remove(ctx.getSlotPermit())) { + taskFinished.countDown(); + } + } + + @Override + public Optional getMaximumSlots() { + return inner.getMaximumSlots(); + } + } +}