Fix AINode status error after ConfigNode Leader restarts #13547

Signed-off-by: OneSizeFitQuorum <tanxinyu@apache.org>

iotdb-core/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);
+        });
   }
 
   /**