Support WorkflowImplementationOptionsCustomizer (#2503)

Support WorkflowImplementationOptionsCustomizer

Author: Quinn-With-Two-Ns

temporal-spring-boot-autoconfigure/src/main/java/io/temporal/spring/boot/WorkflowImplementationOptionsCustomizer.java

@@ -0,0 +1,48 @@
+package io.temporal.spring.boot;
+
+import io.temporal.common.metadata.POJOWorkflowMethodMetadata;
+import io.temporal.worker.Worker;
+import io.temporal.worker.WorkflowImplementationOptions;
+import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
+
+/**
+ * Bean of this class can be added to Spring context to get a fine control over {@link
+ * WorkflowImplementationOptions} objects that are created by Temporal Spring Boot Autoconfigure
+ * module.
+ *
+ * <p>Only one bean of this or {@code
+ * TemporalOptionsCustomizer<WorkflowImplementationOptions.Builder>} type can be added to Spring
+ * context.
+ */
+public interface WorkflowImplementationOptionsCustomizer
+    extends TemporalOptionsCustomizer<WorkflowImplementationOptions.Builder> {
+
+  @Nonnull
+  @Override
+  default WorkflowImplementationOptions.Builder customize(
+      @Nonnull WorkflowImplementationOptions.Builder optionsBuilder) {
+    return optionsBuilder;
+  }
+
+  /**
+   * This method can modify some fields of the provided {@code
+   * WorkflowImplementationOptions.Builder} or create a new builder altogether and return it back to
+   * be used. This method is called after the {@code WorkflowImplementationOptions.Builder} is
+   * initialized by the Temporal Spring Boot module, so changes done by this method will override
+   * Spring Boot configuration values.
+   *
+   * @param optionsBuilder {@code *Options.Builder} to customize
+   * @param worker the worker that the workflow implementation is being registered to.
+   * @param clazz the class of the workflow implementation that is being registered.
+   * @param workflowMethod the metadata of the workflow method on the interface that is being
+   *     registered. null if the class is a {@link io.temporal.workflow.DynamicWorkflow}.
+   * @return modified {@code optionsBuilder} or a new builder instance to be used by the caller code
+   */
+  @Nonnull
+  WorkflowImplementationOptions.Builder customize(
+      @Nonnull WorkflowImplementationOptions.Builder optionsBuilder,
+      @Nonnull Worker worker,
+      @Nonnull Class<?> clazz,
+      @Nullable POJOWorkflowMethodMetadata workflowMethod);
+}

temporal-spring-boot-autoconfigure/src/main/java/io/temporal/spring/boot/autoconfigure/template/WorkersTemplate.java

@@ -501,7 +501,7 @@ private <T> void configureWorkflowImplementation(Worker worker, Class<?> clazz)
                 clazz, Collections.singletonList(executeMethod));
         WorkflowImplementationOptions workflowImplementationOptions =
             new WorkflowImplementationOptionsTemplate(workflowImplementationCustomizer)
-                .createWorkflowImplementationOptions();
+                .createWorkflowImplementationOptions(worker, clazz, null);
         worker.registerWorkflowImplementationFactory(
             DynamicWorkflow.class,
             (encodedValues) -> {
@@ -544,10 +544,6 @@ private <T> void configureWorkflowImplementation(Worker worker, Class<?> clazz)
               + clazz);
     }
 
-    WorkflowImplementationOptions workflowImplementationOptions =
-        new WorkflowImplementationOptionsTemplate(workflowImplementationCustomizer)
-            .createWorkflowImplementationOptions();
-
     WorkerDeploymentOptions deploymentOptions = worker.getWorkerOptions().getDeploymentOptions();
 
     // If the workflow implementation class has a constructor annotated with @WorkflowInit,
@@ -570,6 +566,9 @@ private <T> void configureWorkflowImplementation(Worker worker, Class<?> clazz)
             deploymentOptions.isUsingVersioning());
       }
 
+      WorkflowImplementationOptions workflowImplementationOptions =
+          new WorkflowImplementationOptionsTemplate(workflowImplementationCustomizer)
+              .createWorkflowImplementationOptions(worker, clazz, workflowMethod);
       worker.registerWorkflowImplementationFactory(
           (Class<T>) workflowMethod.getWorkflowInterface(),
           (encodedValues) -> {
@@ -602,6 +601,9 @@ private <T> void configureWorkflowImplementation(Worker worker, Class<?> clazz)
               deploymentOptions.getDefaultVersioningBehavior(),
               deploymentOptions.isUsingVersioning());
         }
