add test for dispatch timeout checker

Author: 苏义超

dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/task/dispatcher/WorkerGroupDispatcherTest.java

@@ -19,22 +19,37 @@
 
 import static org.awaitility.Awaitility.await;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.Mockito.argThat;
+import static org.mockito.Mockito.atLeast;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.doThrow;
 import static org.mockito.Mockito.inOrder;
 import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.never;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import org.apache.dolphinscheduler.dao.entity.TaskInstance;
+import org.apache.dolphinscheduler.dao.entity.WorkflowInstance;
 import org.apache.dolphinscheduler.plugin.task.api.TaskExecutionContext;
 import org.apache.dolphinscheduler.server.master.config.MasterConfig;
+import org.apache.dolphinscheduler.server.master.config.MasterDispatchTimeoutCheckerConfig;
+import org.apache.dolphinscheduler.server.master.engine.WorkflowEventBus;
 import org.apache.dolphinscheduler.server.master.engine.task.client.ITaskExecutorClient;
+import org.apache.dolphinscheduler.server.master.engine.task.lifecycle.event.TaskFailedLifecycleEvent;
+import org.apache.dolphinscheduler.server.master.engine.task.lifecycle.event.TaskFatalLifecycleEvent;
 import org.apache.dolphinscheduler.server.master.engine.task.runnable.ITaskExecutionRunnable;
+import org.apache.dolphinscheduler.server.master.exception.dispatch.NoAvailableWorkerException;
 import org.apache.dolphinscheduler.server.master.exception.dispatch.TaskDispatchException;
+import org.apache.dolphinscheduler.server.master.exception.dispatch.WorkerGroupNotFoundException;
 
 import java.time.Duration;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ThreadLocalRandom;
+import java.util.concurrent.TimeUnit;
 
+import org.junit.jupiter.api.Assertions;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 import org.mockito.InOrder;
@@ -141,4 +156,285 @@ void dispatch_TaskDispatchFails_RetryLogicWorks() throws TaskDispatchException {
                 .untilAsserted(() -> verify(taskExecutorClient, times(2)).dispatch(taskExecutionRunnable));
     }
 
