External Storage Integration: Activity worker, client

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

@@ -110,8 +110,9 @@ public static WorkflowClient newInstance(
             .getMetricsScope()
             .tagged(MetricsTag.defaultTags(options.getNamespace()));
     ExternalStorage externalStorage = options.getExternalStorage();
-    this.externalStorageRunner =
+    ExternalStorageRunner externalStorageRunner =
         externalStorage == null ? null : ExternalStorageRunner.create(externalStorage);
+    this.externalStorageRunner = externalStorageRunner;
     this.genericClient = new GenericWorkflowClientImpl(workflowServiceStubs, metricsScope);
     this.interceptors = options.getInterceptors();
     this.workflowClientCallsInvoker = initializeClientInvoker();
@@ -120,7 +121,8 @@ public static WorkflowClient newInstance(
             workflowServiceStubs,
             options.getNamespace(),
             options.getIdentity(),
-            options.getDataConverter());
+            options.getDataConverter(),
+            externalStorageRunner);
 
     java.time.Duration heartbeatInterval = options.getWorkerHeartbeatInterval();
     if (!heartbeatInterval.isNegative()) {
@@ -133,7 +135,8 @@ public static WorkflowClient newInstance(
 
   private WorkflowClientCallsInterceptor initializeClientInvoker() {
     WorkflowClientCallsInterceptor workflowClientInvoker =
-        new RootWorkflowClientInvoker(genericClient, options, workerFactoryRegistry);
+        new RootWorkflowClientInvoker(
+            genericClient, options, workerFactoryRegistry, externalStorageRunner);
     for (WorkflowClientInterceptor clientInterceptor : interceptors) {
       workflowClientInvoker =
           clientInterceptor.workflowClientCallsInterceptor(workflowClientInvoker);

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

@@ -1,5 +1,6 @@
 package io.temporal.client;
 
+import io.temporal.api.common.v1.Payload;
 import io.temporal.api.workflowservice.v1.DescribeWorkflowExecutionResponse;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.payload.context.WorkflowSerializationContext;
@@ -29,15 +30,13 @@ public String getStaticSummary() {
     if (!response.getExecutionConfig().getUserMetadata().hasSummary()) {
       return null;
     }
+    Payload summary = response.getExecutionConfig().getUserMetadata().getSummary();
     return dataConverter
         .withContext(
             new WorkflowSerializationContext(
                 response.getWorkflowExecutionInfo().getParentNamespaceId(),
                 response.getWorkflowExecutionInfo().getExecution().getWorkflowId()))
-        .fromPayload(
-            response.getExecutionConfig().getUserMetadata().getSummary(),
-            String.class,
-            String.class);
+        .fromPayload(summary, String.class, String.class);
   }
 
   /**
@@ -51,15 +50,13 @@ public String getStaticDetails() {
     if (!response.getExecutionConfig().getUserMetadata().hasDetails()) {
       return null;
     }
+    Payload details = response.getExecutionConfig().getUserMetadata().getDetails();
     return dataConverter
         .withContext(
             new WorkflowSerializationContext(
                 response.getWorkflowExecutionInfo().getParentNamespaceId(),
                 response.getWorkflowExecutionInfo().getExecution().getWorkflowId()))
-        .fromPayload(
-            response.getExecutionConfig().getUserMetadata().getDetails(),
-            String.class,
-            String.class);
+        .fromPayload(details, String.class, String.class);
   }
 
   /** Returns the raw response from the Temporal service. */

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.ExternalStorageRunner;
 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 ExternalStorageRunner 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 ExternalStorageRunner 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,6 +10,7 @@
 import io.temporal.common.CancellationToken;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.internal.client.external.ManualActivityCompletionClientFactory;
+import io.temporal.internal.payload.storage.ExternalStorageRunner;
 import io.temporal.payload.context.ActivitySerializationContext;
 import io.temporal.workflow.Functions;
 import java.lang.reflect.Type;
@@ -18,6 +19,7 @@
 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 +57,8 @@ class ActivityExecutionContextImpl implements InternalActivityExecutionContext {
       String identity,
       Duration maxHeartbeatThrottleInterval,
       Duration defaultHeartbeatThrottleInterval,
-      Functions.Proc closeCallback) {
+      Functions.Proc closeCallback,
+      @Nullable ExternalStorageRunner externalStorage) {
     this.client = client;
     this.activity = activity;
     this.metricsScope = metricsScope;
@@ -73,7 +76,8 @@ class ActivityExecutionContextImpl implements InternalActivityExecutionContext {
             metricsScope,
             identity,
             maxHeartbeatThrottleInterval,
-            defaultHeartbeatThrottleInterval);
+            defaultHeartbeatThrottleInterval,
+            externalStorage);
   }
 
   /**
@@ -155,7 +159,10 @@ public ManualActivityCompletionClient useLocalManualCompletion() {
           new ActivitySerializationContext(info);
       return new CompletionAwareManualCompletionClient(
           manualCompletionClientFactory.getClient(
-              info.getTaskToken(), metricsScope, activitySerializationContext),
+              info.getTaskToken(),
+              metricsScope,
+              activitySerializationContext,
+              HeartbeatContextImpl.storageTargetForActivity(info.getNamespace(), info)),
           completionHandle);
     } finally {
       lock.unlock();

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

@@ -1,29 +1,37 @@
 package io.temporal.internal.activity;
 
+import com.google.common.base.Strings;
+import com.google.protobuf.ByteString;
 import com.uber.m3.tally.Scope;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
 import io.temporal.activity.ActivityExecutionContext;
 import io.temporal.activity.ActivityInfo;
 import io.temporal.api.common.v1.Payloads;
 import io.temporal.api.enums.v1.TimeoutType;
+import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatRequest;
 import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatResponse;
 import io.temporal.client.*;
 import io.temporal.common.CancellationToken;
 import io.temporal.common.converter.DataConverter;
 import io.temporal.failure.TimeoutFailure;
 import io.temporal.internal.client.ActivityClientHelper;
 import io.temporal.internal.concurrent.structured.CancelSource;
+import io.temporal.internal.payload.storage.ActivityStorageTargets;
+import io.temporal.internal.payload.storage.ExternalStorageRunner;
 import io.temporal.payload.context.ActivitySerializationContext;
+import io.temporal.payload.storage.StorageDriverTargetInfo;
 import io.temporal.serviceclient.WorkflowServiceStubs;
 import java.lang.reflect.Type;
 import java.time.Duration;
 import java.util.Optional;
+import java.util.concurrent.CancellationException;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.ScheduledFuture;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import javax.annotation.Nullable;
 import javax.annotation.concurrent.ThreadSafe;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
@@ -58,6 +66,7 @@ static long getLocalHeartbeatTimeoutBufferMillis() {
   private final long heartbeatIntervalMillis;
   private final DataConverter dataConverter;
   private final DataConverter dataConverterWithActivityContext;
+  private final @Nullable ExternalStorageRunner externalStorage;
 
   private final Scope metricsScope;
   private final Optional<Payloads> prevAttemptHeartbeatDetails;
@@ -89,7 +98,8 @@ public HeartbeatContextImpl(
       Scope metricsScope,
       String identity,
       Duration maxHeartbeatThrottleInterval,
-      Duration defaultHeartbeatThrottleInterval) {
+      Duration defaultHeartbeatThrottleInterval,
+      @Nullable ExternalStorageRunner externalStorage) {
     this(
         service,
         namespace,
@@ -100,6 +110,7 @@ public HeartbeatContextImpl(
         identity,
         maxHeartbeatThrottleInterval,
         defaultHeartbeatThrottleInterval,
+        externalStorage,
         getLocalHeartbeatTimeoutBufferMillis());
   }
 
@@ -113,10 +124,12 @@ public HeartbeatContextImpl(
       String identity,
       Duration maxHeartbeatThrottleInterval,
       Duration defaultHeartbeatThrottleInterval,
+      @Nullable ExternalStorageRunner externalStorage,
       long localHeartbeatTimeoutBufferMillis) {
     this.service = service;
     this.metricsScope = metricsScope;
     this.dataConverter = dataConverter;
+    this.externalStorage = externalStorage;
     this.dataConverterWithActivityContext =
         dataConverter.withContext(
             new ActivitySerializationContext(
@@ -330,16 +343,59 @@ private void checkHeartbeatTimeoutDeadlineLocked() {
     }
   }
 
+  private StorageDriverTargetInfo activityStorageTarget() {
+    return storageTargetForActivity(namespace, info);
+  }
+
+  /**
+   * Standalone activities target the activity; workflow activities target their workflow, matching
+   * where {@link io.temporal.internal.worker.ActivityWorker} stores the activity task payloads. A
+   * non-empty {@code activityRunId} marks a standalone activity.
+   */
+  static StorageDriverTargetInfo storageTargetForActivity(String namespace, ActivityInfo info) {
+    return ActivityStorageTargets.newBuilder(namespace)
+        .setActivity(info.getActivityId(), info.getActivityRunId(), info.getActivityType())
+        .setWorkflow(
+            Strings.emptyToNull(info.getWorkflowId()),
+            Strings.emptyToNull(info.getWorkflowRunId()),
+            info.getWorkflowType())
+        .build();
+  }
+
+  /**
+   * Offloads large heartbeat payloads aborting if the store call runs longer than the heartbeat
+   * interval or if the activity is cancelled.
+   */
+  private void offloadHeartbeat(RecordActivityTaskHeartbeatRequest.Builder builder) {
+    CancelSource<CancellationException> offloadCancel =
+        new CancelSource<>(CancellationException::new);
+    ScheduledFuture<?> timeout =
+        heartbeatExecutor.schedule(
+            (Runnable) offloadCancel::cancel, heartbeatIntervalMillis, TimeUnit.MILLISECONDS);
+    CancellationToken.Registration onActivityCancel =
+        cancellationSource.token().onCancel(offloadCancel::cancel);
+    try {
+      externalStorage.store(builder, activityStorageTarget(), null, offloadCancel.token());
+    } finally {
+      timeout.cancel(false);
+      onActivityCancel.close();
+    }
+  }
+
   private void sendHeartbeatRequest(Object details) {
     try {
+      RecordActivityTaskHeartbeatRequest.Builder builder =
+          RecordActivityTaskHeartbeatRequest.newBuilder()
+              .setTaskToken(ByteString.copyFrom(info.getTaskToken()))
+              .setNamespace(namespace)
+              .setIdentity(identity);
+      dataConverterWithActivityContext.toPayloads(details).ifPresent(builder::setDetails);
+      if (externalStorage != null) {
+        offloadHeartbeat(builder);
+      }
+      RecordActivityTaskHeartbeatRequest request = builder.build();
       RecordActivityTaskHeartbeatResponse status =
-          ActivityClientHelper.sendHeartbeatRequest(
-              service,
-              namespace,
-              identity,
-              info.getTaskToken(),
-              dataConverterWithActivityContext.toPayloads(details),
-              metricsScope);
+          ActivityClientHelper.sendHeartbeatRequest(service, request, metricsScope);
       if (status.getCancelRequested()) {
         requestCancelLocked();
       } else if (status.getActivityReset()) {

temporal-sdk/src/main/java/io/temporal/internal/client/ActivityClientHelper.java

@@ -2,19 +2,13 @@
 
 import static io.temporal.serviceclient.MetricsTag.METRICS_TAGS_CALL_OPTIONS_KEY;
 
-import com.google.common.base.Preconditions;
-import com.google.protobuf.ByteString;
 import com.uber.m3.tally.Scope;
 import io.temporal.activity.ManualActivityCompletionClient;
-import io.temporal.api.common.v1.Payloads;
-import io.temporal.api.common.v1.WorkflowExecution;
 import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatByIdRequest;
 import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatByIdResponse;
 import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatRequest;
 import io.temporal.api.workflowservice.v1.RecordActivityTaskHeartbeatResponse;
 import io.temporal.serviceclient.WorkflowServiceStubs;
-import java.util.Optional;
-import javax.annotation.Nonnull;
 
 /**
  * Contains methods that could but didn't become a part of the main {@link
@@ -26,43 +20,21 @@ private ActivityClientHelper() {}
 
   public static RecordActivityTaskHeartbeatResponse sendHeartbeatRequest(
       WorkflowServiceStubs service,
-      String namespace,
-      String identity,
-      byte[] taskToken,
-      Optional<Payloads> payloads,
+      RecordActivityTaskHeartbeatRequest request,
       Scope metricsScope) {
-    RecordActivityTaskHeartbeatRequest.Builder request =
-        RecordActivityTaskHeartbeatRequest.newBuilder()
-            .setTaskToken(ByteString.copyFrom(taskToken))
-            .setNamespace(namespace)
-            .setIdentity(identity);
-    payloads.ifPresent(request::setDetails);
     return service
         .blockingStub()
         .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope)
-        .recordActivityTaskHeartbeat(request.build());
+        .recordActivityTaskHeartbeat(request);
   }
 
   public static RecordActivityTaskHeartbeatByIdResponse recordActivityTaskHeartbeatById(
       WorkflowServiceStubs service,
-      String namespace,
-      String identity,
-      WorkflowExecution execution,
-      @Nonnull String activityId,
-      Optional<Payloads> payloads,
+      RecordActivityTaskHeartbeatByIdRequest request,
       Scope metricsScope) {
-    Preconditions.checkNotNull(activityId, "Either activity id or task token are required");
-    RecordActivityTaskHeartbeatByIdRequest.Builder request =
-        RecordActivityTaskHeartbeatByIdRequest.newBuilder()
-            .setRunId(execution.getRunId())
-            .setWorkflowId(execution.getWorkflowId())
-            .setActivityId(activityId)
-            .setNamespace(namespace)
-            .setIdentity(identity);
-    payloads.ifPresent(request::setDetails);
     return service
         .blockingStub()
         .withOption(METRICS_TAGS_CALL_OPTIONS_KEY, metricsScope)
-        .recordActivityTaskHeartbeatById(request.build());
+        .recordActivityTaskHeartbeatById(request);
   }
 }