[Improvement-16979] Unify PriorityDelayQueue with AbstractDelayEventBus (#17155)

Author: reele

dolphinscheduler-eventbus/src/main/java/org/apache/dolphinscheduler/eventbus/AbstractDelayEventBus.java

@@ -42,6 +42,11 @@ public Optional<T> peek() {
         return Optional.ofNullable(delayEventQueue.peek());
     }
 
+    @Override
+    public T take() throws InterruptedException {
+        return delayEventQueue.take();
+    }
+
     @Override
     public Optional<T> remove() {
         return Optional.ofNullable(delayEventQueue.remove());

dolphinscheduler-eventbus/src/main/java/org/apache/dolphinscheduler/eventbus/IEventBus.java

@@ -54,6 +54,12 @@ public interface IEventBus<T extends IEvent> {
      */
     Optional<T> poll() throws InterruptedException;
 
+    /**
+     * Remove the head event from the bus. This method will block if the event bus is empty.
+     * <p> If the thread is interrupted, an {@link InterruptedException} will be thrown.
+     */
+    T take() throws InterruptedException;
+
     /**
      * peek the head event from the bus. This method will not block if the event bus is empty will return empty optional.
      */

…ter/runner/queue/PriorityDelayQueue.java …dispatcher/TaskDispatchableEventBus.javadolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/runner/queue/PriorityDelayQueue.java renamed to dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/task/dispatcher/TaskDispatchableEventBus.java dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/runner/queue/PriorityDelayQueue.java renamed to dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/task/dispatcher/TaskDispatchableEventBus.java

@@ -15,30 +15,33 @@
  * limitations under the License.
  */
 
-package org.apache.dolphinscheduler.server.master.runner.queue;
+package org.apache.dolphinscheduler.server.master.engine.task.dispatcher;
 
-import java.util.concurrent.DelayQueue;
+import org.apache.dolphinscheduler.eventbus.AbstractDelayEventBus;
+import org.apache.dolphinscheduler.server.master.engine.task.dispatcher.event.TaskDispatchableEvent;
 
 import lombok.SneakyThrows;
 
-public class PriorityDelayQueue<V extends DelayEntry> {
-
-    private final DelayQueue<V> queue = new DelayQueue<>();
+public class TaskDispatchableEventBus<V extends TaskDispatchableEvent<T>, T extends Comparable<T>>
+        extends
+            AbstractDelayEventBus<V> {
 
     public void add(V v) {
-        queue.put(v);
+        super.publish(v);
     }
 
     @SneakyThrows
     public V take() {
-        return queue.take();
+        return super.take();
     }
 
+    // Only use in test
     public int size() {
-        return queue.size();
+        return delayEventQueue.size();
     }
 
+    // Only use in test
     public void clear() {
-        queue.clear();
+        delayEventQueue.clear();
     }
 }

dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/task/dispatcher/WorkerGroupDispatcher.java

@@ -20,9 +20,8 @@
 import org.apache.dolphinscheduler.common.thread.BaseDaemonThread;
 import org.apache.dolphinscheduler.plugin.task.api.utils.LogUtils;
 import org.apache.dolphinscheduler.server.master.engine.task.client.ITaskExecutorClient;
+import org.apache.dolphinscheduler.server.master.engine.task.dispatcher.event.TaskDispatchableEvent;
 import org.apache.dolphinscheduler.server.master.engine.task.runnable.ITaskExecutionRunnable;
-import org.apache.dolphinscheduler.server.master.runner.queue.PriorityAndDelayBasedTaskEntry;
-import org.apache.dolphinscheduler.server.master.runner.queue.PriorityDelayQueue;
 import org.apache.dolphinscheduler.task.executor.log.TaskExecutorMDCUtils;
 
 import java.util.Set;
@@ -34,7 +33,7 @@
 /**
  * WorkerGroupTaskDispatcher is responsible for dispatching tasks from the task queue.
  * The main responsibilities include:
- * 1. Continuously fetching tasks from the {@link PriorityDelayQueue} for dispatch.
+ * 1. Continuously fetching tasks from the {@link TaskDispatchableEvent} for dispatch.
  * 2. Re-queuing tasks that fail to dispatch according to retry logic.
  * 3. Ensuring thread safety and correct state transitions during task processing.
  */
@@ -43,11 +42,7 @@ public class WorkerGroupDispatcher extends BaseDaemonThread {
 
     private final ITaskExecutorClient taskExecutorClient;
 
-    // TODO The current queue is flawed. When a high-priority task fails,
-    // it will be delayed and will not return to the first or second position.
-    // Tasks with the same priority will preempt its position.
-    // If it needs to be placed at the front of the queue, the queue needs to be re-implemented.
-    private final PriorityDelayQueue<PriorityAndDelayBasedTaskEntry<ITaskExecutionRunnable>> workerGroupQueue;
+    private final TaskDispatchableEventBus<TaskDispatchableEvent<ITaskExecutionRunnable>, ITaskExecutionRunnable> workerGroupEventBus;
 
     private final Set<Integer> waitingDispatchTaskIds;
 
@@ -56,7 +51,7 @@ public class WorkerGroupDispatcher extends BaseDaemonThread {
     public WorkerGroupDispatcher(String workerGroupName, ITaskExecutorClient taskExecutorClient) {
         super("WorkerGroupTaskDispatcher-" + workerGroupName);
         this.taskExecutorClient = taskExecutorClient;
-        this.workerGroupQueue = new PriorityDelayQueue<>();
+        this.workerGroupEventBus = new TaskDispatchableEventBus<>();
         this.waitingDispatchTaskIds = ConcurrentHashMap.newKeySet();
         log.info("Initialize WorkerGroupDispatcher: {}", this.getName());
     }
@@ -75,7 +70,7 @@ public synchronized void start() {
     @Override
     public void run() {
         while (runningFlag.get()) {
-            PriorityAndDelayBasedTaskEntry<ITaskExecutionRunnable> taskEntry = workerGroupQueue.take();
+            TaskDispatchableEvent<ITaskExecutionRunnable> taskEntry = workerGroupEventBus.take();
             ITaskExecutionRunnable taskExecutionRunnable = taskEntry.getData();
             try (
                     TaskExecutorMDCUtils.MDCAutoClosable ignore =
@@ -120,7 +115,7 @@ private void doDispatchTask(ITaskExecutionRunnable taskExecutionRunnable) {
      */
     public void dispatchTask(final ITaskExecutionRunnable taskExecutionRunnable, final long delayTimeMills) {
         waitingDispatchTaskIds.add(taskExecutionRunnable.getId());
-        workerGroupQueue.add(new PriorityAndDelayBasedTaskEntry<>(delayTimeMills, taskExecutionRunnable));
+        workerGroupEventBus.add(new TaskDispatchableEvent<>(delayTimeMills, taskExecutionRunnable));
     }
 
     public boolean removeTask(ITaskExecutionRunnable taskExecutionRunnable) {
@@ -141,6 +136,6 @@ public synchronized void close() {
     }
 
     int queueSize() {
-        return this.workerGroupQueue.size();
+        return this.workerGroupEventBus.size();
     }
 }

dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/task/dispatcher/event/TaskDispatchableEvent.java

@@ -0,0 +1,57 @@
+/*
+ * 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.
+ */
+
+package org.apache.dolphinscheduler.server.master.engine.task.dispatcher.event;
+
+import static com.google.common.base.Preconditions.checkNotNull;
+
+import org.apache.dolphinscheduler.eventbus.AbstractDelayEvent;
+
+import java.util.concurrent.Delayed;
+
+import lombok.Getter;
+
+@Getter
+public class TaskDispatchableEvent<V extends Comparable<V>> extends AbstractDelayEvent {
+
+    protected final V data;
+
+    public TaskDispatchableEvent(long delayTimeMills, V data) {
+        super(delayTimeMills);
+        this.data = checkNotNull(data, "data is null");
+    }
+
+    @Override
+    public int compareTo(Delayed other) {
+        if (!(other instanceof TaskDispatchableEvent)) {
+            throw new RuntimeException("The object being compared is not a TaskReadyForDispatchEvent.");
+        }
+
+        @SuppressWarnings("unchecked")
+        final TaskDispatchableEvent<V> otherEvent = (TaskDispatchableEvent<V>) other;
+
+        // there should compare data first for priority
+        if (data != null && otherEvent.data != null) {
+            final int compareResult = data.compareTo(otherEvent.data);
+            if (compareResult != 0) {
+                return compareResult;
+            }
+        }
+
+        return super.compareTo(other);
+    }
+}