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 2b8499bb4812..73c2b66772f2 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 @@ 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 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 1791f3c942aa..cee889bf05c8 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 void deleteByIds(Collection clientIds) { 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 61a934bd148e..8910395360c3 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.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 serverState = + new AtomicReference<>(JdbcRegistryServerState.INIT); private final List 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 JdbcRegistryServer(JdbcRegistryDataRepository jdbcRegistryDataRepository, transactionTemplate, schedulerThreadExecutor); this.jdbcRegistryLockManager = new JdbcRegistryLockManager( jdbcRegistryProperties, jdbcRegistryLockRepository); - this.jdbcRegistryServerState = JdbcRegistryServerState.INIT; lastSuccessHeartbeat = System.currentTimeMillis(); } @Override - public void start() { - if (jdbcRegistryServerState != JdbcRegistryServerState.INIT) { + public synchronized void start() { + if (serverState.get() != JdbcRegistryServerState.INIT) { // The server is already started or stopped, will not start again. return; } @@ -119,7 +120,10 @@ public void start() { 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 void deregisterClient(IJdbcRegistryClient jdbcRegistryClient) { @Override public JdbcRegistryServerState getServerState() { - return jdbcRegistryServerState; + return serverState.get(); } @Override @@ -256,7 +260,22 @@ public void releaseJdbcRegistryLock(Long clientId, String lockKey) { @Override public void close() { - jdbcRegistryServerState = JdbcRegistryServerState.STOPPED; + synchronized (this) { + while (true) { + JdbcRegistryServerState currentState = serverState.get(); + if (currentState == JdbcRegistryServerState.STOPPED) { + log.warn("The JdbcRegistryServer is already STOPPED."); + return; + } + if (serverState.compareAndSet(currentState, JdbcRegistryServerState.STOPPED)) { + break; + } + // A heartbeat can change the state without this monitor. Do not drop the close request. + log.debug("Failed to stop JdbcRegistryServer from state {}, current state is {}, retrying", + currentState, + serverState.get()); + } + } schedulerThreadExecutor.shutdown(); List clientIds = jdbcRegistryClients.stream() .map(IJdbcRegistryClient::getJdbcRegistryClientIdentify) @@ -269,23 +288,26 @@ public void close() { 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 jdbcRegistryClients = jdbcRegistryClientRepository.queryAll(); - final Set deadJdbcRegistryClientIds = jdbcRegistryClients + final Set 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 existJdbcRegistryClientIds = jdbcRegistryClients .stream() .map(JdbcRegistryClientHeartbeatDTO::getId) - .filter(id -> !deadJdbcRegistryClientIds.contains(id)) + .filter(id -> !deletedJdbcRegistryClientIds.contains(id)) .collect(Collectors.toSet()); jdbcRegistryDataManager.getAllJdbcRegistryData() .stream() @@ -321,8 +343,11 @@ private void refreshClientsHeartbeat() { 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 +366,59 @@ private void refreshClientsHeartbeat() { } JdbcRegistryClientHeartbeatDTO clone = jdbcRegistryClientHeartbeatDTO.clone(); clone.setLastHeartbeatTime(now); - jdbcRegistryClientRepository.updateById(jdbcRegistryClientHeartbeatDTO); + if (!jdbcRegistryClientRepository.updateById(clone)) { + log.error("The client heartbeat has expired: {}", 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 = currentState == JdbcRegistryServerState.SUSPENDED; + if (reconnected) { + if (!serverState.compareAndSet(JdbcRegistryServerState.SUSPENDED, JdbcRegistryServerState.STARTED)) { + log.debug("Failed to reconnect JdbcRegistryServer; current state is {}", serverState.get()); + return; + } + } else if (currentState != JdbcRegistryServerState.STARTED) { + return; + } + // Serialize heartbeat side effects with close(), even if close wins after the state transition. + synchronized (this) { + if (serverState.get() != JdbcRegistryServerState.STARTED) { + return; + } + lastSuccessHeartbeat = now; + if (reconnected) { + doTriggerReconnectedListener(); + } } - lastSuccessHeartbeat = now; 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; + 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; + } + synchronized (this) { + if (serverState.get() == JdbcRegistryServerState.DISCONNECTED) { doTriggerOnDisConnectedListener(); } - break; - default: - break; + } + } else if (currentState == JdbcRegistryServerState.STARTED + && !serverState.compareAndSet(JdbcRegistryServerState.STARTED, JdbcRegistryServerState.SUSPENDED)) { + log.debug("Failed to suspend JdbcRegistryServer; current state is {}", serverState.get()); + return; } } } @@ -373,6 +426,9 @@ private void refreshClientsHeartbeat() { private void doTriggerReconnectedListener() { log.info("Trigger:onReconnected listener."); connectionStateListeners.forEach(listener -> { + if (serverState.get() != JdbcRegistryServerState.STARTED) { + return; + } try { listener.onReconnected(); } catch (Exception ex) { @@ -395,6 +451,9 @@ private void doTriggerOnConnectedListener() { private void doTriggerOnDisConnectedListener() { log.info("Trigger:onDisConnected listener."); connectionStateListeners.forEach(listener -> { + if (serverState.get() != JdbcRegistryServerState.DISCONNECTED) { + return; + } try { listener.onDisConnected(); } catch (Exception ex) { 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 000000000000..cd171d448cd0 --- /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,480 @@ +/* + * 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 static org.awaitility.Awaitility.await; + +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.CountDownLatch; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.TimeUnit; +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.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; +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 close_shouldOnlyPurgeClientsOnceWhenCalledConcurrently() throws Exception { + CountDownLatch firstPurgeStarted = new CountDownLatch(1); + CountDownLatch allowFirstPurgeToFinish = new CountDownLatch(1); + AtomicInteger purgeInvocations = new AtomicInteger(); + Mockito.doAnswer(invocation -> { + if (purgeInvocations.incrementAndGet() == 1) { + firstPurgeStarted.countDown(); + allowFirstPurgeToFinish.await(5, TimeUnit.SECONDS); + } + return null; + }).when(jdbcRegistryClientRepository).deleteByIds(Mockito.any()); + ExecutorService closeExecutor = Executors.newFixedThreadPool(2); + Future firstClose = closeExecutor.submit(jdbcRegistryServer::close); + Future secondClose = null; + + try { + Truth.assertThat(firstPurgeStarted.await(5, TimeUnit.SECONDS)).isTrue(); + secondClose = closeExecutor.submit(jdbcRegistryServer::close); + secondClose.get(5, TimeUnit.SECONDS); + + Truth.assertThat(purgeInvocations.get()).isEqualTo(1); + } finally { + allowFirstPurgeToFinish.countDown(); + firstClose.get(5, TimeUnit.SECONDS); + if (secondClose != null) { + secondClose.get(5, TimeUnit.SECONDS); + } + closeExecutor.shutdownNow(); + } + } + + @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()); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void refreshClientsHeartbeat_shouldStopNotifyingWhenListenerClosesServer(boolean heartbeatSucceeds) { + setServerState(JdbcRegistryServerState.SUSPENDED); + ReflectionTestUtils.setField(jdbcRegistryServer, "lastSuccessHeartbeat", 0L); + Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any())).thenReturn(heartbeatSucceeds); + if (heartbeatSucceeds) { + Mockito.doAnswer(invocation -> { + jdbcRegistryServer.close(); + return null; + }).when(connectionStateListener).onReconnected(); + } else { + Mockito.doAnswer(invocation -> { + jdbcRegistryServer.close(); + return null; + }).when(connectionStateListener).onDisConnected(); + } + ConnectionStateListener laterListener = Mockito.mock(ConnectionStateListener.class); + jdbcRegistryServer.subscribeConnectionStateChange(laterListener); + + ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat"); + + Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.STOPPED); + Mockito.verifyNoInteractions(laterListener); + } + + @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 heartbeatMap = + (Map) 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_shouldNotDisconnectWhenCloseWinsRace() throws Exception { + setServerState(JdbcRegistryServerState.STARTED); + CountDownLatch heartbeatUpdateStarted = new CountDownLatch(1); + CountDownLatch allowHeartbeatUpdateToFinish = new CountDownLatch(1); + Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any())).thenAnswer(invocation -> { + heartbeatUpdateStarted.countDown(); + allowHeartbeatUpdateToFinish.await(5, TimeUnit.SECONDS); + return false; + }); + ExecutorService heartbeatExecutor = Executors.newSingleThreadExecutor(); + Future heartbeatFuture = heartbeatExecutor.submit(() -> { + ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat"); + }); + + try { + Truth.assertThat(heartbeatUpdateStarted.await(5, TimeUnit.SECONDS)).isTrue(); + jdbcRegistryServer.close(); + allowHeartbeatUpdateToFinish.countDown(); + heartbeatFuture.get(5, TimeUnit.SECONDS); + } finally { + allowHeartbeatUpdateToFinish.countDown(); + heartbeatExecutor.shutdownNow(); + } + + Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.STOPPED); + Mockito.verify(connectionStateListener, Mockito.never()).onDisConnected(); + } + + @Test + void refreshClientsHeartbeat_shouldNotReconnectWhenCloseWinsSuccessfulHeartbeatRace() throws Exception { + setServerState(JdbcRegistryServerState.SUSPENDED); + ReflectionTestUtils.setField(jdbcRegistryServer, "lastSuccessHeartbeat", 42L); + CountDownLatch heartbeatUpdateStarted = new CountDownLatch(1); + CountDownLatch allowHeartbeatUpdateToFinish = new CountDownLatch(1); + Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any())).thenAnswer(invocation -> { + heartbeatUpdateStarted.countDown(); + allowHeartbeatUpdateToFinish.await(5, TimeUnit.SECONDS); + return true; + }); + ExecutorService heartbeatExecutor = Executors.newSingleThreadExecutor(); + Future heartbeatFuture = heartbeatExecutor.submit(() -> { + ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat"); + }); + + try { + Truth.assertThat(heartbeatUpdateStarted.await(5, TimeUnit.SECONDS)).isTrue(); + jdbcRegistryServer.close(); + allowHeartbeatUpdateToFinish.countDown(); + heartbeatFuture.get(5, TimeUnit.SECONDS); + } finally { + allowHeartbeatUpdateToFinish.countDown(); + heartbeatExecutor.shutdownNow(); + } + + Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.STOPPED); + Truth.assertThat((long) ReflectionTestUtils.getField(jdbcRegistryServer, "lastSuccessHeartbeat")) + .isEqualTo(42L); + Mockito.verify(connectionStateListener, Mockito.never()).onReconnected(); + Mockito.verify(connectionStateListener, Mockito.never()).onDisConnected(); + } + + @Test + void refreshClientsHeartbeat_shouldNotSuspendWhenCloseWinsFailedHeartbeatRace() throws Exception { + setServerState(JdbcRegistryServerState.STARTED); + CountDownLatch heartbeatUpdateStarted = new CountDownLatch(1); + CountDownLatch allowHeartbeatUpdateToFail = new CountDownLatch(1); + Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any())).thenAnswer(invocation -> { + heartbeatUpdateStarted.countDown(); + allowHeartbeatUpdateToFail.await(5, TimeUnit.SECONDS); + throw new RuntimeException("Heartbeat update failed"); + }); + ExecutorService heartbeatExecutor = Executors.newSingleThreadExecutor(); + Future heartbeatFuture = heartbeatExecutor.submit(() -> { + ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat"); + }); + + try { + Truth.assertThat(heartbeatUpdateStarted.await(5, TimeUnit.SECONDS)).isTrue(); + jdbcRegistryServer.close(); + allowHeartbeatUpdateToFail.countDown(); + heartbeatFuture.get(5, TimeUnit.SECONDS); + } finally { + allowHeartbeatUpdateToFail.countDown(); + heartbeatExecutor.shutdownNow(); + } + + Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.STOPPED); + Mockito.verify(connectionStateListener, Mockito.never()).onReconnected(); + Mockito.verify(connectionStateListener, Mockito.never()).onDisConnected(); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void refreshClientsHeartbeat_shouldPreserveCloseWhenFailureTransitionLosesRace(boolean timedOut) throws Exception { + setServerState(JdbcRegistryServerState.STARTED); + ReflectionTestUtils.setField(jdbcRegistryServer, "lastSuccessHeartbeat", + timedOut ? 0L : System.currentTimeMillis()); + JdbcRegistryProperties properties = Mockito.spy((JdbcRegistryProperties) ReflectionTestUtils + .getField(jdbcRegistryServer, "jdbcRegistryProperties")); + ReflectionTestUtils.setField(jdbcRegistryServer, "jdbcRegistryProperties", properties); + CountDownLatch timeoutCheckStarted = new CountDownLatch(1); + CountDownLatch allowTimeoutCheck = new CountDownLatch(1); + Mockito.doAnswer(invocation -> { + timeoutCheckStarted.countDown(); + Truth.assertThat(allowTimeoutCheck.await(5, TimeUnit.SECONDS)).isTrue(); + return timedOut ? Duration.ZERO : Duration.ofDays(1); + }).when(properties).getSessionTimeout(); + Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any())).thenReturn(false); + ExecutorService heartbeatExecutor = Executors.newSingleThreadExecutor(); + Future heartbeatFuture = heartbeatExecutor + .submit(() -> ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat")); + + try { + Truth.assertThat(timeoutCheckStarted.await(5, TimeUnit.SECONDS)).isTrue(); + jdbcRegistryServer.close(); + allowTimeoutCheck.countDown(); + heartbeatFuture.get(5, TimeUnit.SECONDS); + } finally { + allowTimeoutCheck.countDown(); + heartbeatExecutor.shutdownNow(); + } + + Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.STOPPED); + Mockito.verifyNoInteractions(connectionStateListener); + } + + @ParameterizedTest + @ValueSource(booleans = {true, false}) + void refreshClientsHeartbeat_shouldNotNotifyWhenClosedAfterStateTransition(boolean heartbeatSucceeds) throws Exception { + setServerState(JdbcRegistryServerState.SUSPENDED); + ReflectionTestUtils.setField(jdbcRegistryServer, "lastSuccessHeartbeat", 42L); + Mockito.when(jdbcRegistryClientRepository.updateById(Mockito.any())).thenReturn(heartbeatSucceeds); + ExecutorService heartbeatExecutor = Executors.newSingleThreadExecutor(); + Future heartbeatFuture; + try { + synchronized (jdbcRegistryServer) { + heartbeatFuture = heartbeatExecutor + .submit(() -> ReflectionTestUtils.invokeMethod(jdbcRegistryServer, "refreshClientsHeartbeat")); + // Pause notification after the CAS, then let close() win before the heartbeat commits its effects. + await().atMost(Duration.ofSeconds(5)) + .until(() -> jdbcRegistryServer + .getServerState() == (heartbeatSucceeds ? JdbcRegistryServerState.STARTED + : JdbcRegistryServerState.DISCONNECTED)); + jdbcRegistryServer.close(); + } + heartbeatFuture.get(5, TimeUnit.SECONDS); + } finally { + heartbeatExecutor.shutdownNow(); + } + + Truth.assertThat(jdbcRegistryServer.getServerState()).isEqualTo(JdbcRegistryServerState.STOPPED); + Truth.assertThat((long) ReflectionTestUtils.getField(jdbcRegistryServer, "lastSuccessHeartbeat")) + .isEqualTo(42L); + Mockito.verifyNoInteractions(connectionStateListener); + } + + @Test + void refreshClientsHeartbeat_shouldPersistCurrentHeartbeatTimestamp() { + ArgumentCaptor 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) ReflectionTestUtils.getField(jdbcRegistryServer, "serverState")) + .set(state); + } +}