From 449cab71190815646e38d2063a5a2ce34d0edbfb Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Tue, 18 Aug 2026 16:39:59 -0400 Subject: [PATCH 1/5] feature(extstore): integrate into workflow worker pipeline, including replay handler. --- .../replay/ReplayWorkflowTaskHandler.java | 10 +- .../ServiceWorkflowHistoryIterator.java | 21 +++- .../internal/worker/WorkflowWorker.java | 103 +++++++++++++-- .../ServiceWorkflowHistoryIteratorTest.java | 117 ++++++++++++++++++ .../internal/worker/WorkflowWorkerTest.java | 62 ++++++++++ 5 files changed, 300 insertions(+), 13 deletions(-) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java b/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java index f5b7cb0d29..479b08379d 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java @@ -23,6 +23,7 @@ import io.temporal.common.converter.DataConverter; import io.temporal.internal.common.ProtobufTimeUtils; import io.temporal.internal.common.WorkflowExecutionUtils; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.worker.*; import io.temporal.payload.context.WorkflowSerializationContext; import io.temporal.serviceclient.MetricsTag; @@ -77,6 +78,12 @@ public WorkflowTaskHandler.Result handleWorkflowTask(PollWorkflowTaskQueueRespon String workflowType = workflowTask.getWorkflowType().getName(); Scope metricsScope = options.getMetricsScope().tagged(ImmutableMap.of(MetricsTag.WORKFLOW_TYPE, workflowType)); + ExternalStorageRunner externalStorage = options.getExternalStorage(); + if (externalStorage == null) { + ExternalStorageRunner.throwIfContainsReference(workflowTask); + } else { + workflowTask = externalStorage.retrieve(workflowTask); + } return handleWorkflowTaskWithQuery(workflowTask.toBuilder(), metricsScope); } @@ -94,7 +101,8 @@ private Result handleWorkflowTaskWithQuery( logWorkflowTaskToBeProcessed(workflowTask, createdNew); ServiceWorkflowHistoryIterator historyIterator = - new ServiceWorkflowHistoryIterator(service, namespace, workflowTask, metricsScope); + new ServiceWorkflowHistoryIterator( + service, namespace, workflowTask, metricsScope, options.getExternalStorage()); boolean finalCommand; Result result; diff --git a/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java b/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java index 229b66186e..14598af207 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java @@ -12,12 +12,14 @@ import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryRequest; import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse; import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponseOrBuilder; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.retryer.GrpcRetryer; import io.temporal.serviceclient.RpcRetryOptions; import io.temporal.serviceclient.WorkflowServiceStubs; import java.time.Duration; import java.util.Iterator; import java.util.NoSuchElementException; +import javax.annotation.Nullable; /** Supports iteration over history while loading new pages through calls to the service. */ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator { @@ -29,6 +31,7 @@ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator { private final Scope metricsScope; private final PollWorkflowTaskQueueResponseOrBuilder task; private final GrpcRetryer grpcRetryer; + private final @Nullable ExternalStorageRunner externalStorage; private Deadline deadline; private Iterator current; ByteString nextPageToken; @@ -38,10 +41,20 @@ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator { String namespace, PollWorkflowTaskQueueResponseOrBuilder task, Scope metricsScope) { + this(service, namespace, task, metricsScope, null); + } + + ServiceWorkflowHistoryIterator( + WorkflowServiceStubs service, + String namespace, + PollWorkflowTaskQueueResponseOrBuilder task, + Scope metricsScope, + @Nullable ExternalStorageRunner externalStorage) { this.service = service; this.namespace = namespace; this.task = task; this.metricsScope = metricsScope; + this.externalStorage = externalStorage; // TODO Refactor WorkflowHistoryIteratorTest or WorkflowHistoryIterator to remove this check. // `service == null` shouldn't be allowed as it's needed for a normal functioning of this // class. @@ -64,7 +77,13 @@ public boolean hasNext() { // true. GetWorkflowExecutionHistoryResponse response = queryWorkflowExecutionHistory(); - current = response.getHistory().getEventsList().iterator(); + History history = response.getHistory(); + if (externalStorage == null) { + ExternalStorageRunner.throwIfContainsReference(history); + } else { + history = externalStorage.retrieve(history); + } + current = history.getEventsList().iterator(); nextPageToken = response.getNextPageToken(); // Server can return an empty page, but a valid nextPageToken that contains // more events. diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java index 3eed1099d3..c23f4cade5 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java @@ -6,21 +6,31 @@ import com.google.common.base.Preconditions; import com.google.common.base.Strings; import com.google.protobuf.ByteString; +import com.google.protobuf.MessageOrBuilder; import com.uber.m3.tally.Scope; import com.uber.m3.tally.Stopwatch; import com.uber.m3.util.ImmutableMap; import io.grpc.StatusRuntimeException; +import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributesOrBuilder; +import io.temporal.api.command.v1.SignalExternalWorkflowExecutionCommandAttributesOrBuilder; +import io.temporal.api.command.v1.StartChildWorkflowExecutionCommandAttributesOrBuilder; import io.temporal.api.common.v1.WorkflowExecution; import io.temporal.api.enums.v1.QueryResultType; import io.temporal.api.enums.v1.TaskQueueKind; import io.temporal.api.enums.v1.WorkflowTaskFailedCause; import io.temporal.api.failure.v1.Failure; import io.temporal.api.workflowservice.v1.*; +import io.temporal.common.CancellationToken; import io.temporal.failure.ApplicationFailure; import io.temporal.internal.logging.LoggerTag; +import io.temporal.internal.payload.storage.ExternalStorageRunner; +import io.temporal.internal.payload.visitor.MessageVisitor; import io.temporal.internal.retryer.GrpcMessageTooLargeException; import io.temporal.internal.retryer.GrpcRetryer; import io.temporal.payload.context.WorkflowSerializationContext; +import io.temporal.payload.storage.StorageDriverActivityInfo; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import io.temporal.payload.storage.StorageDriverWorkflowInfo; import io.temporal.serviceclient.MetricsTag; import io.temporal.serviceclient.RpcRetryOptions; import io.temporal.serviceclient.WorkflowServiceStubs; @@ -381,6 +391,57 @@ public String toString() { options.getIdentity(), namespace, taskQueue); } + private void storeOutboundPayloads( + com.google.protobuf.Message.Builder builder, @Nullable StorageDriverTargetInfo target) { + ExternalStorageRunner externalStorage = options.getExternalStorage(); + if (externalStorage != null) { + externalStorage.store(builder, target); + } + } + + private void storeOutboundPayloads( + com.google.protobuf.Message.Builder builder, + @Nullable StorageDriverTargetInfo target, + MessageVisitor targetVisitor) { + ExternalStorageRunner externalStorage = options.getExternalStorage(); + if (externalStorage != null) { + externalStorage.store(builder, target, targetVisitor, CancellationToken.none()); + } + } + + @Nullable + private StorageDriverTargetInfo workflowStorageTarget( + WorkflowExecution execution, String workflowType) { + if (options.getExternalStorage() == null) { + return null; + } + return new StorageDriverWorkflowInfo( + namespace, execution.getWorkflowId(), execution.getRunId(), workflowType); + } + + static StorageDriverTargetInfo refineStorageTarget( + String namespace, StorageDriverTargetInfo current, MessageOrBuilder message) { + if (message instanceof ScheduleActivityTaskCommandAttributesOrBuilder) { + ScheduleActivityTaskCommandAttributesOrBuilder attrs = + (ScheduleActivityTaskCommandAttributesOrBuilder) message; + return new StorageDriverActivityInfo( + namespace, attrs.getActivityId(), null, attrs.getActivityType().getName()); + } + if (message instanceof StartChildWorkflowExecutionCommandAttributesOrBuilder) { + StartChildWorkflowExecutionCommandAttributesOrBuilder attrs = + (StartChildWorkflowExecutionCommandAttributesOrBuilder) message; + return new StorageDriverWorkflowInfo( + namespace, attrs.getWorkflowId(), null, attrs.getWorkflowType().getName()); + } + if (message instanceof SignalExternalWorkflowExecutionCommandAttributesOrBuilder) { + WorkflowExecution execution = + ((SignalExternalWorkflowExecutionCommandAttributesOrBuilder) message).getExecution(); + return new StorageDriverWorkflowInfo( + namespace, execution.getWorkflowId(), execution.getRunId(), null); + } + return current; + } + private class TaskHandlerImpl implements PollTaskExecutor.TaskHandler { final WorkflowTaskHandler handler; @@ -453,7 +514,10 @@ public void handle(WorkflowTask task) throws Exception { if (queryCompleted != null) { try { sendDirectQueryCompletedResponse( - currentTask.getTaskToken(), queryCompleted.toBuilder(), workflowTypeScope); + currentTask.getTaskToken(), + queryCompleted.toBuilder(), + workflowTypeScope, + workflowStorageTarget(workflowExecution, workflowType)); } catch (StatusRuntimeException e) { GrpcMessageTooLargeException tooLargeException = GrpcMessageTooLargeException.tryWrap(e); @@ -473,7 +537,10 @@ public void handle(WorkflowTask task) throws Exception { .setErrorMessage(failure.getMessage()) .setFailure(failure); sendDirectQueryCompletedResponse( - currentTask.getTaskToken(), queryFailedBuilder, workflowTypeScope); + currentTask.getTaskToken(), + queryFailedBuilder, + workflowTypeScope, + workflowStorageTarget(workflowExecution, workflowType)); } } else { try { @@ -489,7 +556,8 @@ public void handle(WorkflowTask task) throws Exception { currentTask.getTaskToken(), requestBuilder, result.getRequestRetryOptions(), - workflowTypeScope); + workflowTypeScope, + workflowStorageTarget(workflowExecution, workflowType)); // If we were processing a speculative WFT the server may instruct us that the // task was dropped by resting out event ID. long resetEventId = response.getResetHistoryEventId(); @@ -509,7 +577,8 @@ public void handle(WorkflowTask task) throws Exception { currentTask.getTaskToken(), taskFailed.toBuilder(), result.getRequestRetryOptions(), - workflowTypeScope); + workflowTypeScope, + workflowStorageTarget(workflowExecution, workflowType)); } // Apply post-completion metrics only if runnable present and the above succeeded @@ -546,7 +615,8 @@ public void handle(WorkflowTask task) throws Exception { currentTask.getTaskToken(), taskFailedBuilder, result.getRequestRetryOptions(), - workflowTypeScope); + workflowTypeScope, + workflowStorageTarget(workflowExecution, workflowType)); } } } catch (Exception e) { @@ -651,7 +721,8 @@ private RespondWorkflowTaskCompletedResponse sendTaskCompleted( ByteString taskToken, RespondWorkflowTaskCompletedRequest.Builder taskCompleted, RpcRetryOptions retryOptions, - Scope workflowTypeMetricsScope) { + Scope workflowTypeMetricsScope, + @Nullable StorageDriverTargetInfo storageTarget) { GrpcRetryer.GrpcRetryerOptions grpcRetryOptions = new GrpcRetryer.GrpcRetryerOptions( RpcRetryOptions.newBuilder().buildWithDefaultsFrom(retryOptions), null); @@ -674,12 +745,16 @@ private RespondWorkflowTaskCompletedResponse sendTaskCompleted( taskCompleted.setBinaryChecksum(options.getBuildId()); } + MessageVisitor storageTargetVisitor = + (current, message) -> refineStorageTarget(namespace, current, message); + storeOutboundPayloads(taskCompleted, storageTarget, storageTargetVisitor); + RespondWorkflowTaskCompletedRequest request = taskCompleted.build(); return grpcRetryer.retryWithResult( () -> service .blockingStub() .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, workflowTypeMetricsScope) - .respondWorkflowTaskCompleted(taskCompleted.build()), + .respondWorkflowTaskCompleted(request), grpcRetryOptions); } @@ -688,7 +763,8 @@ private void sendTaskFailed( ByteString taskToken, RespondWorkflowTaskFailedRequest.Builder taskFailed, RpcRetryOptions retryOptions, - Scope workflowTypeMetricsScope) { + Scope workflowTypeMetricsScope, + @Nullable StorageDriverTargetInfo storageTarget) { GrpcRetryer.GrpcRetryerOptions grpcRetryOptions = new GrpcRetryer.GrpcRetryerOptions( RpcRetryOptions.newBuilder().buildWithDefaultsFrom(retryOptions), null); @@ -702,25 +778,30 @@ private void sendTaskFailed( taskFailed.setWorkerVersion(options.workerVersionStamp()); } + storeOutboundPayloads(taskFailed, storageTarget); + RespondWorkflowTaskFailedRequest request = taskFailed.build(); grpcRetryer.retry( () -> service .blockingStub() .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, workflowTypeMetricsScope) - .respondWorkflowTaskFailed(taskFailed.build()), + .respondWorkflowTaskFailed(request), grpcRetryOptions); } private void sendDirectQueryCompletedResponse( ByteString taskToken, RespondQueryTaskCompletedRequest.Builder queryCompleted, - Scope workflowTypeMetricsScope) { + Scope workflowTypeMetricsScope, + @Nullable StorageDriverTargetInfo storageTarget) { queryCompleted.setTaskToken(taskToken).setNamespace(namespace); + storeOutboundPayloads(queryCompleted, storageTarget); + RespondQueryTaskCompletedRequest request = queryCompleted.build(); // Do not retry query response service .blockingStub() .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, workflowTypeMetricsScope) - .respondQueryTaskCompleted(queryCompleted.build()); + .respondQueryTaskCompleted(request); } private void logExceptionDuringResultReporting( diff --git a/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java b/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java index ad0c665800..d302facdca 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java @@ -1,12 +1,29 @@ package io.temporal.internal.replay; import com.google.protobuf.ByteString; +import io.temporal.api.common.v1.Payload; +import io.temporal.api.common.v1.Payloads; import io.temporal.api.history.v1.History; +import io.temporal.api.history.v1.HistoryEvent; +import io.temporal.api.history.v1.WorkflowExecutionStartedEventAttributes; import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse; import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse; +import io.temporal.internal.payload.storage.ExternalStorageNotConfiguredException; +import io.temporal.internal.payload.storage.ExternalStorageRunner; +import io.temporal.payload.storage.ExternalStorage; +import io.temporal.payload.storage.StorageDriver; +import io.temporal.payload.storage.StorageDriverClaim; +import io.temporal.payload.storage.StorageDriverRetrieveContext; +import io.temporal.payload.storage.StorageDriverStoreContext; import io.temporal.testUtils.HistoryUtils; import java.nio.charset.Charset; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashMap; +import java.util.List; +import java.util.Map; import java.util.NoSuchElementException; +import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import org.junit.Assert; import org.junit.Test; @@ -84,4 +101,104 @@ GetWorkflowExecutionHistoryResponse queryWorkflowExecutionHistory() { Assert.assertThrows(NoSuchElementException.class, iterator::next); Assert.assertEquals(4, timesCalledServer.get()); } + + @Test + public void resolvesExternalStorageReferencesInFetchedPages() { + ExternalStorageRunner storage = inMemoryStorage(); + History inline = historyWithInput(payload("big-input")); + History.Builder builder = inline.toBuilder(); + storage.store(builder, null); + History stored = builder.build(); + Assert.assertNotEquals( + "stored history should hold a reference, not the inline payload", inline, stored); + + ServiceWorkflowHistoryIterator iterator = fetchingIterator(stored, storage); + + HistoryEvent event = iterator.next(); + Assert.assertEquals( + payload("big-input"), + event.getWorkflowExecutionStartedEventAttributes().getInput().getPayloads(0)); + } + + @Test + public void failsLoudWhenAFetchedPageHasAReferenceAndStorageIsNotConfigured() { + History.Builder builder = historyWithInput(payload("big-input")).toBuilder(); + inMemoryStorage().store(builder, null); + History stored = builder.build(); + + ServiceWorkflowHistoryIterator iterator = fetchingIterator(stored, null); + + Assert.assertThrows(ExternalStorageNotConfiguredException.class, iterator::hasNext); + } + + private static ServiceWorkflowHistoryIterator fetchingIterator( + History page, ExternalStorageRunner storage) { + PollWorkflowTaskQueueResponse workflowTask = + PollWorkflowTaskQueueResponse.newBuilder().setNextPageToken(NEXT_PAGE_TOKEN).build(); + return new ServiceWorkflowHistoryIterator(null, "default", workflowTask, null, storage) { + @Override + GetWorkflowExecutionHistoryResponse queryWorkflowExecutionHistory() { + return GetWorkflowExecutionHistoryResponse.newBuilder().setHistory(page).build(); + } + }; + } + + private static ExternalStorageRunner inMemoryStorage() { + return ExternalStorageRunner.create( + ExternalStorage.newBuilder() + .setDriver(new InMemoryDriver()) + .setPayloadSizeThreshold(0) + .build()); + } + + private static History historyWithInput(Payload payload) { + return History.newBuilder() + .addEvents( + HistoryEvent.newBuilder() + .setWorkflowExecutionStartedEventAttributes( + WorkflowExecutionStartedEventAttributes.newBuilder() + .setInput(Payloads.newBuilder().addPayloads(payload)))) + .build(); + } + + private static Payload payload(String data) { + return Payload.newBuilder().setData(ByteString.copyFromUtf8(data)).build(); + } + + private static final class InMemoryDriver implements StorageDriver { + private final Map objects = new HashMap<>(); + private int counter = 0; + + @Override + public String getName() { + return "test"; + } + + @Override + public String getType() { + return "test.inmemory"; + } + + @Override + public synchronized CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + List claims = new ArrayList<>(); + for (Payload payload : payloads) { + String key = "k-" + (counter++); + objects.put(key, payload); + claims.add(new StorageDriverClaim(Collections.singletonMap("key", key))); + } + return CompletableFuture.completedFuture(claims); + } + + @Override + public synchronized CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + List payloads = new ArrayList<>(); + for (StorageDriverClaim claim : claims) { + payloads.add(objects.get(claim.getClaimData().get("key"))); + } + return CompletableFuture.completedFuture(payloads); + } + } } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java index 5cd1fc8d3e..c97d85a188 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java @@ -12,6 +12,11 @@ import com.uber.m3.tally.RootScopeBuilder; import com.uber.m3.tally.Scope; import com.uber.m3.util.ImmutableMap; +import io.temporal.api.command.v1.CompleteWorkflowExecutionCommandAttributes; +import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributes; +import io.temporal.api.command.v1.SignalExternalWorkflowExecutionCommandAttributes; +import io.temporal.api.command.v1.StartChildWorkflowExecutionCommandAttributes; +import io.temporal.api.common.v1.ActivityType; import io.temporal.api.common.v1.WorkflowExecution; import io.temporal.api.common.v1.WorkflowType; import io.temporal.api.workflowservice.v1.*; @@ -20,6 +25,9 @@ import io.temporal.internal.replay.ReplayWorkflow; import io.temporal.internal.replay.ReplayWorkflowFactory; import io.temporal.internal.replay.ReplayWorkflowTaskHandler; +import io.temporal.payload.storage.StorageDriverActivityInfo; +import io.temporal.payload.storage.StorageDriverTargetInfo; +import io.temporal.payload.storage.StorageDriverWorkflowInfo; import io.temporal.serviceclient.WorkflowServiceStubs; import io.temporal.testUtils.Eventually; import io.temporal.testUtils.HistoryUtils; @@ -448,4 +456,58 @@ private ReplayWorkflowFactory setUpMockWorkflowFactory() throws Throwable { when(mockWorkflow.eventLoop()).thenReturn(false); return mockFactory; } + + @Test + public void refineStorageTargetPointsActivityCommandsAtTheActivity() { + StorageDriverTargetInfo workflowDefault = + new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "MyWorkflow"); + ScheduleActivityTaskCommandAttributes command = + ScheduleActivityTaskCommandAttributes.newBuilder() + .setActivityId("act-1") + .setActivityType(ActivityType.newBuilder().setName("MyActivity")) + .build(); + + assertEquals( + new StorageDriverActivityInfo("ns", "act-1", null, "MyActivity"), + WorkflowWorker.refineStorageTarget("ns", workflowDefault, command)); + } + + @Test + public void refineStorageTargetPointsChildWorkflowCommandsAtTheChild() { + StorageDriverTargetInfo parent = + new StorageDriverWorkflowInfo("ns", "parent", "parent-run", "Parent"); + StartChildWorkflowExecutionCommandAttributes command = + StartChildWorkflowExecutionCommandAttributes.newBuilder() + .setWorkflowId("child-1") + .setWorkflowType(WorkflowType.newBuilder().setName("Child")) + .build(); + + assertEquals( + new StorageDriverWorkflowInfo("ns", "child-1", null, "Child"), + WorkflowWorker.refineStorageTarget("ns", parent, command)); + } + + @Test + public void refineStorageTargetPointsSignalCommandsAtTheTargetWorkflow() { + StorageDriverTargetInfo self = new StorageDriverWorkflowInfo("ns", "self", "self-run", "Self"); + SignalExternalWorkflowExecutionCommandAttributes command = + SignalExternalWorkflowExecutionCommandAttributes.newBuilder() + .setExecution( + WorkflowExecution.newBuilder().setWorkflowId("other").setRunId("other-run")) + .build(); + + assertEquals( + new StorageDriverWorkflowInfo("ns", "other", "other-run", null), + WorkflowWorker.refineStorageTarget("ns", self, command)); + } + + @Test + public void refineStorageTargetKeepsTheCurrentTargetForOtherCommands() { + StorageDriverTargetInfo current = + new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "MyWorkflow"); + CompleteWorkflowExecutionCommandAttributes command = + CompleteWorkflowExecutionCommandAttributes.newBuilder().build(); + + assertSame(current, WorkflowWorker.refineStorageTarget("ns", current, command)); + } } From d550aeaacc01f54a10a605b76d1334dd9321577f Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Mon, 24 Aug 2026 13:19:06 -0400 Subject: [PATCH 2/5] feat(extstore): make sure sticky cache miss path also retrieves external payloads after fetching history. --- .../replay/ReplayWorkflowTaskHandler.java | 6 + ...orkflowRunTaskHandlerTaskHandlerTests.java | 112 ++++++++++++++++++ 2 files changed, 118 insertions(+) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java b/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java index 479b08379d..351d70d9e2 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java @@ -403,6 +403,12 @@ private WorkflowRunTaskHandler createStatefulHandler( .blockingStub() .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope) .getWorkflowExecutionHistory(getHistoryRequest); + ExternalStorageRunner externalStorage = options.getExternalStorage(); + if (externalStorage == null) { + ExternalStorageRunner.throwIfContainsReference(getHistoryResponse); + } else { + getHistoryResponse = externalStorage.retrieve(getHistoryResponse); + } workflowTask .setHistory(getHistoryResponse.getHistory()) .setNextPageToken(getHistoryResponse.getNextPageToken()); diff --git a/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java b/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java index ed6446678a..cca0db9b02 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java @@ -7,32 +7,46 @@ import static org.junit.Assume.assumeFalse; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import com.google.protobuf.ByteString; import com.google.protobuf.util.Durations; import com.uber.m3.tally.NoopScope; +import io.temporal.api.common.v1.Payload; +import io.temporal.api.common.v1.Payloads; import io.temporal.api.enums.v1.EventType; import io.temporal.api.history.v1.History; import io.temporal.api.history.v1.HistoryEvent; import io.temporal.api.taskqueue.v1.StickyExecutionAttributes; import io.temporal.api.workflowservice.v1.*; import io.temporal.internal.common.InternalUtils; +import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.statemachines.ExecuteLocalActivityParameters; import io.temporal.internal.worker.SingleWorkerOptions; import io.temporal.internal.worker.WorkflowExecutorCache; import io.temporal.internal.worker.WorkflowRunLockManager; import io.temporal.internal.worker.WorkflowTaskHandler; +import io.temporal.payload.storage.ExternalStorage; +import io.temporal.payload.storage.StorageDriver; +import io.temporal.payload.storage.StorageDriverClaim; +import io.temporal.payload.storage.StorageDriverRetrieveContext; +import io.temporal.payload.storage.StorageDriverStoreContext; import io.temporal.serviceclient.Version; import io.temporal.serviceclient.WorkflowServiceStubs; import io.temporal.testUtils.HistoryUtils; import io.temporal.testing.internal.SDKTestWorkflowRule; import java.time.Duration; +import java.util.ArrayList; +import java.util.Collections; import java.util.HashMap; import java.util.List; +import java.util.Map; import java.util.Optional; +import java.util.concurrent.CompletableFuture; import org.junit.Rule; import org.junit.Test; +import org.mockito.ArgumentCaptor; public class ReplayWorkflowRunTaskHandlerTaskHandlerTests { @@ -121,6 +135,67 @@ public void workflowTaskFailOnIncompleteHistory() throws Throwable { result.getTaskFailed().getFailure().getMessage()); } + @Test + public void resolvesExternalStorageReferencesInFetchedFullHistory() throws Throwable { + ExternalStorageRunner externalStorage = + ExternalStorageRunner.create( + ExternalStorage.newBuilder() + .setDriver(new InMemoryStorageDriver()) + .setPayloadSizeThreshold(0) + .build()); + PollWorkflowTaskQueueResponse fullTask = HistoryUtils.generateWorkflowTaskWithInitialHistory(); + HistoryEvent startedEvent = fullTask.getHistory().getEvents(0); + Payload input = Payload.newBuilder().setData(ByteString.copyFromUtf8("input")).build(); + History.Builder storedHistory = + fullTask.getHistory().toBuilder() + .setEvents( + 0, + startedEvent.toBuilder() + .setWorkflowExecutionStartedEventAttributes( + startedEvent.getWorkflowExecutionStartedEventAttributes().toBuilder() + .setInput(Payloads.newBuilder().addPayloads(input)))); + externalStorage.store(storedHistory, null); + + WorkflowServiceStubs client = mock(WorkflowServiceStubs.class); + when(client.getServerCapabilities()) + .thenReturn(() -> GetSystemInfoResponse.Capabilities.newBuilder().build()); + WorkflowServiceGrpc.WorkflowServiceBlockingStub blockingStub = + mock(WorkflowServiceGrpc.WorkflowServiceBlockingStub.class); + when(client.blockingStub()).thenReturn(blockingStub); + when(blockingStub.withOption(any(), any())).thenReturn(blockingStub); + when(blockingStub.getWorkflowExecutionHistory(any())) + .thenReturn( + GetWorkflowExecutionHistoryResponse.newBuilder().setHistory(storedHistory).build()); + + ReplayWorkflow workflow = mock(ReplayWorkflow.class); + when(workflow.eventLoop()).thenReturn(true); + when(workflow.getOutput()).thenReturn(Optional.empty()); + WorkflowContext workflowContext = mock(WorkflowContext.class); + when(workflowContext.getRunningUpdateHandlers()).thenReturn(new HashMap<>()); + when(workflow.getWorkflowContext()).thenReturn(workflowContext); + ReplayWorkflowFactory workflowFactory = mock(ReplayWorkflowFactory.class); + when(workflowFactory.getWorkflow(any(), any())).thenReturn(workflow); + WorkflowTaskHandler taskHandler = + new ReplayWorkflowTaskHandler( + "namespace", + workflowFactory, + new WorkflowExecutorCache(10, new WorkflowRunLockManager(), new NoopScope()), + SingleWorkerOptions.newBuilder().setExternalStorage(externalStorage).build(), + null, + Duration.ofSeconds(5), + client, + null); + + taskHandler.handleWorkflowTask( + fullTask.toBuilder().setHistory(History.getDefaultInstance()).build()); + + ArgumentCaptor event = ArgumentCaptor.forClass(HistoryEvent.class); + verify(workflow).start(event.capture(), any()); + assertEquals( + input, + event.getValue().getWorkflowExecutionStartedEventAttributes().getInput().getPayloads(0)); + } + @Test public void localActivityMeteringHelper() { ReplayWorkflowRunTaskHandler.LocalActivityMeteringHelper laMeteringHelper = @@ -231,4 +306,41 @@ private ReplayWorkflowFactory setUpMockWorkflowFactory() throws Throwable { when(mockWorkflow.getWorkflowContext()).thenReturn(mockWorkflowContext); return mockFactory; } + + private static final class InMemoryStorageDriver implements StorageDriver { + private final Map payloads = new HashMap<>(); + private int nextKey; + + @Override + public String getName() { + return "test"; + } + + @Override + public String getType() { + return "test.in-memory"; + } + + @Override + public synchronized CompletableFuture> store( + StorageDriverStoreContext context, List payloads) { + List claims = new ArrayList<>(); + for (Payload payload : payloads) { + String key = Integer.toString(nextKey++); + this.payloads.put(key, payload); + claims.add(new StorageDriverClaim(Collections.singletonMap("key", key))); + } + return CompletableFuture.completedFuture(claims); + } + + @Override + public synchronized CompletableFuture> retrieve( + StorageDriverRetrieveContext context, List claims) { + List retrieved = new ArrayList<>(); + for (StorageDriverClaim claim : claims) { + retrieved.add(payloads.get(claim.getClaimData().get("key"))); + } + return CompletableFuture.completedFuture(retrieved); + } + } } From b1e4e73346555e9aa25f2a69a8bff6b26c3c03f2 Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Mon, 24 Aug 2026 14:05:52 -0400 Subject: [PATCH 3/5] refactor(extstore): refactor the way we derive storage targets by using command.getAttributesCase() + switch for exhaustiveness checking. --- .../internal/worker/WorkflowWorker.java | 74 +++++++++++----- .../storage/ExternalStorageRunnerTest.java | 23 +++-- .../internal/worker/WorkflowWorkerTest.java | 88 ++++++++++++++----- 3 files changed, 135 insertions(+), 50 deletions(-) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java index c23f4cade5..4d68511bc7 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java @@ -11,9 +11,7 @@ import com.uber.m3.tally.Stopwatch; import com.uber.m3.util.ImmutableMap; import io.grpc.StatusRuntimeException; -import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributesOrBuilder; -import io.temporal.api.command.v1.SignalExternalWorkflowExecutionCommandAttributesOrBuilder; -import io.temporal.api.command.v1.StartChildWorkflowExecutionCommandAttributesOrBuilder; +import io.temporal.api.command.v1.*; import io.temporal.api.common.v1.WorkflowExecution; import io.temporal.api.enums.v1.QueryResultType; import io.temporal.api.enums.v1.TaskQueueKind; @@ -419,27 +417,59 @@ private StorageDriverTargetInfo workflowStorageTarget( namespace, execution.getWorkflowId(), execution.getRunId(), workflowType); } - static StorageDriverTargetInfo refineStorageTarget( + static StorageDriverTargetInfo deriveStorageTarget( String namespace, StorageDriverTargetInfo current, MessageOrBuilder message) { - if (message instanceof ScheduleActivityTaskCommandAttributesOrBuilder) { - ScheduleActivityTaskCommandAttributesOrBuilder attrs = - (ScheduleActivityTaskCommandAttributesOrBuilder) message; - return new StorageDriverActivityInfo( - namespace, attrs.getActivityId(), null, attrs.getActivityType().getName()); + if (!(message instanceof CommandOrBuilder)) { + return current; } - if (message instanceof StartChildWorkflowExecutionCommandAttributesOrBuilder) { - StartChildWorkflowExecutionCommandAttributesOrBuilder attrs = - (StartChildWorkflowExecutionCommandAttributesOrBuilder) message; - return new StorageDriverWorkflowInfo( - namespace, attrs.getWorkflowId(), null, attrs.getWorkflowType().getName()); - } - if (message instanceof SignalExternalWorkflowExecutionCommandAttributesOrBuilder) { - WorkflowExecution execution = - ((SignalExternalWorkflowExecutionCommandAttributesOrBuilder) message).getExecution(); - return new StorageDriverWorkflowInfo( - namespace, execution.getWorkflowId(), execution.getRunId(), null); + CommandOrBuilder command = (CommandOrBuilder) message; + // Keep this exhaustive so new command attributes require an explicit target decision. + switch (command.getAttributesCase()) { + case SCHEDULE_ACTIVITY_TASK_COMMAND_ATTRIBUTES: + ScheduleActivityTaskCommandAttributesOrBuilder activity = + command.getScheduleActivityTaskCommandAttributesOrBuilder(); + return new StorageDriverActivityInfo( + namespace, activity.getActivityId(), null, activity.getActivityType().getName()); + case START_CHILD_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + StartChildWorkflowExecutionCommandAttributesOrBuilder child = + command.getStartChildWorkflowExecutionCommandAttributesOrBuilder(); + return new StorageDriverWorkflowInfo( + namespace, child.getWorkflowId(), null, child.getWorkflowType().getName()); + case SIGNAL_EXTERNAL_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + WorkflowExecution execution = + command.getSignalExternalWorkflowExecutionCommandAttributes().getExecution(); + return new StorageDriverWorkflowInfo( + namespace, execution.getWorkflowId(), execution.getRunId(), null); + case CONTINUE_AS_NEW_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + if (current instanceof StorageDriverWorkflowInfo) { + ContinueAsNewWorkflowExecutionCommandAttributesOrBuilder continueAsNew = + command.getContinueAsNewWorkflowExecutionCommandAttributesOrBuilder(); + StorageDriverWorkflowInfo currentWorkflow = (StorageDriverWorkflowInfo) current; + String workflowType = continueAsNew.getWorkflowType().getName(); + return new StorageDriverWorkflowInfo( + namespace, + currentWorkflow.getId(), + null, + Strings.isNullOrEmpty(workflowType) ? currentWorkflow.getType() : workflowType); + } + return current; + case ATTRIBUTES_NOT_SET: + case START_TIMER_COMMAND_ATTRIBUTES: + case COMPLETE_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + case FAIL_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + case REQUEST_CANCEL_ACTIVITY_TASK_COMMAND_ATTRIBUTES: + case CANCEL_TIMER_COMMAND_ATTRIBUTES: + case CANCEL_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + case REQUEST_CANCEL_EXTERNAL_WORKFLOW_EXECUTION_COMMAND_ATTRIBUTES: + case RECORD_MARKER_COMMAND_ATTRIBUTES: + case UPSERT_WORKFLOW_SEARCH_ATTRIBUTES_COMMAND_ATTRIBUTES: + case PROTOCOL_MESSAGE_COMMAND_ATTRIBUTES: + case MODIFY_WORKFLOW_PROPERTIES_COMMAND_ATTRIBUTES: + case SCHEDULE_NEXUS_OPERATION_COMMAND_ATTRIBUTES: + case REQUEST_CANCEL_NEXUS_OPERATION_COMMAND_ATTRIBUTES: + return current; } - return current; + throw new IllegalStateException("Unhandled command attributes: " + command.getAttributesCase()); } private class TaskHandlerImpl implements PollTaskExecutor.TaskHandler { @@ -746,7 +776,7 @@ private RespondWorkflowTaskCompletedResponse sendTaskCompleted( } MessageVisitor storageTargetVisitor = - (current, message) -> refineStorageTarget(namespace, current, message); + (current, message) -> deriveStorageTarget(namespace, current, message); storeOutboundPayloads(taskCompleted, storageTarget, storageTargetVisitor); RespondWorkflowTaskCompletedRequest request = taskCompleted.build(); return grpcRetryer.retryWithResult( diff --git a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageRunnerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageRunnerTest.java index b5c432ff80..0ab1cd6532 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageRunnerTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageRunnerTest.java @@ -8,6 +8,7 @@ import com.google.protobuf.ByteString; import io.temporal.api.command.v1.Command; +import io.temporal.api.command.v1.CommandOrBuilder; import io.temporal.api.command.v1.CompleteWorkflowExecutionCommandAttributes; import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributes; import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributesOrBuilder; @@ -16,6 +17,7 @@ import io.temporal.api.common.v1.Payload; import io.temporal.api.common.v1.Payloads; import io.temporal.api.common.v1.SearchAttributes; +import io.temporal.api.sdk.v1.UserMetadata; import io.temporal.api.workflowservice.v1.RespondWorkflowTaskCompletedRequest; import io.temporal.common.CancellationToken; import io.temporal.internal.concurrent.structured.CancelSource; @@ -139,7 +141,7 @@ public void throwIfContainsReferenceAllowsInlinePayloads() { } @Test - public void storeAppliesPerCommandTargetFromMessageVisitor() { + public void storeScopesCommandTargetOverAttributesAndMetadata() { TargetCapturingDriver driver = new TargetCapturingDriver("d1"); ExternalStorageRunner storage = transformer(driver, 0); @@ -147,6 +149,8 @@ public void storeAppliesPerCommandTargetFromMessageVisitor() { RespondWorkflowTaskCompletedRequest.newBuilder() .addCommands( Command.newBuilder() + .setUserMetadata( + UserMetadata.newBuilder().setSummary(payload("activity-summary"))) .setScheduleActivityTaskCommandAttributes( ScheduleActivityTaskCommandAttributes.newBuilder() .setActivityId("act-1") @@ -163,11 +167,15 @@ public void storeAppliesPerCommandTargetFromMessageVisitor() { new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "MyWorkflow"); MessageVisitor visitor = (current, message) -> { - if (message instanceof ScheduleActivityTaskCommandAttributesOrBuilder) { - ScheduleActivityTaskCommandAttributesOrBuilder attrs = - (ScheduleActivityTaskCommandAttributesOrBuilder) message; - return new StorageDriverActivityInfo( - "ns", attrs.getActivityId(), null, attrs.getActivityType().getName()); + if (message instanceof CommandOrBuilder) { + CommandOrBuilder command = (CommandOrBuilder) message; + if (command.getAttributesCase() + == Command.AttributesCase.SCHEDULE_ACTIVITY_TASK_COMMAND_ATTRIBUTES) { + ScheduleActivityTaskCommandAttributesOrBuilder attrs = + command.getScheduleActivityTaskCommandAttributesOrBuilder(); + return new StorageDriverActivityInfo( + "ns", attrs.getActivityId(), null, attrs.getActivityType().getName()); + } } return current; }; @@ -177,6 +185,9 @@ public void storeAppliesPerCommandTargetFromMessageVisitor() { assertEquals( new StorageDriverActivityInfo("ns", "act-1", null, "MyActivity"), driver.targetFor("activity-input")); + assertEquals( + new StorageDriverActivityInfo("ns", "act-1", null, "MyActivity"), + driver.targetFor("activity-summary")); assertEquals(workflowTarget, driver.targetFor("wf-result")); } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java index c97d85a188..2554227fcb 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java @@ -12,7 +12,9 @@ import com.uber.m3.tally.RootScopeBuilder; import com.uber.m3.tally.Scope; import com.uber.m3.util.ImmutableMap; +import io.temporal.api.command.v1.Command; import io.temporal.api.command.v1.CompleteWorkflowExecutionCommandAttributes; +import io.temporal.api.command.v1.ContinueAsNewWorkflowExecutionCommandAttributes; import io.temporal.api.command.v1.ScheduleActivityTaskCommandAttributes; import io.temporal.api.command.v1.SignalExternalWorkflowExecutionCommandAttributes; import io.temporal.api.command.v1.StartChildWorkflowExecutionCommandAttributes; @@ -458,56 +460,98 @@ private ReplayWorkflowFactory setUpMockWorkflowFactory() throws Throwable { } @Test - public void refineStorageTargetPointsActivityCommandsAtTheActivity() { + public void deriveStorageTargetPointsActivityCommandsAtTheActivity() { StorageDriverTargetInfo workflowDefault = new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "MyWorkflow"); - ScheduleActivityTaskCommandAttributes command = - ScheduleActivityTaskCommandAttributes.newBuilder() - .setActivityId("act-1") - .setActivityType(ActivityType.newBuilder().setName("MyActivity")) + Command command = + Command.newBuilder() + .setScheduleActivityTaskCommandAttributes( + ScheduleActivityTaskCommandAttributes.newBuilder() + .setActivityId("act-1") + .setActivityType(ActivityType.newBuilder().setName("MyActivity"))) .build(); assertEquals( new StorageDriverActivityInfo("ns", "act-1", null, "MyActivity"), - WorkflowWorker.refineStorageTarget("ns", workflowDefault, command)); + WorkflowWorker.deriveStorageTarget("ns", workflowDefault, command)); } @Test - public void refineStorageTargetPointsChildWorkflowCommandsAtTheChild() { + public void deriveStorageTargetPointsChildWorkflowCommandsAtTheChild() { StorageDriverTargetInfo parent = new StorageDriverWorkflowInfo("ns", "parent", "parent-run", "Parent"); - StartChildWorkflowExecutionCommandAttributes command = - StartChildWorkflowExecutionCommandAttributes.newBuilder() - .setWorkflowId("child-1") - .setWorkflowType(WorkflowType.newBuilder().setName("Child")) + Command command = + Command.newBuilder() + .setStartChildWorkflowExecutionCommandAttributes( + StartChildWorkflowExecutionCommandAttributes.newBuilder() + .setWorkflowId("child-1") + .setWorkflowType(WorkflowType.newBuilder().setName("Child"))) .build(); assertEquals( new StorageDriverWorkflowInfo("ns", "child-1", null, "Child"), - WorkflowWorker.refineStorageTarget("ns", parent, command)); + WorkflowWorker.deriveStorageTarget("ns", parent, command)); } @Test - public void refineStorageTargetPointsSignalCommandsAtTheTargetWorkflow() { + public void deriveStorageTargetPointsSignalCommandsAtTheTargetWorkflow() { StorageDriverTargetInfo self = new StorageDriverWorkflowInfo("ns", "self", "self-run", "Self"); - SignalExternalWorkflowExecutionCommandAttributes command = - SignalExternalWorkflowExecutionCommandAttributes.newBuilder() - .setExecution( - WorkflowExecution.newBuilder().setWorkflowId("other").setRunId("other-run")) + Command command = + Command.newBuilder() + .setSignalExternalWorkflowExecutionCommandAttributes( + SignalExternalWorkflowExecutionCommandAttributes.newBuilder() + .setExecution( + WorkflowExecution.newBuilder() + .setWorkflowId("other") + .setRunId("other-run"))) .build(); assertEquals( new StorageDriverWorkflowInfo("ns", "other", "other-run", null), - WorkflowWorker.refineStorageTarget("ns", self, command)); + WorkflowWorker.deriveStorageTarget("ns", self, command)); } @Test - public void refineStorageTargetKeepsTheCurrentTargetForOtherCommands() { + public void deriveStorageTargetPointsContinueAsNewAtTheNewRun() { + StorageDriverTargetInfo current = + new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "CurrentWorkflow"); + Command command = + Command.newBuilder() + .setContinueAsNewWorkflowExecutionCommandAttributes( + ContinueAsNewWorkflowExecutionCommandAttributes.newBuilder() + .setWorkflowType(WorkflowType.newBuilder().setName("NextWorkflow"))) + .build(); + + assertEquals( + new StorageDriverWorkflowInfo("ns", "wf-1", null, "NextWorkflow"), + WorkflowWorker.deriveStorageTarget("ns", current, command)); + } + + @Test + public void deriveStorageTargetKeepsWorkflowTypeForContinueAsNewWithoutOverride() { + StorageDriverTargetInfo current = + new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "CurrentWorkflow"); + Command command = + Command.newBuilder() + .setContinueAsNewWorkflowExecutionCommandAttributes( + ContinueAsNewWorkflowExecutionCommandAttributes.newBuilder()) + .build(); + + assertEquals( + new StorageDriverWorkflowInfo("ns", "wf-1", null, "CurrentWorkflow"), + WorkflowWorker.deriveStorageTarget("ns", current, command)); + } + + @Test + public void deriveStorageTargetKeepsTheCurrentTargetForOtherCommands() { StorageDriverTargetInfo current = new StorageDriverWorkflowInfo("ns", "wf-1", "run-1", "MyWorkflow"); - CompleteWorkflowExecutionCommandAttributes command = - CompleteWorkflowExecutionCommandAttributes.newBuilder().build(); + Command command = + Command.newBuilder() + .setCompleteWorkflowExecutionCommandAttributes( + CompleteWorkflowExecutionCommandAttributes.newBuilder()) + .build(); - assertSame(current, WorkflowWorker.refineStorageTarget("ns", current, command)); + assertSame(current, WorkflowWorker.deriveStorageTarget("ns", current, command)); } } From 2ad55595a17a4bedc02240c40681efc81ac7706e Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Thu, 27 Aug 2026 14:30:23 -0400 Subject: [PATCH 4/5] Explicitly pass cancellation tokens for external storage methods. --- .../internal/replay/ReplayWorkflowTaskHandler.java | 5 +++-- .../internal/replay/ServiceWorkflowHistoryIterator.java | 3 ++- .../java/io/temporal/internal/worker/WorkflowWorker.java | 7 ++----- .../ReplayWorkflowRunTaskHandlerTaskHandlerTests.java | 3 ++- .../replay/ServiceWorkflowHistoryIteratorTest.java | 5 +++-- 5 files changed, 12 insertions(+), 11 deletions(-) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java b/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java index 351d70d9e2..3a8d66a8e4 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java @@ -20,6 +20,7 @@ import io.temporal.api.taskqueue.v1.StickyExecutionAttributes; import io.temporal.api.taskqueue.v1.TaskQueue; import io.temporal.api.workflowservice.v1.*; +import io.temporal.common.CancellationToken; import io.temporal.common.converter.DataConverter; import io.temporal.internal.common.ProtobufTimeUtils; import io.temporal.internal.common.WorkflowExecutionUtils; @@ -82,7 +83,7 @@ public WorkflowTaskHandler.Result handleWorkflowTask(PollWorkflowTaskQueueRespon if (externalStorage == null) { ExternalStorageRunner.throwIfContainsReference(workflowTask); } else { - workflowTask = externalStorage.retrieve(workflowTask); + workflowTask = externalStorage.retrieve(workflowTask, CancellationToken.none()); } return handleWorkflowTaskWithQuery(workflowTask.toBuilder(), metricsScope); } @@ -407,7 +408,7 @@ private WorkflowRunTaskHandler createStatefulHandler( if (externalStorage == null) { ExternalStorageRunner.throwIfContainsReference(getHistoryResponse); } else { - getHistoryResponse = externalStorage.retrieve(getHistoryResponse); + getHistoryResponse = externalStorage.retrieve(getHistoryResponse, CancellationToken.none()); } workflowTask .setHistory(getHistoryResponse.getHistory()) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java b/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java index 14598af207..fbb33033ba 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java @@ -12,6 +12,7 @@ import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryRequest; import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse; import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponseOrBuilder; +import io.temporal.common.CancellationToken; import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.retryer.GrpcRetryer; import io.temporal.serviceclient.RpcRetryOptions; @@ -81,7 +82,7 @@ public boolean hasNext() { if (externalStorage == null) { ExternalStorageRunner.throwIfContainsReference(history); } else { - history = externalStorage.retrieve(history); + history = externalStorage.retrieve(history, CancellationToken.none()); } current = history.getEventsList().iterator(); nextPageToken = response.getNextPageToken(); diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java index 4d68511bc7..f16fa9644f 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java @@ -391,16 +391,13 @@ public String toString() { private void storeOutboundPayloads( com.google.protobuf.Message.Builder builder, @Nullable StorageDriverTargetInfo target) { - ExternalStorageRunner externalStorage = options.getExternalStorage(); - if (externalStorage != null) { - externalStorage.store(builder, target); - } + storeOutboundPayloads(builder, target, null); } private void storeOutboundPayloads( com.google.protobuf.Message.Builder builder, @Nullable StorageDriverTargetInfo target, - MessageVisitor targetVisitor) { + @Nullable MessageVisitor targetVisitor) { ExternalStorageRunner externalStorage = options.getExternalStorage(); if (externalStorage != null) { externalStorage.store(builder, target, targetVisitor, CancellationToken.none()); diff --git a/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java b/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java index cca0db9b02..afb933c3aa 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandlerTaskHandlerTests.java @@ -20,6 +20,7 @@ import io.temporal.api.history.v1.HistoryEvent; import io.temporal.api.taskqueue.v1.StickyExecutionAttributes; import io.temporal.api.workflowservice.v1.*; +import io.temporal.common.CancellationToken; import io.temporal.internal.common.InternalUtils; import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.internal.statemachines.ExecuteLocalActivityParameters; @@ -154,7 +155,7 @@ public void resolvesExternalStorageReferencesInFetchedFullHistory() throws Throw .setWorkflowExecutionStartedEventAttributes( startedEvent.getWorkflowExecutionStartedEventAttributes().toBuilder() .setInput(Payloads.newBuilder().addPayloads(input)))); - externalStorage.store(storedHistory, null); + externalStorage.store(storedHistory, null, null, CancellationToken.none()); WorkflowServiceStubs client = mock(WorkflowServiceStubs.class); when(client.getServerCapabilities()) diff --git a/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java b/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java index d302facdca..3eb43c3467 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java @@ -8,6 +8,7 @@ import io.temporal.api.history.v1.WorkflowExecutionStartedEventAttributes; import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse; import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse; +import io.temporal.common.CancellationToken; import io.temporal.internal.payload.storage.ExternalStorageNotConfiguredException; import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.payload.storage.ExternalStorage; @@ -107,7 +108,7 @@ public void resolvesExternalStorageReferencesInFetchedPages() { ExternalStorageRunner storage = inMemoryStorage(); History inline = historyWithInput(payload("big-input")); History.Builder builder = inline.toBuilder(); - storage.store(builder, null); + storage.store(builder, null, null, CancellationToken.none()); History stored = builder.build(); Assert.assertNotEquals( "stored history should hold a reference, not the inline payload", inline, stored); @@ -123,7 +124,7 @@ public void resolvesExternalStorageReferencesInFetchedPages() { @Test public void failsLoudWhenAFetchedPageHasAReferenceAndStorageIsNotConfigured() { History.Builder builder = historyWithInput(payload("big-input")).toBuilder(); - inMemoryStorage().store(builder, null); + inMemoryStorage().store(builder, null, null, CancellationToken.none()); History stored = builder.build(); ServiceWorkflowHistoryIterator iterator = fetchingIterator(stored, null); From 6d90f7e97fb978d72829aa121d83499f202d8904 Mon Sep 17 00:00:00 2001 From: Chris Constable Date: Thu, 27 Aug 2026 16:23:43 -0400 Subject: [PATCH 5/5] more cancellation token threading --- .../replay/ReplayWorkflowTaskHandler.java | 36 +++++++++++++++++-- .../ServiceWorkflowHistoryIterator.java | 10 ++++-- .../internal/worker/SyncWorkflowWorker.java | 15 ++++++-- .../internal/worker/WorkflowWorker.java | 8 +++-- .../ServiceWorkflowHistoryIteratorTest.java | 28 ++++++++++++++- .../internal/worker/WorkflowWorkerTest.java | 10 ++++-- 6 files changed, 92 insertions(+), 15 deletions(-) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java b/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java index 3a8d66a8e4..a32e521d99 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowTaskHandler.java @@ -36,6 +36,7 @@ import java.time.Duration; import java.util.List; import java.util.Objects; +import java.util.concurrent.CancellationException; import java.util.concurrent.atomic.AtomicBoolean; import java.util.stream.Collectors; import org.slf4j.Logger; @@ -53,6 +54,7 @@ public final class ReplayWorkflowTaskHandler implements WorkflowTaskHandler { private final WorkflowServiceStubs service; private final TaskQueue stickyTaskQueue; private final LocalActivityDispatcher localActivityDispatcher; + private final CancellationToken storageCancellation; public ReplayWorkflowTaskHandler( String namespace, @@ -63,6 +65,29 @@ public ReplayWorkflowTaskHandler( Duration stickyTaskQueueScheduleToStartTimeout, WorkflowServiceStubs service, LocalActivityDispatcher localActivityDispatcher) { + this( + namespace, + asyncWorkflowFactory, + cache, + options, + stickyTaskQueue, + stickyTaskQueueScheduleToStartTimeout, + service, + localActivityDispatcher, + CancellationToken.none()); + } + + public ReplayWorkflowTaskHandler( + String namespace, + ReplayWorkflowFactory asyncWorkflowFactory, + WorkflowExecutorCache cache, + SingleWorkerOptions options, + TaskQueue stickyTaskQueue, + Duration stickyTaskQueueScheduleToStartTimeout, + WorkflowServiceStubs service, + LocalActivityDispatcher localActivityDispatcher, + CancellationToken storageCancellation) { + this.storageCancellation = storageCancellation; this.namespace = namespace; this.workflowFactory = asyncWorkflowFactory; this.cache = cache; @@ -83,7 +108,7 @@ public WorkflowTaskHandler.Result handleWorkflowTask(PollWorkflowTaskQueueRespon if (externalStorage == null) { ExternalStorageRunner.throwIfContainsReference(workflowTask); } else { - workflowTask = externalStorage.retrieve(workflowTask, CancellationToken.none()); + workflowTask = externalStorage.retrieve(workflowTask, storageCancellation); } return handleWorkflowTaskWithQuery(workflowTask.toBuilder(), metricsScope); } @@ -103,7 +128,12 @@ private Result handleWorkflowTaskWithQuery( ServiceWorkflowHistoryIterator historyIterator = new ServiceWorkflowHistoryIterator( - service, namespace, workflowTask, metricsScope, options.getExternalStorage()); + service, + namespace, + workflowTask, + metricsScope, + options.getExternalStorage(), + storageCancellation); boolean finalCommand; Result result; @@ -408,7 +438,7 @@ private WorkflowRunTaskHandler createStatefulHandler( if (externalStorage == null) { ExternalStorageRunner.throwIfContainsReference(getHistoryResponse); } else { - getHistoryResponse = externalStorage.retrieve(getHistoryResponse, CancellationToken.none()); + getHistoryResponse = externalStorage.retrieve(getHistoryResponse, storageCancellation); } workflowTask .setHistory(getHistoryResponse.getHistory()) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java b/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java index fbb33033ba..6eccdace7a 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java @@ -20,6 +20,7 @@ import java.time.Duration; import java.util.Iterator; import java.util.NoSuchElementException; +import java.util.concurrent.CancellationException; import javax.annotation.Nullable; /** Supports iteration over history while loading new pages through calls to the service. */ @@ -33,6 +34,7 @@ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator { private final PollWorkflowTaskQueueResponseOrBuilder task; private final GrpcRetryer grpcRetryer; private final @Nullable ExternalStorageRunner externalStorage; + private final CancellationToken storageCancellation; private Deadline deadline; private Iterator current; ByteString nextPageToken; @@ -42,7 +44,7 @@ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator { String namespace, PollWorkflowTaskQueueResponseOrBuilder task, Scope metricsScope) { - this(service, namespace, task, metricsScope, null); + this(service, namespace, task, metricsScope, null, CancellationToken.none()); } ServiceWorkflowHistoryIterator( @@ -50,7 +52,9 @@ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator { String namespace, PollWorkflowTaskQueueResponseOrBuilder task, Scope metricsScope, - @Nullable ExternalStorageRunner externalStorage) { + @Nullable ExternalStorageRunner externalStorage, + CancellationToken storageCancellation) { + this.storageCancellation = storageCancellation; this.service = service; this.namespace = namespace; this.task = task; @@ -82,7 +86,7 @@ public boolean hasNext() { if (externalStorage == null) { ExternalStorageRunner.throwIfContainsReference(history); } else { - history = externalStorage.retrieve(history, CancellationToken.none()); + history = externalStorage.retrieve(history, storageCancellation); } current = history.getEventsList().iterator(); nextPageToken = response.getNextPageToken(); diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/SyncWorkflowWorker.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/SyncWorkflowWorker.java index be128a5e62..b58db1c1ec 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/SyncWorkflowWorker.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/SyncWorkflowWorker.java @@ -10,6 +10,7 @@ import io.temporal.internal.activity.ActivityExecutionContextFactory; import io.temporal.internal.activity.ActivityTaskHandlerImpl; import io.temporal.internal.activity.LocalActivityExecutionContextFactoryImpl; +import io.temporal.internal.concurrent.structured.CancelSource; import io.temporal.internal.replay.ReplayWorkflowTaskHandler; import io.temporal.internal.sync.POJOWorkflowImplementationFactory; import io.temporal.internal.sync.WorkflowThreadExecutor; @@ -54,6 +55,8 @@ public class SyncWorkflowWorker implements SuspendableWorker { private final POJOWorkflowImplementationFactory factory; private final DataConverter dataConverter; private final ActivityTaskHandlerImpl laTaskHandler; + private final CancelSource storageCancellation = + new CancelSource<>(() -> new CancellationException("Worker shutdown")); private boolean runningLocalActivityWorker; public SyncWorkflowWorker( @@ -111,7 +114,8 @@ public SyncWorkflowWorker( stickyTaskQueue, singleWorkerOptions.getStickyQueueScheduleToStartTimeout(), client.getWorkflowServiceStubs(), - laWorker.getLocalActivityScheduler()); + laWorker.getLocalActivityScheduler(), + storageCancellation.token()); workflowWorker = new WorkflowWorker( @@ -126,7 +130,8 @@ public SyncWorkflowWorker( eagerActivityDispatcher, maxEagerActivityReservationsPerWorkflowTask, slotSupplier, - namespaceCapabilities); + namespaceCapabilities, + storageCancellation.token()); // Exists to support Worker#replayWorkflowExecution functionality. // This handler has to be non-sticky to avoid evicting actual executions from the cache @@ -139,7 +144,8 @@ public SyncWorkflowWorker( null, Duration.ZERO, client.getWorkflowServiceStubs(), - laWorker.getLocalActivityScheduler()); + laWorker.getLocalActivityScheduler(), + storageCancellation.token()); queryReplayHelper = new QueryReplayHelper(nonStickyReplayTaskHandler); } @@ -175,6 +181,9 @@ public boolean start() { @Override public CompletableFuture shutdown(ShutdownManager shutdownManager, boolean interruptTasks) { + if (interruptTasks) { + storageCancellation.cancel(); + } return workflowWorker .shutdown(shutdownManager, interruptTasks) .thenCompose(ignore -> laWorker.shutdown(shutdownManager, interruptTasks)) diff --git a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java index f16fa9644f..7ff7f4a200 100644 --- a/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java +++ b/temporal-sdk/src/main/java/io/temporal/internal/worker/WorkflowWorker.java @@ -35,6 +35,7 @@ import io.temporal.worker.*; import io.temporal.worker.tuning.*; import java.util.*; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import java.util.concurrent.RejectedExecutionException; import java.util.concurrent.TimeUnit; @@ -67,6 +68,7 @@ final class WorkflowWorker implements SuspendableWorker { private final PollerTracker pollerTracker = new PollerTracker(); private final PollerTracker stickyPollerTracker = new PollerTracker(); private final NamespaceCapabilities namespaceCapabilities; + private final CancellationToken storageCancellation; private PollTaskExecutor pollTaskExecutor; @@ -88,7 +90,8 @@ public WorkflowWorker( @Nonnull EagerActivityDispatcher eagerActivityDispatcher, int maxEagerActivityReservationsPerWorkflowTask, @Nonnull SlotSupplier slotSupplier, - @Nonnull NamespaceCapabilities namespaceCapabilities) { + @Nonnull NamespaceCapabilities namespaceCapabilities, + CancellationToken storageCancellation) { this.service = Objects.requireNonNull(service); this.namespace = Objects.requireNonNull(namespace); this.taskQueue = Objects.requireNonNull(taskQueue); @@ -105,6 +108,7 @@ public WorkflowWorker( this.maxEagerActivityReservationsPerWorkflowTask = maxEagerActivityReservationsPerWorkflowTask; this.slotSupplier = new TrackingSlotSupplier<>(slotSupplier, this.workerMetricsScope); this.namespaceCapabilities = namespaceCapabilities; + this.storageCancellation = storageCancellation; } @Override @@ -400,7 +404,7 @@ private void storeOutboundPayloads( @Nullable MessageVisitor targetVisitor) { ExternalStorageRunner externalStorage = options.getExternalStorage(); if (externalStorage != null) { - externalStorage.store(builder, target, targetVisitor, CancellationToken.none()); + externalStorage.store(builder, target, targetVisitor, storageCancellation); } } diff --git a/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java b/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java index 3eb43c3467..43ebf51275 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/replay/ServiceWorkflowHistoryIteratorTest.java @@ -9,6 +9,7 @@ import io.temporal.api.workflowservice.v1.GetWorkflowExecutionHistoryResponse; import io.temporal.api.workflowservice.v1.PollWorkflowTaskQueueResponse; import io.temporal.common.CancellationToken; +import io.temporal.internal.concurrent.structured.CancelSource; import io.temporal.internal.payload.storage.ExternalStorageNotConfiguredException; import io.temporal.internal.payload.storage.ExternalStorageRunner; import io.temporal.payload.storage.ExternalStorage; @@ -24,6 +25,7 @@ import java.util.List; import java.util.Map; import java.util.NoSuchElementException; +import java.util.concurrent.CancellationException; import java.util.concurrent.CompletableFuture; import java.util.concurrent.atomic.AtomicInteger; import org.junit.Assert; @@ -132,11 +134,35 @@ public void failsLoudWhenAFetchedPageHasAReferenceAndStorageIsNotConfigured() { Assert.assertThrows(ExternalStorageNotConfiguredException.class, iterator::hasNext); } + @Test + public void aCancelledTokenAbortsRetrievalOfAFetchedPage() { + ExternalStorageRunner storage = inMemoryStorage(); + History.Builder builder = historyWithInput(payload("big-input")).toBuilder(); + storage.store(builder, null, null, CancellationToken.none()); + History stored = builder.build(); + + CancelSource source = + new CancelSource<>(() -> new CancellationException("Worker shutdown")); + source.cancel(); + + ServiceWorkflowHistoryIterator iterator = fetchingIterator(stored, storage, source.token()); + + Assert.assertThrows(CancellationException.class, iterator::hasNext); + } + private static ServiceWorkflowHistoryIterator fetchingIterator( History page, ExternalStorageRunner storage) { + return fetchingIterator(page, storage, CancellationToken.none()); + } + + private static ServiceWorkflowHistoryIterator fetchingIterator( + History page, + ExternalStorageRunner storage, + CancellationToken storageCancellation) { PollWorkflowTaskQueueResponse workflowTask = PollWorkflowTaskQueueResponse.newBuilder().setNextPageToken(NEXT_PAGE_TOKEN).build(); - return new ServiceWorkflowHistoryIterator(null, "default", workflowTask, null, storage) { + return new ServiceWorkflowHistoryIterator( + null, "default", workflowTask, null, storage, storageCancellation) { @Override GetWorkflowExecutionHistoryResponse queryWorkflowExecutionHistory() { return GetWorkflowExecutionHistoryResponse.newBuilder().setHistory(page).build(); diff --git a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java index 2554227fcb..687db40e8e 100644 --- a/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java +++ b/temporal-sdk/src/test/java/io/temporal/internal/worker/WorkflowWorkerTest.java @@ -22,6 +22,7 @@ import io.temporal.api.common.v1.WorkflowExecution; import io.temporal.api.common.v1.WorkflowType; import io.temporal.api.workflowservice.v1.*; +import io.temporal.common.CancellationToken; import io.temporal.common.reporter.TestStatsReporter; import io.temporal.internal.common.InternalUtils; import io.temporal.internal.replay.ReplayWorkflow; @@ -96,7 +97,8 @@ public void concurrentPollRequestLockTest() throws Exception { eagerActivityDispatcher, 3, slotSupplier, - new NamespaceCapabilities()); + new NamespaceCapabilities(), + CancellationToken.none()); WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub = mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class); @@ -268,7 +270,8 @@ public void respondWorkflowTaskFailureMetricTest() throws Exception { eagerActivityDispatcher, 3, slotSupplier, - new NamespaceCapabilities()); + new NamespaceCapabilities(), + CancellationToken.none()); WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub = mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class); @@ -413,7 +416,8 @@ public boolean isAnyTypeSupported() { eagerActivityDispatcher, 3, slotSupplier, - new NamespaceCapabilities()); + new NamespaceCapabilities(), + CancellationToken.none()); WorkflowServiceGrpc.WorkflowServiceFutureStub futureStub = mock(WorkflowServiceGrpc.WorkflowServiceFutureStub.class);