Message and concurrent payload visitors

Author: jmaeagle99

temporal-sdk/build.gradle

@@ -65,6 +65,35 @@ dependencies {
     java21Implementation files(sourceSets.main.output.classesDirs) { builtBy compileJava }
 }
 
+// --- Payload visitor code generation ---
+// A build-time generator (compiled in its own source set against the proto classes from
+// temporal-serviceclient) emits GeneratedPayloadVisitor.java, which knows how to walk every
+// payload-bearing Temporal API message. The generated source is added to the main source set.
+sourceSets {
+    payloadVisitorGenerator {
+        java {
+            srcDirs = ['src/payloadVisitorGenerator/java']
+        }
+    }
+}
+
+dependencies {
+    payloadVisitorGeneratorImplementation project(':temporal-serviceclient')
+}
+
+def generatedPayloadVisitorDir = layout.buildDirectory.dir('generated/payloadvisitor/java')
+
+def generatePayloadVisitor = tasks.register('generatePayloadVisitor', JavaExec) {
+    dependsOn 'compilePayloadVisitorGeneratorJava'
+    classpath = sourceSets.payloadVisitorGenerator.runtimeClasspath
+    mainClass = 'io.temporal.internal.payload.visitor.gen.PayloadVisitorGenerator'
+    args generatedPayloadVisitorDir.get().asFile.absolutePath
+    inputs.files(sourceSets.payloadVisitorGenerator.runtimeClasspath)
+    outputs.dir(generatedPayloadVisitorDir)
+}
+
+sourceSets.main.java.srcDir(generatePayloadVisitor)
+
 tasks.named('compileJava17Java') {
     options.release = 17
 }

temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/GeneratedVisitor.java

@@ -0,0 +1,11 @@
+package io.temporal.internal.payload.visitor;
+
+import com.google.protobuf.Message;
+
+/**
+ * Generated traversal for one message type: visits the message's payload fields and recurses into
+ * its child messages. There is one per message type that can contain a payload.
+ */
+interface GeneratedVisitor {
+  void visit(Traversal traversal, Message.Builder builder);
+}

temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/MessageRegistryEntry.java

@@ -0,0 +1,18 @@
+package io.temporal.internal.payload.visitor;
+
+import com.google.protobuf.Message;
+import java.util.function.Supplier;
+
+/**
+ * How to traverse one message type, and how to create an empty builder for it (used to unpack
+ * {@code google.protobuf.Any} values).
+ */
+final class MessageRegistryEntry {
+  final GeneratedVisitor visitor;
+  final Supplier<Message.Builder> newBuilder;
+
+  MessageRegistryEntry(GeneratedVisitor visitor, Supplier<Message.Builder> newBuilder) {
+    this.visitor = visitor;
+    this.newBuilder = newBuilder;
+  }
+}

temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/MessageVisitor.java

@@ -0,0 +1,23 @@
+package io.temporal.internal.payload.visitor;
+
+import com.google.protobuf.MessageOrBuilder;
+
+/**
+ * Callback invoked when traversal enters a proto message. The returned value becomes the contextual
+ * value in scope for that message and everything within it, and is restored to the enclosing value
+ * once traversal leaves the message. The message is provided as a builder and may be inspected or
+ * mutated.
+ *
+ * @param <C> type of the contextual value
+ */
+@FunctionalInterface
+interface MessageVisitor<C> {
+  /**
+   * Handles a message being entered and returns the contextual value for it and its contents.
+   *
+   * @param current the contextual value in scope from the enclosing message
+   * @param message the message being entered
+   * @return the contextual value to use for this message and its contents
+   */
+  C onEnter(C current, MessageOrBuilder message);
+}

temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/MessageVisitorOptions.java

