This is an automated email from the ASF dual-hosted git repository.
Aias00 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/shenyu.git
The following commit(s) were added to refs/heads/master by this push:
new 58dd5a517e fix: improve websocket client closure safety, timer
shutdown robustness, and lifecycle management (#6982)
58dd5a517e is described below
commit 58dd5a517ee32d1bf8b58e479881894439c54537
Author: Limbo <[email protected]>
AuthorDate: Wed Aug 26 16:41:31 2026 +0800
fix: improve websocket client closure safety, timer shutdown robustness,
and lifecycle management (#6982)
Co-authored-by: aias00 <[email protected]>
---
.../common/timer/HierarchicalWheelTimer.java | 24 +++++-
.../common/timer/HierarchicalWheelTimerTest.java | 41 +++++++++
.../data/websocket/WebsocketSyncDataService.java | 34 +++++---
.../websocket/client/ShenyuWebsocketClient.java | 22 ++++-
.../websocket/WebsocketSyncDataServiceTest.java | 98 +++++++++++++++++++---
.../client/ShenyuWebsocketClientTest.java | 58 +++++++++++++
6 files changed, 249 insertions(+), 28 deletions(-)
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 bb6f4f6f6b..1abffc5095 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 @@ public class HierarchicalWheelTimer implements Timer {
@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 class HierarchicalWheelTimer implements Timer {
}
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 class HierarchicalWheelTimer implements Timer {
@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 755d42af34..a910ea3abf 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 class HierarchicalWheelTimerTest {
timer = WheelTimerFactory.newWheelTimer();
timerTaskList = new TimerTaskList(taskCount);
}
+
+ /**
+ * Tear down.
+ */
+ @AfterEach
+ public void tearDown() {
+ timer.shutdown();
+ }
/**
* Test timer.
@@ -86,6 +99,34 @@ public class HierarchicalWheelTimerTest {
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 2b5abaa950..481ad9854b 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 class WebsocketSyncDataService implements
SyncDataService {
final
List<org.apache.shenyu.sync.data.api.AiProxyApiKeyDataSubscriber>
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 class WebsocketSyncDataService implements
SyncDataService {
});
}
- private void masterCheck() {
+ private synchronized void masterCheck() {
+ if (closed) {
+ return;
+ }
if (LOG.isDebugEnabled()) {
LOG.debug("master checking task start...");
}
@@ -165,18 +170,25 @@ public class WebsocketSyncDataService implements
SyncDataService {
}
@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 ddf5878055..76917229d0 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 final class ShenyuWebsocketClient extends
WebSocketClient {
* 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 @@ public final class ShenyuWebsocketClient extends
WebSocketClient {
}
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 @@ public final class ShenyuWebsocketClient extends
WebSocketClient {
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 a34ef976d6..65da2272a7 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.MetaDataSubscriber;
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.lang.reflect.Method;
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.<MetaDataSubscriber>emptyList(),
- Collections.<AuthDataSubscriber>emptyList(),
- Collections.<ProxySelectorDataSubscriber>emptyList(),
- Collections.<DiscoveryUpstreamDataSubscriber>emptyList(),
- Collections.<AiProxyApiKeyDataSubscriber>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 final class WebsocketSyncDataServiceTest {
websocketSyncDataService.close();
}
}
+
+ @Test
+ @SuppressWarnings("unchecked")
+ public void testCloseShutsDownPrivateTimer() throws Exception {
+ Timer sharedTimer = mock(Timer.class);
+ Timer privateTimer = mock(Timer.class);
+ try (MockedStatic<WheelTimerFactory> 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<ShenyuWebsocketClient> clients =
(List<ShenyuWebsocketClient>) 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> 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<ShenyuWebsocketClient> clients =
(List<ShenyuWebsocketClient>) 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.<MetaDataSubscriber>emptyList(),
+ Collections.<AuthDataSubscriber>emptyList(),
+ Collections.<ProxySelectorDataSubscriber>emptyList(),
+ Collections.<DiscoveryUpstreamDataSubscriber>emptyList(),
+ Collections.<AiProxyApiKeyDataSubscriber>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 dcb42b7db6..db730837c0 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.dto.WebsocketData;
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 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.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 class ShenyuWebsocketClientTest {
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 @@ public class ShenyuWebsocketClientTest {
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.
