Jackie-Jiang commented on code in PR #11315:
URL: https://github.com/apache/pinot/pull/11315#discussion_r1290771989


##########
pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/minion/PinotHelixTaskResourceManager.java:
##########
@@ -326,27 +324,35 @@ public synchronized Set<String> getTasks(String taskType) 
{
 
   /**
    * Get all task states for the given task type.
+   * NOTE: For tasks just submitted without the context created, count them as 
NOT_STARTED.
    *
    * @param taskType Task type
    * @return Map from task name to task state
    */
   public synchronized Map<String, TaskState> getTaskStates(String taskType) {
-    Map<String, TaskState> helixJobStates = new HashMap<>();
+    WorkflowConfig workflowConfig = 
_taskDriver.getWorkflowConfig(getHelixJobQueueName(taskType));
+    if (workflowConfig == null) {
+      return Collections.emptyMap();
+    }
+    Set<String> helixJobs = workflowConfig.getJobDag().getAllNodes();
+    if (helixJobs.isEmpty()) {
+      return Collections.emptyMap();
+    }
     WorkflowContext workflowContext = 
_taskDriver.getWorkflowContext(getHelixJobQueueName(taskType));
-
     if (workflowContext == null) {
-      return helixJobStates;
-    }
-    helixJobStates = workflowContext.getJobStates();
-    Map<String, TaskState> taskStates = new HashMap<>(helixJobStates.size());
-    for (Map.Entry<String, TaskState> entry : helixJobStates.entrySet()) {
-      taskStates.put(getPinotTaskName(entry.getKey()), entry.getValue());
+      return helixJobs.stream()
+          
.collect(Collectors.toMap(PinotHelixTaskResourceManager::getPinotTaskName, 
ignored -> TaskState.NOT_STARTED));
+    } else {
+      Map<String, TaskState> helixJobStates = workflowContext.getJobStates();
+      return 
helixJobs.stream().collect(Collectors.toMap(PinotHelixTaskResourceManager::getPinotTaskName,
+          helixJobName -> helixJobStates.getOrDefault(helixJobName, 
TaskState.NOT_STARTED)));
     }
-    return taskStates;
   }
 
   /**
    * This method returns a count of sub-tasks in various states, given the 
top-level task name.
+   * TODO: It doesn't count tasks just submitted without the context created.

Review Comment:
   Good point. Let me fix all of them



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]

Reply via email to