@@ -0,0 +1,59 @@
+package io.temporal.internal.payload.visitor;
+
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Options for visiting the messages of a proto message, without visiting individual payloads.
+ *
+ * @param <C> type of the contextual value supplied to the visitor
+ */
+final class MessageVisitorOptions<C> {
+  private final @Nonnull MessageVisitor<C> messageVisitor;
+  private final @Nullable C initialContext;
+
+  private MessageVisitorOptions(Builder<C> b) {
+    this.messageVisitor = b.messageVisitor;
+    this.initialContext = b.initialContext;
+  }
+
+  public static <C> Builder<C> newBuilder() {
+    return new Builder<>();
+  }
+
+  @Nonnull
+  public MessageVisitor<C> getMessageVisitor() {
+    return messageVisitor;
+  }
+
+  @Nullable
+  public C getInitialContext() {
+    return initialContext;
+  }
+
+  public static final class Builder<C> {
+    private MessageVisitor<C> messageVisitor;
+    private C initialContext;
+
+    private Builder() {}
+
+    /** Required. The message visitor. */
+    public Builder<C> setMessageVisitor(@Nonnull MessageVisitor<C> messageVisitor) {
+      this.messageVisitor = messageVisitor;
+      return this;
+    }
+
+    /** Optional. The contextual value in scope before any message is entered. */
+    public Builder<C> setInitialContext(@Nullable C initialContext) {
+      this.initialContext = initialContext;
+      return this;
+    }
+
+    public MessageVisitorOptions<C> build() {
+      if (messageVisitor == null) {
+        throw new IllegalArgumentException("messageVisitor is required");
+      }
+      return new MessageVisitorOptions<>(this);
+    }
+  }
+}

temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/MessageVisitors.java

@@ -0,0 +1,43 @@
+package io.temporal.internal.payload.visitor;
+
+import com.google.protobuf.Message;
+import javax.annotation.Nonnull;
+
+/**
+ * Visits the messages within a proto message, invoking the message visitor on each, without
+ * visiting individual payloads. Only messages that can contain a payload are visited.
+ *
+ * <p>This is an SDK-internal utility; it is not part of the public API.
+ */
+final class MessageVisitors {
+  private MessageVisitors() {}
+
+  /** Visits the messages in {@code builder} in place. */
+  public static <C> void visit(
+      @Nonnull Message.Builder builder, @Nonnull MessageVisitorOptions<C> options) {
+    Traversal traversal =
+        new Traversal(
+            null,
+            options.getMessageVisitor(),
+            options.getInitialContext(),
+            /* skipSearchAttributes= */ false,
+            /* skipHeaders= */ false,
+            1,
+            null,
+            GeneratedPayloadVisitor.REGISTRY);
+    traversal.dispatch(builder);
+    traversal.execute();
+  }
+
+  /**
+   * Visits the messages in {@code message}, returning a copy with any changes applied; the input is
+   * unchanged.
+   */
+  @SuppressWarnings("unchecked")
+  public static <C, T extends Message> T visit(
+      @Nonnull T message, @Nonnull MessageVisitorOptions<C> options) {
+    Message.Builder builder = message.toBuilder();
+    visit(builder, options);
+    return (T) builder.build();
+  }
+}

temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitor.java

@@ -0,0 +1,26 @@
+package io.temporal.internal.payload.visitor;
+
+import io.temporal.api.common.v1.Payload;
+import java.util.List;
+
+/**
+ * Callback for a sequence of payloads found in a proto message. The returned list replaces those
+ * payloads; return the same list to leave them unchanged.
+ *
+ * <p>When the visited field holds a single payload the list has one element and the visitor must
+ * return exactly one payload. With a concurrency limit greater than one, visits may run on multiple
+ * threads, so implementations must be thread-safe.
+ *
+ * @param <C> type of the contextual value supplied to each visit
+ */
+@FunctionalInterface
+interface PayloadVisitor<C> {
+  /**
+   * Visits a sequence of payloads and returns their replacements.
+   *
+   * @param context the location of these payloads and the contextual value in scope
+   * @param payloads the payloads found at this location
+   * @return the replacement payloads
+   */
+  List<Payload> visit(PayloadVisitorContext<C> context, List<Payload> payloads);
+}

temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitorContext.java

@@ -0,0 +1,33 @@
+package io.temporal.internal.payload.visitor;
+
+import com.google.protobuf.MessageOrBuilder;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * The context for one payload visitor call: the contextual value in scope and the message that
+ * contains the payloads being visited.
+ *
+ * @param <C> type of the contextual value
+ */
+final class PayloadVisitorContext<C> {
+  private final @Nullable C context;
+  private final @Nonnull MessageOrBuilder parent;
+
+  PayloadVisitorContext(@Nullable C context, @Nonnull MessageOrBuilder parent) {
+    this.context = context;
+    this.parent = parent;
+  }
+
+  /** The contextual value in scope at this location, or {@code null} if none. */
+  @Nullable
+  public C getContext() {
+    return context;
+  }
+
+  /** The message that directly contains the payloads being visited. */
+  @Nonnull
+  public MessageOrBuilder getParent() {
+    return parent;
+  }
+}

temporal-sdk/src/main/java/io/temporal/internal/payload/visitor/PayloadVisitorOptions.java

