External Storage Integration: Lazy resolving references, general refactoring

temporal-sdk/src/main/java/io/temporal/common/converter/PayloadAndFailureDataConverter.java

@@ -8,6 +8,8 @@
 import io.temporal.api.common.v1.Payloads;
 import io.temporal.api.failure.v1.Failure;
 import io.temporal.failure.DefaultFailureConverter;
+import io.temporal.internal.payload.storage.ExternalStorageNotConfiguredException;
+import io.temporal.internal.payload.storage.ExternalStorageReferences;
 import io.temporal.payload.context.SerializationContext;
 import java.lang.reflect.Type;
 import java.util.*;
@@ -71,6 +73,10 @@ public <T> T fromPayload(Payload payload, Class<T> valueClass, Type valueType)
       return (T) new RawValue(payload);
     }
 
+    if (ExternalStorageReferences.isReference(payload)) {
+      throw new ExternalStorageNotConfiguredException();
+    }
+
     try {
       String encoding =
           payload.getMetadataOrThrow(EncodingKeys.METADATA_ENCODING_KEY).toString(UTF_8);

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

@@ -0,0 +1,162 @@
+package io.temporal.internal.payload.storage;
+
+import com.google.common.base.Throwables;
+import com.google.protobuf.Message;
+import io.temporal.api.common.v1.Payload;
+import io.temporal.api.sdk.v1.ExternalStorageReference;
+import io.temporal.common.CancellationToken;
+import io.temporal.internal.payload.visitor.MessageVisitor;
+import io.temporal.internal.payload.visitor.PayloadVisitorOptions;
+import io.temporal.internal.payload.visitor.PayloadVisitors;
+import io.temporal.payload.storage.ExternalStorageOptions;
+import io.temporal.payload.storage.StorageDriver;
+import io.temporal.payload.storage.StorageDriverTargetInfo;
+import java.util.List;
+import java.util.concurrent.CancellationException;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionException;
+import java.util.concurrent.ExecutionException;
+import javax.annotation.Nullable;
+
+/**
+ * External storage offloads large payloads via {@link StorageDriver}s. It walks messages using
+ * {@link PayloadVisitors} transforming payloads to and from {@link ExternalStorageReference} using
+ * {@link ExternalStoragePayloadTransformer}. Use {@link ExternalStorageOptions} via {@link#create}
+ * to configure external storage.
+ */
+public final class ExternalStorage {
+  private final ExternalStoragePayloadTransformer payloadTransformer;
+  private final int payloadVisitConcurrency;
+
+  public static ExternalStorage create(ExternalStorageOptions options) {
+    return new ExternalStorage(
+        ExternalStoragePayloadTransformer.fromOptions(options),
+        options.getMaxConcurrentPayloadVisits());
+  }
+
+  ExternalStorage(
+      ExternalStoragePayloadTransformer payloadTransformer, int payloadVisitConcurrency) {
+    this.payloadTransformer = payloadTransformer;
+    this.payloadVisitConcurrency = payloadVisitConcurrency;
+  }
+
+  public <T extends Message> T storeBlocking(T message, @Nullable StorageDriverTargetInfo target) {
+    return storeBlocking(message, target, CancellationToken.none());
+  }
+
+  public <T extends Message> T storeBlocking(
+      T message,
+      @Nullable StorageDriverTargetInfo target,
+      CancellationToken<CancellationException> cancellationToken) {
+    return getOrThrowIfCancelled(store(message, target, cancellationToken), cancellationToken);
+  }
+
+  public <T extends Message> T storeBlocking(
+      T message,
+      @Nullable StorageDriverTargetInfo target,
+      @Nullable MessageVisitor<StorageDriverTargetInfo> targetVisitor) {
+    CancellationToken<CancellationException> cancellationToken = CancellationToken.none();
+    return getOrThrowIfCancelled(
+        PayloadVisitors.visit(message, storeOptions(target, targetVisitor, cancellationToken)),
+        cancellationToken);
+  }
+
+  public <T extends Message> T retrieveBlocking(T message) {
+    CancellationToken<CancellationException> cancellationToken = CancellationToken.none();
+    return getOrThrowIfCancelled(retrieve(message, cancellationToken), cancellationToken);
+  }
+
+  public <T extends Message> CompletableFuture<T> retrieveAsync(T message) {
+    return retrieve(message, CancellationToken.none());
+  }
+
+  /**
+   * Throws {@link ExternalStorageNotConfiguredException} if {@code message} contains any reference
+   * payload. Used at inbound task boundaries when external storage is not configured.
+   */
+  public static void throwIfContainsReference(Message message) {
+    PayloadVisitorOptions<Void> options =
+        PayloadVisitorOptions.<Void>newBuilder(
+                (context, payloads) -> {
+                  for (Payload payload : payloads) {
+                    if (ExternalStorageReferences.isReference(payload)) {
+                      CompletableFuture<List<Payload>> found = new CompletableFuture<>();
+                      found.completeExceptionally(new ExternalStorageNotConfiguredException());
+                      return found;
+                    }
+                  }
+                  return CompletableFuture.completedFuture(payloads);
+                })
+            .setSkipSearchAttributes(true)
+            .build();
+    try {
+      PayloadVisitors.visit(message.toBuilder(), options).join();
+    } catch (CompletionException e) {
+      Throwable cause = e.getCause() != null ? e.getCause() : e;
+      Throwables.throwIfUnchecked(cause);
+      throw e;
+    }
+  }
+
+  private static <T> T getOrThrowIfCancelled(
+      CompletableFuture<T> future, CancellationToken<CancellationException> cancellationToken) {
+    try {
+      CompletableFuture.anyOf(future, cancellationToken.getCancellationFuture()).get();
+      return future.get();
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      throw new CancellationException("External storage operation interrupted");
+    } catch (ExecutionException e) {
+      Throwable cause = e.getCause() != null ? e.getCause() : e;
+      Throwables.throwIfUnchecked(cause);
+      throw new CompletionException(cause);
+    }
+  }
+
+  <T extends Message> CompletableFuture<T> store(
+      T message,
+      @Nullable StorageDriverTargetInfo target,
+      CancellationToken<CancellationException> cancellationToken) {
+    return PayloadVisitors.visit(message, storeOptions(target, null, cancellationToken));
+  }
+
+  CompletableFuture<Void> store(
+      Message.Builder builder,
+      @Nullable StorageDriverTargetInfo target,
+      CancellationToken<CancellationException> cancellationToken) {
+    return PayloadVisitors.visit(builder, storeOptions(target, null, cancellationToken));
+  }
+
+  <T extends Message> CompletableFuture<T> retrieve(
+      T message, CancellationToken<CancellationException> cancellationToken) {
+    return PayloadVisitors.visit(message, retrieveOptions(cancellationToken));
+  }
+
+  CompletableFuture<Void> retrieve(
+      Message.Builder builder, CancellationToken<CancellationException> cancellationToken) {
+    return PayloadVisitors.visit(builder, retrieveOptions(cancellationToken));
+  }
+
+  private PayloadVisitorOptions<StorageDriverTargetInfo> storeOptions(
+      @Nullable StorageDriverTargetInfo target,
+      @Nullable MessageVisitor<StorageDriverTargetInfo> targetVisitor,
+      CancellationToken<CancellationException> cancellationToken) {
+    return PayloadVisitorOptions.<StorageDriverTargetInfo>newBuilder(
+            (visitedTarget, payloads) ->
+                payloadTransformer.store(payloads, visitedTarget, cancellationToken))
+        .setInitialContext(target)
+        .setMessageVisitor(targetVisitor)
+        .setConcurrency(payloadVisitConcurrency)
+        .setSkipSearchAttributes(true)
+        .build();
+  }
+
+  private PayloadVisitorOptions<Void> retrieveOptions(
+      CancellationToken<CancellationException> cancellationToken) {
+    return PayloadVisitorOptions.<Void>newBuilder(
+            (context, payloads) -> payloadTransformer.retrieve(payloads, cancellationToken))
+        .setConcurrency(payloadVisitConcurrency)
+        .setSkipSearchAttributes(true)
+        .build();
+  }
+}

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

@@ -0,0 +1,16 @@
+package io.temporal.internal.payload.storage;
+
+import io.temporal.common.converter.DataConverterException;
+
+/**
+ * Signals that an external storage reference reached a data converter without storage configured.
+ * Logs a TMPRL1105 error.
+ */
+public final class ExternalStorageNotConfiguredException extends DataConverterException {
+  public ExternalStorageNotConfiguredException() {
+    super(
+        "[TMPRL1105] Encountered an external-storage reference payload but external storage is not "
+            + "configured. Configure WorkflowClientOptions.Builder.setExternalStorage(...) with a "
+            + "driver able to retrieve it.");
+  }
+}

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

@@ -9,7 +9,7 @@
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 
-final class ExternalStorageReferences {
+public final class ExternalStorageReferences {
   private static final String ENCODING_PROTOBUF_JSON = "json/protobuf";
   private static final String REFERENCE_MESSAGE_TYPE =
       ExternalStorageReference.getDescriptor().getFullName();
@@ -64,8 +64,7 @@ static Payload toReferencePayload(
    * producer that omits it still yields a readable reference.
    */
   static @Nullable ParsedReference tryParseReference(@Nonnull Payload payload) {
-    if (!hasMetadata(payload, EncodingKeys.METADATA_ENCODING_KEY, ENCODING_PROTOBUF_JSON)
-        || !hasMetadata(payload, EncodingKeys.METADATA_MESSAGE_TYPE_KEY, REFERENCE_MESSAGE_TYPE)) {
+    if (!isReference(payload)) {
       return null;
     }
     ExternalStorageReference.Builder builder = ExternalStorageReference.newBuilder();
@@ -79,6 +78,12 @@ static Payload toReferencePayload(
         reference.getDriverName(), new StorageDriverClaim(reference.getClaimDataMap()));
   }
 
+  /** True if {@code payload} has an external storage reference encoding and message type. */
+  public static boolean isReference(Payload payload) {
+    return hasMetadata(payload, EncodingKeys.METADATA_ENCODING_KEY, ENCODING_PROTOBUF_JSON)
+        && hasMetadata(payload, EncodingKeys.METADATA_MESSAGE_TYPE_KEY, REFERENCE_MESSAGE_TYPE);
+  }
+
   private static boolean hasMetadata(Payload payload, String key, String expected) {
     ByteString value = payload.getMetadataMap().get(key);
     return value != null && expected.equals(value.toStringUtf8());