From 120d9a4f214bbd93d0796fb5d7782907bd5a502c Mon Sep 17 00:00:00 2001 From: jmaeagle99 <44687433+jmaeagle99@users.noreply.github.com> Date: Thu, 29 Oct 2026 10:36:25 -0700 Subject: [PATCH] Test sending of paginated workflow task completions Cover the paginated send path in WorkflowWorker, which previously had no unit tests: intermediate pages are sent in order before the final page, a completion is sent whole when the namespace does not support pagination, buffer loss on any page resends every page from the first, other errors and worker shutdown end the resend loop, and the gRPC retryer never retries a lost buffer on its own. --- ...orkflowWorkerCompletionPaginationTest.java | 389 ++++++++++++++++++ 1 file changed, 389 insertions(+) create mode 100644 temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerCompletionPaginationTest.java 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(); + } + } +}