+        WorkflowImplementationOptions workflowImplementationOptions =
+            new WorkflowImplementationOptionsTemplate(workflowImplementationCustomizer)
+                .createWorkflowImplementationOptions(worker, clazz, workflowMethod);
         worker.registerWorkflowImplementationFactory(
             (Class<T>) workflowMethod.getWorkflowInterface(),
             () -> (T) beanFactory.createBean(clazz),

temporal-spring-boot-autoconfigure/src/main/java/io/temporal/spring/boot/autoconfigure/template/WorkflowImplementationOptionsTemplate.java

@@ -1,6 +1,9 @@
 package io.temporal.spring.boot.autoconfigure.template;
 
+import io.temporal.common.metadata.POJOWorkflowMethodMetadata;
 import io.temporal.spring.boot.TemporalOptionsCustomizer;
+import io.temporal.spring.boot.WorkflowImplementationOptionsCustomizer;
+import io.temporal.worker.Worker;
 import io.temporal.worker.WorkflowImplementationOptions;
 import javax.annotation.Nullable;
 
@@ -13,11 +16,17 @@ public WorkflowImplementationOptionsTemplate(
     this.customizer = customizer;
   }
 
-  public WorkflowImplementationOptions createWorkflowImplementationOptions() {
+  public WorkflowImplementationOptions createWorkflowImplementationOptions(
+      Worker worker, Class<?> clazz, POJOWorkflowMethodMetadata workflowMethod) {
     WorkflowImplementationOptions.Builder options = WorkflowImplementationOptions.newBuilder();
 
     if (customizer != null) {
       options = customizer.customize(options);
+      if (customizer instanceof WorkflowImplementationOptionsCustomizer) {
+        options =
+            ((WorkflowImplementationOptionsCustomizer) customizer)
+                .customize(options, worker, clazz, workflowMethod);
+      }
     }
 
     return options.build();

temporal-spring-boot-autoconfigure/src/test/java/io/temporal/spring/boot/autoconfigure/OptionsCustomizersTest.java

@@ -7,9 +7,10 @@
 import io.temporal.client.WorkflowClientOptions;
 import io.temporal.spring.boot.TemporalOptionsCustomizer;
 import io.temporal.spring.boot.WorkerOptionsCustomizer;
+import io.temporal.spring.boot.WorkflowImplementationOptionsCustomizer;
+import io.temporal.spring.boot.autoconfigure.bytaskqueue.TestWorkflowImpl;
 import io.temporal.testing.TestEnvironmentOptions;
 import io.temporal.worker.WorkerFactoryOptions;
-import io.temporal.worker.WorkflowImplementationOptions;
 import java.util.List;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
@@ -31,6 +32,7 @@ public class OptionsCustomizersTest {
 
   @Autowired List<TemporalOptionsCustomizer<?>> customizers;
   @Autowired WorkerOptionsCustomizer workerCustomizer;
+  @Autowired WorkflowImplementationOptionsCustomizer workflowImplementationOptionsCustomizer;
 
   @BeforeEach
   void setUp() {
@@ -43,6 +45,8 @@ public void testCustomizersGotCalled() {
     assertEquals(5, customizers.size());
     customizers.forEach(c -> verify(c).customize(any()));
     verify(workerCustomizer).customize(any(), eq("UnitTest"), eq("UnitTest"));
+    verify(workflowImplementationOptionsCustomizer)
+        .customize(any(), any(), eq(TestWorkflowImpl.class), any());
   }
 
   @ComponentScan(
@@ -68,9 +72,14 @@ public TemporalOptionsCustomizer<WorkerFactoryOptions.Builder> workerFactoryCust
     }
 
     @Bean
-    public TemporalOptionsCustomizer<WorkflowImplementationOptions.Builder>
-        WorkflowImplementationCustomizer() {
-      return getReturningMock();
+    public WorkflowImplementationOptionsCustomizer WorkflowImplementationCustomizer() {
+      WorkflowImplementationOptionsCustomizer mock =
+          mock(WorkflowImplementationOptionsCustomizer.class);
+      when(mock.customize(any())).thenAnswer(invocation -> invocation.getArgument(0)).getMock();
+      when(mock.customize(any(), any(), any(), any()))
+          .thenAnswer(invocation -> invocation.getArgument(0))
+          .getMock();
+      return mock;
     }
 
     @Bean