[Fix-18274][Registry] Handle expired JDBC heartbeat sessions

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<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);
+
 }

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<Long> 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);

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<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 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<Long> 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<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 +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,38 +366,69 @@ 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;
             }
         }
     }
 
     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) {