[DS-6961][MasterServer] batch dispatch (#6962)

caishunfeng · caishunfeng · web-flow · commit e6fe39ea47f3 · 2021-11-25T12:55:37.000+08:00
* [DS-6961][MasterServer] batch dispatch

* fix test

Co-authored-by: caishunfeng &lt;534328519@qq.com&gt;
diff --git a/dolphinscheduler-server/src/main/java/org/apache/dolphinscheduler/server/master/config/MasterConfig.java b/dolphinscheduler-server/src/main/java/org/apache/dolphinscheduler/server/master/config/MasterConfig.java
@@ -31,7 +31,6 @@ public class MasterConfig {
     private int fetchCommandNum;
     private int preExecThreads;
     private int execThreads;
-    private int execTaskNum;
     private int dispatchTaskNumber;
     private HostSelector hostSelector;
     private int heartbeatInterval;
@@ -74,14 +73,6 @@ public void setExecThreads(int execThreads) {
         this.execThreads = execThreads;
     }
 
-    public int getExecTaskNum() {
-        return execTaskNum;
-    }
-
-    public void setExecTaskNum(int execTaskNum) {
-        this.execTaskNum = execTaskNum;
-    }
-
     public int getDispatchTaskNumber() {
         return dispatchTaskNumber;
     }
diff --git a/dolphinscheduler-server/src/main/java/org/apache/dolphinscheduler/server/master/consumer/TaskPriorityQueueConsumer.java b/dolphinscheduler-server/src/main/java/org/apache/dolphinscheduler/server/master/consumer/TaskPriorityQueueConsumer.java
@@ -19,12 +19,14 @@
 
 import org.apache.dolphinscheduler.common.Constants;
 import org.apache.dolphinscheduler.common.thread.Stopper;
+import org.apache.dolphinscheduler.common.thread.ThreadUtils;
 import org.apache.dolphinscheduler.dao.entity.TaskInstance;
 import org.apache.dolphinscheduler.server.master.config.MasterConfig;
 import org.apache.dolphinscheduler.server.master.dispatch.ExecutorDispatcher;
 import org.apache.dolphinscheduler.server.master.dispatch.context.ExecutionContext;
 import org.apache.dolphinscheduler.server.master.dispatch.enums.ExecutorType;
 import org.apache.dolphinscheduler.server.master.dispatch.exceptions.ExecuteException;
+import org.apache.dolphinscheduler.service.exceptions.TaskPriorityQueueException;
 import org.apache.dolphinscheduler.service.process.ProcessService;
 import org.apache.dolphinscheduler.service.queue.TaskPriority;
 import org.apache.dolphinscheduler.service.queue.TaskPriorityQueue;
@@ -33,6 +35,8 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
 
 import javax.annotation.PostConstruct;
@@ -78,30 +82,24 @@ public class TaskPriorityQueueConsumer extends Thread {
     @Autowired
     private MasterConfig masterConfig;
 
+    /**
+     * consumer thread pool
+     */
+    private ThreadPoolExecutor consumerThreadPoolExecutor;
+
     @PostConstruct
     public void init() {
-        super.setName("TaskUpdateQueueConsumerThread");
+        this.consumerThreadPoolExecutor = (ThreadPoolExecutor) ThreadUtils.newDaemonFixedThreadExecutor("TaskUpdateQueueConsumerThread", masterConfig.getDispatchTaskNumber());
         super.start();
     }
 
     @Override
     public void run() {
-        List<TaskPriority> failedDispatchTasks = new ArrayList<>();
+        int fetchTaskNum = masterConfig.getDispatchTaskNumber();
         while (Stopper.isRunning()) {
             try {
-                int fetchTaskNum = masterConfig.getDispatchTaskNumber();
-                failedDispatchTasks.clear();
-                for (int i = 0; i < fetchTaskNum; i++) {
-                    TaskPriority taskPriority = taskPriorityQueue.poll(Constants.SLEEP_TIME_MILLIS, TimeUnit.MILLISECONDS);
-                    if (Objects.isNull(taskPriority)) {
-                        continue;
-                    }
+                List<TaskPriority> failedDispatchTasks = this.batchDispatch(fetchTaskNum);
 
-                    boolean dispatchResult = dispatch(taskPriority);
-                    if (!dispatchResult) {
-                        failedDispatchTasks.add(taskPriority);
-                    }
-                }
                 if (!failedDispatchTasks.isEmpty()) {
                     for (TaskPriority dispatchFailedTask : failedDispatchTasks) {
                         taskPriorityQueue.put(dispatchFailedTask);
@@ -118,13 +116,41 @@ public void run() {
         }
     }
 
+    /**
+     * batch dispatch with thread pool
+     */
+    private List<TaskPriority> batchDispatch(int fetchTaskNum) throws TaskPriorityQueueException, InterruptedException {
+        List<TaskPriority> failedDispatchTasks = new ArrayList<>();
+        CountDownLatch latch = new CountDownLatch(fetchTaskNum);
+
+        for (int i = 0; i < fetchTaskNum; i++) {
+            TaskPriority taskPriority = taskPriorityQueue.poll(Constants.SLEEP_TIME_MILLIS, TimeUnit.MILLISECONDS);
+            if (Objects.isNull(taskPriority)) {
+                latch.countDown();
+                continue;
+            }
+
+            consumerThreadPoolExecutor.submit(() -> {
+                boolean dispatchResult = this.dispatchTask(taskPriority);
+                if (!dispatchResult) {
+                    failedDispatchTasks.add(taskPriority);
+                }
+                latch.countDown();
+            });
+        }
+
+        latch.await();
+
+        return failedDispatchTasks;
+    }
+
     /**
      * dispatch task
      *
      * @param taskPriority taskPriority
      * @return result
      */
-    protected boolean dispatch(TaskPriority taskPriority) {
+    protected boolean dispatchTask(TaskPriority taskPriority) {
         boolean result = false;
         try {
             TaskExecutionContext context = taskPriority.getTaskExecutionContext();
@@ -158,8 +184,6 @@ public Boolean taskInstanceIsFinalState(int taskInstanceId) {
 
     /**
      * check if task need to check state, if true, refresh the checkpoint
-     * @param taskPriority
-     * @return
      */
     private boolean isTaskNeedToCheck(TaskPriority taskPriority) {
         long now = System.currentTimeMillis();
diff --git a/dolphinscheduler-server/src/test/java/org/apache/dolphinscheduler/server/master/WorkflowExecuteThreadTest.java b/dolphinscheduler-server/src/test/java/org/apache/dolphinscheduler/server/master/WorkflowExecuteThreadTest.java
@@ -88,7 +88,6 @@ public void init() throws Exception {
 
         applicationContext = mock(ApplicationContext.class);
         config = new MasterConfig();
-        config.setExecTaskNum(1);
         Mockito.when(applicationContext.getBean(MasterConfig.class)).thenReturn(config);
 
         processInstance = mock(ProcessInstance.class);
diff --git a/dolphinscheduler-server/src/test/java/org/apache/dolphinscheduler/server/master/consumer/TaskPriorityQueueConsumerTest.java b/dolphinscheduler-server/src/test/java/org/apache/dolphinscheduler/server/master/consumer/TaskPriorityQueueConsumerTest.java
@@ -307,7 +307,7 @@ public void testDispatch() {
 
         TaskPriority taskPriority = new TaskPriority();
         taskPriority.setTaskId(1);
-        boolean res = taskPriorityQueueConsumer.dispatch(taskPriority);
+        boolean res = taskPriorityQueueConsumer.dispatchTask(taskPriority);
 
         Assert.assertFalse(res);
     }

Original file line number	Diff line number	Diff line change
`@@ -307,7 +307,7 @@ public void testDispatch() {`
`307`	`307`
`308`	`308`	`TaskPriority taskPriority = new TaskPriority();`
`309`	`309`	`taskPriority.setTaskId(1);`
`310`		`- boolean res = taskPriorityQueueConsumer.dispatch(taskPriority);`
	`310`	`+ boolean res = taskPriorityQueueConsumer.dispatchTask(taskPriority);`
`311`	`311`
`312`	`312`	`Assert.assertFalse(res);`
`313`	`313`	`}`