@@ -0,0 +1,141 @@
+package io.temporal.internal.payload.visitor;
+
+import java.util.concurrent.Executor;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Options for visiting the payloads of a proto message.
+ *
+ * @param <C> type of the contextual value supplied to the visitor
+ */
+final class PayloadVisitorOptions<C> {
+  private final @Nonnull PayloadVisitor<C> payloadVisitor;
+  private final @Nullable MessageVisitor<C> messageVisitor;
+  private final @Nullable C initialContext;
+  private final boolean skipSearchAttributes;
+  private final boolean skipHeaders;
+  private final int concurrency;
+  private final @Nullable Executor executor;
+
+  private PayloadVisitorOptions(Builder<C> b) {
+    this.payloadVisitor = b.payloadVisitor;
+    this.messageVisitor = b.messageVisitor;
+    this.initialContext = b.initialContext;
+    this.skipSearchAttributes = b.skipSearchAttributes;
+    this.skipHeaders = b.skipHeaders;
+    this.concurrency = b.concurrency;
+    this.executor = b.executor;
+  }
+
+  public static <C> Builder<C> newBuilder() {
+    return new Builder<>();
+  }
+
+  @Nonnull
+  public PayloadVisitor<C> getPayloadVisitor() {
+    return payloadVisitor;
+  }
+
+  @Nullable
+  public MessageVisitor<C> getMessageVisitor() {
+    return messageVisitor;
+  }
+
+  @Nullable
+  public C getInitialContext() {
+    return initialContext;
+  }
+
+  /** Whether search attribute payloads are skipped. */
+  public boolean isSkipSearchAttributes() {
+    return skipSearchAttributes;
+  }
+
+  /** Whether header payloads are skipped. */
+  public boolean isSkipHeaders() {
+    return skipHeaders;
+  }
+
+  /** Maximum number of visits that may run concurrently; {@code 1} is sequential. */
+  public int getConcurrency() {
+    return concurrency;
+  }
+
+  /** Executor for concurrent visits; {@code null} when concurrency is {@code 1}. */
+  @Nullable
+  public Executor getExecutor() {
+    return executor;
+  }
+
+  public static final class Builder<C> {
+    private PayloadVisitor<C> payloadVisitor;
+    private MessageVisitor<C> messageVisitor;
+    private C initialContext;
+    private boolean skipSearchAttributes;
+    private boolean skipHeaders;
+    private int concurrency = 1;
+    private Executor executor;
+
+    private Builder() {}
+
+    /** Required. The payload visitor. */
+    public Builder<C> setPayloadVisitor(@Nonnull PayloadVisitor<C> payloadVisitor) {
+      this.payloadVisitor = payloadVisitor;
+      return this;
+    }
+
+    /** Optional. A callback invoked when entering each message. */
+    public Builder<C> setMessageVisitor(@Nullable MessageVisitor<C> messageVisitor) {
+      this.messageVisitor = messageVisitor;
+      return this;
+    }
+
+    /** Optional. The contextual value in scope before any message is entered. */
+    public Builder<C> setInitialContext(@Nullable C initialContext) {
+      this.initialContext = initialContext;
+      return this;
+    }
+
+    /** Whether to skip search attribute payloads. */
+    public Builder<C> setSkipSearchAttributes(boolean skipSearchAttributes) {
+      this.skipSearchAttributes = skipSearchAttributes;
+      return this;
+    }
+
+    /** Whether to skip header payloads. */
+    public Builder<C> setSkipHeaders(boolean skipHeaders) {
+      this.skipHeaders = skipHeaders;
+      return this;
+    }
+
+    /**
+     * Maximum number of concurrent visits; must be at least {@code 1} (the default, sequential). A
+     * value greater than {@code 1} requires an executor (see {@link #setExecutor}).
+     */
+    public Builder<C> setConcurrency(int concurrency) {
+      this.concurrency = concurrency;
+      return this;
+    }
+
+    /** Executor for concurrent visits. Required when concurrency is greater than {@code 1}. */
+    public Builder<C> setExecutor(@Nullable Executor executor) {
+      this.executor = executor;
+      return this;
+    }
+
+    public PayloadVisitorOptions<C> build() {
+      if (payloadVisitor == null) {
+        throw new IllegalArgumentException("payloadVisitor is required");
+      }
+      if (concurrency < 1) {
+        throw new IllegalArgumentException("concurrency must be at least 1, got " + concurrency);
+      }
+      if (concurrency > 1 && executor == null) {
+        throw new IllegalArgumentException(
+            "executor is required when concurrency is greater than 1");
+      }
+      return new PayloadVisitorOptions<>(this);
+    }
+  }
+}