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/BasePoller.java

@@ -9,6 +9,7 @@
 import java.time.Duration;
 import java.util.Objects;
 import java.util.concurrent.*;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.atomic.AtomicReference;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -176,6 +177,8 @@ static boolean shouldIgnoreDuringShutdown(Throwable ex) {
     ex instanceof RejectedExecutionException
         // if the worker thread gets InterruptedException - it's normal during shutdown
         || ex instanceof InterruptedException
+        || ex instanceof CancellationException
+        || ex.getCause() instanceof CancellationException
         // if we get wrapped InterruptedException like what PollTask or GRPC clients do with
         // setting Thread.interrupted() on - it's normal during shutdown too. See PollTask
         // javadoc.

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() {
@@ -414,6 +429,10 @@ public ExternalStorageRunner getExternalStorageRunner() {
     return externalStorageRunner;
   }
 
+  public CancellationToken<CancellationException> getStorageCancellation() {
+    return storageCancellation;
+  }
+
   public WorkerVersioningOptions getWorkerVersioningOptions() {
     return new WorkerVersioningOptions(
         this.getBuildId(), this.isUsingBuildIdForVersioning(), this.getDeploymentOptions());