+    @Test
+    void dispatchTask_WorkerGroupNotFound_TimeoutDisabled_ShouldKeepRetrying() throws TaskDispatchException {
+        // Given
+        ITaskExecutionRunnable task = mockTaskExecutionRunnableWithFirstDispatchTime(System.currentTimeMillis());
+        WorkerGroupNotFoundException ex = new WorkerGroupNotFoundException("no worker group");
+        doThrow(ex).when(taskExecutorClient).dispatch(task);
+
+        dispatcher.start();
+        dispatcher.dispatchTask(task, 0);
+
+        // When & Then
+        await().atMost(Duration.ofSeconds(3))
+                .untilAsserted(() -> {
+                    // Ensure it's retrying
+                    verify(taskExecutorClient, atLeast(2)).dispatch(task);
+
+                    // Ensure NO event has been published during this time
+                    WorkflowEventBus eventBus = task.getWorkflowEventBus();
+                    verify(eventBus, never()).publish(any(TaskFailedLifecycleEvent.class));
+                    verify(eventBus, never()).publish(any(TaskFatalLifecycleEvent.class));
+                });
+    }
+
+    @Test
+    void dispatchTask_WorkerGroupNotFound_TimeoutEnabledAndExceeded_ShouldPublishFailedEvent() throws TaskDispatchException {
+        // Given
+        MasterDispatchTimeoutCheckerConfig dispatchTimeoutCheckerConfig = new MasterDispatchTimeoutCheckerConfig();
+        dispatchTimeoutCheckerConfig.setEnabled(true);
+        dispatchTimeoutCheckerConfig.setMaxTaskDispatchDuration(Duration.ofMillis(200));
+
+        dispatcher = new WorkerGroupDispatcher("TestGroup", taskExecutorClient, dispatchTimeoutCheckerConfig);
+
+        ITaskExecutionRunnable taskExecutionRunnable = mockTaskExecutionRunnableWithFirstDispatchTime(
+                System.currentTimeMillis() - 500);
+
+        WorkerGroupNotFoundException ex = new WorkerGroupNotFoundException("worker group not found");
+        doThrow(ex).when(taskExecutorClient).dispatch(taskExecutionRunnable);
+
+        dispatcher.start();
+        dispatcher.dispatchTask(taskExecutionRunnable, 0);
+
+        // Then
+        await().atMost(Duration.ofSeconds(2))
+                .untilAsserted(() -> {
+                    verify(taskExecutorClient, times(1)).dispatch(taskExecutionRunnable);
+                    WorkflowEventBus eventBus = taskExecutionRunnable.getWorkflowEventBus();
+                    verify(eventBus).publish(argThat(evt -> evt instanceof TaskFatalLifecycleEvent &&
+                            ((TaskFatalLifecycleEvent) evt).getTaskExecutionRunnable() == taskExecutionRunnable));
+                    verify(eventBus, never()).publish(any(TaskFailedLifecycleEvent.class));
+                });
+    }
+
+    @Test
+    void dispatchTask_WorkerGroupNotFound_TimeoutEnabledButNotExceeded_ShouldNotPublishAnyFailureEvent() throws TaskDispatchException, InterruptedException {
+        // Given: Dispatcher configured with a 5-minute timeout (enabled)
+        MasterDispatchTimeoutCheckerConfig dispatchTimeoutCheckerConfig = new MasterDispatchTimeoutCheckerConfig();
+        dispatchTimeoutCheckerConfig.setEnabled(true);
+        dispatchTimeoutCheckerConfig.setMaxTaskDispatchDuration(Duration.ofMinutes(5));
+
+        dispatcher = new WorkerGroupDispatcher("TestGroup", taskExecutorClient, dispatchTimeoutCheckerConfig);
+
+        // Mock task with first dispatch time set to 100ms ago → well within timeout window
+        ITaskExecutionRunnable taskExecutionRunnable = mockTaskExecutionRunnableWithFirstDispatchTime(
+                System.currentTimeMillis() - 100);
+
+        // Use CountDownLatch to reliably detect actual dispatch invocation
+        CountDownLatch dispatchCalled = new CountDownLatch(1);
+
+        // Stub client to throw WorkerGroupNotFoundException and signal the latch
+        doAnswer(invocation -> {
+            dispatchCalled.countDown(); // Confirm dispatch was attempted
+            throw new WorkerGroupNotFoundException("Worker group 'TestGroup' does not exist");
+        }).when(taskExecutorClient).dispatch(taskExecutionRunnable);
+
+        // When: Start dispatcher and dispatch the task
+        dispatcher.start();
+        dispatcher.dispatchTask(taskExecutionRunnable, 0);
+
+        // Wait up to 1 second for the dispatch attempt to complete
+        boolean dispatched = dispatchCalled.await(1000, TimeUnit.MILLISECONDS);
+        Assertions.assertTrue(dispatched, "Expected dispatch() to be called within 1 second");
+
+        // Then: Verify NO failure events are published because timeout has NOT been exceeded
+        WorkflowEventBus eventBus = taskExecutionRunnable.getWorkflowEventBus();
+        verify(eventBus, never()).publish(any(TaskFailedLifecycleEvent.class));
+        verify(eventBus, never()).publish(any(TaskFatalLifecycleEvent.class));
+    }
+
+    @Test
+    void dispatchTask_NoAvailableWorker_TimeoutDisabled_ShouldKeepRetrying() throws TaskDispatchException {
+        // Given
+        ITaskExecutionRunnable task = mockTaskExecutionRunnableWithFirstDispatchTime(System.currentTimeMillis());
+        NoAvailableWorkerException ex = new NoAvailableWorkerException("no worker");
+        doThrow(ex).when(taskExecutorClient).dispatch(task);
+
+        dispatcher.start();
+        dispatcher.dispatchTask(task, 0);
+
+        // When & Then
+        await().atMost(Duration.ofSeconds(3))
+                .untilAsserted(() -> {
+                    // Ensure it's retrying
+                    verify(taskExecutorClient, atLeast(2)).dispatch(task);
+
+                    // Ensure NO event has been published during this time
+                    WorkflowEventBus eventBus = task.getWorkflowEventBus();
+                    verify(eventBus, never()).publish(any(TaskFailedLifecycleEvent.class));
+                    verify(eventBus, never()).publish(any(TaskFatalLifecycleEvent.class));
+                });
+    }
+
+    @Test
+    void dispatchTask_NoAvailableWorker_TimeoutEnabledAndExceeded_ShouldPublishFailedEvent() throws TaskDispatchException {
+        // Given: enable timeout (200ms), task already waited 500ms
+        MasterDispatchTimeoutCheckerConfig dispatchTimeoutCheckerConfig = new MasterDispatchTimeoutCheckerConfig();
+        dispatchTimeoutCheckerConfig.setEnabled(true);
+        dispatchTimeoutCheckerConfig.setMaxTaskDispatchDuration(Duration.ofMillis(200));
+
+        dispatcher = new WorkerGroupDispatcher("TestGroup", taskExecutorClient, dispatchTimeoutCheckerConfig);
+
+        ITaskExecutionRunnable taskExecutionRunnable =
+                mockTaskExecutionRunnableWithFirstDispatchTime(System.currentTimeMillis() - 500);
+
+        NoAvailableWorkerException ex = new NoAvailableWorkerException("no worker");
+        doThrow(ex).when(taskExecutorClient).dispatch(taskExecutionRunnable);
+
+        dispatcher.start();
+        dispatcher.dispatchTask(taskExecutionRunnable, 0);
+
+        // Then
+        await().atMost(Duration.ofSeconds(2))
+                .untilAsserted(() -> {
+                    verify(taskExecutorClient, times(1)).dispatch(taskExecutionRunnable);
+                    WorkflowEventBus eventBus = taskExecutionRunnable.getWorkflowEventBus();
+                    verify(eventBus).publish(argThat(evt -> evt instanceof TaskFailedLifecycleEvent &&
+                            ((TaskFailedLifecycleEvent) evt).getTaskExecutionRunnable() == taskExecutionRunnable));
+                    verify(eventBus, never()).publish(any(TaskFatalLifecycleEvent.class));
+                });
+    }
+
+    @Test
+    void dispatchTask_NoAvailableWorker_TimeoutEnabledButNotExceeded_ShouldNotPublishAnyFailureEvent() throws TaskDispatchException, InterruptedException {
+        // Given: Configure dispatcher with a 5-minute dispatch timeout (enabled)
+        MasterDispatchTimeoutCheckerConfig dispatchTimeoutCheckerConfig = new MasterDispatchTimeoutCheckerConfig();
+        dispatchTimeoutCheckerConfig.setEnabled(true);
+        dispatchTimeoutCheckerConfig.setMaxTaskDispatchDuration(Duration.ofMinutes(5));
+
+        dispatcher = new WorkerGroupDispatcher("TestGroup", taskExecutorClient, dispatchTimeoutCheckerConfig);
+
+        // Mock task with first dispatch time set to 100ms ago → ensures it's NOT timed out yet
+        ITaskExecutionRunnable taskExecutionRunnable = mockTaskExecutionRunnableWithFirstDispatchTime(
+                System.currentTimeMillis() - 100);
+
+        // Use CountDownLatch to reliably detect when dispatch is actually invoked (avoids timing flakiness)
+        CountDownLatch dispatchCalled = new CountDownLatch(1);
+
+        // Stub the client to throw NoAvailableWorkerException on dispatch and signal the latch
+        doAnswer(invocation -> {
+            dispatchCalled.countDown(); // Signal that dispatch was attempted
+            throw new NoAvailableWorkerException("no worker");
+        }).when(taskExecutorClient).dispatch(taskExecutionRunnable);
+
+        // When: Start dispatcher and trigger task dispatch
+        dispatcher.start();
+        dispatcher.dispatchTask(taskExecutionRunnable, 0);
+
+        // Wait up to 1 second for the dispatch attempt to occur (ensures async execution completes)
+        boolean dispatched = dispatchCalled.await(1000, TimeUnit.MILLISECONDS);
+        Assertions.assertTrue(dispatched, "Expected dispatch() to be called within 1 second");
+
+        // Then: Verify NO failure events are published since timeout has NOT been exceeded
+        WorkflowEventBus eventBus = taskExecutionRunnable.getWorkflowEventBus();
+        verify(eventBus, never()).publish(any(TaskFailedLifecycleEvent.class));
+        verify(eventBus, never()).publish(any(TaskFatalLifecycleEvent.class));
+    }
+
+    @Test
+    void dispatchTask_GenericTaskDispatchException_TimeoutDisabled_ShouldKeepRetrying() throws TaskDispatchException {
+        // Given
+        ITaskExecutionRunnable task = mockTaskExecutionRunnableWithFirstDispatchTime(System.currentTimeMillis());
+        TaskDispatchException ex = new TaskDispatchException("generic dispatch error");
+        doThrow(ex).when(taskExecutorClient).dispatch(task);
+
+        dispatcher.start();
+        dispatcher.dispatchTask(task, 0);
+
+        // When & Then
+        await().atMost(Duration.ofSeconds(3))
+                .untilAsserted(() -> {
+                    // Ensure it's retrying
+                    verify(taskExecutorClient, atLeast(2)).dispatch(task);
+
+                    // Ensure NO event has been published during this time
+                    WorkflowEventBus eventBus = task.getWorkflowEventBus();
+                    verify(eventBus, never()).publish(any(TaskFailedLifecycleEvent.class));
+                    verify(eventBus, never()).publish(any(TaskFatalLifecycleEvent.class));
+                });
+    }
+
+    @Test
+    void dispatchTask_GenericTaskDispatchException_TimeoutEnabledAndExceeded_ShouldPublishFailedEvent() throws TaskDispatchException {
+        // Given
+        MasterDispatchTimeoutCheckerConfig dispatchTimeoutCheckerConfig = new MasterDispatchTimeoutCheckerConfig();
+        dispatchTimeoutCheckerConfig.setEnabled(true);
+        dispatchTimeoutCheckerConfig.setMaxTaskDispatchDuration(Duration.ofMillis(200));
+
+        dispatcher = new WorkerGroupDispatcher("TestGroup", taskExecutorClient, dispatchTimeoutCheckerConfig);
+
+        ITaskExecutionRunnable taskExecutionRunnable = mockTaskExecutionRunnableWithFirstDispatchTime(
+                System.currentTimeMillis() - 500);
+
+        TaskDispatchException ex = new TaskDispatchException("generic dispatch error");
+        doThrow(ex).when(taskExecutorClient).dispatch(taskExecutionRunnable);
+
+        dispatcher.start();
+        dispatcher.dispatchTask(taskExecutionRunnable, 0);
+
+        // Then
+        await().atMost(Duration.ofSeconds(2))
+                .untilAsserted(() -> {
+                    verify(taskExecutorClient, times(1)).dispatch(taskExecutionRunnable);
+                    WorkflowEventBus eventBus = taskExecutionRunnable.getWorkflowEventBus();
+                    verify(eventBus).publish(argThat(evt -> evt instanceof TaskFailedLifecycleEvent &&
+                            ((TaskFailedLifecycleEvent) evt).getTaskExecutionRunnable() == taskExecutionRunnable));
+                    verify(eventBus, never()).publish(any(TaskFatalLifecycleEvent.class));
+                });
+    }
+
+    @Test
+    void dispatchTask_GenericTaskDispatchException_TimeoutEnabledButNotExceeded_ShouldNotPublishAnyFailureEvent() throws TaskDispatchException, InterruptedException {
+        // Given: Dispatcher configured with a 5-minute dispatch timeout (enabled)
+        MasterDispatchTimeoutCheckerConfig config = new MasterDispatchTimeoutCheckerConfig();
+        config.setEnabled(true);
+        config.setMaxTaskDispatchDuration(Duration.ofMinutes(5));
+
+        dispatcher = new WorkerGroupDispatcher("TestGroup", taskExecutorClient, config);
+
+        // Mock task with first dispatch time set to 100ms ago → well within timeout window
+        ITaskExecutionRunnable task = mockTaskExecutionRunnableWithFirstDispatchTime(
+                System.currentTimeMillis() - 100);
+
+        // Use CountDownLatch to reliably detect when dispatch is actually invoked
+        CountDownLatch dispatchCalled = new CountDownLatch(1);
+
+        // Stub client to throw a generic TaskDispatchException and signal the latch
+        doAnswer(invocation -> {
+            dispatchCalled.countDown(); // Confirm dispatch attempt occurred
+            throw new TaskDispatchException("Generic dispatch error");
+        }).when(taskExecutorClient).dispatch(task);
+
+        // When: Start dispatcher and trigger task dispatch
+        dispatcher.start();
+        dispatcher.dispatchTask(task, 0);
+
+        // Wait up to 1 second for the dispatch attempt to complete (handles async execution)
+        boolean dispatched = dispatchCalled.await(1000, TimeUnit.MILLISECONDS);
+        Assertions.assertTrue(dispatched, "Expected dispatch() to be called within 1 second");
+
+        // Then: Verify NO failure events are published because timeout has NOT been exceeded
+        WorkflowEventBus eventBus = task.getWorkflowEventBus();
+        verify(eventBus, never()).publish(any(TaskFailedLifecycleEvent.class));
+        verify(eventBus, never()).publish(any(TaskFatalLifecycleEvent.class));
+    }
+
+    private ITaskExecutionRunnable mockTaskExecutionRunnableWithFirstDispatchTime(long firstDispatchTime) {
+        ITaskExecutionRunnable taskExecutionRunnable = mock(ITaskExecutionRunnable.class);
+        TaskInstance taskInstance = mock(TaskInstance.class);
+        WorkflowInstance workflowInstance = mock(WorkflowInstance.class);
+        WorkflowEventBus eventBus = mock(WorkflowEventBus.class);
+
+        TaskExecutionContext context = mock(TaskExecutionContext.class);
+        when(context.getFirstDispatchTime()).thenReturn(firstDispatchTime);
+
+        when(taskExecutionRunnable.getTaskInstance()).thenReturn(taskInstance);
+        when(taskExecutionRunnable.getWorkflowInstance()).thenReturn(workflowInstance);
+        when(taskExecutionRunnable.getWorkflowEventBus()).thenReturn(eventBus);
+        when(taskExecutionRunnable.getId()).thenReturn(ThreadLocalRandom.current().nextInt(1000, 9999));
+        when(taskExecutionRunnable.getTaskExecutionContext()).thenReturn(context);
+
+        return taskExecutionRunnable;
+    }
 }