Add preferred version provider to worker options (#2942)

Author: Sushisource

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

@@ -15,7 +15,10 @@
 import io.temporal.internal.common.ProtobufTimeUtils;
 import io.temporal.internal.common.SdkFlag;
 import io.temporal.internal.statemachines.*;
+import io.temporal.internal.sync.WorkflowInternal;
 import io.temporal.internal.worker.SingleWorkerOptions;
+import io.temporal.worker.PreferredVersionProvider;
+import io.temporal.worker.PreferredVersionProviderInput;
 import io.temporal.workflow.Functions;
 import io.temporal.workflow.Functions.Func;
 import io.temporal.workflow.Functions.Func1;
@@ -337,7 +340,20 @@ public Integer getVersion(
       int minSupported,
       int maxSupported,
       Functions.Proc2<Integer, RuntimeException> callback) {
-    return workflowStateMachines.getVersion(changeId, minSupported, maxSupported, callback);
+    PreferredVersionProvider preferredVersionProvider = workerOptions.getPreferredVersionProvider();
+    return workflowStateMachines.getVersion(
+        changeId,
+        minSupported,
+        maxSupported,
+        preferredVersionProvider == null
+            ? null
+            : (min, max) ->
+                WorkflowInternal.readOnly(
+                    () ->
+                        preferredVersionProvider.getPreferredVersion(
+                            new PreferredVersionProviderInput(
+                                WorkflowInternal.getWorkflowInfo(), changeId, min, max))),
+        callback);
   }
 
   @Override

temporal-sdk/src/main/java/io/temporal/internal/statemachines/VersionStateMachine.java

@@ -14,8 +14,10 @@
 import io.temporal.api.history.v1.HistoryEvent;
 import io.temporal.internal.history.VersionMarkerUtils;
 import io.temporal.worker.NonDeterministicException;
+import io.temporal.worker.VersionPreference;
 import io.temporal.workflow.Functions;
 import java.util.Objects;
+import java.util.function.BiFunction;
 import javax.annotation.Nullable;
 
 final class VersionStateMachine {
@@ -137,17 +139,22 @@ class InvocationStateMachine
     private final Functions.Func1<Integer, SearchAttributes> upsertSearchAttributeCallback;
     private final Functions.Proc2<Integer, RuntimeException> resultCallback;
 
+    @Nullable
+    private final BiFunction<Integer, Integer, VersionPreference> preferredVersionProvider;
+
     InvocationStateMachine(
         int minSupported,
         int maxSupported,
         boolean waitForMarkerRecordedReplaying,
         Functions.Func1<Integer, SearchAttributes> upsertSearchAttributeCallback,
+        @Nullable BiFunction<Integer, Integer, VersionPreference> preferredVersionProvider,
         Functions.Proc2<Integer, RuntimeException> callback) {
       super(STATE_MACHINE_DEFINITION, VersionStateMachine.this.commandSink, stateMachineSink);
       this.minSupported = minSupported;
       this.maxSupported = maxSupported;
       this.waitForMarkerRecordedReplaying = waitForMarkerRecordedReplaying;
       this.upsertSearchAttributeCallback = upsertSearchAttributeCallback;
+      this.preferredVersionProvider = preferredVersionProvider;
       this.resultCallback = Objects.requireNonNull(callback);
     }
 
@@ -210,13 +217,17 @@ private void validateVersionAndThrow(boolean preloaded) {
         throw new IllegalStateException((preloaded ? "preloaded " : "") + " version not set");
       }
       if (versionToUse < minSupported || versionToUse > maxSupported) {
-        throw new UnsupportedVersion.UnsupportedVersionException(
-            String.format(
-                "Version %d of changeId %s is not supported. Supported v is between %d and %d.",
-                versionToUse, changeId, minSupported, maxSupported));
+        throwUnsupportedVersion(versionToUse);
       }
     }
 
+    private void throwUnsupportedVersion(int unsupportedVersion) {
+      throw new UnsupportedVersion.UnsupportedVersionException(
+          String.format(
+              "Version %d of changeId %s is not supported. Supported v is between %d and %d.",
+              unsupportedVersion, changeId, minSupported, maxSupported));
+    }
+
     void notifyFromVersion(boolean preloaded) {
       Integer versionToUse = preloaded ? preloadedVersion : version;
       resultCallback.apply(versionToUse, null);
@@ -227,8 +238,8 @@ void notifyFromException(RuntimeException ex) {
     }
 
     void notifyFromVersionExecuting() {
-      // the only case when we don't need to validate before notification because
-      // we just initialized the version with maxVersion
+      // No validation needed before notification here: resolveVersionExecuting() already guarantees
+      // the version it produces is within [minSupported, maxSupported].
       notifyFromVersion(false);
     }
 
@@ -238,7 +249,7 @@ State createMarkerExecuting() {
         addCommand(StateMachineCommandUtils.RECORD_MARKER_FAKE_COMMAND);
         return State.SKIPPED;
       } else {
-        version = maxSupported;
+        version = resolveVersionExecuting();
         SearchAttributes sa = upsertSearchAttributeCallback.apply(version);
         writeVersionChangeSA = sa != null;
         RecordMarkerCommandAttributes markerAttributes =
@@ -253,6 +264,27 @@ State createMarkerExecuting() {
       }
     }
 
+    private int resolveVersionExecuting() {
+      if (preferredVersionProvider == null) {
+        return maxSupported;
+      }
+      VersionPreference versionPreference =
+          preferredVersionProvider.apply(minSupported, maxSupported);
+      if (versionPreference == null) {
+        return maxSupported;
+      }
+
+      int preferredVersion = versionPreference.getVersion();
+      if (preferredVersion >= minSupported && preferredVersion <= maxSupported) {
+        return preferredVersion;
+      }
+      if (versionPreference.isClampToSupportedRange()) {
+        return Math.min(Math.max(preferredVersion, minSupported), maxSupported);
+      }
+      throwUnsupportedVersion(preferredVersion);
+      throw new IllegalStateException("unreachable");
+    }
+
     void notifySkippedExecuting() {
       cancelCommand();
       try {
@@ -411,13 +443,15 @@ public Integer getVersion(
       int maxSupported,
       boolean waitForMarkerRecordedReplaying,
       Functions.Func1<Integer, SearchAttributes> upsertSearchAttributeCallback,
+      @Nullable BiFunction<Integer, Integer, VersionPreference> preferredVersionProvider,
       Functions.Proc2<Integer, RuntimeException> callback) {
     InvocationStateMachine ism =
         new InvocationStateMachine(
             minSupported,
             maxSupported,
             waitForMarkerRecordedReplaying,
             upsertSearchAttributeCallback,
+            preferredVersionProvider,
             callback);
     ism.explicitEvent(ExplicitEvent.CHECK_EXECUTION_STATE);
     ism.explicitEvent(ExplicitEvent.SCHEDULE);

temporal-sdk/src/main/java/io/temporal/internal/statemachines/WorkflowStateMachines.java

@@ -30,12 +30,14 @@
 import io.temporal.serviceclient.Version;
 import io.temporal.worker.MetricsType;
 import io.temporal.worker.NonDeterministicException;
+import io.temporal.worker.VersionPreference;
 import io.temporal.worker.WorkflowImplementationOptions;
 import io.temporal.workflow.ChildWorkflowCancellationType;
 import io.temporal.workflow.Functions;
 import io.temporal.workflow.NexusOperationCancellationType;
 import java.nio.charset.StandardCharsets;
 import java.util.*;
+import java.util.function.BiFunction;
 import javax.annotation.Nonnull;
 import javax.annotation.Nullable;
 import org.slf4j.Logger;
@@ -1244,6 +1246,15 @@ public Integer getVersion(
       int minSupported,
       int maxSupported,
       Functions.Proc2<Integer, RuntimeException> callback) {
+    return getVersion(changeId, minSupported, maxSupported, null, callback);
+  }
+
+  public Integer getVersion(
+      String changeId,
+      int minSupported,
+      int maxSupported,
+      @Nullable BiFunction<Integer, Integer, VersionPreference> preferredVersionProvider,
+      Functions.Proc2<Integer, RuntimeException> callback) {
     VersionStateMachine stateMachine =
         versions.computeIfAbsent(
             changeId,
@@ -1275,6 +1286,7 @@ public Integer getVersion(
           }
           return sa;
         },
+        preferredVersionProvider,
         (v, e) -> {
           callback.apply(v, e);
           // without this getVersion call will trigger the end of WFT,

temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflowContext.java

@@ -1143,21 +1143,26 @@ private <R> R mutableSideEffectImpl(
   @Override
   public int getVersion(String changeId, int minSupported, int maxSupported) {
     CompletablePromise<Integer> result = Workflow.newPromise();
-    Integer versionToUse =
-        replayContext.getVersion(
-            changeId,
-            minSupported,
-            maxSupported,
-            (v, e) ->
-                runner.executeInWorkflowThread(
-                    "version-callback",
-                    () -> {
-                      if (v != null) {
-                        result.complete(v);
-                      } else {
-                        result.completeExceptionally(e);
-                      }
-                    }));
+    Integer versionToUse;
+    try {
+      versionToUse =
+          replayContext.getVersion(
+              changeId,
+              minSupported,
+              maxSupported,
+              (v, e) ->
+                  runner.executeInWorkflowThread(
+                      "version-callback",
+                      () -> {
+                        if (v != null) {
+                          result.complete(v);
+                        } else {
+                          result.completeExceptionally(e);
+                        }
+                      }));
+    } catch (UnsupportedVersion.UnsupportedVersionException ex) {
+      throw new UnsupportedVersion(ex);
+    }
     /*
      * If we are replaying a workflow and encounter a getVersion call it is possible that this call did not exist
      * on the original execution. If the call did not exist on the original execution then we cannot block on results

temporal-sdk/src/main/java/io/temporal/internal/sync/WorkflowInternal.java

@@ -674,6 +674,17 @@ public static <T> T deadlockDetectorOff(Functions.Func<T> func) {
     }
   }
 
+  public static <T> T readOnly(Functions.Func<T> func) {
+    SyncWorkflowContext workflowContext = getRootWorkflowContext();
+    boolean previousReadOnly = workflowContext.isReadOnly();
+    workflowContext.setReadOnly(true);
+    try {
+      return func.apply();
+    } finally {
+      workflowContext.setReadOnly(previousReadOnly);
+    }
+  }
+
   public static WorkflowInfo getWorkflowInfo() {
     return new WorkflowInfoImpl(getRootWorkflowContext().getReplayContext());
   }

temporal-sdk/src/main/java/io/temporal/internal/worker/SingleWorkerOptions.java

@@ -7,6 +7,7 @@
 import io.temporal.common.converter.DataConverter;
 import io.temporal.common.converter.GlobalDataConverter;
 import io.temporal.common.interceptors.WorkerInterceptor;
+import io.temporal.worker.PreferredVersionProvider;
 import io.temporal.worker.WorkerDeploymentOptions;
 import java.time.Duration;
 import java.util.List;
@@ -43,6 +44,7 @@ public static final class Builder {
     private String workerInstanceKey;
     private boolean allowActivityHeartbeatDuringShutdown;
     private String workerControlTaskQueue;
+    private PreferredVersionProvider preferredVersionProvider;
 
     private Builder() {}
 
@@ -70,6 +72,7 @@ private Builder(SingleWorkerOptions options) {
       this.workerInstanceKey = options.getWorkerInstanceKey();
       this.allowActivityHeartbeatDuringShutdown = options.getAllowActivityHeartbeatDuringShutdown();
       this.workerControlTaskQueue = options.getWorkerControlTaskQueue();
+      this.preferredVersionProvider = options.getPreferredVersionProvider();
     }
 
     public Builder setIdentity(String identity) {
@@ -177,6 +180,11 @@ public Builder setWorkerControlTaskQueue(String workerControlTaskQueue) {
       return this;
     }
 
+    public Builder setPreferredVersionProvider(PreferredVersionProvider preferredVersionProvider) {
+      this.preferredVersionProvider = preferredVersionProvider;
+      return this;
+    }
+
     public SingleWorkerOptions build() {
       PollerOptions pollerOptions = this.pollerOptions;
       if (pollerOptions == null) {
@@ -218,7 +226,8 @@ public SingleWorkerOptions build() {
           this.deploymentOptions,
           this.workerInstanceKey,
           this.allowActivityHeartbeatDuringShutdown,
-          this.workerControlTaskQueue);
+          this.workerControlTaskQueue,
+          this.preferredVersionProvider);
     }
   }
 
@@ -242,6 +251,7 @@ public SingleWorkerOptions build() {
   private final String workerInstanceKey;
   private final boolean allowActivityHeartbeatDuringShutdown;
   private final String workerControlTaskQueue;
+  private final PreferredVersionProvider preferredVersionProvider;
 
   private SingleWorkerOptions(
       String identity,
@@ -263,7 +273,8 @@ private SingleWorkerOptions(
       WorkerDeploymentOptions deploymentOptions,
       String workerInstanceKey,
       boolean allowActivityHeartbeatDuringShutdown,
-      String workerControlTaskQueue) {
+      String workerControlTaskQueue,
+      PreferredVersionProvider preferredVersionProvider) {
     this.identity = identity;
     this.binaryChecksum = binaryChecksum;
     this.buildId = buildId;
@@ -284,6 +295,7 @@ private SingleWorkerOptions(
     this.workerInstanceKey = workerInstanceKey;
     this.allowActivityHeartbeatDuringShutdown = allowActivityHeartbeatDuringShutdown;
     this.workerControlTaskQueue = workerControlTaskQueue;
+    this.preferredVersionProvider = preferredVersionProvider;
   }
 
   public String getIdentity() {
@@ -377,6 +389,10 @@ public String getWorkerControlTaskQueue() {
     return workerControlTaskQueue;
   }
 
+  public PreferredVersionProvider getPreferredVersionProvider() {
+    return preferredVersionProvider;
+  }
+
   public WorkerVersioningOptions getWorkerVersioningOptions() {
     return new WorkerVersioningOptions(
         this.getBuildId(), this.isUsingBuildIdForVersioning(), this.getDeploymentOptions());

temporal-sdk/src/main/java/io/temporal/worker/PreferredVersionProvider.java

@@ -0,0 +1,30 @@
+package io.temporal.worker;
+
+import io.temporal.common.Experimental;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Provides the version to record for a {@link io.temporal.workflow.Workflow#getVersion(String, int,
+ * int)} call when the version marker is created for the first time.
+ *
+ * <p>This provider is called only during non-replay workflow task execution, before the SDK records
+ * a version marker for the change ID. It is not called during replay or after a version was already
+ * memoized for the same change ID.
+ */
+@Experimental
+@FunctionalInterface
+public interface PreferredVersionProvider {
+
+  /**
+   * Returns the preferred version for the supplied {@code getVersion} call, or {@code null} to use
+   * the SDK default of {@code maxSupported}.
+   *
+   * <p>This method is invoked on the workflow thread. Any exception it throws propagates out of the
+   * {@code getVersion} call and fails the current workflow task, which will be retried; a provider
+   * that always throws will therefore block the workflow. Implementations should be deterministic
+   * and side-effect free.
+   */
+  @Nullable
+  VersionPreference getPreferredVersion(@Nonnull PreferredVersionProviderInput input);
+}