Skip external storage downloads the workflow worker never uses

CHANGELOG.md

@@ -69,6 +69,8 @@ to docs, or any other relevant information.
 - `TemporalDevServer`, `TemporalDevServerOptions`, `TestWorkflowEnvironment.startLocal`, and the `useDevServer` methods
   of `TestWorkflowRule` and `TestWorkflowExtension` are no longer marked experimental. `TemporalDevServerOptions` remain
   specific to the Temporal CLI dev server implementation and may no longer be valid if that implementation changes.
+- With external storage, large activity arguments are no longer downloaded twice. Workflow workers now skip
+  downloading the inputs they sent to activities, child workflows, signals and Nexus operations.
 
 ### Fixed
 - `ActivityClient` now sends the `temporal-namespace` gRPC header like `WorkflowClient` and `ScheduleClient` do,

temporal-sdk/src/main/java/io/temporal/internal/payload/storage/ExternalStorageRunner.java

@@ -3,6 +3,8 @@
 import com.google.common.base.Throwables;
 import com.google.protobuf.Message;
 import io.temporal.api.common.v1.Payload;
+import io.temporal.api.enums.v1.EventType;
+import io.temporal.api.history.v1.HistoryEventOrBuilder;
 import io.temporal.api.sdk.v1.ExternalStorageReference;
 import io.temporal.common.CancellationToken;
 import io.temporal.internal.payload.visitor.MessageVisitor;
@@ -11,6 +13,8 @@
 import io.temporal.payload.storage.ExternalStorage;
 import io.temporal.payload.storage.StorageDriver;
 import io.temporal.payload.storage.StorageDriverTargetInfo;
+import java.util.EnumSet;
+import java.util.Set;
 import java.util.concurrent.CancellationException;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CompletionException;
