From 1257ea83d43916fee9ce071fd6b9317d5bf1c077 Mon Sep 17 00:00:00 2001 From: Dani Khan Date: Mon, 26 Oct 2026 23:42:46 +0000 Subject: [PATCH 1/2] Fix queries buffered during Continue-As-New Treat Continue-As-New as terminal when deciding whether to schedule a new workflow task on the current run. This lets the old history commit before buffered queries dispatch. Register each query future before dispatch so an immediate worker response can resolve it. Add a regression test that buffers a query during Continue-As-New, responds during dispatch, and verifies the query stays on the closed run. --- .../TestWorkflowMutableStateImpl.java | 8 +- .../TestWorkflowMutableStateImplTest.java | 197 ++++++++++++++++++ 2 files changed, 200 insertions(+), 5 deletions(-) create mode 100644 temporal-test-server/src/test/java/io/temporal/internal/testservice/TestWorkflowMutableStateImplTest.java diff --git a/temporal-test-server/src/main/java/io/temporal/internal/testservice/TestWorkflowMutableStateImpl.java b/temporal-test-server/src/main/java/io/temporal/internal/testservice/TestWorkflowMutableStateImpl.java index 804cf7906c..e36bf92e19 100644 --- a/temporal-test-server/src/main/java/io/temporal/internal/testservice/TestWorkflowMutableStateImpl.java +++ b/temporal-test-server/src/main/java/io/temporal/internal/testservice/TestWorkflowMutableStateImpl.java @@ -536,10 +536,7 @@ public void completeWorkflowTask( } WorkflowTaskData data = this.workflowTaskStateMachine.getData(); - boolean completed = - workflow.getState() == StateMachines.State.COMPLETED - || workflow.getState() == StateMachines.State.FAILED - || workflow.getState() == StateMachines.State.CANCELED; + boolean completed = isTerminalState(); if (!completed && ((ctx.isNeedWorkflowTask() || !workflowTaskStateMachine.getData().bufferedEvents.isEmpty()) @@ -603,9 +600,10 @@ public void completeWorkflowTask( stickyExecutionAttributes == null ? startRequest.getTaskQueue().getName() : stickyExecutionAttributes.getWorkerTaskQueue().getName()); + // Register before dispatch because a worker can respond immediately. + this.queries.put(queryId.getQueryId(), consistent.getResult()); store.sendQueryTask( executionId, taskQueueId, task, getStartRequest().getPriority()); - this.queries.put(queryId.getQueryId(), consistent.getResult()); } } data.queryBuffer.clear(); diff --git a/temporal-test-server/src/test/java/io/temporal/internal/testservice/TestWorkflowMutableStateImplTest.java b/temporal-test-server/src/test/java/io/temporal/internal/testservice/TestWorkflowMutableStateImplTest.java new file mode 100644 index 0000000000..db93bc0243 --- /dev/null +++ b/temporal-test-server/src/test/java/io/temporal/internal/testservice/TestWorkflowMutableStateImplTest.java @@ -0,0 +1,197 @@ +package io.temporal.internal.testservice; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNotEquals; +import static org.junit.Assert.assertNotNull; +import static org.junit.Assert.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.spy; + +import com.google.common.util.concurrent.ListenableFuture; +import com.google.protobuf.ByteString; +import com.google.protobuf.util.Durations; +import io.temporal.api.command.v1.Command; +import io.temporal.api.command.v1.ContinueAsNewWorkflowExecutionCommandAttributes; +import io.temporal.api.common.v1.Payload; +import io.temporal.api.common.v1.Payloads; +import io.temporal.api.common.v1.Priority; +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.enums.v1.EventType; +import io.temporal.api.enums.v1.QueryResultType; +import io.temporal.api.history.v1.HistoryEvent; +import io.temporal.api.query.v1.WorkflowQuery; +import io.temporal.api.taskqueue.v1.TaskQueue; +import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueRequest; +import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse; +import io.temporal.api.workflowservice.v1.QueryWorkflowRequest; +import io.temporal.api.workflowservice.v1.QueryWorkflowResponse; +import io.temporal.api.workflowservice.v1.RespondQueryTaskCompletedRequest; +import io.temporal.api.workflowservice.v1.RespondWorkflowTaskCompletedRequest; +import io.temporal.api.workflowservice.v1.StartWorkflowExecutionRequest; +import io.temporal.api.workflowservice.v1.StartWorkflowExecutionResponse; +import io.temporal.api.workflowservice.v1.WorkflowServiceGrpc; +import java.time.Clock; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.UUID; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; +import org.junit.Test; + +public class TestWorkflowMutableStateImplTest { + + @Test(timeout = 60000) + public void queryBufferedDuringContinueAsNewIsDispatchedToClosedRun() throws Exception { + SelfAdvancingTimerImpl timer = new SelfAdvancingTimerImpl(0, Clock.systemUTC()); + timer.lockTimeSkipping("query test"); + TestWorkflowStoreImpl store = spy(new TestWorkflowStoreImpl(timer)); + TestWorkflowService service = + new TestWorkflowService( + store, new TestVisibilityStoreImpl(), new TestNexusEndpointStoreImpl(), timer); + InProcessGRPCServer server = new InProcessGRPCServer(Collections.singletonList(service)); + ListenableFuture queryFuture = null; + try { + WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub = + WorkflowServiceGrpc.newBlockingStub(Objects.requireNonNull(server.getChannel())); + WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub = + WorkflowServiceGrpc.newFutureStub(Objects.requireNonNull(server.getChannel())); + String namespace = "query-test"; + String workflowId = "query-during-continue-as-new"; + TaskQueue originalTaskQueue = TaskQueue.newBuilder().setName("original-run").build(); + TaskQueue continuedTaskQueue = TaskQueue.newBuilder().setName("continued-run").build(); + StartWorkflowExecutionResponse startResponse = + blockingStub + .withDeadlineAfter(5, TimeUnit.SECONDS) + .startWorkflowExecution( + StartWorkflowExecutionRequest.newBuilder() + .setNamespace(namespace) + .setWorkflowId(workflowId) + .setWorkflowType(WorkflowType.newBuilder().setName("QueryWorkflow")) + .setTaskQueue(originalTaskQueue) + .setRequestId(UUID.randomUUID().toString()) + .setWorkflowTaskTimeout(Durations.fromSeconds(60)) + .build()); + WorkflowExecution originalExecution = + WorkflowExecution.newBuilder() + .setWorkflowId(workflowId) + .setRunId(startResponse.getRunId()) + .build(); + PollWorkflowTaskQueueResponse workflowTask = + blockingStub + .withDeadlineAfter(5, TimeUnit.SECONDS) + .pollWorkflowTaskQueue( + PollWorkflowTaskQueueRequest.newBuilder() + .setNamespace(namespace) + .setTaskQueue(originalTaskQueue) + .build()); + assertEquals(originalExecution, workflowTask.getWorkflowExecution()); + + CountDownLatch queryBuffered = new CountDownLatch(1); + // No other work is submitted before this latch fires, so this observes the query buffer. + doAnswer( + invocation -> { + invocation.callRealMethod(); + queryBuffered.countDown(); + return null; + }) + .when(store) + .applyTimersAndLocks(any(RequestContext.class)); + + Payloads queryAnswer = + Payloads.newBuilder() + .addPayloads(Payload.newBuilder().setData(ByteString.copyFromUtf8("original-run"))) + .build(); + AtomicReference dispatchedQuery = new AtomicReference<>(); + // Complete the query before dispatch returns to exercise the fastest possible worker reply. + doAnswer( + invocation -> { + invocation.callRealMethod(); + PollWorkflowTaskQueueResponse.Builder queryTask = invocation.getArgument(2); + dispatchedQuery.set(queryTask.build()); + blockingStub + .withDeadlineAfter(5, TimeUnit.SECONDS) + .respondQueryTaskCompleted( + RespondQueryTaskCompletedRequest.newBuilder() + .setTaskToken(queryTask.getTaskToken()) + .setCompletedType(QueryResultType.QUERY_RESULT_TYPE_ANSWERED) + .setQueryResult(queryAnswer) + .build()); + return null; + }) + .when(store) + .sendQueryTask( + any(ExecutionId.class), + any(TestWorkflowStore.TaskQueueId.class), + any(PollWorkflowTaskQueueResponse.Builder.class), + any(Priority.class)); + + queryFuture = + futureStub + .withDeadlineAfter(20, TimeUnit.SECONDS) + .queryWorkflow( + QueryWorkflowRequest.newBuilder() + .setNamespace(namespace) + .setExecution(WorkflowExecution.newBuilder().setWorkflowId(workflowId)) + .setQuery(WorkflowQuery.newBuilder().setQueryType("state")) + .build()); + assertTrue("Query was not buffered", queryBuffered.await(5, TimeUnit.SECONDS)); + + blockingStub + .withDeadlineAfter(5, TimeUnit.SECONDS) + .respondWorkflowTaskCompleted( + RespondWorkflowTaskCompletedRequest.newBuilder() + .setTaskToken(workflowTask.getTaskToken()) + .addCommands( + Command.newBuilder() + .setCommandType( + CommandType.COMMAND_TYPE_CONTINUE_AS_NEW_WORKFLOW_EXECUTION) + .setContinueAsNewWorkflowExecutionCommandAttributes( + ContinueAsNewWorkflowExecutionCommandAttributes.newBuilder() + .setTaskQueue(continuedTaskQueue))) + .build()); + + assertEquals(queryAnswer, queryFuture.get(5, TimeUnit.SECONDS).getQueryResult()); + PollWorkflowTaskQueueResponse queryTask = dispatchedQuery.get(); + assertNotNull("Query was not dispatched", queryTask); + assertEquals(originalExecution, queryTask.getWorkflowExecution()); + assertTrue(queryTask.hasQuery()); + List history = queryTask.getHistory().getEventsList(); + HistoryEvent closeEvent = history.get(history.size() - 1); + assertEquals( + EventType.EVENT_TYPE_WORKFLOW_EXECUTION_CONTINUED_AS_NEW, closeEvent.getEventType()); + assertEquals( + 1, + history.stream() + .filter(event -> event.getEventType() == EventType.EVENT_TYPE_WORKFLOW_TASK_SCHEDULED) + .count()); + + String continuedRunId = + closeEvent.getWorkflowExecutionContinuedAsNewEventAttributes().getNewExecutionRunId(); + assertFalse(continuedRunId.isEmpty()); + assertNotEquals(originalExecution.getRunId(), continuedRunId); + PollWorkflowTaskQueueResponse continuedTask = + blockingStub + .withDeadlineAfter(5, TimeUnit.SECONDS) + .pollWorkflowTaskQueue( + PollWorkflowTaskQueueRequest.newBuilder() + .setNamespace(namespace) + .setTaskQueue(continuedTaskQueue) + .build()); + assertEquals(continuedRunId, continuedTask.getWorkflowExecution().getRunId()); + assertFalse(continuedTask.hasQuery()); + } finally { + if (queryFuture != null) { + queryFuture.cancel(true); + } + server.shutdownNow(); + service.close(); + server.awaitTermination(5, TimeUnit.SECONDS); + } + } +} From fc42be88a12b2db4130588e2f986672338b005fb Mon Sep 17 00:00:00 2001 From: Dani Khan Date: Tue, 27 Oct 2026 00:14:32 +0000 Subject: [PATCH 2/2] Add changelog entry for buffered query fix --- CHANGELOG.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d426a29297..1303848f91 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -79,6 +79,8 @@ to docs, or any other relevant information. behavior during replay. - Test server now honors retry expiration deadlines that fall exactly on a whole second. Previously such deadlines were ignored and retries were scheduled past them instead of failing with `RETRY_STATE_TIMEOUT`. +- Fixed test server queries timing out when they are buffered during Continue-As-New. The query remains bound to the + original run and uses its committed history. - Worker heartbeats now list each plugin name once. Previously a plugin registered on both the `WorkflowClient` and the `WorkerFactory` was listed twice.