From f25c19664ab4d518a81c30853d8540f9a60ba565 Mon Sep 17 00:00:00 2001
From: Maros Marsalek <mmarsalek@frinx.io>
Date: Wed, 28 Oct 2020 11:14:25 +0100
Subject: [PATCH] Preserve task domains for tasks inside do_while loop

On every iteration, apply task domains as set in the workflow execution.
This feature was missing and made do wile behave incorrectly when using
task domains (since the domains were reset for tasks inside the loop)

Signed-off-by: Maros Marsalek <mmarsalek@frinx.io>
---
 .../core/execution/WorkflowExecutor.java      |  1 +
 .../core/execution/TestWorkflowExecutor.java  | 23 +++++++++++++++++++
 2 files changed, 24 insertions(+)

diff --git a/core/src/main/java/com/netflix/conductor/core/execution/WorkflowExecutor.java b/core/src/main/java/com/netflix/conductor/core/execution/WorkflowExecutor.java
index 4012c1e57b..d42101a5c1 100644
--- a/core/src/main/java/com/netflix/conductor/core/execution/WorkflowExecutor.java
+++ b/core/src/main/java/com/netflix/conductor/core/execution/WorkflowExecutor.java
@@ -1638,6 +1638,7 @@ private boolean rerunWF(String workflowId, String taskId, Map<String, Object> ta
     public void scheduleNextIteration(Task loopTask, Workflow workflow) {
         //Schedule only first loop over task. Rest will be taken care in Decider Service when this task will get completed.
         List<Task> scheduledLoopOverTasks = deciderService.getTasksToBeScheduled(workflow, loopTask.getWorkflowTask().getLoopOver().get(0), loopTask.getRetryCount(), null);
+        setTaskDomains(scheduledLoopOverTasks, workflow);
         scheduledLoopOverTasks.stream().forEach(t -> {
             t.setReferenceTaskName(TaskUtils.appendIteration(t.getReferenceTaskName(), loopTask.getIteration()));
             t.setIteration(loopTask.getIteration());
diff --git a/core/src/test/java/com/netflix/conductor/core/execution/TestWorkflowExecutor.java b/core/src/test/java/com/netflix/conductor/core/execution/TestWorkflowExecutor.java
index d78d8eb7ab..7f9e1bffab 100644
--- a/core/src/test/java/com/netflix/conductor/core/execution/TestWorkflowExecutor.java
+++ b/core/src/test/java/com/netflix/conductor/core/execution/TestWorkflowExecutor.java
@@ -1469,6 +1469,29 @@ public void testStartWorkflow() {
         verify(executionDAOFacade, times(1)).createWorkflow(any(Workflow.class));
     }
 
+    @Test
+    public void testScheduleNextIteration() {
+        Workflow workflow = generateSampleWorkflow();
+        workflow.setTaskToDomain(new HashMap<String, String>() {{
+            put("TEST", "domain1");
+        }});
+        Task loopTask = mock(Task.class);
+        WorkflowTask loopWfTask = mock(WorkflowTask.class);
+        when(loopTask.getWorkflowTask()).thenReturn(loopWfTask);
+        List<WorkflowTask> loopOver = new ArrayList<WorkflowTask>(){{
+            WorkflowTask e = new WorkflowTask();
+            e.setType(TaskType.TASK_TYPE_SIMPLE);
+            e.setName("TEST");
+            e.setTaskDefinition(new TaskDef());
+            add(e);
+        }};
+        when(loopWfTask.getLoopOver()).thenReturn(loopOver);
+
+        workflowExecutor.scheduleNextIteration(loopTask, workflow);
+
+        verify(executionDAOFacade).getTaskPollDataByDomain("TEST", "domain1");
+    }
+
     private Workflow generateSampleWorkflow() {
         //setup
         Workflow workflow = new Workflow();