diff --git a/shenyu-common/src/main/java/org/apache/shenyu/common/timer/HierarchicalWheelTimer.java b/shenyu-common/src/main/java/org/apache/shenyu/common/timer/HierarchicalWheelTimer.java index bb6f4f6f6b1e..1abffc50958b 100644 --- a/shenyu-common/src/main/java/org/apache/shenyu/common/timer/HierarchicalWheelTimer.java +++ b/shenyu-common/src/main/java/org/apache/shenyu/common/timer/HierarchicalWheelTimer.java @@ -113,10 +113,16 @@ private void addTimerTaskEntry(final TimerTaskList.TimerTaskEntry timerTaskEntry @Override public void advanceClock(final long timeoutMs) throws InterruptedException { + if (taskExecutor.isShutdown()) { + return; + } TimerTaskList bucket = delayQueue.poll(timeoutMs, TimeUnit.MILLISECONDS); if (Objects.nonNull(bucket)) { writeLock.lock(); try { + if (taskExecutor.isShutdown()) { + return; + } while (Objects.nonNull(bucket)) { timingWheel.advanceClock(bucket.getExpiration()); bucket.flush(this::addTimerTaskEntry); @@ -129,6 +135,9 @@ public void advanceClock(final long timeoutMs) throws InterruptedException { } private void start() { + if (taskExecutor.isShutdown()) { + throw new IllegalStateException("Timer already shutdown"); + } int state = WORKER_STATE_UPDATER.get(this); if (state == 0) { if (WORKER_STATE_UPDATER.compareAndSet(this, 0, 1)) { @@ -144,28 +153,35 @@ public int size() { @Override public void shutdown() { - taskExecutor.shutdown(); + writeLock.lock(); + try { + workerThread.interrupt(); + taskExecutor.shutdown(); + } finally { + writeLock.unlock(); + } } private static class Worker implements Runnable { - private final Timer timer; + private final HierarchicalWheelTimer timer; /** * Instantiates a new Worker. * * @param timer the timer */ - Worker(final Timer timer) { + Worker(final HierarchicalWheelTimer timer) { this.timer = timer; } @Override public void run() { - while (true) { + while (!Thread.currentThread().isInterrupted()) { try { timer.advanceClock(100L); } catch (InterruptedException ignored) { + Thread.currentThread().interrupt(); } } } diff --git a/shenyu-common/src/test/java/org/apache/shenyu/common/timer/HierarchicalWheelTimerTest.java b/shenyu-common/src/test/java/org/apache/shenyu/common/timer/HierarchicalWheelTimerTest.java index 755d42af34e1..a910ea3abf8e 100644 --- a/shenyu-common/src/test/java/org/apache/shenyu/common/timer/HierarchicalWheelTimerTest.java +++ b/shenyu-common/src/test/java/org/apache/shenyu/common/timer/HierarchicalWheelTimerTest.java @@ -17,14 +17,19 @@ package org.apache.shenyu.common.timer; +import org.junit.jupiter.api.AfterEach; import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; +import java.lang.reflect.Field; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicInteger; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertSame; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; /** * HierarchicalWheelTimerTest . @@ -54,6 +59,14 @@ public void setUp() { timer = WheelTimerFactory.newWheelTimer(); timerTaskList = new TimerTaskList(taskCount); } + + /** + * Tear down. + */ + @AfterEach + public void tearDown() { + timer.shutdown(); + } /** * Test timer. @@ -86,6 +99,34 @@ public void run(final TaskEntity taskEntity) { timerTask.cancel(); assertEquals(timer.size(), 0); } + + /** + * Test shutdown. + * + * @throws Exception reflection exception + */ + @Test + public void testShutdownStopsWorkerAndRejectsNewTasks() throws Exception { + timer.add(new TimerTask(TimeUnit.MINUTES.toMillis(1)) { + @Override + public void run(final TaskEntity taskEntity) { + } + }); + Field workerThreadField = HierarchicalWheelTimer.class.getDeclaredField("workerThread"); + workerThreadField.setAccessible(true); + Thread workerThread = (Thread) workerThreadField.get(timer); + assertTrue(workerThread.isAlive()); + + timer.shutdown(); + + workerThread.join(TimeUnit.SECONDS.toMillis(1)); + assertFalse(workerThread.isAlive()); + assertThrows(IllegalStateException.class, () -> timer.add(new TimerTask(1) { + @Override + public void run(final TaskEntity taskEntity) { + } + })); + } /** * Test list foreach. diff --git a/shenyu-sync-data-center/shenyu-sync-data-websocket/src/main/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataService.java b/shenyu-sync-data-center/shenyu-sync-data-websocket/src/main/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataService.java index 2b5abaa950c0..481ad9854b20 100644 --- a/shenyu-sync-data-center/shenyu-sync-data-websocket/src/main/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataService.java +++ b/shenyu-sync-data-center/shenyu-sync-data-websocket/src/main/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataService.java @@ -82,6 +82,8 @@ public class WebsocketSyncDataService implements SyncDataService { private TimerTask timerTask; + private boolean closed; + private final ServerProperties serverProperties; /** @@ -107,7 +109,7 @@ public WebsocketSyncDataService( final List aiProxyApiKeyDataSubscribers, final ServerProperties serverProperties) { - this.timer = WheelTimerFactory.getSharedTimer(); + this.timer = WheelTimerFactory.newWheelTimer(); this.websocketConfig = websocketConfig; this.pluginDataSubscriber = pluginDataSubscriber; this.metaDataSubscribers = metaDataSubscribers; @@ -131,7 +133,10 @@ public void doRun(final String key, final TimerTask timerTask) { }); } - private void masterCheck() { + private synchronized void masterCheck() { + if (closed) { + return; + } if (LOG.isDebugEnabled()) { LOG.debug("master checking task start..."); } @@ -165,18 +170,25 @@ private void masterCheck() { } @Override - public void close() { - if (CollectionUtils.isNotEmpty(clients)) { - for (ShenyuWebsocketClient client : clients) { - if (Objects.nonNull(client)) { - client.close(); + public synchronized void close() { + if (closed) { + return; + } + closed = true; + try { + if (Objects.nonNull(timerTask)) { + timerTask.cancel(); + } + if (CollectionUtils.isNotEmpty(clients)) { + for (ShenyuWebsocketClient client : clients) { + if (Objects.nonNull(client)) { + client.nowClose(); + } } } + } finally { + timer.shutdown(); } - if (Objects.nonNull(timerTask)) { - timerTask.cancel(); - } - timer.shutdown(); } private ShenyuWebsocketClient createClient(final String url) { diff --git a/shenyu-sync-data-center/shenyu-sync-data-websocket/src/main/java/org/apache/shenyu/plugin/sync/data/websocket/client/ShenyuWebsocketClient.java b/shenyu-sync-data-center/shenyu-sync-data-websocket/src/main/java/org/apache/shenyu/plugin/sync/data/websocket/client/ShenyuWebsocketClient.java index ddf5878055b2..76917229d059 100644 --- a/shenyu-sync-data-center/shenyu-sync-data-websocket/src/main/java/org/apache/shenyu/plugin/sync/data/websocket/client/ShenyuWebsocketClient.java +++ b/shenyu-sync-data-center/shenyu-sync-data-websocket/src/main/java/org/apache/shenyu/plugin/sync/data/websocket/client/ShenyuWebsocketClient.java @@ -103,12 +103,16 @@ public final class ShenyuWebsocketClient extends WebSocketClient { private final String namespaceId; + private final AtomicBoolean manuallyClosed = new AtomicBoolean(false); + private final AtomicBoolean reconnecting = new AtomicBoolean(false); private volatile long lastReconnectAttemptTime; private final AtomicInteger reconnectBackoff = new AtomicInteger(0); + private volatile Thread reconnectThread; + /** * Instantiates a new shenyu websocket client. * @@ -263,14 +267,22 @@ public void close() { * now close. will cancel the task execution. */ public void nowClose() { - this.close(); + this.manuallyClosed.set(true); if (Objects.nonNull(timerTask)) { timerTask.cancel(); } + Thread currentReconnectThread = this.reconnectThread; + if (Objects.nonNull(currentReconnectThread)) { + currentReconnectThread.interrupt(); + } + this.close(); } private void healthCheck() { try { + if (this.manuallyClosed.get()) { + return; + } if (!this.isOpen()) { if (this.reconnecting.compareAndSet(false, true)) { RECONNECT_EXECUTOR.submit(this::doReconnect); @@ -287,7 +299,11 @@ private void healthCheck() { } private void doReconnect() { + this.reconnectThread = Thread.currentThread(); try { + if (this.manuallyClosed.get()) { + return; + } long backoff = calculateBackoff(); long since = System.currentTimeMillis() - lastReconnectAttemptTime; long waitMs = backoff - since; @@ -305,7 +321,11 @@ private void doReconnect() { reconnectBackoff.set(Math.min(reconnectBackoff.get() + 1, 10)); LOG.error("websocket reconnect server[{}] error", this.getURI(), e); } finally { + this.reconnectThread = null; this.reconnecting.set(false); + if (this.manuallyClosed.get()) { + this.close(); + } } } diff --git a/shenyu-sync-data-center/shenyu-sync-data-websocket/src/test/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataServiceTest.java b/shenyu-sync-data-center/shenyu-sync-data-websocket/src/test/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataServiceTest.java index a34ef976d69d..65da2272a7b7 100644 --- a/shenyu-sync-data-center/shenyu-sync-data-websocket/src/test/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataServiceTest.java +++ b/shenyu-sync-data-center/shenyu-sync-data-websocket/src/test/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataServiceTest.java @@ -18,6 +18,9 @@ package org.apache.shenyu.plugin.sync.data.websocket; import org.apache.shenyu.common.config.ShenyuConfig; +import org.apache.shenyu.common.timer.Timer; +import org.apache.shenyu.common.timer.TimerTask; +import org.apache.shenyu.common.timer.WheelTimerFactory; import org.apache.shenyu.plugin.sync.data.websocket.client.ShenyuWebsocketClient; import org.apache.shenyu.plugin.sync.data.websocket.config.WebsocketConfig; import org.apache.shenyu.sync.data.api.AiProxyApiKeyDataSubscriber; @@ -27,6 +30,8 @@ import org.apache.shenyu.sync.data.api.PluginDataSubscriber; import org.apache.shenyu.sync.data.api.ProxySelectorDataSubscriber; import org.junit.jupiter.api.Test; +import org.mockito.InOrder; +import org.mockito.MockedStatic; import org.springframework.boot.autoconfigure.web.ServerProperties; import java.lang.reflect.Field; @@ -34,8 +39,14 @@ import java.util.Collections; import java.util.List; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.inOrder; 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; @@ -44,18 +55,7 @@ public final class WebsocketSyncDataServiceTest { @Test @SuppressWarnings("unchecked") public void testMasterCheckClosesRemovedClient() throws Exception { - WebsocketConfig websocketConfig = new WebsocketConfig(); - websocketConfig.setUrls(Collections.emptyList()); - WebsocketSyncDataService websocketSyncDataService = new WebsocketSyncDataService( - websocketConfig, - new ShenyuConfig(), - mock(PluginDataSubscriber.class), - Collections.emptyList(), - Collections.emptyList(), - Collections.emptyList(), - Collections.emptyList(), - Collections.emptyList(), - mock(ServerProperties.class)); + WebsocketSyncDataService websocketSyncDataService = createWebsocketSyncDataService(); ShenyuWebsocketClient websocketClient = mock(ShenyuWebsocketClient.class); when(websocketClient.isOpen()).thenReturn(false); Field clientsField = WebsocketSyncDataService.class.getDeclaredField("clients"); @@ -74,4 +74,78 @@ public void testMasterCheckClosesRemovedClient() throws Exception { websocketSyncDataService.close(); } } + + @Test + @SuppressWarnings("unchecked") + public void testCloseShutsDownPrivateTimer() throws Exception { + Timer sharedTimer = mock(Timer.class); + Timer privateTimer = mock(Timer.class); + try (MockedStatic wheelTimerFactory = mockStatic(WheelTimerFactory.class)) { + wheelTimerFactory.when(WheelTimerFactory::getSharedTimer).thenReturn(sharedTimer); + wheelTimerFactory.when(WheelTimerFactory::newWheelTimer).thenReturn(privateTimer); + final WebsocketSyncDataService websocketSyncDataService = createWebsocketSyncDataService(); + ShenyuWebsocketClient websocketClient = mock(ShenyuWebsocketClient.class); + Field clientsField = WebsocketSyncDataService.class.getDeclaredField("clients"); + clientsField.setAccessible(true); + List clients = (List) clientsField + .get(websocketSyncDataService); + clients.add(websocketClient); + TimerTask timerTask = mock(TimerTask.class); + Field timerTaskField = WebsocketSyncDataService.class.getDeclaredField("timerTask"); + timerTaskField.setAccessible(true); + timerTaskField.set(websocketSyncDataService, timerTask); + + websocketSyncDataService.close(); + Method masterCheck = WebsocketSyncDataService.class.getDeclaredMethod("masterCheck"); + masterCheck.setAccessible(true); + masterCheck.invoke(websocketSyncDataService); + websocketSyncDataService.close(); + + InOrder closeOrder = inOrder(timerTask, websocketClient); + closeOrder.verify(timerTask).cancel(); + closeOrder.verify(websocketClient).nowClose(); + verify(websocketClient, times(1)).nowClose(); + verify(timerTask, times(1)).cancel(); + verify(privateTimer, times(1)).shutdown(); + verify(sharedTimer, never()).shutdown(); + wheelTimerFactory.verify(WheelTimerFactory::getSharedTimer, never()); + } + } + + @Test + @SuppressWarnings("unchecked") + public void testCloseShutsDownPrivateTimerWhenClientCloseFails() throws Exception { + final Timer privateTimer = mock(Timer.class); + try (MockedStatic wheelTimerFactory = mockStatic(WheelTimerFactory.class)) { + wheelTimerFactory.when(WheelTimerFactory::newWheelTimer).thenReturn(privateTimer); + final WebsocketSyncDataService websocketSyncDataService = createWebsocketSyncDataService(); + final ShenyuWebsocketClient websocketClient = mock(ShenyuWebsocketClient.class); + final IllegalStateException clientCloseException = new IllegalStateException("client close failed"); + doThrow(clientCloseException).when(websocketClient).nowClose(); + final Field clientsField = WebsocketSyncDataService.class.getDeclaredField("clients"); + clientsField.setAccessible(true); + final List clients = (List) clientsField + .get(websocketSyncDataService); + clients.add(websocketClient); + + assertThrows(IllegalStateException.class, websocketSyncDataService::close); + + verify(privateTimer).shutdown(); + } + } + + private WebsocketSyncDataService createWebsocketSyncDataService() { + WebsocketConfig websocketConfig = new WebsocketConfig(); + websocketConfig.setUrls(Collections.emptyList()); + return new WebsocketSyncDataService( + websocketConfig, + new ShenyuConfig(), + mock(PluginDataSubscriber.class), + Collections.emptyList(), + Collections.emptyList(), + Collections.emptyList(), + Collections.emptyList(), + Collections.emptyList(), + mock(ServerProperties.class)); + } } diff --git a/shenyu-sync-data-center/shenyu-sync-data-websocket/src/test/java/org/apache/shenyu/plugin/sync/data/websocket/client/ShenyuWebsocketClientTest.java b/shenyu-sync-data-center/shenyu-sync-data-websocket/src/test/java/org/apache/shenyu/plugin/sync/data/websocket/client/ShenyuWebsocketClientTest.java index dcb42b7db6c0..db730837c0ff 100644 --- a/shenyu-sync-data-center/shenyu-sync-data-websocket/src/test/java/org/apache/shenyu/plugin/sync/data/websocket/client/ShenyuWebsocketClientTest.java +++ b/shenyu-sync-data-center/shenyu-sync-data-websocket/src/test/java/org/apache/shenyu/plugin/sync/data/websocket/client/ShenyuWebsocketClientTest.java @@ -22,6 +22,7 @@ import org.apache.shenyu.common.enums.ConfigGroupEnum; import org.apache.shenyu.common.enums.DataEventTypeEnum; import org.apache.shenyu.common.exception.ShenyuException; +import org.apache.shenyu.common.timer.TimerTask; import org.apache.shenyu.common.utils.GsonUtils; import org.apache.shenyu.sync.data.api.AuthDataSubscriber; import org.apache.shenyu.sync.data.api.MetaDataSubscriber; @@ -47,14 +48,17 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import org.mockito.Answers; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.spy; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.withSettings; @@ -135,6 +139,7 @@ public void testOnError() { private ShenyuWebsocketClient createMockClient() { ShenyuWebsocketClient client = mock(ShenyuWebsocketClient.class, withSettings().defaultAnswer(Answers.CALLS_REAL_METHODS)); + setField(client, "manuallyClosed", new AtomicBoolean(false)); setField(client, "reconnecting", new AtomicBoolean(false)); setField(client, "reconnectBackoff", new AtomicInteger(0)); setField(client, "lastReconnectAttemptTime", 0L); @@ -264,6 +269,59 @@ void testHealthCheckResetsBackoffAndSendsPingWhenOpen() { verify(client).sendPing(); } + @Test + void testNowClosePreventsPendingReconnect() throws InterruptedException { + ShenyuWebsocketClient client = createMockClient(); + TimerTask timerTask = mock(TimerTask.class); + setField(client, "timerTask", timerTask); + doReturn(false).when(client).isOpen(); + + client.nowClose(); + invokePrivate(client, "doReconnect"); + + verify(timerTask).cancel(); + verify(client, never()).reconnectBlocking(); + } + + @Test + void testNowCloseInterruptsRunningReconnect() { + ShenyuWebsocketClient client = createMockClient(); + Thread reconnectThread = mock(Thread.class); + setField(client, "reconnectThread", reconnectThread); + doReturn(false).when(client).isOpen(); + + client.nowClose(); + + verify(reconnectThread).interrupt(); + } + + @Test + void testNowCloseCancelsTimerTaskWhenSocketCloseFails() { + ShenyuWebsocketClient client = createMockClient(); + TimerTask timerTask = mock(TimerTask.class); + setField(client, "timerTask", timerTask); + IllegalStateException closeException = new IllegalStateException("socket close failed"); + doThrow(closeException).when(client).close(); + + assertThrows(IllegalStateException.class, client::nowClose); + + verify(timerTask).cancel(); + } + + @Test + void testReconnectClosesConnectionWhenNowClosedDuringAttempt() throws InterruptedException { + ShenyuWebsocketClient client = createMockClient(); + doAnswer(invocation -> { + ((AtomicBoolean) getField(client, "manuallyClosed")).set(true); + return true; + }).when(client).reconnectBlocking(); + doNothing().when(client).close(); + + invokePrivate(client, "doReconnect"); + + verify(client).close(); + } + // ---------- doReconnect tests ---------- // reconnectBlocking() is stubbed to throw, so no real socket connection is attempted.