External Storage Integration: WorkflowWorker, replay, history

temporal-sdk/src/main/java/io/temporal/internal/replay/ReplayWorkflowRunTaskHandler.java

@@ -175,7 +175,9 @@ public WorkflowTaskResult handleWorkflowTask(
               .setForceWorkflowTask(
                   localActivityTaskCount > 0 && !context.isWorkflowMethodCompleted())
               .setNonfirstLocalActivityAttempts(localActivityMeteringHelper.getNonfirstAttempts())
-              .setSdkFlags(newSdkFlags);
+              .setSdkFlags(newSdkFlags)
+              .setParentWorkflowExecution(context.getParentWorkflowExecution())
+              .setContinuedAsNew(context.getContinuedExecutionRunId().isPresent());
       if (workflowStateMachines.sdkNameToWrite() != null) {
         result.setWriteSdkName(workflowStateMachines.sdkNameToWrite());
       }

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;
@@ -34,6 +35,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;
@@ -89,12 +91,19 @@ private Result handleWorkflowTaskWithQuery(
     boolean useCache = stickyTaskQueue != null;
 
     try {
+      workflowTask = retrieveStoredPayloads(workflowTask);
       workflowRunTaskHandler =
           getOrCreateWorkflowExecutor(useCache, workflowTask, metricsScope, createdNew);
       logWorkflowTaskToBeProcessed(workflowTask, createdNew);
 
       ServiceWorkflowHistoryIterator historyIterator =
-          new ServiceWorkflowHistoryIterator(service, namespace, workflowTask, metricsScope);
+          new ServiceWorkflowHistoryIterator(
+              service,
+              namespace,
+              workflowTask,
+              metricsScope,
+              options.getExternalStorageRunner(),
+              options.getStorageCancellation());
       boolean finalCommand;
       Result result;
 
@@ -132,7 +141,7 @@ private Result handleWorkflowTaskWithQuery(
       }
 
       return result;
-    } catch (InterruptedException e) {
+    } catch (InterruptedException | CancellationException e) {
       throw e;
     } catch (Throwable e) {
       // Note here that the executor might not be in the cache, even when the caching is on. In that
@@ -170,6 +179,18 @@ private Result handleWorkflowTaskWithQuery(
     }
   }
 
+  private PollWorkflowTaskQueueResponse.Builder retrieveStoredPayloads(
+      PollWorkflowTaskQueueResponse.Builder workflowTask) {
+    ExternalStorageRunner externalStorageRunner = options.getExternalStorageRunner();
+    if (externalStorageRunner == null) {
+      ExternalStorageRunner.throwIfContainsReference(workflowTask.build());
+      return workflowTask;
+    }
+    return externalStorageRunner
+        .retrieve(workflowTask.build(), options.getStorageCancellation())
+        .toBuilder();
+  }
+
   private Result createCompletedWFTRequest(
       String workflowType,
       PollWorkflowTaskQueueResponseOrBuilder workflowTask,
@@ -253,7 +274,8 @@ private Result createCompletedWFTRequest(
         null,
         result.isFinalCommand(),
         eventIdSetHandle,
-        result.getApplyPostCompletionMetrics());
+        result.getApplyPostCompletionMetrics(),
+        result.isContinuedAsNew() ? null : result.getParentWorkflowExecution());
   }
 
   private Result failureToWFTResult(
@@ -395,6 +417,13 @@ private WorkflowRunTaskHandler createStatefulHandler(
               .blockingStub()
               .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope)
               .getWorkflowExecutionHistory(getHistoryRequest);
+      ExternalStorageRunner externalStorageRunner = options.getExternalStorageRunner();
+      if (externalStorageRunner == null) {
+        ExternalStorageRunner.throwIfContainsReference(getHistoryResponse);
+      } else {
+        getHistoryResponse =
+            externalStorageRunner.retrieve(getHistoryResponse, options.getStorageCancellation());
+      }
       workflowTask
           .setHistory(getHistoryResponse.getHistory())
           .setNextPageToken(getHistoryResponse.getNextPageToken());

temporal-sdk/src/main/java/io/temporal/internal/replay/ServiceWorkflowHistoryIterator.java

@@ -12,12 +12,16 @@
 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;
 import io.temporal.serviceclient.WorkflowServiceStubs;
 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. */
 class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator {
@@ -29,6 +33,8 @@ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator {
   private final Scope metricsScope;
   private final PollWorkflowTaskQueueResponseOrBuilder task;
   private final GrpcRetryer grpcRetryer;
+  private final @Nullable ExternalStorageRunner externalStorageRunner;
+  private final CancellationToken<CancellationException> storageCancellation;
   private Deadline deadline;
   private Iterator<HistoryEvent> current;
   ByteString nextPageToken;
@@ -38,10 +44,22 @@ class ServiceWorkflowHistoryIterator implements WorkflowHistoryIterator {
       String namespace,
       PollWorkflowTaskQueueResponseOrBuilder task,
       Scope metricsScope) {
+    this(service, namespace, task, metricsScope, null, CancellationToken.none());
+  }
+
+  ServiceWorkflowHistoryIterator(
+      WorkflowServiceStubs service,
+      String namespace,
+      PollWorkflowTaskQueueResponseOrBuilder task,
+      Scope metricsScope,
+      @Nullable ExternalStorageRunner externalStorageRunner,
+      CancellationToken<CancellationException> storageCancellation) {
+    this.storageCancellation = storageCancellation;
     this.service = service;
     this.namespace = namespace;
     this.task = task;
     this.metricsScope = metricsScope;
+    this.externalStorageRunner = externalStorageRunner;
     // 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 +82,13 @@ public boolean hasNext() {
       // true.
       GetWorkflowExecutionHistoryResponse response = queryWorkflowExecutionHistory();
 
-      current = response.getHistory().getEventsList().iterator();
+      History history = response.getHistory();
+      if (externalStorageRunner == null) {
+        ExternalStorageRunner.throwIfContainsReference(history);
+      } else {
+        history = externalStorageRunner.retrieve(history, storageCancellation);
+      }
+      current = history.getEventsList().iterator();
       nextPageToken = response.getNextPageToken();
       // Server can return an empty page, but a valid nextPageToken that contains
       // more events.

temporal-sdk/src/main/java/io/temporal/internal/replay/WorkflowTaskResult.java

@@ -1,12 +1,14 @@
 package io.temporal.internal.replay;
 
 import io.temporal.api.command.v1.Command;
+import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.api.protocol.v1.Message;
 import io.temporal.api.query.v1.WorkflowQueryResult;
 import io.temporal.common.VersioningBehavior;
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import javax.annotation.Nullable;
 
 public final class WorkflowTaskResult {
 
@@ -26,6 +28,8 @@ public static final class Builder {
     private String writeSdkVersion;
     private VersioningBehavior versioningBehavior;
     private Runnable applyPostCompletionMetrics;
+    private @Nullable WorkflowExecution parentWorkflowExecution;
+    private boolean continuedAsNew;
 
     public Builder setCommands(List<Command> commands) {
       this.commands = commands;
@@ -77,6 +81,16 @@ public Builder setVersioningBehavior(VersioningBehavior versioningBehavior) {
       return this;
     }
 
+    public Builder setParentWorkflowExecution(@Nullable WorkflowExecution parentWorkflowExecution) {
+      this.parentWorkflowExecution = parentWorkflowExecution;
+      return this;
+    }
+
+    public Builder setContinuedAsNew(boolean continuedAsNew) {
+      this.continuedAsNew = continuedAsNew;
+      return this;
+    }
+
     public Builder setApplyPostCompletionMetrics(Runnable applyPostCompletionMetrics) {
       this.applyPostCompletionMetrics = applyPostCompletionMetrics;
       return this;
@@ -94,7 +108,9 @@ public WorkflowTaskResult build() {
           writeSdkName,
           writeSdkVersion,
           versioningBehavior == null ? VersioningBehavior.UNSPECIFIED : versioningBehavior,
-          applyPostCompletionMetrics);
+          applyPostCompletionMetrics,
+          parentWorkflowExecution,
+          continuedAsNew);
     }
   }
 
@@ -109,6 +125,8 @@ public WorkflowTaskResult build() {
   private final String writeSdkVersion;
   private final VersioningBehavior versioningBehavior;
   private final Runnable applyPostCompletionMetrics;
+  private final @Nullable WorkflowExecution parentWorkflowExecution;
+  private final boolean continuedAsNew;
 
   private WorkflowTaskResult(
       List<Command> commands,
@@ -121,7 +139,9 @@ private WorkflowTaskResult(
       String writeSdkName,
       String writeSdkVersion,
       VersioningBehavior versioningBehavior,
-      Runnable applyPostCompletionMetrics) {
+      Runnable applyPostCompletionMetrics,
+      @Nullable WorkflowExecution parentWorkflowExecution,
+      boolean continuedAsNew) {
     this.commands = commands;
     this.messages = messages;
     this.nonfirstLocalActivityAttempts = nonfirstLocalActivityAttempts;
@@ -136,6 +156,19 @@ private WorkflowTaskResult(
     this.writeSdkVersion = writeSdkVersion;
     this.versioningBehavior = versioningBehavior;
     this.applyPostCompletionMetrics = applyPostCompletionMetrics;
+    this.parentWorkflowExecution = parentWorkflowExecution;
+    this.continuedAsNew = continuedAsNew;
+  }
+
+  /** The workflow that started this one as a child, or {@code null} if it has no parent. */
+  @Nullable
+  public WorkflowExecution getParentWorkflowExecution() {
+    return parentWorkflowExecution;
+  }
+
+  /** Whether this run was created by a continue-as-new rather than started directly. */
+  public boolean isContinuedAsNew() {
+    return continuedAsNew;
   }
 
   public List<Command> getCommands() {

temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java

@@ -3,6 +3,7 @@
 import com.uber.m3.tally.NoopScope;
 import com.uber.m3.tally.Scope;
 import io.temporal.api.common.v1.WorkerVersionStamp;
+import io.temporal.common.CancellationToken;
 import io.temporal.common.context.ContextPropagator;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.common.converter.GlobalDataConverter;
@@ -12,6 +13,7 @@
 import io.temporal.worker.WorkerDeploymentOptions;
 import java.time.Duration;
 import java.util.List;
+import java.util.concurrent.CancellationException;
 import javax.annotation.Nullable;
 
 public final class SingleWorkerOptions {
@@ -48,6 +50,7 @@ public static final class Builder {
     private String workerControlTaskQueue;
     private PreferredVersionProvider preferredVersionProvider;
     private @Nullable ExternalStorageRunner externalStorageRunner;
+    private CancellationToken<CancellationException> storageCancellation = CancellationToken.none();
 
     private Builder() {}
 
@@ -77,6 +80,7 @@ private Builder(SingleWorkerOptions options) {
       this.workerControlTaskQueue = options.getWorkerControlTaskQueue();
       this.preferredVersionProvider = options.getPreferredVersionProvider();
       this.externalStorageRunner = options.getExternalStorageRunner();
+      this.storageCancellation = options.getStorageCancellation();
     }
 
     public Builder setIdentity(String identity) {
@@ -189,6 +193,13 @@ public Builder setPreferredVersionProvider(PreferredVersionProvider preferredVer
       return this;
     }
 
+    /** Cancelled when this worker stops, to abandon its in-flight external storage work. */
+    public Builder setStorageCancellation(
+        CancellationToken<CancellationException> storageCancellation) {
+      this.storageCancellation = storageCancellation;
+      return this;
+    }
+
     public Builder setExternalStorageRunner(@Nullable ExternalStorageRunner externalStorageRunner) {
       this.externalStorageRunner = externalStorageRunner;
       return this;
@@ -237,7 +248,8 @@ public SingleWorkerOptions build() {
           this.allowActivityHeartbeatDuringShutdown,
           this.workerControlTaskQueue,
           this.preferredVersionProvider,
-          this.externalStorageRunner);
+          this.externalStorageRunner,
+          this.storageCancellation);
     }
   }
 
@@ -263,6 +275,7 @@ public SingleWorkerOptions build() {
   private final String workerControlTaskQueue;
   private final PreferredVersionProvider preferredVersionProvider;
   private final @Nullable ExternalStorageRunner externalStorageRunner;
+  private final CancellationToken<CancellationException> storageCancellation;
 
   private SingleWorkerOptions(
       String identity,
@@ -286,7 +299,8 @@ private SingleWorkerOptions(
       boolean allowActivityHeartbeatDuringShutdown,
       String workerControlTaskQueue,
       PreferredVersionProvider preferredVersionProvider,
-      @Nullable ExternalStorageRunner externalStorageRunner) {
+      @Nullable ExternalStorageRunner externalStorageRunner,
+      CancellationToken<CancellationException> storageCancellation) {
     this.identity = identity;
     this.binaryChecksum = binaryChecksum;
     this.buildId = buildId;
@@ -309,6 +323,7 @@ private SingleWorkerOptions(
     this.workerControlTaskQueue = workerControlTaskQueue;
     this.preferredVersionProvider = preferredVersionProvider;
     this.externalStorageRunner = externalStorageRunner;
+    this.storageCancellation = storageCancellation;
   }
 
   public String getIdentity() {
@@ -411,6 +426,10 @@ public ExternalStorageRunner getExternalStorageRunner() {
     return externalStorageRunner;
   }
 
+  public CancellationToken<CancellationException> getStorageCancellation() {
+    return storageCancellation;
+  }
+
   public WorkerVersioningOptions getWorkerVersioningOptions() {
     return new WorkerVersioningOptions(
         this.getBuildId(), this.isUsingBuildIdForVersioning(), this.getDeploymentOptions());

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<CancellationException> storageCancellation =
+      new CancelSource<>(() -> new CancellationException("Worker shutdown"));
   private boolean runningLocalActivityWorker;
 
   public SyncWorkflowWorker(
@@ -71,6 +74,10 @@ public SyncWorkflowWorker(
       @Nonnull SlotSupplier<WorkflowSlotInfo> slotSupplier,
       @Nonnull SlotSupplier<LocalActivitySlotInfo> laSlotSupplier,
       @Nonnull NamespaceCapabilities namespaceCapabilities) {
+    singleWorkerOptions =
+        SingleWorkerOptions.newBuilder(singleWorkerOptions)
+            .setStorageCancellation(storageCancellation.token())
+            .build();
     this.identity = singleWorkerOptions.getIdentity();
     this.namespace = namespace;
     this.taskQueue = taskQueue;
@@ -175,8 +182,12 @@ public boolean start() {
 
   @Override
   public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean interruptTasks) {
-    return workflowWorker
-        .shutdown(shutdownManager, interruptTasks)
+    CompletableFuture<Void> workflowWorkerShutdown =
+        workflowWorker.shutdown(shutdownManager, interruptTasks);
+    if (interruptTasks) {
+      storageCancellation.cancel();
+    }
+    return workflowWorkerShutdown
         .thenCompose(ignore -> laWorker.shutdown(shutdownManager, interruptTasks))
         .exceptionally(
             e -> {