feat(extstore): integrate into activity worker, heartbeats, and client pipelines.

Author: cconstable

temporal-sdk/src/main/java/io/temporal/client/WorkflowClientInternalImpl.java

@@ -18,12 +18,15 @@
 import io.temporal.internal.WorkflowThreadMarker;
 import io.temporal.internal.client.*;
 import io.temporal.internal.client.NexusStartWorkflowResponse;
+import io.temporal.internal.client.external.ExternalStorageGenericWorkflowClient;
 import io.temporal.internal.client.external.GenericWorkflowClient;
 import io.temporal.internal.client.external.GenericWorkflowClientImpl;
 import io.temporal.internal.client.external.ManualActivityCompletionClientFactory;
 import io.temporal.internal.common.PluginUtils;
+import io.temporal.internal.payload.storage.ExternalStorage;
 import io.temporal.internal.sync.StubMarker;
 import io.temporal.internal.worker.HeartbeatManager;
+import io.temporal.payload.storage.ExternalStorageOptions;
 import io.temporal.serviceclient.MetricsTag;
 import io.temporal.serviceclient.WorkflowServiceStubs;
 import io.temporal.serviceclient.WorkflowServiceStubsPlugin;
@@ -56,6 +59,7 @@ final class WorkflowClientInternalImpl implements WorkflowClient, WorkflowClient
   private final WorkerFactoryRegistry workerFactoryRegistry = new WorkerFactoryRegistry();
   private final String workerGroupingKey = java.util.UUID.randomUUID().toString();
   private final @Nullable HeartbeatManager heartbeatManager;
