diff --git a/dolphinscheduler-alert/dolphinscheduler-alert-server/src/main/java/org/apache/dolphinscheduler/alert/service/AlertHAServer.java b/dolphinscheduler-alert/dolphinscheduler-alert-server/src/main/java/org/apache/dolphinscheduler/alert/service/AlertHAServer.java index 67382ac470a7..cc5a920cd1c1 100644 --- a/dolphinscheduler-alert/dolphinscheduler-alert-server/src/main/java/org/apache/dolphinscheduler/alert/service/AlertHAServer.java +++ b/dolphinscheduler-alert/dolphinscheduler-alert-server/src/main/java/org/apache/dolphinscheduler/alert/service/AlertHAServer.java @@ -42,6 +42,7 @@ public void start() { @Override public void close() { + super.close(); log.info("AlertHAServer shutdown..."); } } diff --git a/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/ITaskGroupCoordinator.java b/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/ITaskGroupCoordinator.java index 854097b02124..7d841aa15d9b 100644 --- a/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/ITaskGroupCoordinator.java +++ b/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/ITaskGroupCoordinator.java @@ -52,6 +52,11 @@ public interface ITaskGroupCoordinator extends AutoCloseable { */ void start(); + /** + * Request the worker to stop without joining the worker. Call {@link #close()} to wait before restarting. + */ + void requestStop(); + /** * If the {@link TaskInstance#getTaskGroupId()} > 0, and the TaskGroup flag is {@link Flag#YES} then the task instance need to use task group. * diff --git a/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/IWorkflowSerialCoordinator.java b/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/IWorkflowSerialCoordinator.java index bbd9c248c6de..12d57143fc65 100644 --- a/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/IWorkflowSerialCoordinator.java +++ b/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/IWorkflowSerialCoordinator.java @@ -21,6 +21,11 @@ public interface IWorkflowSerialCoordinator extends AutoCloseable { void start(); + /** + * Request the worker to stop without joining the worker. Call {@link #close()} to wait before restarting. + */ + void requestStop(); + @Override void close(); diff --git a/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/MasterCoordinator.java b/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/MasterCoordinator.java index eb2a9dda02cd..81d5533b8edd 100644 --- a/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/MasterCoordinator.java +++ b/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/MasterCoordinator.java @@ -41,11 +41,7 @@ @Component public class MasterCoordinator extends AbstractHAServer { - private final ITaskGroupCoordinator taskGroupCoordinator; - - private final IFailoverCoordinator failoverCoordinator; - - private final IWorkflowSerialCoordinator workflowSerialCoordinator; + private final MasterCoordinatorListener masterCoordinatorListener; public MasterCoordinator(final Registry registry, final MasterConfig masterConfig, @@ -56,11 +52,9 @@ public MasterCoordinator(final Registry registry, registry, RegistryNodeType.MASTER_COORDINATOR.getRegistryPath(), masterConfig.getMasterAddress()); - this.taskGroupCoordinator = taskGroupCoordinator; - this.failoverCoordinator = failoverCoordinator; - this.workflowSerialCoordinator = workflowSerialCoordinator; - addServerStatusChangeListener( - new MasterCoordinatorListener(taskGroupCoordinator, failoverCoordinator, workflowSerialCoordinator)); + this.masterCoordinatorListener = + new MasterCoordinatorListener(taskGroupCoordinator, failoverCoordinator, workflowSerialCoordinator); + addServerStatusChangeListener(masterCoordinatorListener); } @Override @@ -71,7 +65,8 @@ public void start() { @Override public void close() { - taskGroupCoordinator.close(); + super.close(); + masterCoordinatorListener.changeToStandBy(); log.info("MasterCoordinator shutdown..."); } @@ -108,11 +103,14 @@ public void changeToActive() { @Override public void changeToStandBy() { - taskGroupCoordinator.close(); - workflowSerialCoordinator.close(); + // Stop both workers before waiting: either may be blocked in a database call. + taskGroupCoordinator.requestStop(); + workflowSerialCoordinator.requestStop(); if (failoverCoordinatorFuture != null) { failoverCoordinatorFuture.cancel(true); } + taskGroupCoordinator.close(); + workflowSerialCoordinator.close(); } } diff --git a/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/TaskGroupCoordinator.java b/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/TaskGroupCoordinator.java index 3a6d81c24c6b..8fc25a300029 100644 --- a/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/TaskGroupCoordinator.java +++ b/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/TaskGroupCoordinator.java @@ -77,7 +77,7 @@ public class TaskGroupCoordinator implements ITaskGroupCoordinator, AutoCloseabl @Autowired private TransactionTemplate transactionTemplate; - private boolean flag = false; + private volatile boolean flag = false; private Thread internalThread; @@ -113,9 +113,19 @@ private void doStart() { try { final StopWatch taskGroupCoordinatorRoundCost = StopWatch.createStarted(); + // A database call may outlive a stop request; check flag before starting the next phase. amendTaskGroupUseSize(); + if (!flag) { + return; + } amendTaskGroupQueueStatus(); + if (!flag) { + return; + } dealWithForceStartTaskGroupQueue(); + if (!flag) { + return; + } dealWithWaitingTaskGroupQueue(); taskGroupCoordinatorRoundCost.stop(); @@ -124,7 +134,9 @@ private void doStart() { log.error("TaskGroupCoordinator error", e); } finally { // sleep 5s - ThreadUtils.sleep(Constants.SLEEP_TIME_MILLIS * 5); + if (flag) { + ThreadUtils.sleep(Constants.SLEEP_TIME_MILLIS * 5); + } } } } @@ -141,7 +153,14 @@ private void amendTaskGroupUseSize() { StopWatch taskGroupCoordinatorRoundTimeCost = StopWatch.createStarted(); for (TaskGroup taskGroup : taskGroups) { + if (!flag) { + return; + } int actualUseSize = taskGroupQueueDao.countUsingTaskGroupQueueByGroupId(taskGroup.getId()); + // The query may finish after a stop request; do not start the update in that case. + if (!flag) { + return; + } if (taskGroup.getUseSize() == actualUseSize) { continue; } @@ -161,10 +180,10 @@ private void amendTaskGroupQueueStatus() { int minTaskGroupQueueId = -1; int limit = DEFAULT_LIMIT; StopWatch taskGroupCoordinatorRoundTimeCost = StopWatch.createStarted(); - while (true) { + while (flag) { List taskGroupQueues = taskGroupQueueDao.queryInQueueTaskGroupQueue(minTaskGroupQueueId, limit); - if (CollectionUtils.isEmpty(taskGroupQueues)) { + if (!flag || CollectionUtils.isEmpty(taskGroupQueues)) { break; } amendTaskGroupQueueStatus(taskGroupQueues); @@ -188,6 +207,9 @@ private void amendTaskGroupQueueStatus(List taskGroupQueues) { .collect(Collectors.toMap(TaskInstance::getId, Function.identity())); for (TaskGroupQueue taskGroupQueue : taskGroupQueues) { + if (!flag) { + return; + } int taskId = taskGroupQueue.getTaskId(); final TaskInstance taskInstance = taskInstanceMap.get(taskId); @@ -214,10 +236,10 @@ private void dealWithForceStartTaskGroupQueue() { int minTaskGroupQueueId = -1; int limit = DEFAULT_LIMIT; StopWatch taskGroupCoordinatorRoundTimeCost = StopWatch.createStarted(); - while (true) { + while (flag) { final List taskGroupQueues = taskGroupQueueDao.queryWaitNotifyForceStartTaskGroupQueue(minTaskGroupQueueId, limit); - if (CollectionUtils.isEmpty(taskGroupQueues)) { + if (!flag || CollectionUtils.isEmpty(taskGroupQueues)) { break; } dealWithForceStartTaskGroupQueue(taskGroupQueues); @@ -235,6 +257,9 @@ private void dealWithForceStartTaskGroupQueue(List taskGroupQueu // Notify the related waiting task instance // Set the taskGroupQueue status to RELEASE and remove it from queue for (final TaskGroupQueue taskGroupQueue : taskGroupQueues) { + if (!flag) { + return; + } try { LogUtils.setTaskInstanceIdMDC(taskGroupQueue.getTaskId()); if (!notifyForceStartTaskGroupQueue(taskGroupQueue)) { @@ -290,6 +315,9 @@ private void dealWithWaitingTaskGroupQueue() { return; } for (TaskGroup taskGroup : taskGroups) { + if (!flag) { + return; + } int availableSize = taskGroup.getGroupSize() - taskGroup.getUseSize(); if (availableSize <= 0) { log.info("TaskGroup {} is full, available size is {}", taskGroup, availableSize); @@ -307,6 +335,9 @@ private void dealWithWaitingTaskGroupQueue() { continue; } for (TaskGroupQueue taskGroupQueue : taskGroupQueues) { + if (!flag) { + return; + } try { LogUtils.setTaskInstanceIdMDC(taskGroupQueue.getTaskId()); if (!acquireTaskGroupSlotAndNotify(taskGroupQueue)) { @@ -516,21 +547,36 @@ private void deleteTaskGroupQueueSlot(TaskGroupQueue taskGroupQueue) { log.debug("TaskGroupQueue has already been released: {}", taskGroupQueue.getId()); } + @Override + public synchronized void requestStop() { + flag = false; + if (internalThread != null) { + internalThread.interrupt(); + } + } + @Override public synchronized void close() { - if (!flag) { - log.warn("TaskGroupCoordinator is already closed"); - return; + if (Thread.currentThread() == internalThread) { + throw new IllegalStateException("TaskGroupCoordinator cannot close its own worker thread"); } - flag = false; - try { - if (internalThread != null) { - internalThread.interrupt(); + // A prior stop request does not mean the worker has finished. + requestStop(); + boolean interrupted = false; + if (internalThread != null) { + // Keep start() waiting until the old worker has finished, including any in-flight JDBC call. + while (internalThread.isAlive()) { + try { + internalThread.join(); + } catch (InterruptedException ex) { + interrupted = true; + } } - } catch (Exception ex) { - log.error("Close internalThread failed", ex); + internalThread = null; + } + if (interrupted) { + Thread.currentThread().interrupt(); } - internalThread = null; log.info("TaskGroupCoordinator closed"); } } diff --git a/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/workflow/serial/WorkflowSerialCoordinator.java b/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/workflow/serial/WorkflowSerialCoordinator.java index a11223c45b2e..f743c7b0e06c 100644 --- a/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/workflow/serial/WorkflowSerialCoordinator.java +++ b/dolphinscheduler-master/src/main/java/org/apache/dolphinscheduler/server/master/engine/workflow/serial/WorkflowSerialCoordinator.java @@ -95,7 +95,12 @@ private void doStart() { try { final StopWatch workflowSerialCoordinatorRoundCost = StopWatch.createStarted(); final List serialCommandsGroups = fetchSerialCommands(); - serialCommandsGroups.forEach(this::handleSerialCommand); + // Fetching or handling a group may outlive a stop request; check before handling the next group. + serialCommandsGroups.forEach(serialCommandsGroup -> { + if (flag) { + handleSerialCommand(serialCommandsGroup); + } + }); log.debug("WorkflowSerialCoordinator handled SerialCommandsGroup size: {}, cost: {}/ms ", serialCommandsGroups.size(), workflowSerialCoordinatorRoundCost.getDuration().toMillis()); @@ -103,7 +108,9 @@ private void doStart() { log.error("WorkflowSerialCoordinator error", e); } finally { // sleep 5s - ThreadUtils.sleep(TimeUnit.SECONDS.toMillis(DEFAULT_FETCH_INTERVAL_SECONDS)); + if (flag) { + ThreadUtils.sleep(TimeUnit.SECONDS.toMillis(DEFAULT_FETCH_INTERVAL_SECONDS)); + } } } } @@ -168,21 +175,36 @@ private SerialCommandsGroup createSerialCommandsGroup(SerialCommandDto serialCom .build(); } + @Override + public synchronized void requestStop() { + flag = false; + if (internalThread != null) { + internalThread.interrupt(); + } + } + @Override public synchronized void close() { - if (!flag) { - log.warn("WorkflowSerialCoordinator is already closed"); - return; + if (Thread.currentThread() == internalThread) { + throw new IllegalStateException("WorkflowSerialCoordinator cannot close its own worker thread"); } - flag = false; - try { - if (internalThread != null) { - internalThread.interrupt(); + // A prior stop request does not mean the worker has finished. + requestStop(); + boolean interrupted = false; + if (internalThread != null) { + // Keep start() waiting until the old worker has finished, including any in-flight JDBC call. + while (internalThread.isAlive()) { + try { + internalThread.join(); + } catch (InterruptedException ex) { + interrupted = true; + } } - } catch (Exception ex) { - log.error("Close internalThread failed", ex); + internalThread = null; + } + if (interrupted) { + Thread.currentThread().interrupt(); } - internalThread = null; log.info("WorkflowSerialCoordinator closed"); } } diff --git a/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/MasterCoordinatorTest.java b/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/MasterCoordinatorTest.java new file mode 100644 index 000000000000..a4c3dd7b5dd3 --- /dev/null +++ b/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/MasterCoordinatorTest.java @@ -0,0 +1,206 @@ +/* + * 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; + +import static org.junit.jupiter.api.Assertions.assertThrows; + +import org.apache.dolphinscheduler.common.enums.WorkflowExecutionTypeEnum; +import org.apache.dolphinscheduler.common.thread.ThreadUtils; +import org.apache.dolphinscheduler.dao.entity.WorkflowDefinitionLog; +import org.apache.dolphinscheduler.dao.model.SerialCommandDto; +import org.apache.dolphinscheduler.dao.repository.SerialCommandDao; +import org.apache.dolphinscheduler.dao.repository.TaskGroupDao; +import org.apache.dolphinscheduler.dao.repository.WorkflowDefinitionLogDao; +import org.apache.dolphinscheduler.registry.api.Registry; +import org.apache.dolphinscheduler.server.master.config.MasterConfig; +import org.apache.dolphinscheduler.server.master.engine.workflow.serial.SerialCommandDiscardHandler; +import org.apache.dolphinscheduler.server.master.engine.workflow.serial.SerialCommandPriorityHandler; +import org.apache.dolphinscheduler.server.master.engine.workflow.serial.SerialCommandWaitHandler; +import org.apache.dolphinscheduler.server.master.engine.workflow.serial.WorkflowSerialCoordinator; +import org.apache.dolphinscheduler.server.master.failover.IFailoverCoordinator; +import org.apache.dolphinscheduler.server.master.utils.MasterThreadFactory; + +import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.InjectMocks; +import org.mockito.Mock; +import org.mockito.MockedStatic; +import org.mockito.Mockito; +import org.mockito.junit.jupiter.MockitoExtension; +import org.mockito.junit.jupiter.MockitoSettings; +import org.mockito.quality.Strictness; +import org.springframework.test.util.ReflectionTestUtils; + +@ExtendWith(MockitoExtension.class) +@MockitoSettings(strictness = Strictness.LENIENT) +class MasterCoordinatorTest { + + @InjectMocks + private WorkflowSerialCoordinator workflowSerialCoordinator; + + @Mock + private SerialCommandDao serialCommandDao; + + @Mock + private WorkflowDefinitionLogDao workflowDefinitionLogDao; + + @Mock + private SerialCommandWaitHandler serialCommandWaitHandler; + + @Mock + private SerialCommandDiscardHandler serialCommandDiscardHandler; + + @Mock + private SerialCommandPriorityHandler serialCommandPriorityHandler; + + @Test + void closeShouldStopBothWorkersBeforeWaitingForEither() throws Exception { + TaskGroupCoordinator taskGroupCoordinator = new TaskGroupCoordinator(); + TaskGroupDao taskGroupDao = Mockito.mock(TaskGroupDao.class); + ReflectionTestUtils.setField(taskGroupCoordinator, "taskGroupDao", taskGroupDao); + ScheduledExecutorService scheduler = Mockito.mock(ScheduledExecutorService.class); + ScheduledFuture scheduled = Mockito.mock(ScheduledFuture.class); + Mockito.doReturn(scheduled).when(scheduler).scheduleWithFixedDelay( + Mockito.any(Runnable.class), Mockito.anyLong(), Mockito.anyLong(), Mockito.any(TimeUnit.class)); + ExecutorService electionExecutor = Executors.newFixedThreadPool(1, runnable -> { + Thread worker = new Thread(() -> { + // The listener schedules failover work on the election worker, not the test thread. + try (MockedStatic factory = Mockito.mockStatic(MasterThreadFactory.class)) { + factory.when(MasterThreadFactory::getDefaultSchedulerThreadExecutor).thenReturn(scheduler); + runnable.run(); + } + }, "test-master-election"); + worker.setDaemon(true); + return worker; + }); + Registry registry = Mockito.mock(Registry.class); + Mockito.when(registry.acquireLock(Mockito.anyString())).thenReturn(true); + MasterConfig config = new MasterConfig(); + config.setMasterAddress("master-0:5678"); + MasterCoordinator masterCoordinator; + try ( + MockedStatic threadUtils = + Mockito.mockStatic(ThreadUtils.class, Mockito.CALLS_REAL_METHODS)) { + threadUtils.when(() -> ThreadUtils.newDaemonFixedThreadExecutor("HA-Election-%d", 1)) + .thenReturn(electionExecutor); + masterCoordinator = new MasterCoordinator(registry, config, taskGroupCoordinator, + Mockito.mock(IFailoverCoordinator.class), workflowSerialCoordinator); + } + CountDownLatch taskGroupFetching = new CountDownLatch(1); + CountDownLatch serialFetching = new CountDownLatch(1); + CountDownLatch taskGroupCanceled = new CountDownLatch(1); + CountDownLatch serialCanceled = new CountDownLatch(1); + CountDownLatch releaseTaskGroup = new CountDownLatch(1); + CountDownLatch releaseSerial = new CountDownLatch(1); + CountDownLatch closed = new CountDownLatch(1); + CountDownLatch failoverCanceled = new CountDownLatch(1); + Mockito.when(scheduled.cancel(true)).thenAnswer(invocation -> { + failoverCanceled.countDown(); + return true; + }); + AtomicReference taskGroupWorker = new AtomicReference<>(); + AtomicReference serialWorker = new AtomicReference<>(); + AtomicReference failure = new AtomicReference<>(); + Mockito.when(taskGroupDao.queryAllTaskGroups()).thenAnswer(invocation -> { + taskGroupWorker.set(Thread.currentThread()); + taskGroupFetching.countDown(); + awaitDatabaseResponse(releaseTaskGroup, taskGroupCanceled, failure); + return Collections.emptyList(); + }); + Mockito.when(serialCommandDao.fetchSerialCommands(Mockito.anyInt())).thenAnswer(invocation -> { + serialWorker.set(Thread.currentThread()); + serialFetching.countDown(); + awaitDatabaseResponse(releaseSerial, serialCanceled, failure); + return Collections.singletonList(SerialCommandDto.builder() + .workflowDefinitionCode(1L).workflowDefinitionVersion(1).build()); + }); + WorkflowDefinitionLog definition = new WorkflowDefinitionLog(); + definition.setExecutionType(WorkflowExecutionTypeEnum.SERIAL_WAIT); + Mockito.when(workflowDefinitionLogDao.queryByDefinitionCodeAndVersion(1L, 1)).thenReturn(definition); + Thread closer = new Thread(() -> { + try { + masterCoordinator.close(); + } catch (Throwable ex) { + failure.compareAndSet(null, ex); + } finally { + closed.countDown(); + } + }); + closer.setDaemon(true); + try { + masterCoordinator.start(); + Assertions.assertTrue(serialFetching.await(5, TimeUnit.SECONDS)); + // Keep the real one-minute TaskGroup startup delay; both lifecycle implementations participate. + Assertions.assertTrue(taskGroupFetching.await(90, TimeUnit.SECONDS)); + closer.start(); + Assertions.assertTrue(taskGroupCanceled.await(5, TimeUnit.SECONDS)); + Assertions.assertTrue(serialCanceled.await(5, TimeUnit.SECONDS)); + releaseSerial.countDown(); + serialWorker.get().join(5000); + Assertions.assertFalse(serialWorker.get().isAlive()); + // A stop request alone must not make a new generation eligible to start. + assertThrows(IllegalStateException.class, workflowSerialCoordinator::start); + Mockito.verifyNoInteractions(serialCommandWaitHandler); + // Serial has stopped while TaskGroup is still blocked; close must still join TaskGroup. + Assertions.assertTrue(taskGroupWorker.get().isAlive()); + Assertions.assertEquals(1L, closed.getCount()); + Assertions.assertTrue(failoverCanceled.await(5, TimeUnit.SECONDS)); + Mockito.verify(scheduled).cancel(true); + releaseTaskGroup.countDown(); + Assertions.assertTrue(closed.await(5, TimeUnit.SECONDS)); + Assertions.assertFalse(taskGroupWorker.get().isAlive()); + Assertions.assertNull(failure.get()); + } finally { + releaseSerial.countDown(); + releaseTaskGroup.countDown(); + closer.join(5000); + masterCoordinator.close(); + // Also release Serial if a regression in the Master entry point omitted it. + workflowSerialCoordinator.close(); + Assertions.assertTrue(electionExecutor.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + private static void awaitDatabaseResponse(CountDownLatch release, CountDownLatch canceled, + AtomicReference failure) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(100); + while (release.getCount() != 0) { + try { + long remaining = deadline - System.nanoTime(); + if (remaining <= 0 || !release.await(remaining, TimeUnit.NANOSECONDS)) { + failure.compareAndSet(null, new AssertionError("Blocked DAO was not released")); + return; + } + } catch (InterruptedException ignored) { + // JDBC may ignore cancellation; only the test-controlled response releases the call. + canceled.countDown(); + } + } + } + +} diff --git a/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/TaskGroupCoordinatorTest.java b/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/TaskGroupCoordinatorTest.java index 6f4b7d516ca4..878742853e46 100644 --- a/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/TaskGroupCoordinatorTest.java +++ b/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/TaskGroupCoordinatorTest.java @@ -35,7 +35,11 @@ import org.apache.dolphinscheduler.dao.repository.TaskInstanceDao; import org.apache.dolphinscheduler.dao.repository.WorkflowInstanceDao; +import java.util.Collections; import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; @@ -171,4 +175,144 @@ void releaseTaskGroupSlot() { verify(taskGroupQueueDao, Mockito.times(1)).deleteById(taskGroupQueue); } + @Test + void closeShouldWaitForPreviousFetchBeforeRestart() throws Exception { + verifyCloseDuringFetch(false); + } + + @Test + void interruptedCloseShouldFinishWaitingAndRestoreInterrupt() throws Exception { + verifyCloseDuringFetch(true); + } + + private void verifyCloseDuringFetch(boolean interruptCloser) throws Exception { + CountDownLatch fetching = new CountDownLatch(1); + CountDownLatch canceled = new CountDownLatch(1); + CountDownLatch releaseFetch = new CountDownLatch(1); + CountDownLatch restarted = new CountDownLatch(1); + CountDownLatch closed = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + AtomicReference firstThread = new AtomicReference<>(); + Mockito.when(taskGroupDao.queryAllTaskGroups()).thenAnswer(invocation -> { + if (firstThread.compareAndSet(null, Thread.currentThread())) { + fetching.countDown(); + // Model a JDBC request that returns only after the server responds, despite interruption. + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (releaseFetch.getCount() != 0) { + try { + long remaining = deadline - System.nanoTime(); + if (remaining <= 0 || !releaseFetch.await(remaining, TimeUnit.NANOSECONDS)) { + failure.compareAndSet(null, new AssertionError("Blocked DAO was not released")); + return Collections.emptyList(); + } + } catch (InterruptedException ignored) { + canceled.countDown(); + } + } + TaskGroup staleGroup = new TaskGroup(); + staleGroup.setId(1); + staleGroup.setUseSize(1); + return Collections.singletonList(staleGroup); + } + if (Thread.currentThread() == firstThread.get()) { + failure.compareAndSet(null, new AssertionError("Canceled worker resumed polling")); + } + restarted.countDown(); + return Collections.emptyList(); + }); + Thread closer = new Thread(() -> { + try { + taskGroupCoordinator.close(); + Assertions.assertFalse(firstThread.get().isAlive(), "close must finish the old worker"); + Assertions.assertEquals(interruptCloser, Thread.currentThread().isInterrupted()); + } catch (Throwable ex) { + failure.compareAndSet(null, ex); + } finally { + closed.countDown(); + } + }); + Thread starter = new Thread(() -> { + try { + taskGroupCoordinator.start(); + } catch (Throwable ex) { + failure.compareAndSet(null, ex); + } + }); + closer.setDaemon(true); + starter.setDaemon(true); + try { + taskGroupCoordinator.start(); + // Allow the coordinator's normal one-minute startup delay before observing the DAO call. + Assertions.assertTrue(fetching.await(90, TimeUnit.SECONDS)); + closer.start(); + Assertions.assertTrue(canceled.await(5, TimeUnit.SECONDS)); + if (interruptCloser) { + closer.interrupt(); + } + starter.start(); + // The start call must wait on the lifecycle monitor while close drains the old request. + awaitBlocked(starter); + Assertions.assertEquals(1L, closed.getCount()); + Assertions.assertEquals(1L, restarted.getCount()); + releaseFetch.countDown(); + Assertions.assertTrue(closed.await(5, TimeUnit.SECONDS)); + // A new worker also observes the normal startup delay. + Assertions.assertTrue(restarted.await(90, TimeUnit.SECONDS)); + closer.join(5000); + starter.join(5000); + Assertions.assertFalse(closer.isAlive()); + Assertions.assertFalse(starter.isAlive()); + Assertions.assertNull(failure.get()); + // A nonempty batch fetched before cancellation must not be processed after it returns. + Mockito.verify(taskGroupQueueDao, Mockito.never()).countUsingTaskGroupQueueByGroupId(Mockito.anyInt()); + taskGroupCoordinator.close(); + taskGroupCoordinator.close(); + taskGroupCoordinator.start(); + taskGroupCoordinator.close(); + } finally { + // Release external work before waiting for close; otherwise cleanup itself would deadlock. + releaseFetch.countDown(); + closer.join(5000); + starter.join(5000); + Assertions.assertFalse(closer.isAlive()); + Assertions.assertFalse(starter.isAlive()); + taskGroupCoordinator.close(); + } + } + + @Test + void workerShouldNotCloseItself() throws Exception { + CountDownLatch checked = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + Mockito.when(taskGroupDao.queryAllTaskGroups()).thenAnswer(invocation -> { + try { + Assertions.assertThrows(IllegalStateException.class, () -> taskGroupCoordinator.close()); + } catch (Throwable ex) { + failure.compareAndSet(null, ex); + } finally { + checked.countDown(); + } + return Collections.emptyList(); + }); + try { + taskGroupCoordinator.start(); + // Exercise the actual worker after its normal one-minute startup delay. + Assertions.assertTrue(checked.await(90, TimeUnit.SECONDS)); + Assertions.assertNull(failure.get()); + // Rejection must leave lifecycle state unchanged so the owner can still close the worker. + Assertions.assertThrows(IllegalStateException.class, () -> taskGroupCoordinator.start()); + } finally { + taskGroupCoordinator.close(); + } + } + + private static void awaitBlocked(Thread thread) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (thread.isAlive() && thread.getState() != Thread.State.BLOCKED + && System.nanoTime() < deadline) { + Thread.yield(); + } + Assertions.assertEquals(Thread.State.BLOCKED, thread.getState()); + } + } diff --git a/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/workflow/serial/WorkflowSerialCoordinatorTest.java b/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/workflow/serial/WorkflowSerialCoordinatorTest.java index cac58146efcd..86119a4f4447 100644 --- a/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/workflow/serial/WorkflowSerialCoordinatorTest.java +++ b/dolphinscheduler-master/src/test/java/org/apache/dolphinscheduler/server/master/engine/workflow/serial/WorkflowSerialCoordinatorTest.java @@ -20,16 +20,40 @@ import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; import static org.junit.jupiter.api.Assertions.assertThrows; +import org.apache.dolphinscheduler.common.enums.WorkflowExecutionTypeEnum; +import org.apache.dolphinscheduler.dao.entity.WorkflowDefinitionLog; +import org.apache.dolphinscheduler.dao.model.SerialCommandDto; import org.apache.dolphinscheduler.dao.repository.SerialCommandDao; +import org.apache.dolphinscheduler.dao.repository.TaskGroupDao; import org.apache.dolphinscheduler.dao.repository.WorkflowDefinitionLogDao; +import org.apache.dolphinscheduler.registry.api.Event; +import org.apache.dolphinscheduler.registry.api.Registry; +import org.apache.dolphinscheduler.registry.api.SubscribeListener; +import org.apache.dolphinscheduler.registry.api.ha.AbstractHAServer; +import org.apache.dolphinscheduler.server.master.engine.ITaskGroupCoordinator; +import org.apache.dolphinscheduler.server.master.engine.MasterCoordinator; +import org.apache.dolphinscheduler.server.master.engine.TaskGroupCoordinator; +import org.apache.dolphinscheduler.server.master.failover.IFailoverCoordinator; +import org.apache.dolphinscheduler.server.master.utils.MasterThreadFactory; +import java.util.Collections; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.InjectMocks; import org.mockito.Mock; +import org.mockito.MockedStatic; +import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; import org.mockito.junit.jupiter.MockitoSettings; import org.mockito.quality.Strictness; +import org.springframework.test.util.ReflectionTestUtils; @ExtendWith(MockitoExtension.class) @MockitoSettings(strictness = Strictness.LENIENT) @@ -81,4 +105,351 @@ void closeShouldBeIdempotent() { assertDoesNotThrow(() -> workflowSerialCoordinator.close()); } + @Test + void closeShouldWaitForPreviousFetchBeforeRestart() throws Exception { + verifyCloseDuringFetch(false); + } + + @Test + void interruptedCloseShouldFinishWaitingAndRestoreInterrupt() throws Exception { + verifyCloseDuringFetch(true); + } + + private void verifyCloseDuringFetch(boolean interruptCloser) throws Exception { + CountDownLatch fetching = new CountDownLatch(1); + CountDownLatch canceled = new CountDownLatch(1); + CountDownLatch releaseFetch = new CountDownLatch(1); + CountDownLatch restarted = new CountDownLatch(1); + CountDownLatch closed = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + AtomicReference firstThread = new AtomicReference<>(); + Mockito.when(serialCommandDao.fetchSerialCommands(Mockito.anyInt())).thenAnswer(invocation -> { + if (firstThread.compareAndSet(null, Thread.currentThread())) { + fetching.countDown(); + // Model a JDBC request that returns only after the server responds, despite interruption. + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (releaseFetch.getCount() != 0) { + try { + long remaining = deadline - System.nanoTime(); + if (remaining <= 0 || !releaseFetch.await(remaining, TimeUnit.NANOSECONDS)) { + failure.compareAndSet(null, new AssertionError("Blocked DAO was not released")); + return Collections.emptyList(); + } + } catch (InterruptedException ignored) { + canceled.countDown(); + } + } + return Collections.singletonList(SerialCommandDto.builder() + .workflowDefinitionCode(1L).workflowDefinitionVersion(1).build()); + } + if (Thread.currentThread() == firstThread.get()) { + failure.compareAndSet(null, new AssertionError("Canceled worker resumed polling")); + } + restarted.countDown(); + return Collections.emptyList(); + }); + WorkflowDefinitionLog definition = new WorkflowDefinitionLog(); + definition.setExecutionType(WorkflowExecutionTypeEnum.SERIAL_WAIT); + Mockito.when(workflowDefinitionLogDao.queryByDefinitionCodeAndVersion(1L, 1)).thenReturn(definition); + Thread closer = new Thread(() -> { + try { + workflowSerialCoordinator.close(); + Assertions.assertFalse(firstThread.get().isAlive(), "close must finish the old worker"); + Assertions.assertEquals(interruptCloser, Thread.currentThread().isInterrupted()); + } catch (Throwable ex) { + failure.compareAndSet(null, ex); + } finally { + closed.countDown(); + } + }); + Thread starter = new Thread(() -> { + try { + workflowSerialCoordinator.start(); + } catch (Throwable ex) { + failure.compareAndSet(null, ex); + } + }); + closer.setDaemon(true); + starter.setDaemon(true); + try { + workflowSerialCoordinator.start(); + Assertions.assertTrue(fetching.await(5, TimeUnit.SECONDS)); + closer.start(); + Assertions.assertTrue(canceled.await(5, TimeUnit.SECONDS)); + if (interruptCloser) { + closer.interrupt(); + } + starter.start(); + // The start call must wait on the lifecycle monitor while close drains the old request. + awaitBlocked(starter); + Assertions.assertEquals(1L, closed.getCount()); + Assertions.assertEquals(1L, restarted.getCount()); + releaseFetch.countDown(); + Assertions.assertTrue(closed.await(5, TimeUnit.SECONDS)); + Assertions.assertTrue(restarted.await(5, TimeUnit.SECONDS)); + closer.join(5000); + starter.join(5000); + Assertions.assertFalse(closer.isAlive()); + Assertions.assertFalse(starter.isAlive()); + Assertions.assertNull(failure.get()); + // A nonempty batch fetched before cancellation must not be processed after it returns. + Mockito.verifyNoInteractions(serialCommandWaitHandler); + workflowSerialCoordinator.close(); + workflowSerialCoordinator.close(); + workflowSerialCoordinator.start(); + workflowSerialCoordinator.close(); + } finally { + // Release external work before waiting for close; otherwise cleanup itself would deadlock. + releaseFetch.countDown(); + closer.join(5000); + starter.join(5000); + Assertions.assertFalse(closer.isAlive()); + Assertions.assertFalse(starter.isAlive()); + workflowSerialCoordinator.close(); + } + } + + @Test + void workerShouldNotCloseItself() throws Exception { + CountDownLatch checked = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + Mockito.when(serialCommandDao.fetchSerialCommands(Mockito.anyInt())).thenAnswer(invocation -> { + try { + Assertions.assertThrows(IllegalStateException.class, () -> workflowSerialCoordinator.close()); + } catch (Throwable ex) { + failure.compareAndSet(null, ex); + } finally { + checked.countDown(); + } + return Collections.emptyList(); + }); + try { + workflowSerialCoordinator.start(); + Assertions.assertTrue(checked.await(5, TimeUnit.SECONDS)); + Assertions.assertNull(failure.get()); + // Rejection must leave lifecycle state unchanged so the owner can still close the worker. + Assertions.assertThrows(IllegalStateException.class, () -> workflowSerialCoordinator.start()); + } finally { + workflowSerialCoordinator.close(); + } + } + + @Test + void queuedRoleChangesShouldDrainOldWorkerBeforeReactivation() throws Exception { + Registry registry = Mockito.mock(Registry.class); + AtomicReference owner = new AtomicReference<>(); + AtomicReference subscriber = new AtomicReference<>(); + AbstractHAServer server = new AbstractHAServer(registry, "/coordinator", "master-0:5678") { + }; + MasterCoordinator.MasterCoordinatorListener listener = new MasterCoordinator.MasterCoordinatorListener( + Mockito.mock(ITaskGroupCoordinator.class), Mockito.mock(IFailoverCoordinator.class), + workflowSerialCoordinator) { + + @Override + public void changeToActive() { + // Static mocks are thread-local: intercept the factory on the HA worker itself. + try (MockedStatic factory = Mockito.mockStatic(MasterThreadFactory.class)) { + factory.when(MasterThreadFactory::getDefaultSchedulerThreadExecutor) + .thenReturn(Mockito.mock(ScheduledExecutorService.class)); + super.changeToActive(); + } + } + }; + server.addServerStatusChangeListener(listener); + Mockito.when(registry.acquireLock("/coordinator-lock")).thenReturn(true); + Mockito.when(registry.exists("/coordinator")).thenAnswer(invocation -> owner.get() != null); + Mockito.when(registry.get("/coordinator")).thenAnswer(invocation -> owner.get()); + Mockito.doAnswer(invocation -> { + owner.set(invocation.getArgument(1)); + return null; + }).when(registry).put(Mockito.eq("/coordinator"), Mockito.anyString(), Mockito.eq(true)); + Mockito.doAnswer(invocation -> { + subscriber.set(invocation.getArgument(1)); + return null; + }).when(registry).subscribe(Mockito.eq("/coordinator"), Mockito.any()); + + CountDownLatch fetching = new CountDownLatch(1); + CountDownLatch canceled = new CountDownLatch(1); + CountDownLatch releaseFetch = new CountDownLatch(1); + CountDownLatch restarted = new CountDownLatch(1); + CountDownLatch callbacksReturned = new CountDownLatch(1); + AtomicReference oldWorker = new AtomicReference<>(); + AtomicReference failure = new AtomicReference<>(); + WorkflowDefinitionLog definition = new WorkflowDefinitionLog(); + definition.setExecutionType(WorkflowExecutionTypeEnum.SERIAL_WAIT); + Mockito.when(workflowDefinitionLogDao.queryByDefinitionCodeAndVersion(1L, 1)).thenReturn(definition); + Mockito.when(serialCommandDao.fetchSerialCommands(Mockito.anyInt())).thenAnswer(invocation -> { + if (oldWorker.compareAndSet(null, Thread.currentThread())) { + fetching.countDown(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (releaseFetch.getCount() != 0) { + try { + long remaining = deadline - System.nanoTime(); + if (remaining <= 0 || !releaseFetch.await(remaining, TimeUnit.NANOSECONDS)) { + failure.compareAndSet(null, new AssertionError("Blocked DAO was not released")); + return Collections.emptyList(); + } + } catch (InterruptedException ignored) { + canceled.countDown(); + } + } + return Collections.singletonList(SerialCommandDto.builder() + .workflowDefinitionCode(1L).workflowDefinitionVersion(1).build()); + } + if (Thread.currentThread() == oldWorker.get()) { + failure.compareAndSet(null, new AssertionError("Canceled worker resumed polling")); + } + restarted.countDown(); + return Collections.emptyList(); + }); + Thread callbacks = new Thread(() -> { + try { + owner.set("master-1:5678#peer"); + subscriber.get().notify(new Event("/coordinator", "/coordinator", "", Event.Type.REMOVE)); + Assertions.assertTrue(canceled.await(5, TimeUnit.SECONDS)); + // The HA worker is draining the old DAO call. A second notification must still return. + owner.set(null); + subscriber.get().notify(new Event("/coordinator", "/coordinator", "", Event.Type.REMOVE)); + } catch (Throwable ex) { + failure.compareAndSet(null, ex); + } finally { + callbacksReturned.countDown(); + } + }); + callbacks.setDaemon(true); + try { + server.start(); + Assertions.assertTrue(fetching.await(5, TimeUnit.SECONDS)); + callbacks.start(); + Assertions.assertTrue(callbacksReturned.await(5, TimeUnit.SECONDS)); + Assertions.assertNull(failure.get()); + Assertions.assertFalse(server.isActive()); + Assertions.assertTrue(oldWorker.get().isAlive()); + Assertions.assertEquals(1L, restarted.getCount()); + releaseFetch.countDown(); + Assertions.assertTrue(restarted.await(5, TimeUnit.SECONDS)); + Assertions.assertFalse(oldWorker.get().isAlive()); + Assertions.assertTrue(server.isActive()); + Assertions.assertNull(failure.get()); + Mockito.verifyNoInteractions(serialCommandWaitHandler); + } finally { + releaseFetch.countDown(); + callbacks.join(5000); + Assertions.assertFalse(callbacks.isAlive()); + server.close(); + // Also cancel the listener's scheduled failover task, including when an assertion fails. + listener.changeToStandBy(); + } + } + + @Test + void demotionShouldStopBothWorkersBeforeWaitingForEither() throws Exception { + TaskGroupCoordinator taskGroupCoordinator = new TaskGroupCoordinator(); + TaskGroupDao taskGroupDao = Mockito.mock(TaskGroupDao.class); + ReflectionTestUtils.setField(taskGroupCoordinator, "taskGroupDao", taskGroupDao); + MasterCoordinator.MasterCoordinatorListener listener = new MasterCoordinator.MasterCoordinatorListener( + taskGroupCoordinator, Mockito.mock(IFailoverCoordinator.class), workflowSerialCoordinator); + ScheduledExecutorService scheduler = Mockito.mock(ScheduledExecutorService.class); + ScheduledFuture scheduled = Mockito.mock(ScheduledFuture.class); + Mockito.doReturn(scheduled).when(scheduler).scheduleWithFixedDelay( + Mockito.any(Runnable.class), Mockito.anyLong(), Mockito.anyLong(), Mockito.any(TimeUnit.class)); + CountDownLatch taskGroupFetching = new CountDownLatch(1); + CountDownLatch serialFetching = new CountDownLatch(1); + CountDownLatch taskGroupCanceled = new CountDownLatch(1); + CountDownLatch serialCanceled = new CountDownLatch(1); + CountDownLatch releaseTaskGroup = new CountDownLatch(1); + CountDownLatch releaseSerial = new CountDownLatch(1); + CountDownLatch closed = new CountDownLatch(1); + CountDownLatch failoverCanceled = new CountDownLatch(1); + Mockito.when(scheduled.cancel(true)).thenAnswer(invocation -> { + failoverCanceled.countDown(); + return true; + }); + AtomicReference taskGroupWorker = new AtomicReference<>(); + AtomicReference serialWorker = new AtomicReference<>(); + AtomicReference failure = new AtomicReference<>(); + Mockito.when(taskGroupDao.queryAllTaskGroups()).thenAnswer(invocation -> { + taskGroupWorker.set(Thread.currentThread()); + taskGroupFetching.countDown(); + awaitDatabaseResponse(releaseTaskGroup, taskGroupCanceled, failure); + return Collections.emptyList(); + }); + Mockito.when(serialCommandDao.fetchSerialCommands(Mockito.anyInt())).thenAnswer(invocation -> { + serialWorker.set(Thread.currentThread()); + serialFetching.countDown(); + awaitDatabaseResponse(releaseSerial, serialCanceled, failure); + return Collections.singletonList(SerialCommandDto.builder() + .workflowDefinitionCode(1L).workflowDefinitionVersion(1).build()); + }); + WorkflowDefinitionLog definition = new WorkflowDefinitionLog(); + definition.setExecutionType(WorkflowExecutionTypeEnum.SERIAL_WAIT); + Mockito.when(workflowDefinitionLogDao.queryByDefinitionCodeAndVersion(1L, 1)).thenReturn(definition); + Thread closer = new Thread(() -> { + try { + listener.changeToStandBy(); + } catch (Throwable ex) { + failure.compareAndSet(null, ex); + } finally { + closed.countDown(); + } + }); + closer.setDaemon(true); + try (MockedStatic factory = Mockito.mockStatic(MasterThreadFactory.class)) { + factory.when(MasterThreadFactory::getDefaultSchedulerThreadExecutor).thenReturn(scheduler); + listener.changeToActive(); + Assertions.assertTrue(serialFetching.await(5, TimeUnit.SECONDS)); + // Keep the real one-minute TaskGroup startup delay; both lifecycle implementations participate. + Assertions.assertTrue(taskGroupFetching.await(90, TimeUnit.SECONDS)); + closer.start(); + Assertions.assertTrue(taskGroupCanceled.await(5, TimeUnit.SECONDS)); + Assertions.assertTrue(serialCanceled.await(5, TimeUnit.SECONDS)); + releaseSerial.countDown(); + serialWorker.get().join(5000); + Assertions.assertFalse(serialWorker.get().isAlive()); + // A stop request alone must not make a new generation eligible to start. + assertThrows(IllegalStateException.class, workflowSerialCoordinator::start); + Mockito.verifyNoInteractions(serialCommandWaitHandler); + // Serial has stopped while TaskGroup is still blocked; close must still join TaskGroup. + Assertions.assertTrue(taskGroupWorker.get().isAlive()); + Assertions.assertEquals(1L, closed.getCount()); + Assertions.assertTrue(failoverCanceled.await(5, TimeUnit.SECONDS)); + Mockito.verify(scheduled).cancel(true); + releaseTaskGroup.countDown(); + Assertions.assertTrue(closed.await(5, TimeUnit.SECONDS)); + Assertions.assertFalse(taskGroupWorker.get().isAlive()); + Assertions.assertNull(failure.get()); + } finally { + releaseSerial.countDown(); + releaseTaskGroup.countDown(); + closer.join(5000); + listener.changeToStandBy(); + } + } + + private static void awaitDatabaseResponse(CountDownLatch release, CountDownLatch canceled, + AtomicReference failure) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(100); + while (release.getCount() != 0) { + try { + long remaining = deadline - System.nanoTime(); + if (remaining <= 0 || !release.await(remaining, TimeUnit.NANOSECONDS)) { + failure.compareAndSet(null, new AssertionError("Blocked DAO was not released")); + return; + } + } catch (InterruptedException ignored) { + // JDBC may ignore cancellation; only the test-controlled response releases the call. + canceled.countDown(); + } + } + } + + private static void awaitBlocked(Thread thread) { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (thread.isAlive() && thread.getState() != Thread.State.BLOCKED + && System.nanoTime() < deadline) { + Thread.yield(); + } + Assertions.assertEquals(Thread.State.BLOCKED, thread.getState()); + } + } diff --git a/dolphinscheduler-registry/dolphinscheduler-registry-api/src/main/java/org/apache/dolphinscheduler/registry/api/ha/AbstractHAServer.java b/dolphinscheduler-registry/dolphinscheduler-registry-api/src/main/java/org/apache/dolphinscheduler/registry/api/ha/AbstractHAServer.java index a79b14d99d5c..9cb762fcf70f 100644 --- a/dolphinscheduler-registry/dolphinscheduler-registry-api/src/main/java/org/apache/dolphinscheduler/registry/api/ha/AbstractHAServer.java +++ b/dolphinscheduler-registry/dolphinscheduler-registry-api/src/main/java/org/apache/dolphinscheduler/registry/api/ha/AbstractHAServer.java @@ -25,9 +25,13 @@ import org.apache.dolphinscheduler.registry.api.SubscribeListener; import java.util.List; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.RejectedExecutionException; import lombok.extern.slf4j.Slf4j; +import com.google.common.annotations.VisibleForTesting; import com.google.common.collect.Lists; @Slf4j @@ -39,7 +43,9 @@ public abstract class AbstractHAServer implements HAServer { private final String serverIdentify; - private ServerStatus serverStatus; + private volatile ServerStatus serverStatus; + + private final ExecutorService electionExecutor; private final List serverStatusChangeListeners; @@ -48,43 +54,101 @@ public abstract class AbstractHAServer implements HAServer { private static final int DEFAULT_MAX_RETRY_TIMES = 20; public AbstractHAServer(final Registry registry, final String selectorPath, final String serverIdentify) { + this(registry, selectorPath, serverIdentify, + ThreadUtils.newDaemonFixedThreadExecutor("HA-Election-%d", 1)); + } + + @VisibleForTesting + AbstractHAServer(final Registry registry, final String selectorPath, final String serverIdentify, + final ExecutorService electionExecutor) { + this.electionExecutor = electionExecutor; this.registry = registry; this.selectorPath = checkNotNull(selectorPath); - this.serverIdentify = checkNotNull(serverIdentify); + // Include the creation time to distinguish restarts at the same address. + this.serverIdentify = checkNotNull(serverIdentify) + "#" + System.currentTimeMillis(); this.serverStatus = ServerStatus.STAND_BY; this.serverStatusChangeListeners = Lists.newArrayList(new DefaultServerStatusChangeListener()); } @Override public void start() { - registry.subscribe(selectorPath, new SubscribeListener() { - - @Override - public void notify(Event event) { - if (Event.Type.REMOVE.equals(event.getType())) { - if (serverIdentify.equals(event.getEventData())) { - statusChange(ServerStatus.STAND_BY); - } else { - if (participateElection()) { - statusChange(ServerStatus.ACTIVE); - } + try { + registry.subscribe(selectorPath, new SubscribeListener() { + + @Override + public void notify(Event event) { + if (Event.Type.REMOVE.equals(event.getType())) { + enqueueElection(); } } + + @Override + public SubscribeScope getSubscribeScope() { + return SubscribeScope.PATH_ONLY; + } + }); + // Preserve startup completion and failure semantics; callbacks only enqueue work. + electionExecutor.submit(this::reconcileElection).get(); + } catch (InterruptedException e) { + close(); + Thread.currentThread().interrupt(); + throw new IllegalStateException("Interrupted while starting HA server", e); + } catch (ExecutionException e) { + close(); + if (e.getCause() instanceof RuntimeException) { + throw (RuntimeException) e.getCause(); } + throw new IllegalStateException("Failed to start HA server", e.getCause()); + } catch (RuntimeException e) { + close(); + throw e; + } + } - @Override - public SubscribeScope getSubscribeScope() { - return SubscribeScope.PATH_ONLY; + private void enqueueElection() { + if (electionExecutor.isShutdown()) { + return; + } + try { + electionExecutor.execute(() -> { + try { + reconcileElection(); + } catch (Exception e) { + log.error("Failed to reconcile HA ownership for {}", serverIdentify, e); + } + }); + } catch (RejectedExecutionException e) { + // The Registry can deliver callbacks concurrently with close(). + if (!electionExecutor.isShutdown()) { + throw e; } - }); + } + } - if (participateElection()) { + private void reconcileElection() { + // REMOVE is a request to recheck ownership, not a role decision to replay later. + boolean elected = participateElection(); + if (electionExecutor.isShutdown()) { + return; + } + if (elected) { statusChange(ServerStatus.ACTIVE); } else { + statusChange(ServerStatus.STAND_BY); log.info("Server {} is standby", serverIdentify); } } + @Override + public void close() { + // Let queued startup futures finish as no-ops. shutdownNow would strand their callers. + electionExecutor.shutdown(); + synchronized (this) { + // Wait for any publication already in progress, without joining the event worker: + // an Alert listener can close this server from that worker itself. + } + } + @Override public boolean isActive() { return ServerStatus.ACTIVE.equals(getServerStatus()); @@ -96,22 +160,35 @@ public boolean participateElection() { // If meet exception during participate election, will retry. // This can avoid the situation that the server is not elected as leader due to network jitter. for (int i = 0; i < DEFAULT_MAX_RETRY_TIMES; i++) { + if (electionExecutor.isShutdown()) { + return false; + } try { + if (!registry.acquireLock(electionLock)) { + return false; + } try { - if (registry.acquireLock(electionLock)) { - if (!registry.exists(selectorPath)) { - registry.put(selectorPath, serverIdentify, true); - return true; - } + if (electionExecutor.isShutdown()) { + return false; + } + boolean selectorExists = registry.exists(selectorPath); + if (electionExecutor.isShutdown()) { + return false; + } + if (selectorExists) { return serverIdentify.equals(registry.get(selectorPath)); } - return false; + registry.put(selectorPath, serverIdentify, true); + return true; } finally { registry.releaseLock(electionLock); } } catch (Exception e) { log.error("Participate election error, meet an exception, will retry after {}ms", DEFAULT_RETRY_INTERVAL, e); + if (electionExecutor.isShutdown()) { + return false; + } ThreadUtils.sleep(DEFAULT_RETRY_INTERVAL); } } @@ -129,15 +206,22 @@ public ServerStatus getServerStatus() { return serverStatus; } - private void statusChange(ServerStatus targetStatus) { + // Use the same monitor as external close() so it waits for ongoing status updates and listener calls. + private synchronized void statusChange(ServerStatus targetStatus) { + if (electionExecutor.isShutdown()) { + return; + } final ServerStatus originStatus = serverStatus; serverStatus = targetStatus; - synchronized (this) { - try { - serverStatusChangeListeners.forEach(listener -> listener.change(originStatus, serverStatus)); - } catch (Exception ex) { - log.error("Trigger ServerStatusChangeListener from {} -> {} error", originStatus, targetStatus, ex); - } + try { + serverStatusChangeListeners.forEach(listener -> { + // A listener may close this server; do not invoke subsequent listeners after that. + if (!electionExecutor.isShutdown()) { + listener.change(originStatus, targetStatus); + } + }); + } catch (Exception ex) { + log.error("Trigger ServerStatusChangeListener from {} -> {} error", originStatus, targetStatus, ex); } } } diff --git a/dolphinscheduler-registry/dolphinscheduler-registry-api/src/test/java/org/apache/dolphinscheduler/registry/api/ha/AbstractHAServerTest.java b/dolphinscheduler-registry/dolphinscheduler-registry-api/src/test/java/org/apache/dolphinscheduler/registry/api/ha/AbstractHAServerTest.java new file mode 100644 index 000000000000..bb39e4ec857e --- /dev/null +++ b/dolphinscheduler-registry/dolphinscheduler-registry-api/src/test/java/org/apache/dolphinscheduler/registry/api/ha/AbstractHAServerTest.java @@ -0,0 +1,557 @@ +/* + * 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.registry.api.ha; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.CALLS_REAL_METHODS; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.mockStatic; +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.common.thread.ThreadUtils; +import org.apache.dolphinscheduler.registry.api.Event; +import org.apache.dolphinscheduler.registry.api.Registry; +import org.apache.dolphinscheduler.registry.api.SubscribeListener; + +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.LinkedBlockingQueue; +import java.util.concurrent.ThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicReference; + +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.mockito.MockedStatic; + +class AbstractHAServerTest { + + private static final String SELECTOR_PATH = "/coordinator"; + private static final String ELECTION_LOCK = SELECTOR_PATH + "-lock"; + private static final String ADDRESS = "master-0:5678"; + + private final List servers = new ArrayList<>(); + private final List electionExecutors = new ArrayList<>(); + private final AtomicReference retryAction = new AtomicReference<>(() -> { + }); + private ExecutorService electionExecutor; + private final AtomicReference workerFailure = new AtomicReference<>(); + + private Registry registry; + private AtomicReference owner; + private AtomicReference subscriber; + private AbstractServerStatusChangeListener statusListener; + private AbstractHAServer server; + + @BeforeEach + void setUp() { + registry = mock(Registry.class); + owner = new AtomicReference<>(); + subscriber = new AtomicReference<>(); + statusListener = mock(AbstractServerStatusChangeListener.class, CALLS_REAL_METHODS); + server = newServer(); + electionExecutor = electionExecutors.get(0); + server.addServerStatusChangeListener(statusListener); + when(registry.acquireLock(ELECTION_LOCK)).thenReturn(true); + when(registry.exists(SELECTOR_PATH)).thenAnswer(invocation -> owner.get() != null); + when(registry.get(SELECTOR_PATH)).thenAnswer(invocation -> owner.get()); + doAnswer(invocation -> { + owner.set(invocation.getArgument(1)); + return null; + }).when(registry).put(eq(SELECTOR_PATH), anyString(), eq(true)); + doAnswer(invocation -> { + subscriber.set(invocation.getArgument(1)); + return null; + }).when(registry).subscribe(eq(SELECTOR_PATH), org.mockito.ArgumentMatchers.any()); + } + + @AfterEach + void tearDown() throws Exception { + servers.forEach(AbstractHAServer::close); + for (ExecutorService executor : electionExecutors) { + assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS)); + } + assertNoWorkerFailure(); + } + + @Test + void testInitialLeaderAndFollower() { + server.start(); + assertTrue(server.isActive()); + verify(statusListener).changeToActive(); + assertEquals(SubscribeListener.SubscribeScope.PATH_ONLY, subscriber.get().getSubscribeScope()); + + AbstractHAServer follower = newServer("master-1:5678"); + AbstractServerStatusChangeListener followerListener = + mock(AbstractServerStatusChangeListener.class, CALLS_REAL_METHODS); + follower.addServerStatusChangeListener(followerListener); + follower.start(); + assertFalse(follower.isActive()); + verify(followerListener, never()).changeToActive(); + verify(registry, times(1)).put(eq(SELECTOR_PATH), anyString(), eq(true)); + } + + @Test + void testDoesNotAdoptPreviousProcessWithSameAddress() { + // A legacy selector may survive its process until its lease/session expires. + owner.set(ADDRESS); + server.start(); + assertFalse(server.isActive()); + verify(statusListener, never()).changeToActive(); + verify(registry, never()).put(eq(SELECTOR_PATH), anyString(), eq(true)); + + owner.set(null); + remove(ADDRESS); + assertTrue(server.isActive()); + assertNotEquals(ADDRESS, owner.get()); + verify(statusListener).changeToActive(); + } + + @Test + void testDoesNotAdoptPredecessorWithEarlierTimestamp() { + // Seed an earlier incarnation explicitly. Millisecond timestamps do not guarantee + // different identities for two same-address instances created in the same millisecond. + String previousOwner = ADDRESS + "#1"; + owner.set(previousOwner); + server.start(); + assertFalse(server.isActive()); + verify(statusListener, never()).changeToActive(); + verify(registry, never()).put(eq(SELECTOR_PATH), anyString(), eq(true)); + + owner.set(null); + remove(previousOwner); + assertTrue(server.isActive()); + assertNotEquals(previousOwner, owner.get()); + verify(statusListener).changeToActive(); + } + + @Test + void testEmptyRemoveDemotesFormerLeaderWhenPeerOwnsSelector() { + server.start(); + owner.set("master-1:5678#peer-instance"); + remove(""); + assertFalse(server.isActive()); + verify(statusListener).changeToActive(); + verify(statusListener).changeToStandBy(); + } + + @Test + void testDelayedRemoveDoesNotRestartCurrentOwner() { + server.start(); + // The old notification arrives after this instance has already acquired the current key. + remove(owner.get()); + remove(""); + assertTrue(server.isActive()); + verify(statusListener, times(1)).changeToActive(); + verify(statusListener, never()).changeToStandBy(); + } + + @Test + void testOwnRemovalCanReelectWithoutAnotherPeer() { + server.start(); + String previousOwner = owner.get(); + owner.set(null); + remove(previousOwner); + assertTrue(server.isActive()); + verify(registry, times(2)).put(eq(SELECTOR_PATH), anyString(), eq(true)); + verify(statusListener, times(1)).changeToActive(); + verify(statusListener, never()).changeToStandBy(); + } + + @Test + void testAddAndUpdateDoNotTriggerElection() { + owner.set("master-1:5678#peer-instance"); + server.start(); + owner.set(null); + subscriber.get().notify(new Event(SELECTOR_PATH, SELECTOR_PATH, "", Event.Type.ADD)); + subscriber.get().notify(new Event(SELECTOR_PATH, SELECTOR_PATH, "", Event.Type.UPDATE)); + drain(); + assertFalse(server.isActive()); + verify(registry, times(1)).acquireLock(ELECTION_LOCK); + verify(statusListener, never()).changeToActive(); + } + + @Test + void testUnacquiredLockIsNotReleased() { + when(registry.acquireLock(ELECTION_LOCK)).thenReturn(false); + server.start(); + assertFalse(server.isActive()); + verify(registry, never()).releaseLock(ELECTION_LOCK); + verify(statusListener, never()).changeToActive(); + } + + @Test + void testActiveServerDemotesWhenLockCannotBeAcquired() { + server.start(); + when(registry.acquireLock(ELECTION_LOCK)).thenReturn(false); + remove(""); + assertFalse(server.isActive()); + verify(statusListener).changeToStandBy(); + verify(registry, times(1)).releaseLock(ELECTION_LOCK); + } + + @Test + void testAcquisitionErrorDoesNotReleaseUnacquiredLock() { + when(registry.acquireLock(ELECTION_LOCK)).thenThrow(new IllegalStateException("lock unavailable")); + assertThrows(IllegalStateException.class, server::start); + assertTrue(electionExecutor.isShutdown()); + assertFalse(server.isActive()); + verify(registry, never()).releaseLock(ELECTION_LOCK); + verify(statusListener, never()).changeToActive(); + } + + @Test + void testTransientElectionErrorKeepsRoleUntilOwnershipDecision() { + server.start(); + when(registry.exists(SELECTOR_PATH)) + .thenThrow(new IllegalStateException("temporary registry failure")) + .thenReturn(true); + retryAction.set(() -> { + assertTrue(server.isActive()); + verify(statusListener, never()).changeToStandBy(); + }); + remove(""); + assertTrue(server.isActive()); + verify(statusListener, times(1)).changeToActive(); + verify(statusListener, never()).changeToStandBy(); + } + + @Test + void testExhaustedRetriesPreserveOriginalRole() { + server.start(); + when(registry.exists(SELECTOR_PATH)).thenThrow(new IllegalStateException("registry unavailable")); + // Asynchronous callback failures are logged by the worker, not thrown on the Registry thread. + remove(""); + // Exception-driven demotion is deliberately outside this minimal candidate. + assertTrue(server.isActive()); + verify(statusListener, never()).changeToStandBy(); + verify(registry, times(21)).releaseLock(ELECTION_LOCK); + } + + @Test + void testRemoveCannotBeOverwrittenByEarlierStartupElection() throws Exception { + CountDownLatch startupElectionFinished = new CountDownLatch(1); + CountDownLatch allowStartupToReturn = new CountDownLatch(1); + AtomicBoolean firstRelease = new AtomicBoolean(true); + when(registry.releaseLock(ELECTION_LOCK)).thenAnswer(invocation -> { + if (firstRelease.getAndSet(false)) { + // Pause after election, before publication; the notification must queue behind it. + startupElectionFinished.countDown(); + assertTrue(allowStartupToReturn.await(5, TimeUnit.SECONDS)); + } + return true; + }); + ExecutorService caller = Executors.newSingleThreadExecutor(); + try { + Future startup = caller.submit(server::start); + assertTrue(startupElectionFinished.await(5, TimeUnit.SECONDS)); + owner.set("master-1:5678#peer-instance"); + notifyRemove(""); + allowStartupToReturn.countDown(); + startup.get(5, TimeUnit.SECONDS); + drain(); + assertFalse(server.isActive()); + verify(statusListener).changeToActive(); + verify(statusListener).changeToStandBy(); + } finally { + allowStartupToReturn.countDown(); + caller.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void testCallbacksReturnWhileRoleListenerIsBlockedAndRecheckCurrentOwner() throws Exception { + server.start(); + String ownIdentity = owner.get(); + CountDownLatch stopping = new CountDownLatch(1); + CountDownLatch allowStop = new CountDownLatch(1); + AtomicReference listenerThread = new AtomicReference<>(); + server.addServerStatusChangeListener(new AbstractServerStatusChangeListener() { + + @Override + public void changeToActive() { + } + + @Override + public void changeToStandBy() { + listenerThread.set(Thread.currentThread()); + stopping.countDown(); + await(allowStop); + } + }); + try { + owner.set("master-1:5678#peer-instance"); + notifyRemove(""); + assertTrue(stopping.await(5, TimeUnit.SECONDS)); + assertNotEquals(Thread.currentThread(), listenerThread.get()); + // Queue while this instance appears to own the key, then change the owner again. + // The queued request must re-read the Registry, not publish a captured ACTIVE decision. + owner.set(ownIdentity); + notifyRemove(""); + owner.set("master-2:5678#peer-instance"); + allowStop.countDown(); + drain(); + assertFalse(server.isActive()); + verify(statusListener, times(1)).changeToActive(); + } finally { + allowStop.countDown(); + } + } + + @Test + void testElectionAndPublicationRunOnSameWorker() { + AtomicReference electionThread = new AtomicReference<>(); + AtomicReference publicationThread = new AtomicReference<>(); + when(registry.acquireLock(ELECTION_LOCK)).thenAnswer(invocation -> { + electionThread.set(Thread.currentThread()); + return true; + }); + server.addServerStatusChangeListener((origin, target) -> publicationThread.set(Thread.currentThread())); + server.start(); + assertNotEquals(Thread.currentThread(), electionThread.get()); + assertEquals(electionThread.get(), publicationThread.get()); + owner.set("master-1:5678#peer-instance"); + remove(""); + assertEquals(electionThread.get(), publicationThread.get()); + } + + @Test + void testCloseDiscardsPendingRequestAndLateElectionResult() throws Exception { + CountDownLatch acquired = new CountDownLatch(1); + CountDownLatch allowElection = new CountDownLatch(1); + when(registry.acquireLock(ELECTION_LOCK)).thenAnswer(invocation -> { + acquired.countDown(); + assertTrue(allowElection.await(5, TimeUnit.SECONDS)); + return true; + }); + ExecutorService caller = Executors.newSingleThreadExecutor(); + try { + Future startup = caller.submit(server::start); + assertTrue(acquired.await(5, TimeUnit.SECONDS)); + notifyRemove(""); + server.close(); + // A notification after shutdown is ignored, including the submit/shutdown race. + notifyRemove(""); + allowElection.countDown(); + startup.get(5, TimeUnit.SECONDS); + assertTrue(electionExecutor.awaitTermination(5, TimeUnit.SECONDS)); + assertFalse(server.isActive()); + verify(statusListener, never()).changeToActive(); + verify(registry, times(1)).acquireLock(ELECTION_LOCK); + } finally { + allowElection.countDown(); + caller.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void testCloseFromListenerDoesNotDeadlockOrActivateQueuedWork() throws Exception { + server.start(); + CountDownLatch closed = new CountDownLatch(1); + server.addServerStatusChangeListener(new AbstractServerStatusChangeListener() { + + @Override + public void changeToActive() { + } + + @Override + public void changeToStandBy() { + // Alert closes its HA server from the demotion listener on the election worker. + server.close(); + owner.set(null); + notifyRemove(""); + closed.countDown(); + } + }); + AbstractServerStatusChangeListener subsequentListener = mock(AbstractServerStatusChangeListener.class); + server.addServerStatusChangeListener(subsequentListener); + owner.set("master-1:5678#peer-instance"); + notifyRemove(""); + assertTrue(closed.await(5, TimeUnit.SECONDS)); + assertTrue(electionExecutor.awaitTermination(5, TimeUnit.SECONDS)); + assertFalse(server.isActive()); + verify(statusListener, times(1)).changeToActive(); + org.mockito.Mockito.verifyNoInteractions(subsequentListener); + } + + @Test + void testCloseDoesNotStrandQueuedStartupFuture() throws Exception { + CountDownLatch workerBlocked = new CountDownLatch(1); + CountDownLatch releaseWorker = new CountDownLatch(1); + electionExecutor.execute(() -> { + workerBlocked.countDown(); + await(releaseWorker); + }); + ExecutorService caller = Executors.newSingleThreadExecutor(); + try { + assertTrue(workerBlocked.await(5, TimeUnit.SECONDS)); + Future startup = caller.submit(server::start); + // Observe the submitted Future in the queue, not just subscribe() before submission. + ThreadPoolExecutor executor = (ThreadPoolExecutor) electionExecutor; + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (executor.getQueue().isEmpty() && System.nanoTime() < deadline) { + Thread.yield(); + } + assertEquals(1, executor.getQueue().size()); + server.close(); + releaseWorker.countDown(); + startup.get(5, TimeUnit.SECONDS); + assertFalse(server.isActive()); + verify(registry, never()).acquireLock(ELECTION_LOCK); + } finally { + releaseWorker.countDown(); + caller.shutdownNow(); + assertTrue(caller.awaitTermination(5, TimeUnit.SECONDS)); + } + } + + @Test + void testCloseBeforeStartupSubmissionRejectsWithoutWaiting() { + doAnswer(invocation -> { + // Close after subscription but before the startup Future can be submitted. + server.close(); + return null; + }).when(registry).subscribe(eq(SELECTOR_PATH), org.mockito.ArgumentMatchers.any()); + assertThrows(java.util.concurrent.RejectedExecutionException.class, server::start); + verify(registry, never()).acquireLock(ELECTION_LOCK); + } + + @Test + void testExternalCloseWaitsForEnteredListener() throws Exception { + server.start(); + CountDownLatch entered = new CountDownLatch(1); + CountDownLatch releaseListener = new CountDownLatch(1); + AtomicReference failure = new AtomicReference<>(); + server.addServerStatusChangeListener((origin, target) -> { + entered.countDown(); + await(releaseListener); + }); + Thread closer = new Thread(() -> { + try { + server.close(); + } catch (Throwable ex) { + failure.set(ex); + } + }); + closer.setDaemon(true); + try { + owner.set("master-1:5678#peer-instance"); + notifyRemove(""); + assertTrue(entered.await(5, TimeUnit.SECONDS)); + closer.start(); + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5); + while (closer.isAlive() && closer.getState() != Thread.State.BLOCKED + && System.nanoTime() < deadline) { + Thread.yield(); + } + // The listener holds the publication monitor; external close must wait for it. + assertEquals(Thread.State.BLOCKED, closer.getState()); + releaseListener.countDown(); + closer.join(5000); + assertFalse(closer.isAlive()); + assertNull(failure.get()); + } finally { + releaseListener.countDown(); + closer.join(5000); + } + } + + private void remove(String previousOwner) { + notifyRemove(previousOwner); + drain(); + } + + private void notifyRemove(String previousOwner) { + subscriber.get().notify(new Event(SELECTOR_PATH, SELECTOR_PATH, previousOwner, Event.Type.REMOVE)); + } + + private void drain() { + try { + electionExecutor.submit(() -> { + }).get(5, TimeUnit.SECONDS); + assertNoWorkerFailure(); + } catch (Exception e) { + throw new AssertionError("Election worker did not finish", e); + } + } + + private void assertNoWorkerFailure() { + if (workerFailure.get() != null) { + throw new AssertionError("Election worker failed", workerFailure.get()); + } + } + + private void await(CountDownLatch latch) { + try { + assertTrue(latch.await(5, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new AssertionError(e); + } + } + + private AbstractHAServer newServer() { + return newServer(ADDRESS); + } + + private AbstractHAServer newServer(String address) { + ExecutorService executor = new ThreadPoolExecutor(1, 1, 0L, TimeUnit.MILLISECONDS, + new LinkedBlockingQueue<>(), runnable -> { + Thread thread = new Thread(() -> { + // Static mocks are thread-local, so install the retry delay stub on the election worker. + try ( + MockedStatic threadUtils = + mockStatic(ThreadUtils.class, CALLS_REAL_METHODS)) { + threadUtils.when(() -> ThreadUtils.sleep(anyLong())).thenAnswer(invocation -> { + retryAction.get().run(); + return null; + }); + runnable.run(); + } + }, "test-ha-election"); + thread.setUncaughtExceptionHandler( + (failedThread, failure) -> workerFailure.compareAndSet(null, failure)); + return thread; + }); + electionExecutors.add(executor); + AbstractHAServer result = new AbstractHAServer(registry, SELECTOR_PATH, address, executor) { + }; + servers.add(result); + return result; + } +}