Fix AINode status error after ConfigNode Leader restarts #13547

Signed-off-by: OneSizeFitQuorum <[email protected]>
apache · Sep 20, 2024 · b8deec2 · b8deec2
1 parent d4dcbc9
commit b8deec2
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/...re/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/cache/LoadCache.java b/...re/confignode/src/main/java/org/apache/iotdb/confignode/manager/load/cache/LoadCache.java
@@ -19,6 +19,7 @@
 
 package org.apache.iotdb.confignode.manager.load.cache;
 
+import org.apache.iotdb.common.rpc.thrift.TAINodeConfiguration;
 import org.apache.iotdb.common.rpc.thrift.TConfigNodeLocation;
 import org.apache.iotdb.common.rpc.thrift.TConsensusGroupId;
 import org.apache.iotdb.common.rpc.thrift.TConsensusGroupType;
@@ -98,7 +99,8 @@ public LoadCache() {
   public void initHeartbeatCache(IManager configManager) {
     initNodeHeartbeatCache(
         configManager.getNodeManager().getRegisteredConfigNodes(),
-        configManager.getNodeManager().getRegisteredDataNodes());
+        configManager.getNodeManager().getRegisteredDataNodes(),
+        configManager.getNodeManager().getRegisteredAINodes());
     initRegionGroupHeartbeatCache(
         configManager.getClusterSchemaManager().getDatabaseNames().stream()
             .collect(
@@ -110,7 +112,8 @@ public void initHeartbeatCache(IManager configManager) {
   /** Initialize the nodeCacheMap when the ConfigNode-Leader is switched. */
   private void initNodeHeartbeatCache(
       List<TConfigNodeLocation> registeredConfigNodes,
-      List<TDataNodeConfiguration> registeredDataNodes) {
+      List<TDataNodeConfiguration> registeredDataNodes,
+      List<TAINodeConfiguration> registeredAINodes) {
 
     final int CURRENT_NODE_ID = ConfigNodeHeartbeatCache.CURRENT_NODE_ID;
     nodeCacheMap.clear();
@@ -136,6 +139,13 @@ private void initNodeHeartbeatCache(
           int dataNodeId = dataNodeConfiguration.getLocation().getDataNodeId();
           createNodeHeartbeatCache(NodeType.DataNode, dataNodeId);
         });
+
+    // Init AiNodeHeartbeatCache
+    registeredAINodes.forEach(
+        aiNodeConfiguration -> {
+          int aiNodeId = aiNodeConfiguration.getLocation().getAiNodeId();
+          createNodeHeartbeatCache(NodeType.AINode, aiNodeId);
+        });
   }
 
   /**