Adding ability to override activity options. (#772)

* Adding ability to override activity options.

Author: mindaugasbarcauskas

src/main/java/com/uber/cadence/internal/sync/ActivityInvocationHandler.java

@@ -24,6 +24,7 @@
 import com.uber.cadence.workflow.WorkflowInterceptor;
 import java.lang.reflect.InvocationHandler;
 import java.lang.reflect.Method;
+import java.util.Map;
 import java.util.function.Function;
 
 class ActivityInvocationHandler extends ActivityInvocationHandlerBase {
@@ -43,9 +44,31 @@ private ActivityInvocationHandler(ActivityOptions options, WorkflowInterceptor a
   @Override
   protected Function<Object[], Object> getActivityFunc(
       Method method, MethodRetry methodRetry, ActivityMethod activityMethod, String activityName) {
+
+    ActivityOptions activityOptionsOverride = null;
+    WorkflowThread workflowThread = WorkflowInternal.getRootWorkflowContext();
+    if (workflowThread.getDecisionContext() != null
+        && workflowThread.getDecisionContext().getWorkflowImplementationOptions() != null) {
+      Map<String, ActivityOptions> activityOptionsMap =
+          workflowThread
+              .getDecisionContext()
+              .getWorkflowImplementationOptions()
+              .getActivityOptions();
+
+      if (activityOptionsMap.containsKey(activityName)) {
+        activityOptionsOverride = activityOptionsMap.get(activityName);
+      }
+    }
+
     Function<Object[], Object> function;
     ActivityOptions mergedOptions = ActivityOptions.merge(activityMethod, methodRetry, options);
-    ActivityStub stub = ActivityStubImpl.newInstance(mergedOptions, activityExecutor);
+    ActivityStub stub;
+    if (activityOptionsOverride == null) {
+      stub = ActivityStubImpl.newInstance(mergedOptions, activityExecutor);
+    } else {
+      ActivityOptions mergedOverrideOptions = ActivityOptions.merge(activityMethod, methodRetry, activityOptionsOverride);
+      stub = ActivityStubImpl.newInstance(mergedOverrideOptions, activityExecutor);
+    }
 
     function =
         (a) -> stub.execute(activityName, method.getReturnType(), method.getGenericReturnType(), a);

src/main/java/com/uber/cadence/internal/sync/DeterministicRunnerImpl.java

@@ -187,7 +187,12 @@ public Thread newThread(Runnable r) {
 
   private static SyncDecisionContext newDummySyncDecisionContext() {
     return new SyncDecisionContext(
-        new DummyDecisionContext(), JsonDataConverter.getInstance(), null, (next) -> next, null);
+        new DummyDecisionContext(),
+        JsonDataConverter.getInstance(),
+        null,
+        (next) -> next,
+        null,
+        null);
   }
 
   SyncDecisionContext getDecisionContext() {

src/main/java/com/uber/cadence/internal/sync/SyncDecisionContext.java

@@ -40,6 +40,7 @@
 import com.uber.cadence.internal.replay.ExecuteLocalActivityParameters;
 import com.uber.cadence.internal.replay.SignalExternalWorkflowParameters;
 import com.uber.cadence.internal.replay.StartChildWorkflowExecutionParameters;
+import com.uber.cadence.worker.WorkflowImplementationOptions;
 import com.uber.cadence.workflow.ActivityException;
 import com.uber.cadence.workflow.ActivityFailureException;
 import com.uber.cadence.workflow.ActivityTimeoutException;
@@ -88,13 +89,15 @@ final class SyncDecisionContext implements WorkflowInterceptor {
   private final WorkflowTimers timers = new WorkflowTimers();
   private final Map<String, Functions.Func1<byte[], byte[]>> queryCallbacks = new HashMap<>();
   private final byte[] lastCompletionResult;
+  private final WorkflowImplementationOptions workflowImplementationOptions;
 
   public SyncDecisionContext(
       DecisionContext context,
       DataConverter converter,
       List<ContextPropagator> contextPropagators,
       Function<WorkflowInterceptor, WorkflowInterceptor> interceptorFactory,
-      byte[] lastCompletionResult) {
+      byte[] lastCompletionResult,
+      WorkflowImplementationOptions workflowImplementationOptions) {
     this.context = context;
     this.converter = converter;
     this.contextPropagators = contextPropagators;
@@ -105,6 +108,7 @@ public SyncDecisionContext(
     }
     this.headInterceptor = interceptor;
     this.lastCompletionResult = lastCompletionResult;
+    this.workflowImplementationOptions = workflowImplementationOptions;
   }
 
   /**
@@ -764,4 +768,8 @@ public void upsertSearchAttributes(Map<String, Object> searchAttributes) {
     SearchAttributes attr = InternalUtils.convertMapToSearchAttributes(searchAttributes);
     context.upsertSearchAttributes(attr);
   }
+
+  public WorkflowImplementationOptions getWorkflowImplementationOptions() {
+    return workflowImplementationOptions;
+  }
 }

src/main/java/com/uber/cadence/internal/sync/SyncWorkflow.java

@@ -96,7 +96,8 @@ public void start(HistoryEvent event, DecisionContext context) {
             dataConverter,
             contextPropagators,
             interceptorFactory,
-            event.getWorkflowExecutionStartedEventAttributes().getLastCompletionResult());
+            event.getWorkflowExecutionStartedEventAttributes().getLastCompletionResult(),
+            workflowImplementationOptions);
 
     workflowProc =
         new WorkflowRunnable(

src/main/java/com/uber/cadence/internal/sync/WorkflowInternal.java

@@ -390,4 +390,8 @@ public static <R> R getLastCompletionResult(Class<R> resultClass, Type resultTyp
   public static void upsertSearchAttributes(Map<String, Object> searchAttributes) {
     getWorkflowInterceptor().upsertSearchAttributes(searchAttributes);
   }
+
+  public static WorkflowThread getRootWorkflowContext() {
+    return DeterministicRunnerImpl.currentThreadInternal();
+  }
 }

src/main/java/com/uber/cadence/worker/WorkflowImplementationOptions.java

@@ -19,13 +19,17 @@
 
 import static com.uber.cadence.worker.NonDeterministicWorkflowPolicy.BlockWorkflow;
 
+import com.uber.cadence.activity.ActivityOptions;
+import java.util.HashMap;
+import java.util.Map;
 import java.util.Objects;
 
 public final class WorkflowImplementationOptions {
 
   public static final class Builder {
 
     private NonDeterministicWorkflowPolicy nonDeterministicWorkflowPolicy = BlockWorkflow;
+    private Map<String, ActivityOptions> activityOptionOverrides = new HashMap<>();
 
     /**
      * Optional: Sets how decision worker deals with non-deterministic history events (presumably
@@ -39,27 +43,51 @@ public Builder setNonDeterministicWorkflowPolicy(
       return this;
     }
 
+    /**
+     * Set overrides for a specific workflow implementation for activity options.
+     *
+     * @param activityOptionOverrides a map where the key is the activity name and the value is the activity
+     *     options that should override existing activity configuration that comes from @ActivityMethod annotation.
+     */
+    public Builder setActivityOptionOverrides(Map<String, ActivityOptions> activityOptionOverrides) {
+      this.activityOptionOverrides = activityOptionOverrides;
+      return this;
+    }
+
     public WorkflowImplementationOptions build() {
-      return new WorkflowImplementationOptions(nonDeterministicWorkflowPolicy);
+      return new WorkflowImplementationOptions(nonDeterministicWorkflowPolicy, activityOptionOverrides);
     }
   }
 
   private final NonDeterministicWorkflowPolicy nonDeterministicWorkflowPolicy;
+  private Map<String, ActivityOptions> activityOptions;
 
   public WorkflowImplementationOptions(
       NonDeterministicWorkflowPolicy nonDeterministicWorkflowPolicy) {
     this.nonDeterministicWorkflowPolicy = nonDeterministicWorkflowPolicy;
   }
 
+  public WorkflowImplementationOptions(
+      NonDeterministicWorkflowPolicy nonDeterministicWorkflowPolicy,
+      Map<String, ActivityOptions> activityOptions) {
+    this.nonDeterministicWorkflowPolicy = nonDeterministicWorkflowPolicy;
+    this.activityOptions = activityOptions;
+  }
+
   public NonDeterministicWorkflowPolicy getNonDeterministicWorkflowPolicy() {
     return nonDeterministicWorkflowPolicy;
   }
 
+  public Map<String, ActivityOptions> getActivityOptions() {
+    return activityOptions;
+  }
+
   @Override
   public String toString() {
     return "WorkflowImplementationOptions{"
         + "nonDeterministicWorkflowPolicy="
         + nonDeterministicWorkflowPolicy
+        + activityOptions
         + '}';
   }
 
@@ -68,11 +96,12 @@ public boolean equals(Object o) {
     if (this == o) return true;
     if (o == null || getClass() != o.getClass()) return false;
     WorkflowImplementationOptions that = (WorkflowImplementationOptions) o;
-    return nonDeterministicWorkflowPolicy == that.nonDeterministicWorkflowPolicy;
+    return nonDeterministicWorkflowPolicy == that.nonDeterministicWorkflowPolicy
+        && Objects.equals(activityOptions, that.activityOptions);
   }
 
   @Override
   public int hashCode() {
-    return Objects.hash(nonDeterministicWorkflowPolicy);
+    return Objects.hash(nonDeterministicWorkflowPolicy, activityOptions);
   }
 }

src/test/java/com/uber/cadence/internal/sync/DeterministicRunnerTest.java

@@ -727,7 +727,7 @@ public void workflowThreadsWillEvictCacheWhenMaxThreadCountIsHit() throws Throwa
         new DeterministicRunnerImpl(
             threadPool,
             new SyncDecisionContext(
-                decisionContext, JsonDataConverter.getInstance(), null, next -> next, null),
+                decisionContext, JsonDataConverter.getInstance(), null, next -> next, null, null),
             () -> 0L, // clock override
             () -> {
               Promise<Void> thread =
@@ -752,7 +752,7 @@ public void workflowThreadsWillEvictCacheWhenMaxThreadCountIsHit() throws Throwa
         new DeterministicRunnerImpl(
             threadPool,
             new SyncDecisionContext(
-                decisionContext, JsonDataConverter.getInstance(), null, next -> next, null),
+                decisionContext, JsonDataConverter.getInstance(), null, next -> next, null, null),
             () -> 0L, // clock override
             () -> {
               Promise<Void> thread =

src/test/java/com/uber/cadence/internal/sync/SyncDecisionContextTest.java

@@ -38,7 +38,7 @@ public class SyncDecisionContextTest {
   public void setUp() {
     this.context =
         new SyncDecisionContext(
-            mockDecisionContext, JsonDataConverter.getInstance(), null, (next) -> next, null);
+            mockDecisionContext, JsonDataConverter.getInstance(), null, (next) -> next, null, null);
   }
 
   @Test

src/test/java/com/uber/cadence/workflow/WorkflowTest.java

@@ -358,6 +358,15 @@ private void startWorkerFor(Class<?>... workflowTypes) {
     }
   }
 
+  private void startWorkerFor(WorkflowImplementationOptions options, Class<?>... workflowTypes) {
+    worker.registerWorkflowImplementationTypes(options, workflowTypes);
+    if (useExternalService) {
+      workerFactory.start();
+    } else {
+      testEnvironment.start();
+    }
+  }
+
   // TODO: Refactor testEnvironment to support testing through real service to avoid this
   // conditional switches
   void registerDelayedCallback(Duration delay, Runnable r) {
@@ -460,6 +469,16 @@ public void signal1(String arg) {
     }
   }
 
+  public static class TestWorkflowActivityOptionOverride implements TestWorkflow1 {
+
+    @Override
+    public String execute(String taskList) {
+      TestActivities activities = Workflow.newActivityStub(TestActivities.class);
+
+      return activities.timeOutActivity();
+    }
+  }
+
   public static class TestSyncWorkflowImpl implements TestWorkflow1 {
 
     @Override
@@ -3971,6 +3990,9 @@ public interface TestActivities {
 
     void neverComplete();
 
+    @ActivityMethod(startToCloseTimeoutSeconds = 1, scheduleToCloseTimeoutSeconds = 1)
+    String timeOutActivity();
+
     @ActivityMethod(
       scheduleToStartTimeoutSeconds = 5,
       scheduleToCloseTimeoutSeconds = 5,
@@ -4051,6 +4073,7 @@ public String sleepActivity(long milliseconds, int input) {
 
     @Override
     public String activity() {
+
       invocations.add("activity");
       return "activity";
     }
@@ -4184,6 +4207,17 @@ public void neverComplete() {
       Activity.doNotCompleteOnReturn(); // Simulate activity timeout
     }
 
+    @Override
+    public String timeOutActivity() {
+      try {
+        Thread.sleep(2000);
+      } catch (InterruptedException e) {
+        return "activityWasInterrupted";
+      }
+
+      return "activityCompletedSuccessfully";
+    }
+
     @Override
     public void throwIOAnnotated() {
       invocations.add("throwIOAnnotated");
@@ -6280,4 +6314,40 @@ public void testGetVersionWithRetryReplay() throws Exception {
     WorkflowReplayer.replayWorkflowExecutionFromResource(
         "testGetVersionWithRetryHistory.json", TestGetVersionWorkflowRetryImpl.class);
   }
+
+  @Test(expected = ExecutionException.class)
+  public void testWorkflowTimesOutWhenNoOverridesProvided() throws Exception {
+    startWorkerFor(TestWorkflowActivityOptionOverride.class);
+
+    TestWorkflow1 workflowStub =
+        workflowClient.newWorkflowStub(
+            TestWorkflow1.class, newWorkflowOptionsBuilder(taskList).build());
+
+    CompletableFuture<String> result = WorkflowClient.execute(workflowStub::execute, taskList);
+    // Activity doesn't timeout because we overrode the timeouts to be longer.
+    assertEquals("activityCompletedSuccessfully", result.get());
+  }
+
+  @Test
+  public void testActivityDoesntTimeoutWhenTimeoutsAreLonger() throws Exception {
+    Map<String, ActivityOptions> activityOptionsMap = new HashMap<>();
+    activityOptionsMap.put(
+        "TestActivities::timeOutActivity",
+        new ActivityOptions.Builder()
+            .setScheduleToCloseTimeout(Duration.ofSeconds(3))
+            .setStartToCloseTimeout(Duration.ofSeconds(3))
+            .build());
+
+    startWorkerFor(
+        new WorkflowImplementationOptions.Builder().setActivityOptionOverrides(activityOptionsMap).build(),
+        TestWorkflowActivityOptionOverride.class);
+
+    TestWorkflow1 workflowStub =
+        workflowClient.newWorkflowStub(
+            TestWorkflow1.class, newWorkflowOptionsBuilder(taskList).build());
+
+    CompletableFuture<String> result = WorkflowClient.execute(workflowStub::execute, taskList);
+    // Activity doesn't timeout because we overrode the timeouts to be longer.
+    assertEquals("activityCompletedSuccessfully", result.get());
+  }
 }