Skip to content

fix: prevent WebsocketSyncDataService from shutting down shared timer (#6785) #6982

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
yykaue wants to merge 1 commit into
base: master
Choose a base branch
from
Open

fix: prevent WebsocketSyncDataService from shutting down shared timer (#6785) #6982

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 20 additions & 4 deletions shenyu-common/src/main/java/org/apache/shenyu/common/timer/HierarchicalWheelTimer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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)) {
Expand All @@ -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();
}
}
}
Expand Down
41 changes: 41 additions & 0 deletions shenyu-common/src/test/java/org/apache/shenyu/common/timer/HierarchicalWheelTimerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 .
Expand Down Expand Up @@ -54,6 +59,14 @@ public void setUp() {
timer = WheelTimerFactory.newWheelTimer();
timerTaskList = new TimerTaskList(taskCount);
}

/**
* Tear down.
*/
@AfterEach
public void tearDown() {
timer.shutdown();
}

/**
* Test timer.
Expand Down Expand Up @@ -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.
Expand Down
34 changes: 23 additions & 11 deletions .../src/main/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataService.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,8 @@ public class WebsocketSyncDataService implements SyncDataService {

private TimerTask timerTask;

private boolean closed;

private final ServerProperties serverProperties;

/**
Expand All @@ -107,7 +109,7 @@ public WebsocketSyncDataService(
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;
Expand All @@ -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...");
}
Expand Down Expand Up @@ -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) {
Expand Down
22 changes: 21 additions & 1 deletion .../main/java/org/apache/shenyu/plugin/sync/data/websocket/client/ShenyuWebsocketClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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.
*
Expand Down Expand Up @@ -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);
Expand All @@ -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;
Expand All @@ -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();
}
}
}

Expand Down
98 changes: 86 additions & 12 deletions .../test/java/org/apache/shenyu/plugin/sync/data/websocket/WebsocketSyncDataServiceTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -27,15 +30,23 @@
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;
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;

Expand All @@ -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");
Expand All @@ -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> 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));
}
}
Loading
Loading