You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@dolphinscheduler.apache.org by zi...@apache.org on 2026/10/20 06:03:56 UTC
(dolphinscheduler) branch dev updated: [Fix-18274][Registry] Handle expired JDBC heartbeat sessions (#18416)
This is an automated email from the ASF dual-hosted git repository.
SbloodyS pushed a commit to branch dev
in repository https://gitbox.apache.org/repos/asf/dolphinscheduler.git
The following commit(s) were added to refs/heads/dev by this push:
new ba5443ed7b [Fix-18274][Registry] Handle expired JDBC heartbeat sessions (#18416)
ba5443ed7b is described below
commit ba5443ed7b654a8ae5032c0731edd278602e1591
Author: Yanjun Qiu <15...@users.noreply.github.com>
AuthorDate: Tue Oct 20 14:03:50 2026 +0800
[Fix-18274][Registry] Handle expired JDBC heartbeat sessions (#18416)
---
.../mapper/JdbcRegistryClientHeartbeatMapper.java | 6 +
.../repository/JdbcRegistryClientRepository.java | 4 +
.../registry/jdbc/server/JdbcRegistryServer.java | 93 +++++---
.../jdbc/server/JdbcRegistryServerTest.java | 256 +++++++++++++++++++++
4 files changed, 329 insertions(+), 30 deletions(-)
diff --git a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/mapper/JdbcRegistryClientHeartbeatMapper.java b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/mapper/JdbcRegistryClientHeartbeatMapper.java
index 2b8499bb48..73c2b66772 100644
--- a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/mapper/JdbcRegistryClientHeartbeatMapper.java
+++ b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/mapper/JdbcRegistryClientHeartbeatMapper.java
@@ -19,6 +19,8 @@ package org.apache.dolphinscheduler.plugin.registry.jdbc.mapper;
import org.apache.dolphinscheduler.plugin.registry.jdbc.model.DO.JdbcRegistryClientHeartbeat;
+import org.apache.ibatis.annotations.Delete;
+import org.apache.ibatis.annotations.Param;
import org.apache.ibatis.annotations.Select;
import java.util.List;
@@ -30,4 +32,8 @@ public interface JdbcRegistryClientHeartbeatMapper extends BaseMapper<JdbcRegist
@Select("select * from t_ds_jdbc_registry_client_heartbeat")
List<JdbcRegistryClientHeartbeat> selectAll();
+ @Delete("delete from t_ds_jdbc_registry_client_heartbeat "
+ + "where id = #{id} and last_heartbeat_time = #{lastHeartbeatTime}")
+ int deleteByIdAndLastHeartbeatTime(@Param("id") Long id, @Param("lastHeartbeatTime") Long lastHeartbeatTime);
+
}
diff --git a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/repository/JdbcRegistryClientRepository.java b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/repository/JdbcRegistryClientRepository.java
index 1791f3c942..cee889bf05 100644
--- a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/repository/JdbcRegistryClientRepository.java
+++ b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/repository/JdbcRegistryClientRepository.java
@@ -52,6 +52,10 @@ public class JdbcRegistryClientRepository {
jdbcRegistryClientHeartbeatMapper.deleteBatchIds(clientIds);
}
+ public boolean deleteByIdAndLastHeartbeatTime(Long id, Long lastHeartbeatTime) {
+ return jdbcRegistryClientHeartbeatMapper.deleteByIdAndLastHeartbeatTime(id, lastHeartbeatTime) == 1;
+ }
+
public boolean updateById(JdbcRegistryClientHeartbeatDTO jdbcRegistryClientHeartbeatDTO) {
JdbcRegistryClientHeartbeat jdbcRegistryClientHeartbeat =
JdbcRegistryClientHeartbeatDTO.toJdbcRegistryClientHeartbeat(jdbcRegistryClientHeartbeatDTO);
diff --git a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/server/JdbcRegistryServer.java b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/server/JdbcRegistryServer.java
index 61a934bd14..ccd43af77f 100644
--- a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/server/JdbcRegistryServer.java
+++ b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/main/java/org/apache/dolphinscheduler/plugin/registry/jdbc/server/JdbcRegistryServer.java
@@ -45,6 +45,7 @@ import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicReference;
import java.util.stream.Collectors;
import lombok.SneakyThrows;
@@ -70,7 +71,8 @@ public class JdbcRegistryServer implements IJdbcRegistryServer {
private final JdbcRegistryLockManager jdbcRegistryLockManager;
- private JdbcRegistryServerState jdbcRegistryServerState;
+ private final AtomicReference<JdbcRegistryServerState> serverState =
+ new AtomicReference<>(JdbcRegistryServerState.INIT);
private final List<IJdbcRegistryClient> jdbcRegistryClients = new CopyOnWriteArrayList<>();
@@ -81,7 +83,7 @@ public class JdbcRegistryServer implements IJdbcRegistryServer {
private final ScheduledExecutorService schedulerThreadExecutor;
- private Long lastSuccessHeartbeat;
+ private volatile long lastSuccessHeartbeat;
public JdbcRegistryServer(JdbcRegistryDataRepository jdbcRegistryDataRepository,
JdbcRegistryLockRepository jdbcRegistryLockRepository,
@@ -100,13 +102,12 @@ public class JdbcRegistryServer implements IJdbcRegistryServer {
transactionTemplate, schedulerThreadExecutor);
this.jdbcRegistryLockManager = new JdbcRegistryLockManager(
jdbcRegistryProperties, jdbcRegistryLockRepository);
- this.jdbcRegistryServerState = JdbcRegistryServerState.INIT;
lastSuccessHeartbeat = System.currentTimeMillis();
}
@Override
public void start() {
- if (jdbcRegistryServerState != JdbcRegistryServerState.INIT) {
+ if (serverState.get() != JdbcRegistryServerState.INIT) {
// The server is already started or stopped, will not start again.
return;
}
@@ -119,7 +120,10 @@ public class JdbcRegistryServer implements IJdbcRegistryServer {
jdbcRegistryProperties.getSessionTimeout().toMillis(),
TimeUnit.MILLISECONDS);
jdbcRegistryDataManager.start();
- jdbcRegistryServerState = JdbcRegistryServerState.STARTED;
+ if (!serverState.compareAndSet(JdbcRegistryServerState.INIT, JdbcRegistryServerState.STARTED)) {
+ log.warn("The JdbcRegistryServer state changed before startup completed: {}", serverState.get());
+ return;
+ }
doTriggerOnConnectedListener();
schedulerThreadExecutor.scheduleWithFixedDelay(
this::refreshClientsHeartbeat,
@@ -168,7 +172,7 @@ public class JdbcRegistryServer implements IJdbcRegistryServer {
@Override
public JdbcRegistryServerState getServerState() {
- return jdbcRegistryServerState;
+ return serverState.get();
}
@Override
@@ -256,7 +260,10 @@ public class JdbcRegistryServer implements IJdbcRegistryServer {
@Override
public void close() {
- jdbcRegistryServerState = JdbcRegistryServerState.STOPPED;
+ if (serverState.getAndSet(JdbcRegistryServerState.STOPPED) == JdbcRegistryServerState.STOPPED) {
+ log.warn("The JdbcRegistryServer is already STOPPED.");
+ return;
+ }
schedulerThreadExecutor.shutdown();
List<Long> clientIds = jdbcRegistryClients.stream()
.map(IJdbcRegistryClient::getJdbcRegistryClientIdentify)
@@ -269,23 +276,26 @@ public class JdbcRegistryServer implements IJdbcRegistryServer {
private void purgeInvalidJdbcRegistryMetadata() {
final StopWatch stopWatch = StopWatch.createStarted();
- if (jdbcRegistryServerState == JdbcRegistryServerState.STOPPED) {
+ JdbcRegistryServerState currentState = getServerState();
+ if (currentState == JdbcRegistryServerState.STOPPED
+ || currentState == JdbcRegistryServerState.DISCONNECTED) {
return;
}
// remove the client which is already dead from the registry, and remove it's related data and lock.
final List<JdbcRegistryClientHeartbeatDTO> jdbcRegistryClients = jdbcRegistryClientRepository.queryAll();
- final Set<Long> deadJdbcRegistryClientIds = jdbcRegistryClients
+ final Set<Long> deletedJdbcRegistryClientIds = jdbcRegistryClients
.stream()
.filter(JdbcRegistryClientHeartbeatDTO::isDead)
+ .filter(jdbcRegistryClient -> jdbcRegistryClientRepository.deleteByIdAndLastHeartbeatTime(
+ jdbcRegistryClient.getId(), jdbcRegistryClient.getLastHeartbeatTime()))
.map(JdbcRegistryClientHeartbeatDTO::getId)
.collect(Collectors.toSet());
- doPurgeJdbcRegistryClientInDB(deadJdbcRegistryClientIds);
// remove the data and lock which client is not exist.
final Set<Long> existJdbcRegistryClientIds = jdbcRegistryClients
.stream()
.map(JdbcRegistryClientHeartbeatDTO::getId)
- .filter(id -> !deadJdbcRegistryClientIds.contains(id))
+ .filter(id -> !deletedJdbcRegistryClientIds.contains(id))
.collect(Collectors.toSet());
jdbcRegistryDataManager.getAllJdbcRegistryData()
.stream()
@@ -321,8 +331,11 @@ public class JdbcRegistryServer implements IJdbcRegistryServer {
if (CollectionUtils.isEmpty(jdbcRegistryClients)) {
return;
}
- if (jdbcRegistryServerState == JdbcRegistryServerState.STOPPED) {
- log.warn("The JdbcRegistryServer is STOPPED, will not refresh clients: {} heartbeat.",
+ JdbcRegistryServerState currentState = getServerState();
+ if (currentState == JdbcRegistryServerState.STOPPED
+ || currentState == JdbcRegistryServerState.DISCONNECTED) {
+ log.warn("The JdbcRegistryServer is {}, will not refresh clients: {} heartbeat.",
+ currentState,
CollectionUtils.collect(jdbcRegistryClients, IJdbcRegistryClient::getJdbcRegistryClientIdentify));
return;
}
@@ -341,31 +354,51 @@ public class JdbcRegistryServer implements IJdbcRegistryServer {
}
JdbcRegistryClientHeartbeatDTO clone = jdbcRegistryClientHeartbeatDTO.clone();
clone.setLastHeartbeatTime(now);
- jdbcRegistryClientRepository.updateById(jdbcRegistryClientHeartbeatDTO);
+ if (!jdbcRegistryClientRepository.updateById(clone)) {
+ log.error("Refresh client: {} heartbeat failed, the client might have been deleted",
+ jdbcRegistryClientHeartbeatDTO.getId());
+ throw new IllegalStateException(
+ "The client heartbeat record no longer exists: " + jdbcRegistryClientHeartbeatDTO.getId());
+ }
jdbcRegistryClientHeartbeatDTO.setLastHeartbeatTime(clone.getLastHeartbeatTime());
}
- if (jdbcRegistryServerState == JdbcRegistryServerState.SUSPENDED) {
- jdbcRegistryServerState = JdbcRegistryServerState.STARTED;
- doTriggerReconnectedListener();
+ currentState = serverState.get();
+ boolean reconnected = false;
+ if (currentState == JdbcRegistryServerState.SUSPENDED) {
+ if (!serverState.compareAndSet(JdbcRegistryServerState.SUSPENDED, JdbcRegistryServerState.STARTED)) {
+ log.debug("Failed to reconnect JdbcRegistryServer; current state is {}", serverState.get());
+ return;
+ }
+ reconnected = true;
+ } else if (currentState != JdbcRegistryServerState.STARTED) {
+ return;
}
lastSuccessHeartbeat = now;
+ if (reconnected) {
+ doTriggerReconnectedListener();
+ }
log.debug("Success refresh clients: {} heartbeat.",
CollectionUtils.collect(jdbcRegistryClients, IJdbcRegistryClient::getJdbcRegistryClientIdentify));
} catch (Exception ex) {
log.error("Failed to refresh the client's term", ex);
- switch (jdbcRegistryServerState) {
- case STARTED:
- jdbcRegistryServerState = JdbcRegistryServerState.SUSPENDED;
- break;
- case SUSPENDED:
- if (System.currentTimeMillis() - lastSuccessHeartbeat > jdbcRegistryProperties.getSessionTimeout()
- .toMillis()) {
- jdbcRegistryServerState = JdbcRegistryServerState.DISCONNECTED;
- doTriggerOnDisConnectedListener();
- }
- break;
- default:
- break;
+ currentState = serverState.get();
+ if (currentState != JdbcRegistryServerState.STARTED
+ && currentState != JdbcRegistryServerState.SUSPENDED) {
+ return;
+ }
+ long sessionTimeoutMillis = jdbcRegistryProperties.getSessionTimeout().toMillis();
+ if (System.currentTimeMillis() - lastSuccessHeartbeat > sessionTimeoutMillis) {
+ if (!serverState.compareAndSet(currentState, JdbcRegistryServerState.DISCONNECTED)) {
+ log.debug("Failed to disconnect JdbcRegistryServer from state {}, current state is {}",
+ currentState,
+ serverState.get());
+ return;
+ }
+ doTriggerOnDisConnectedListener();
+ } else if (currentState == JdbcRegistryServerState.STARTED
+ && !serverState.compareAndSet(JdbcRegistryServerState.STARTED, JdbcRegistryServerState.SUSPENDED)) {
+ log.debug("Failed to suspend JdbcRegistryServer; current state is {}", serverState.get());
+ return;
}
}
}
diff --git a/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/test/java/org/apache/dolphinscheduler/plugin/registry/jdbc/server/JdbcRegistryServerTest.java b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/test/java/org/apache/dolphinscheduler/plugin/registry/jdbc/server/JdbcRegistryServerTest.java
new file mode 100644
index 0000000000..3e20d4946b
--- /dev/null
+++ b/dolphinscheduler-registry/dolphinscheduler-registry-plugins/dolphinscheduler-registry-jdbc/src/test/java/org/apache/dolphinscheduler/plugin/registry/jdbc/server/JdbcRegistryServerTest.java
@@ -0,0 +1,256 @@
+/*
+ * 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.plugin.registry.jdbc.server;
+
+import org.apache.dolphinscheduler.plugin.registry.jdbc.JdbcRegistryProperties;
+import org.apache.dolphinscheduler.plugin.registry.jdbc.client.IJdbcRegistryClient;
+import org.apache.dolphinscheduler.plugin.registry.jdbc.client.JdbcRegistryClientIdentify;
+import org.apache.dolphinscheduler.plugin.registry.jdbc.model.DTO.JdbcRegistryClientHeartbeatDTO;
+import org.apache.dolphinscheduler.plugin.registry.jdbc.model.DTO.JdbcRegistryLockDTO;
+import org.apache.dolphinscheduler.plugin.registry.jdbc.repository.JdbcRegistryClientRepository;
+import org.apache.dolphinscheduler.plugin.registry.jdbc.repository.JdbcRegistryDataChangeEventRepository;
+import org.apache.dolphinscheduler.plugin.registry.jdbc.repository.JdbcRegistryDataRepository;
+import org.apache.dolphinscheduler.plugin.registry.jdbc.repository.JdbcRegistryLockRepository;
+
+import java.time.Duration;
+import java.util.Collections;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+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.junit.jupiter.api.extension.ExtendWith;
+import org.mockito.ArgumentCaptor;
+import org.mockito.Mock;
+import org.mockito.Mockito;
+import org.mockito.junit.jupiter.MockitoExtension;
+import org.springframework.test.util.ReflectionTestUtils;
+import org.springframework.transaction.support.TransactionTemplate;
+
+import com.google.common.truth.Truth;
+
+@ExtendWith(MockitoExtension.class)
+class JdbcRegistryServerTest {
+
+ private static final JdbcRegistryClientIdentify CLIENT_IDENTIFY =
+ new JdbcRegistryClientIdentify(1L, "test-client");
+
+ @Mock
+ private JdbcRegistryDataRepository jdbcRegistryDataRepository;
+
+ @Mock
+ private JdbcRegistryLockRepository jdbcRegistryLockRepository;
+
+ @Mock
+ private JdbcRegistryClientRepository jdbcRegistryClientRepository;
+
+ @Mock
+ private JdbcRegistryDataChangeEventRepository jdbcRegistryDataChangeEventRepository;
+
+ @Mock
+ private TransactionTemplate transactionTemplate;
+
+ @Mock
+ private IJdbcRegistryClient jdbcRegistryClient;
+
+ @Mock
+ private ConnectionStateListener connectionStateListener;
+
+ private JdbcRegistryServer jdbcRegistryServer;
+
+ @BeforeEach
+ void setUp() {
+ JdbcRegistryProperties jdbcRegistryProperties = new JdbcRegistryProperties();
+ jdbcRegistryProperties.setSessionTimeout(Duration.ofSeconds(1));
+ jdbcRegistryServer = new JdbcRegistryServer(
+ jdbcRegistryDataRepository,
+ jdbcRegistryLockRepository,
+ jdbcRegistryClientRepository,
+ jdbcRegistryDataChangeEventRepository,
+ jdbcRegistryProperties,
+ transactionTemplate);
+ Mockito.when(jdbcRegistryClient.getJdbcRegistryClientIdentify()).thenReturn(CLIENT_IDENTIFY);
+ jdbcRegistryServer.registerClient(jdbcRegistryClient);
+ jdbcRegistryServer.subscribeConnectionStateChange(connectionStateListener);
+ }
+
+ @AfterEach
+ void tearDown() {
+ jdbcRegistryServer.close();
+ }
+
+ @Test
+ void refreshClientsHeartbeat_shouldDisconnectWhenHeartbeatRecordWasPurged() {
+ setServerState(JdbcRegistryServerState.SUSPENDED);
+ ReflectionTestUtils.setField(jdbcRegistryServer, "lastSuccessHeartbeat", 0L);
+ Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any())).thenReturn(false);
+
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+
+ Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.DISCONNECTED);
+ Mockito.verify(connectionStateListener).onDisConnected();
+ }
+
+ @Test
+ void refreshClientsHeartbeat_shouldDisconnectImmediatelyWhenStartedHeartbeatRecordWasPurgedAfterTimeout() {
+ setServerState(JdbcRegistryServerState.STARTED);
+ ReflectionTestUtils.setField(jdbcRegistryServer, "lastSuccessHeartbeat", 0L);
+ Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any())).thenReturn(false);
+
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+
+ Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.DISCONNECTED);
+ Mockito.verify(jdbcRegistryClientRepository).updateById(Mockito.any());
+ Mockito.verify(connectionStateListener).onDisConnected();
+ }
+
+ @Test
+ void refreshClientsHeartbeat_shouldSuspendUntilMissingHeartbeatTimesOut() {
+ setServerState(JdbcRegistryServerState.STARTED);
+ long lastSuccessHeartbeat = System.currentTimeMillis();
+ ReflectionTestUtils.setField(jdbcRegistryServer, "lastSuccessHeartbeat", lastSuccessHeartbeat);
+ Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any())).thenReturn(false);
+
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+
+ Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.SUSPENDED);
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+ Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.SUSPENDED);
+ Truth.assertThat((long) ReflectionTestUtils.getField(jdbcRegistryServer, "lastSuccessHeartbeat"))
+ .isEqualTo(lastSuccessHeartbeat);
+ Mockito.verify(connectionStateListener, Mockito.never()).onDisConnected();
+
+ ReflectionTestUtils.setField(jdbcRegistryServer, "lastSuccessHeartbeat", 0L);
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+
+ Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.DISCONNECTED);
+ Mockito.verify(connectionStateListener).onDisConnected();
+ Mockito.verify(connectionStateListener, Mockito.never()).onReconnected();
+ Mockito.verify(jdbcRegistryClientRepository, Mockito.times(3)).updateById(Mockito.any());
+ }
+
+ @Test
+ void refreshClientsHeartbeat_shouldReconnectOnceAfterTransientFailure() {
+ setServerState(JdbcRegistryServerState.STARTED);
+ ReflectionTestUtils.setField(jdbcRegistryServer, "lastSuccessHeartbeat", System.currentTimeMillis());
+ Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any()))
+ .thenThrow(new IllegalStateException("Database unavailable"))
+ .thenReturn(true);
+
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+ Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.SUSPENDED);
+ Mockito.verifyNoInteractions(connectionStateListener);
+
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+
+ Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.STARTED);
+ Mockito.verify(connectionStateListener).onReconnected();
+ Mockito.verify(connectionStateListener, Mockito.never()).onDisConnected();
+ Mockito.verify(jdbcRegistryClientRepository, Mockito.times(3)).updateById(Mockito.any());
+ }
+
+ @Test
+ void refreshClientsHeartbeat_shouldKeepPartialHeartbeatUpdateWhenLaterClientFails() {
+ IJdbcRegistryClient secondClient = Mockito.mock(IJdbcRegistryClient.class);
+ JdbcRegistryClientIdentify secondClientIdentify = new JdbcRegistryClientIdentify(2L, "second-client");
+ Mockito.when(secondClient.getJdbcRegistryClientIdentify()).thenReturn(secondClientIdentify);
+ jdbcRegistryServer.registerClient(secondClient);
+ setServerState(JdbcRegistryServerState.STARTED);
+ ReflectionTestUtils.setField(jdbcRegistryServer, "lastSuccessHeartbeat", System.currentTimeMillis());
+ AtomicInteger updateInvocations = new AtomicInteger();
+ Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any()))
+ .thenAnswer(invocation -> updateInvocations.incrementAndGet() == 1);
+
+ @SuppressWarnings("unchecked")
+ Map<JdbcRegistryClientIdentify, JdbcRegistryClientHeartbeatDTO> heartbeatMap =
+ (Map<JdbcRegistryClientIdentify, JdbcRegistryClientHeartbeatDTO>) ReflectionTestUtils
+ .getField(jdbcRegistryServer, "jdbcRegistryClientDTOMap");
+ heartbeatMap.get(CLIENT_IDENTIFY).setLastHeartbeatTime(0L);
+
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+
+ Truth.assertThat(updateInvocations.get()).isEqualTo(2);
+ Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.SUSPENDED);
+ Truth.assertThat(heartbeatMap.get(CLIENT_IDENTIFY).getLastHeartbeatTime()).isGreaterThan(0L);
+ Mockito.verify(connectionStateListener, Mockito.never()).onReconnected();
+ Mockito.verify(connectionStateListener, Mockito.never()).onDisConnected();
+ }
+
+ @Test
+ void refreshClientsHeartbeat_shouldPersistCurrentHeartbeatTimestamp() {
+ ArgumentCaptor<JdbcRegistryClientHeartbeatDTO> registeredHeartbeat =
+ ArgumentCaptor.forClass(JdbcRegistryClientHeartbeatDTO.class);
+ Mockito.verify(jdbcRegistryClientRepository).insert(registeredHeartbeat.capture());
+ registeredHeartbeat.getValue().setLastHeartbeatTime(0L);
+ AtomicLong persistedHeartbeatTimestamp = new AtomicLong(-1L);
+ Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any())).thenAnswer(invocation -> {
+ JdbcRegistryClientHeartbeatDTO heartbeat = invocation.getArgument(0);
+ persistedHeartbeatTimestamp.set(heartbeat.getLastHeartbeatTime());
+ return true;
+ });
+
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+
+ Truth.assertThat(persistedHeartbeatTimestamp.get()).isGreaterThan(0L);
+ }
+
+ @Test
+ void refreshClientsHeartbeat_shouldNotRefreshAfterDisconnected() {
+ setServerState(JdbcRegistryServerState.DISCONNECTED);
+
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat");
+
+ Mockito.verify(jdbcRegistryClientRepository, Mockito.never()).updateById(Mockito.any());
+ }
+
+ @Test
+ void purgeInvalidJdbcRegistryMetadata_shouldKeepMetadataWhenHeartbeatWasUpdatedAfterSnapshot() {
+ JdbcRegistryClientHeartbeatDTO staleHeartbeat = JdbcRegistryClientHeartbeatDTO.builder()
+ .id(CLIENT_IDENTIFY.getClientId())
+ .clientName(CLIENT_IDENTIFY.getClientName())
+ .lastHeartbeatTime(System.currentTimeMillis() - Duration.ofSeconds(2).toMillis())
+ .clientConfig(new JdbcRegistryClientHeartbeatDTO.ClientConfig(Duration.ofSeconds(1).toMillis()))
+ .build();
+ JdbcRegistryLockDTO clientLock = JdbcRegistryLockDTO.builder()
+ .id(1L)
+ .clientId(CLIENT_IDENTIFY.getClientId())
+ .build();
+ Mockito.when(jdbcRegistryClientRepository.queryAll()).thenReturn(Collections.singletonList(staleHeartbeat));
+ Mockito.when(jdbcRegistryClientRepository.deleteByIdAndLastHeartbeatTime(
+ staleHeartbeat.getId(), staleHeartbeat.getLastHeartbeatTime())).thenReturn(false);
+ Mockito.when(jdbcRegistryDataRepository.selectAll()).thenReturn(Collections.emptyList());
+ Mockito.when(jdbcRegistryLockRepository.queryAll()).thenReturn(Collections.singletonList(clientLock));
+
+ ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "purgeInvalidJdbcRegistryMetadata");
+
+ Mockito.verify(jdbcRegistryClientRepository).deleteByIdAndLastHeartbeatTime(
+ staleHeartbeat.getId(), staleHeartbeat.getLastHeartbeatTime());
+ Mockito.verify(jdbcRegistryLockRepository, Mockito.never()).deleteById(clientLock.getId());
+ }
+
+ @SuppressWarnings("unchecked")
+ private void setServerState(JdbcRegistryServerState state) {
+ ((AtomicReference<JdbcRegistryServerState>) ReflectionTestUtils.getField(jdbcRegistryServer, "serverState"))
+ .set(state);
+ }
+}