optimization: 安装Gse2.0 查询Agent状态优化 (closed TencentBlueKing#1883)

apps/backend/components/collections/agent_new/base.py

@@ -99,6 +99,7 @@ def get_common_data(cls, data):
             subscription_instance_ids=common_data.subscription_instance_ids,
             sub_inst_id__host_id_map=common_data.sub_inst_id__host_id_map,
             host_id__sub_inst_id_map=common_data.host_id__sub_inst_id_map,
+            sub_inst_id__sub_inst_obj_map=common_data.sub_inst_id__sub_inst_obj_map,
             # Agent 新增的公共数据
             default_ap=default_ap,
             host_id__ap_map=host_id__ap_map,

apps/backend/components/collections/agent_new/get_agent_status.py

@@ -9,7 +9,7 @@
 specific language governing permissions and limitations under the License.
 """
 from collections import defaultdict
-from typing import Dict, List, Set, Union
+from typing import Dict, List, Set, Union, Optional
 
 from django.utils.translation import ugettext_lazy as _
 
@@ -67,15 +67,21 @@ def _schedule(self, data, parent_data, callback_data=None):
 
         # 构造 gse 请求参数
         hosts: List[Dict[str, Union[int, str]]] = []
+        need_to_quey_host_agent_id_map: Dict[int, int] = {}
         for host_id in host_ids_need_to_query:
+            sub_inst_id = common_data.host_id__sub_inst_id_map[host_id]
+            sub_inst = common_data.sub_inst_id__sub_inst_obj_map[sub_inst_id]
             host_obj = common_data.host_id_obj_map[host_id]
-            hosts.append(
-                {
-                    "ip": host_obj.inner_ip or host_obj.inner_ipv6,
-                    "bk_cloud_id": host_obj.bk_cloud_id,
-                    "bk_agent_id": host_obj.bk_agent_id,
-                }
-            )
+
+            bk_agent_id: Optional[str] = sub_inst.instance_info["host"].get("bk_agent_id") or host_obj.bk_agent_id
+            host: Dict[str, Union[int, str]] = {
+                "ip": host_obj.inner_ip or host_obj.inner_ipv6,
+                "bk_cloud_id": host_obj.bk_cloud_id,
+                "bk_agent_id": bk_agent_id,
+            }
+            need_to_quey_host_agent_id_map[host_id] = common_data.gse_api_helper.get_agent_id(host)
+
+            hosts.append(host)
 
         agent_id__agent_state_info_map: Dict[str, Dict] = common_data.gse_api_helper.list_agent_state(hosts)
 
@@ -93,7 +99,7 @@ def _schedule(self, data, parent_data, callback_data=None):
         }
         for host_id in host_ids_need_to_query:
             host_obj = common_data.host_id_obj_map[host_id]
-            agent_id = common_data.gse_api_helper.get_agent_id(host_obj)
+            agent_id = need_to_quey_host_agent_id_map[host_id]
             # bk_agent_alive 默认值为 None 的背景：expect_status 是可传入的，不能设定一个明确状态作为默认值，不然可能与 expect_status 一致
             # 误判为当前 Agent 状态已符合预期
             agent_state_info = agent_id__agent_state_info_map.get(agent_id, {"version": "", "bk_agent_alive": None})

apps/backend/components/collections/agent_new/upgrade_to_agent_id.py

@@ -9,7 +9,7 @@
 specific language governing permissions and limitations under the License.
 """
 
-from typing import Dict, List, Union
+from typing import Dict, List, Optional, Union
 
 from django.utils.translation import ugettext_lazy as _
 
@@ -18,21 +18,24 @@
 
 class UpgradeToAgentIDService(AgentBaseService):
     def _execute(self, data, parent_data, common_data: AgentCommonData):
-        upgrade_hosts: List[Dict[str, Union[int, str]]] = []
         cloud_ip__sub_inst_id_map: Dict[str, int] = {}
-        # 如果主机有AgentID，则调用 upgrade_to_agent_id 将基于 Host IP 的配置升级到基于 Agent-ID 的配置
-        for host_id, sub_inst_id in common_data.host_id__sub_inst_id_map.items():
-            host = common_data.host_id_obj_map[host_id]
-            upgrade_hosts.append(
-                {"ip": host.inner_ip, "bk_cloud_id": host.bk_cloud_id, "bk_agent_id": host.bk_agent_id}
-            )
+        # 如果主机有AgentID，则调用 upgrade_to_agent_id 将基于 Host IP 的配置升级到基于 Agent-ID 的配
+        upgrade_hosts: List[Dict[str, Union[int, str]]] = []
+
+        for host in common_data.host_id_obj_map.values():
+            sub_inst_id = common_data.host_id__sub_inst_id_map[host.bk_host_id]
+            sub_inst = common_data.sub_inst_id__sub_inst_obj_map[sub_inst_id]
+            bk_agent_id: Optional[str] = sub_inst.instance_info["host"].get("bk_agent_id") or host.bk_agent_id
             cloud_ip__sub_inst_id_map[f"{host.bk_cloud_id}:{host.inner_ip}"] = sub_inst_id
 
-        upgrade_hosts = [
-            {"ip": host.inner_ip, "bk_cloud_id": host.bk_cloud_id, "bk_agent_id": host.bk_agent_id}
-            for host in common_data.host_id_obj_map.values()
-            if host.bk_agent_id
-        ]
+            if bk_agent_id:
+                upgrade_hosts.append(
+                    {
+                        "ip": host.inner_ip,
+                        "bk_cloud_id": host.bk_cloud_id,
+                        "bk_agent_id": bk_agent_id,
+                    }
+                )
 
         if not upgrade_hosts:
             return True

apps/backend/components/collections/base.py

@@ -230,6 +230,7 @@ class CommonData:
     sub_inst_id__host_id_map: Dict[int, int]
     host_id__sub_inst_id_map: Dict[int, int]
     ap_id_obj_map: Dict[int, models.AccessPoint]
+    sub_inst_id__sub_inst_obj_map: Dict[int, models.SubscriptionInstanceRecord]
     gse_api_helper: GseApiBaseHelper
     subscription: models.Subscription
     subscription_step: models.SubscriptionStep
@@ -360,6 +361,7 @@ def get_common_data(cls, data):
         subscription_instance_ids = set()
         sub_inst_id__host_id_map = {}
         host_id__sub_inst_id_map = {}
+        sub_inst_id__sub_inst_obj_map = {}
         for subscription_instance in subscription_instances:
             subscription_instance_ids.add(subscription_instance.id)
             # 兼容新安装Agent主机无bk_host_id的场景
@@ -368,6 +370,7 @@ def get_common_data(cls, data):
                 bk_host_ids.add(bk_host_id)
                 sub_inst_id__host_id_map[subscription_instance.id] = bk_host_id
                 host_id__sub_inst_id_map[bk_host_id] = subscription_instance.id
+                sub_inst_id__sub_inst_obj_map[subscription_instance.id] = subscription_instance
 
         host_id_obj_map: Dict[int, models.Host] = models.Host.host_id_obj_map(bk_host_id__in=bk_host_ids)
         ap_id_obj_map = models.AccessPoint.ap_id_obj_map()
@@ -376,6 +379,7 @@ def get_common_data(cls, data):
             host_id_obj_map=host_id_obj_map,
             sub_inst_id__host_id_map=sub_inst_id__host_id_map,
             host_id__sub_inst_id_map=host_id__sub_inst_id_map,
+            sub_inst_id__sub_inst_obj_map=sub_inst_id__sub_inst_obj_map,
             ap_id_obj_map=ap_id_obj_map,
             gse_api_helper=get_gse_api_helper(gse_version=cls.get_meta(data).get("GSE_VERSION")),
             subscription=subscription,

apps/backend/components/collections/plugin.py

@@ -145,6 +145,7 @@ def get_common_data(cls, data):
             subscription_instance_ids=common_data.subscription_instance_ids,
             sub_inst_id__host_id_map=common_data.sub_inst_id__host_id_map,
             host_id__sub_inst_id_map=common_data.host_id__sub_inst_id_map,
+            sub_inst_id__sub_inst_obj_map=common_data.sub_inst_id__sub_inst_obj_map,
             # Plugin 新增的公共数据
             process_statuses=process_statuses,
             target_host_objs=target_host_objs,