@@ -24,6 +28,26 @@
  * configure external storage.
  */
 public final class ExternalStorageRunner {
+  /**
+   * Events that echo the inputs of a command the workflow issued. Replay matches these events by
+   * identifiers only, so their payloads never need to be downloaded by the workflow worker. If the
+   * workflow worker ever needs to read one of these payloads, remove its event type from this set.
+   */
+  private static final Set<EventType> COMMAND_INPUT_EVENTS =
+      EnumSet.of(
+          EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED,
+          EventType.EVENT_TYPE_START_CHILD_WORKFLOW_EXECUTION_INITIATED,
+          EventType.EVENT_TYPE_SIGNAL_EXTERNAL_WORKFLOW_EXECUTION_INITIATED,
+          EventType.EVENT_TYPE_NEXUS_OPERATION_SCHEDULED);
+
+  /** Marks payloads inside {@link #COMMAND_INPUT_EVENTS} so that retrieval leaves them alone. */
+  private static final MessageVisitor<Boolean> SKIP_COMMAND_INPUTS =
+      (skip, message) ->
+          Boolean.TRUE.equals(skip)
+              || (message instanceof HistoryEventOrBuilder
+                  && COMMAND_INPUT_EVENTS.contains(
+                      ((HistoryEventOrBuilder) message).getEventType()));
+
   private final ExternalStoragePayloadTransformer payloadTransformer;
   private final int payloadVisitConcurrency;
 
@@ -92,7 +116,23 @@ public <T extends Message> CompletableFuture<T> retrieveAsync(
       T message,
       CancellationToken<CancellationException> cancellationToken,
       @Nullable StorageOperationMetrics metrics) {
-    return PayloadVisitors.visit(message, retrieveOptions(cancellationToken, metrics));
+    return PayloadVisitors.visit(message, retrieveOptions(null, cancellationToken, metrics));
+  }
+
+  /**
+   * Like {@link #retrieve}, for workflow history processed by the workflow worker. References in
+   * events that echo the workflow's own command inputs are left in place, because replay does not
+   * read them. The workflow's data converter does not download references, so code that reads these
+   * payloads would see the reference instead of the data.
+   */
+  public <T extends Message> T retrieveForReplay(
+      T message,
+      CancellationToken<CancellationException> cancellationToken,
+      @Nullable StorageOperationMetrics metrics) {
+    return getOrThrowIfCancelled(
+        PayloadVisitors.visit(
+            message, retrieveOptions(SKIP_COMMAND_INPUTS, cancellationToken, metrics)),
+        cancellationToken);
   }
 
   /**
@@ -157,12 +197,16 @@ private PayloadVisitorOptions<StorageDriverTargetInfo> storeOptions(
         .build();
   }
 
-  private PayloadVisitorOptions<Void> retrieveOptions(
+  private PayloadVisitorOptions<Boolean> retrieveOptions(
+      @Nullable MessageVisitor<Boolean> skipVisitor,
       CancellationToken<CancellationException> cancellationToken,
       @Nullable StorageOperationMetrics metrics) {
-    return PayloadVisitorOptions.<Void>newBuilder(
-            (context, payloads) ->
-                payloadTransformer.retrieve(payloads, cancellationToken, metrics))
+    return PayloadVisitorOptions.<Boolean>newBuilder(
+            (skip, payloads) ->
+                Boolean.TRUE.equals(skip)
+                    ? CompletableFuture.completedFuture(payloads)
+                    : payloadTransformer.retrieve(payloads, cancellationToken, metrics))
+        .setMessageVisitor(skipVisitor)
         .setConcurrency(payloadVisitConcurrency)
         .setSkipSearchAttributes(true)
         .build();

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

@@ -196,7 +196,7 @@ private PollWorkflowTaskQueueResponse.Builder retrieveStoredPayloads(
       return workflowTask;
     }
     return externalStorageRunner
-        .retrieve(workflowTask.build(), options.getStorageCancellation(), downloadMetrics)
+        .retrieveForReplay(workflowTask.build(), options.getStorageCancellation(), downloadMetrics)
         .toBuilder();
   }
 
@@ -435,7 +435,7 @@ private WorkflowRunTaskHandler createStatefulHandler(
         ExternalStorageRunner.throwIfContainsReference(getHistoryResponse);
       } else {
         getHistoryResponse =
-            externalStorageRunner.retrieve(
+            externalStorageRunner.retrieveForReplay(
                 getHistoryResponse, options.getStorageCancellation(), downloadMetrics);
       }
       workflowTask

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

@@ -90,7 +90,8 @@ public boolean hasNext() {
       if (externalStorageRunner == null) {
         ExternalStorageRunner.throwIfContainsReference(history);
       } else {
-        history = externalStorageRunner.retrieve(history, storageCancellation, downloadMetrics);
+        history =
+            externalStorageRunner.retrieveForReplay(history, storageCancellation, downloadMetrics);
       }
       current = history.getEventsList().iterator();
       nextPageToken = response.getNextPageToken();

temporal-sdk/src/test/java/io/temporal/internal/payload/storage/ExternalStorageRunnerTest.java

@@ -1,6 +1,7 @@
 package io.temporal.internal.payload.storage;
 
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertFalse;
 import static org.junit.Assert.assertNotNull;
 import static org.junit.Assert.assertNull;
 import static org.junit.Assert.assertThrows;
@@ -18,6 +19,15 @@
 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.enums.v1.EventType;
+import io.temporal.api.history.v1.ActivityTaskCompletedEventAttributes;
+import io.temporal.api.history.v1.ActivityTaskScheduledEventAttributes;
+import io.temporal.api.history.v1.History;
+import io.temporal.api.history.v1.HistoryEvent;
+import io.temporal.api.history.v1.NexusOperationScheduledEventAttributes;
+import io.temporal.api.history.v1.SignalExternalWorkflowExecutionInitiatedEventAttributes;
+import io.temporal.api.history.v1.StartChildWorkflowExecutionInitiatedEventAttributes;
+import io.temporal.api.history.v1.WorkflowExecutionStartedEventAttributes;
 import io.temporal.api.sdk.v1.UserMetadata;
 import io.temporal.api.workflowservice.v1.RespondWorkflowTaskCompletedRequest;
 import io.temporal.api.workflowservice.v1.SignalWorkflowExecutionRequest;
@@ -117,6 +127,72 @@ public void walksNestedPayloads() throws Exception {
     assertEquals(command, transformer.retrieve(stored, CancellationToken.none()));
   }
 
+  @Test
+  public void retrieveForReplayLeavesCommandInputEventsAsReferences() throws Exception {
+    TestStorageDriver driver = TestStorageDriver.named("d1");
+    ExternalStorageRunner transformer = transformer(driver, 0);
+    Payloads input = Payloads.newBuilder().addPayloads(payload("x")).build();
+    History history =
+        History.newBuilder()
+            .addEvents(
+                HistoryEvent.newBuilder()
+                    .setEventType(EventType.EVENT_TYPE_WORKFLOW_EXECUTION_STARTED)
+                    .setWorkflowExecutionStartedEventAttributes(
+                        WorkflowExecutionStartedEventAttributes.newBuilder().setInput(input)))
+            .addEvents(
+                HistoryEvent.newBuilder()
+                    .setEventType(EventType.EVENT_TYPE_ACTIVITY_TASK_SCHEDULED)
+                    .setActivityTaskScheduledEventAttributes(
+                        ActivityTaskScheduledEventAttributes.newBuilder().setInput(input)))
+            .addEvents(
+                HistoryEvent.newBuilder()
+                    .setEventType(EventType.EVENT_TYPE_START_CHILD_WORKFLOW_EXECUTION_INITIATED)
+                    .setStartChildWorkflowExecutionInitiatedEventAttributes(
+                        StartChildWorkflowExecutionInitiatedEventAttributes.newBuilder()
+                            .setInput(input)))
+            .addEvents(
+                HistoryEvent.newBuilder()
+                    .setEventType(EventType.EVENT_TYPE_SIGNAL_EXTERNAL_WORKFLOW_EXECUTION_INITIATED)
+                    .setSignalExternalWorkflowExecutionInitiatedEventAttributes(
+                        SignalExternalWorkflowExecutionInitiatedEventAttributes.newBuilder()
+                            .setInput(input)))
+            .addEvents(
+                HistoryEvent.newBuilder()
+                    .setEventType(EventType.EVENT_TYPE_NEXUS_OPERATION_SCHEDULED)
+                    .setNexusOperationScheduledEventAttributes(
+                        NexusOperationScheduledEventAttributes.newBuilder().setInput(payload("x"))))
+            .addEvents(
+                HistoryEvent.newBuilder()
+                    .setEventType(EventType.EVENT_TYPE_ACTIVITY_TASK_COMPLETED)
+                    .setActivityTaskCompletedEventAttributes(
+                        ActivityTaskCompletedEventAttributes.newBuilder().setResult(input)))
+            .build();
+    History.Builder builder = history.toBuilder();
+    transformer.store(builder, null, null, CancellationToken.none());
+    History stored = builder.build();
+
+    History replayed = transformer.retrieveForReplay(stored, CancellationToken.none(), null);
+
+    assertEquals(
+        "the started event's input is needed by replay",
+        history.getEvents(0),
+        replayed.getEvents(0));
+    for (int i = 1; i <= 4; i++) {
+      assertEquals(
+          "a command's echoed input must be left as stored",
+          stored.getEvents(i),
+          replayed.getEvents(i));
+      assertTrue(containsReference(replayed.getEvents(i)));
+    }
+    assertEquals(
+        "the activity result is needed by replay", history.getEvents(5), replayed.getEvents(5));
+    assertFalse(containsReference(replayed.getEvents(5)));
+    assertEquals(
+        "plain retrieval still resolves everything",
+        history,
+        transformer.retrieve(stored, CancellationToken.none()));
+  }
+
   @Test
   public void payloadBelowThresholdLeavesMessageUnchanged() throws Exception {
     TestStorageDriver driver = TestStorageDriver.named("d1");
@@ -284,6 +360,15 @@ private static ExternalStorageRunner transformer(StorageDriver driver, int thres
     return new ExternalStorageRunner(payloadTransformer, 4);
   }
 
+  private static boolean containsReference(HistoryEvent event) {
+    try {
+      ExternalStorageRunner.throwIfContainsReference(event);
+      return false;
+    } catch (ExternalStorageNotConfiguredException e) {
+      return true;
+    }
+  }
+
   private static Payload payload(String data) {
     return Payload.newBuilder().setData(ByteString.copyFromUtf8(data)).build();
   }