+  private final @Nullable ExternalStorage externalStorage;
 
   /**
    * Creates client that connects to an instance of the Temporal Service. Cannot be used from within
@@ -106,15 +110,27 @@ public static WorkflowClient newInstance(
             .getOptions()
             .getMetricsScope()
             .tagged(MetricsTag.defaultTags(options.getNamespace()));
-    this.genericClient = new GenericWorkflowClientImpl(workflowServiceStubs, metricsScope);
+    ExternalStorageOptions externalStorageOptions = options.getExternalStorage();
+    ExternalStorage externalStorage =
+        externalStorageOptions == null ? null : ExternalStorage.create(externalStorageOptions);
+    this.externalStorage = externalStorage;
+    GenericWorkflowClient genericClient =
+        new GenericWorkflowClientImpl(workflowServiceStubs, metricsScope);
+    if (externalStorage != null) {
+      genericClient =
+          new ExternalStorageGenericWorkflowClient(
+              genericClient, externalStorage, options.getNamespace());
+    }
+    this.genericClient = genericClient;
     this.interceptors = options.getInterceptors();
     this.workflowClientCallsInvoker = initializeClientInvoker();
     this.manualActivityCompletionClientFactory =
         ManualActivityCompletionClientFactory.newFactory(
             workflowServiceStubs,
             options.getNamespace(),
             options.getIdentity(),
-            options.getDataConverter());
+            options.getDataConverter(),
+            externalStorage);
 
     java.time.Duration heartbeatInterval = options.getWorkerHeartbeatInterval();
     if (!heartbeatInterval.isNegative()) {
@@ -127,7 +143,8 @@ public static WorkflowClient newInstance(
 
   private WorkflowClientCallsInterceptor initializeClientInvoker() {
     WorkflowClientCallsInterceptor workflowClientInvoker =
-        new RootWorkflowClientInvoker(genericClient, options, workerFactoryRegistry);
+        new RootWorkflowClientInvoker(
+            genericClient, options, workerFactoryRegistry, externalStorage);
     for (WorkflowClientInterceptor clientInterceptor : interceptors) {
       workflowClientInvoker =
           clientInterceptor.workflowClientCallsInterceptor(workflowClientInvoker);
@@ -815,6 +832,12 @@ public HeartbeatManager getHeartbeatManager() {
     return heartbeatManager;
   }
 
+  @Override
+  @Nullable
+  public ExternalStorage getExternalStorage() {
+    return externalStorage;
+  }
+
   @Override
   public NexusStartWorkflowResponse startNexus(
       NexusStartWorkflowRequest request, Functions.Proc workflow) {

temporal-sdk/src/main/java/io/temporal/client/WorkflowClientOptions.java

@@ -7,12 +7,14 @@
 import io.temporal.common.converter.DataConverter;
 import io.temporal.common.converter.GlobalDataConverter;
 import io.temporal.common.interceptors.WorkflowClientInterceptor;
+import io.temporal.payload.storage.ExternalStorageOptions;
 import java.lang.management.ManagementFactory;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.List;
 import java.util.Objects;
+import javax.annotation.Nullable;
 
 /** Options for WorkflowClient configuration. */
 public final class WorkflowClientOptions {
@@ -52,6 +54,7 @@ public static final class Builder {
     private QueryRejectCondition queryRejectCondition;
     private WorkflowClientPlugin[] plugins;
     private Duration workerHeartbeatInterval;
+    private @Nullable ExternalStorageOptions externalStorage;
 
     private Builder() {}
 
@@ -68,6 +71,7 @@ private Builder(WorkflowClientOptions options) {
       queryRejectCondition = options.queryRejectCondition;
       plugins = options.plugins;
       workerHeartbeatInterval = options.workerHeartbeatInterval;
+      externalStorage = options.externalStorage;
     }
 
     public Builder setNamespace(String namespace) {
@@ -170,6 +174,16 @@ public Builder setWorkerHeartbeatInterval(Duration workerHeartbeatInterval) {
       return this;
     }
 
+    /**
+     * Configures offloading of large payloads to external storage for workflows and activities
+     * created through this client and its workers.
+     */
+    @Experimental
+    public Builder setExternalStorage(@Nullable ExternalStorageOptions externalStorage) {
+      this.externalStorage = externalStorage;
+      return this;
+    }
+
     public WorkflowClientOptions build() {
       return new WorkflowClientOptions(
           namespace,
@@ -180,7 +194,8 @@ public WorkflowClientOptions build() {
           contextPropagators,
           queryRejectCondition,
           plugins == null ? EMPTY_PLUGINS : plugins,
-          resolveHeartbeatInterval(workerHeartbeatInterval));
+          resolveHeartbeatInterval(workerHeartbeatInterval),
+          externalStorage);
     }
 
     /**
@@ -207,7 +222,8 @@ public WorkflowClientOptions validateAndBuildWithDefaults() {
               ? QueryRejectCondition.QUERY_REJECT_CONDITION_UNSPECIFIED
               : queryRejectCondition,
           plugins == null ? EMPTY_PLUGINS : plugins,
-          resolveHeartbeatInterval(workerHeartbeatInterval));
+          resolveHeartbeatInterval(workerHeartbeatInterval),
+          externalStorage);
     }
 
     private static Duration resolveHeartbeatInterval(Duration raw) {
@@ -250,6 +266,8 @@ private static Duration resolveHeartbeatInterval(Duration raw) {
 
   private final Duration workerHeartbeatInterval;
 
+  private final @Nullable ExternalStorageOptions externalStorage;
+
   private WorkflowClientOptions(
       String namespace,
       DataConverter dataConverter,
@@ -259,7 +277,8 @@ private WorkflowClientOptions(
       List<ContextPropagator> contextPropagators,
       QueryRejectCondition queryRejectCondition,
       WorkflowClientPlugin[] plugins,
-      Duration workerHeartbeatInterval) {
+      Duration workerHeartbeatInterval,
+      @Nullable ExternalStorageOptions externalStorage) {
     this.namespace = namespace;
     this.dataConverter = dataConverter;
     this.interceptors = interceptors;
@@ -269,6 +288,7 @@ private WorkflowClientOptions(
     this.queryRejectCondition = queryRejectCondition;
     this.plugins = plugins;
     this.workerHeartbeatInterval = workerHeartbeatInterval;
+    this.externalStorage = externalStorage;
   }
 
   /**
@@ -335,6 +355,12 @@ public Duration getWorkerHeartbeatInterval() {
     return workerHeartbeatInterval;
   }
 
+  /** External storage configuration, or null when external storage is disabled. */
+  @Experimental
+  public @Nullable ExternalStorageOptions getExternalStorage() {
+    return externalStorage;
+  }
+
   @Override
   public String toString() {
     return "WorkflowClientOptions{"
@@ -359,6 +385,8 @@ public String toString() {
         + Arrays.toString(plugins)
         + ", workerHeartbeatInterval="
         + workerHeartbeatInterval
+        + ", externalStorage="
+        + externalStorage
         + '}';
   }
 
@@ -376,7 +404,8 @@ public boolean equals(Object o) {
         && queryRejectCondition == that.queryRejectCondition
         && Arrays.equals(plugins, that.plugins)
         && com.google.common.base.Objects.equal(
-            workerHeartbeatInterval, that.workerHeartbeatInterval);
+            workerHeartbeatInterval, that.workerHeartbeatInterval)
+        && com.google.common.base.Objects.equal(externalStorage, that.externalStorage);
   }
 
   @Override
@@ -390,6 +419,7 @@ public int hashCode() {
         contextPropagators,
         queryRejectCondition,
         Arrays.hashCode(plugins),
-        workerHeartbeatInterval);
+        workerHeartbeatInterval,
+        externalStorage);
   }
 }

temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityExecutionContextFactoryImpl.java

@@ -4,13 +4,15 @@
 import io.temporal.client.WorkflowClient;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.internal.client.external.ManualActivityCompletionClientFactory;
+import io.temporal.internal.payload.storage.ExternalStorage;
 import java.nio.ByteBuffer;
 import java.time.Duration;
 import java.util.Arrays;
 import java.util.Objects;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.ScheduledExecutorService;
+import javax.annotation.Nullable;
 
 public class ActivityExecutionContextFactoryImpl implements ActivityExecutionContextFactory {
   private final WorkflowClient client;
@@ -21,6 +23,7 @@ public class ActivityExecutionContextFactoryImpl implements ActivityExecutionCon
   private final DataConverter dataConverter;
   private final ScheduledExecutorService heartbeatExecutor;
   private final ManualActivityCompletionClientFactory manualCompletionClientFactory;
+  private final @Nullable ExternalStorage externalStorage;
   private final ConcurrentMap<ByteBuffer, ActivityExecutionContextImpl> activeContexts =
       new ConcurrentHashMap<>();
 
@@ -31,7 +34,8 @@ public ActivityExecutionContextFactoryImpl(
       Duration maxHeartbeatThrottleInterval,
       Duration defaultHeartbeatThrottleInterval,
       DataConverter dataConverter,
-      ScheduledExecutorService heartbeatExecutor) {
+      ScheduledExecutorService heartbeatExecutor,
+      @Nullable ExternalStorage externalStorage) {
     this.client = Objects.requireNonNull(client);
     this.identity = identity;
     this.namespace = Objects.requireNonNull(namespace);
@@ -40,9 +44,10 @@ public ActivityExecutionContextFactoryImpl(
         Objects.requireNonNull(defaultHeartbeatThrottleInterval);
     this.dataConverter = Objects.requireNonNull(dataConverter);
     this.heartbeatExecutor = Objects.requireNonNull(heartbeatExecutor);
+    this.externalStorage = externalStorage;
     this.manualCompletionClientFactory =
         ManualActivityCompletionClientFactory.newFactory(
-            client.getWorkflowServiceStubs(), namespace, identity, dataConverter);
+            client.getWorkflowServiceStubs(), namespace, identity, dataConverter, externalStorage);
   }
 
   @Override
@@ -63,7 +68,8 @@ public InternalActivityExecutionContext createContext(
             identity,
             maxHeartbeatThrottleInterval,
             defaultHeartbeatThrottleInterval,
-            () -> cleanupContext(info.getTaskToken(), false));
+            () -> cleanupContext(info.getTaskToken(), false),
+            externalStorage);
     activeContexts.put(taskToken, context);
     return context;
   }

temporal-sdk/src/main/java/io/temporal/internal/activity/ActivityExecutionContextImpl.java

@@ -10,14 +10,17 @@
 import io.temporal.common.CancellationToken;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.internal.client.external.ManualActivityCompletionClientFactory;
+import io.temporal.internal.payload.storage.ExternalStorage;
 import io.temporal.payload.context.ActivitySerializationContext;
+import io.temporal.payload.storage.StorageDriverActivityInfo;
 import io.temporal.workflow.Functions;
 import java.lang.reflect.Type;
 import java.time.Duration;
 import java.util.Optional;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import javax.annotation.Nullable;
 import javax.annotation.concurrent.ThreadSafe;
 
 /**
@@ -55,7 +58,8 @@ class ActivityExecutionContextImpl implements InternalActivityExecutionContext {
       String identity,
       Duration maxHeartbeatThrottleInterval,
       Duration defaultHeartbeatThrottleInterval,
-      Functions.Proc closeCallback) {
+      Functions.Proc closeCallback,
+      @Nullable ExternalStorage externalStorage) {
     this.client = client;
     this.activity = activity;
     this.metricsScope = metricsScope;
@@ -73,7 +77,8 @@ class ActivityExecutionContextImpl implements InternalActivityExecutionContext {
             metricsScope,
             identity,
             maxHeartbeatThrottleInterval,
-            defaultHeartbeatThrottleInterval);
+            defaultHeartbeatThrottleInterval,
+            externalStorage);
   }
 
   /**
@@ -155,7 +160,14 @@ public ManualActivityCompletionClient useLocalManualCompletion() {
           new ActivitySerializationContext(info);
       return new CompletionAwareManualCompletionClient(
           manualCompletionClientFactory.getClient(
-              info.getTaskToken(), metricsScope, activitySerializationContext),
+              info.getTaskToken(),
+              metricsScope,
+              activitySerializationContext,
+              new StorageDriverActivityInfo(
+                  info.getNamespace(),
+                  info.getActivityId(),
+                  info.getActivityRunId(),
+                  info.getActivityType())),
           completionHandle);
     } finally {
       lock.unlock();