[helix] branch master updated: Include pending messages in getCurrentInstanceToTaskAssignments method (#1390)

[al...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

alizamus pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/helix.git


The following commit(s) were added to refs/heads/master by this push:
     new 3c16cb3  Include pending messages in getCurrentInstanceToTaskAssignments method (#1390)
3c16cb3 is described below

commit 3c16cb3f8762eb6ff12ede041583adcf2af7a3b6
Author: Ali Reza Zamani Zadeh Najari <an...@linkedin.com>
AuthorDate: Thu Sep 24 17:05:28 2020 -0700

    Include pending messages in getCurrentInstanceToTaskAssignments method (#1390)
    
    This commit includes pending messages in getCurrentInstanceToTaskAssignments
    to calculate previously assigned tasks. It would help to avoid duplicate messages
    being sent to the participants. Also, it would avoid message with mismatch CurrentState
    being sent through the Task Framework pipeline.
---
 .../java/org/apache/helix/task/JobDispatcher.java  |  17 ++-
 ...nstanceToTaskAssignmentsWithPendingMessage.java | 114 +++++++++++++++++++++
 2 files changed, 130 insertions(+), 1 deletion(-)

diff --git a/helix-core/src/main/java/org/apache/helix/task/JobDispatcher.java b/helix-core/src/main/java/org/apache/helix/task/JobDispatcher.java
index 9a48339..e3d7ecd 100644
--- a/helix-core/src/main/java/org/apache/helix/task/JobDispatcher.java
+++ b/helix-core/src/main/java/org/apache/helix/task/JobDispatcher.java
@@ -381,7 +381,8 @@ public class JobDispatcher extends AbstractTaskDispatcher {
    *          are accounted for
    * @param jobName job name
    * @param tasksToDrop instance -> pId's, to gather all pIds that need to be dropped
-   * @return instance -> partitionIds from currentState, if the instance is still live
+   * @return instance -> partitionIds from currentState and pending messages, if the instance is
+   *         still live
    */
   protected static Map<String, SortedSet<Integer>> getCurrentInstanceToTaskAssignments(
       Iterable<String> liveInstances, CurrentStateOutput currStateOutput, String jobName,
@@ -419,6 +420,20 @@ public class JobDispatcher extends AbstractTaskDispatcher {
         }
       }
     }
+
+    Map<Partition, Map<String, Message>> pendingMessageMap =
+        currStateOutput.getPendingMessageMap(jobName);
+    for (Map.Entry<Partition, Map<String, Message>> pendingMessageMapEntry : pendingMessageMap
+        .entrySet()) {
+      for (Map.Entry<String, Message> instancePendingMessageEntry : pendingMessageMapEntry
+          .getValue().entrySet()) {
+        String instance = instancePendingMessageEntry.getKey();
+        int pId = TaskUtil.getPartitionId(pendingMessageMapEntry.getKey().getPartitionName());
+        if (result.containsKey(instance)) {
+          result.get(instance).add(pId);
+        }
+      }
+    }
     return result;
   }
 
diff --git a/helix-core/src/test/java/org/apache/helix/task/TestCurrentInstanceToTaskAssignmentsWithPendingMessage.java b/helix-core/src/test/java/org/apache/helix/task/TestCurrentInstanceToTaskAssignmentsWithPendingMessage.java
new file mode 100644
index 0000000..b8458ec
--- /dev/null
+++ b/helix-core/src/test/java/org/apache/helix/task/TestCurrentInstanceToTaskAssignmentsWithPendingMessage.java
@@ -0,0 +1,114 @@
+package org.apache.helix.task;
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+import static org.mockito.Mockito.*;
+
+import java.util.Arrays;
+import java.util.Collection;
+import java.util.HashMap;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Map;
+import java.util.Random;
+import java.util.Set;
+
+import java.util.SortedSet;
+import org.apache.helix.controller.stages.CurrentStateOutput;
+import org.apache.helix.model.Message;
+import org.apache.helix.model.Partition;
+import org.testng.Assert;
+import org.testng.annotations.Test;
+
+public class TestCurrentInstanceToTaskAssignmentsWithPendingMessage {
+
+  /**
+   * This is a unit test that tests that all task partitions with currentState or pending message is
+   * added to the result of getCurrentInstanceToTaskAssignments method.
+   */
+  @Test
+  public void testCurrentInstanceToTaskAssignmentsWithPendingMessage() {
+    Random random = new Random();
+    String jobName = "job";
+    String nodeName = "localhost";
+    int numTasks = 100;
+
+    // Create an Iterable of LiveInstances
+    Collection<String> liveInstances = new HashSet<>();
+    liveInstances.add("localhost");
+
+    // Create allTaskPartitions
+    Set<Integer> allTaskPartitions = new HashSet<>();
+
+    // Create a mock CurrentStateOutput
+    CurrentStateOutput currentStateOutput = mock(CurrentStateOutput.class);
+
+    // Generate a CurrentStateMap and PendingMessageMap
+    Map<Partition, Map<String, String>> currentStateMap = new HashMap<>();
+    Map<Partition, Map<String, Message>> pendingMessageMap = new HashMap<>();
+
+    int tasksWithCurrentStateOnly = 0;
+    int tasksWithCurrentStateAndPendingMessage = 0;
+    int tasksWithPendingMessageOnly = 0;
+
+    List<String> states =
+        Arrays.asList(TaskPartitionState.INIT.name(), TaskPartitionState.RUNNING.name(),
+            TaskPartitionState.TIMED_OUT.name(), TaskPartitionState.TASK_ERROR.name(),
+            TaskPartitionState.COMPLETED.name(), TaskPartitionState.STOPPED.name(),
+            TaskPartitionState.TASK_ABORTED.name(), TaskPartitionState.DROPPED.name());
+
+    for (int i = 0; i < numTasks; i++) {
+      allTaskPartitions.add(i);
+      Partition task = new Partition(jobName + "_" + i);
+      currentStateMap.put(task, new HashMap<>());
+      pendingMessageMap.put(task, new HashMap<>());
+
+      String currentState = states.get(random.nextInt(states.size()));
+      Message message = new Message(Message.MessageType.STATE_TRANSITION, "12345");
+      message.setToState(states.get(random.nextInt(states.size())));
+      message.setFromState(states.get(random.nextInt(states.size())));
+
+      int randInt = random.nextInt(4);
+      if (randInt == 0) {
+        tasksWithCurrentStateOnly = tasksWithCurrentStateOnly + 1;
+        currentStateMap.get(task).put(nodeName, currentState);
+      } else if (randInt == 1) {
+        tasksWithCurrentStateAndPendingMessage = tasksWithCurrentStateAndPendingMessage + 1;
+        currentStateMap.get(task).put(nodeName, currentState);
+        pendingMessageMap.get(task).put(nodeName, message);
+      } else if (randInt == 2) {
+        tasksWithPendingMessageOnly = tasksWithPendingMessageOnly + 1;
+        pendingMessageMap.get(task).put(nodeName, message);
+      }
+    }
+
+    when(currentStateOutput.getCurrentStateMap(jobName)).thenReturn(currentStateMap);
+    when(currentStateOutput.getPendingMessageMap(jobName)).thenReturn(pendingMessageMap);
+
+    // Create an empty tasksToDrop
+    Map<String, Set<Integer>> tasksToDrop = new HashMap<>();
+
+    // Call the static method we are testing
+    Map<String, SortedSet<Integer>> result = JobDispatcher.getCurrentInstanceToTaskAssignments(
+        liveInstances, currentStateOutput, jobName, tasksToDrop);
+    Assert.assertEquals(result.get(nodeName).size(), (tasksWithCurrentStateOnly
+        + tasksWithCurrentStateAndPendingMessage + tasksWithPendingMessageOnly));
+  }
+}