From fe4ef05053c513f8b8601ac2627e2333953ff892 Mon Sep 17 00:00:00 2001
From: Abhishek Kumar <abhishek.mrt22@gmail.com>
Date: Fri, 13 Sep 2024 16:17:47 +0530
Subject: [PATCH] server,engine-schema: use single query to to list host
 capacities while host capacity update

Signed-off-by: Abhishek Kumar <abhishek.mrt22@gmail.com>
---
 .../com/cloud/capacity/dao/CapacityDao.java   |  2 ++
 .../cloud/capacity/dao/CapacityDaoImpl.java   | 12 ++++++++++
 .../cloud/capacity/CapacityManagerImpl.java   | 22 ++++++++++++++++---
 3 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/engine/schema/src/main/java/com/cloud/capacity/dao/CapacityDao.java b/engine/schema/src/main/java/com/cloud/capacity/dao/CapacityDao.java
index 9616f31d0c5..1bb79ce417a 100644
--- a/engine/schema/src/main/java/com/cloud/capacity/dao/CapacityDao.java
+++ b/engine/schema/src/main/java/com/cloud/capacity/dao/CapacityDao.java
@@ -28,6 +28,8 @@ import com.cloud.utils.db.GenericDao;
 public interface CapacityDao extends GenericDao<CapacityVO, Long> {
     CapacityVO findByHostIdType(Long hostId, short capacityType);
 
+    List<CapacityVO> listByHostIdTypes(Long hostId, List<Short> capacityTypes);
+
     List<Long> listClustersInZoneOrPodByHostCapacities(long id, long vmId, int requiredCpu, long requiredRam, short capacityTypeForOrdering, boolean isZone);
 
     List<Long> listHostsWithEnoughCapacity(int requiredCpu, long requiredRam, Long clusterId, String hostType);
diff --git a/engine/schema/src/main/java/com/cloud/capacity/dao/CapacityDaoImpl.java b/engine/schema/src/main/java/com/cloud/capacity/dao/CapacityDaoImpl.java
index ae75ca49f33..5f6cf9c5083 100644
--- a/engine/schema/src/main/java/com/cloud/capacity/dao/CapacityDaoImpl.java
+++ b/engine/schema/src/main/java/com/cloud/capacity/dao/CapacityDaoImpl.java
@@ -674,6 +674,18 @@ public class CapacityDaoImpl extends GenericDaoBase<CapacityVO, Long> implements
         return findOneBy(sc);
     }
 
+    @Override
+    public List<CapacityVO> listByHostIdTypes(Long hostId, List<Short> capacityTypes) {
+        SearchBuilder<CapacityVO> sb = createSearchBuilder();
+        sb.and("hostId", sb.entity().getHostOrPoolId(), SearchCriteria.Op.EQ);
+        sb.and("type", sb.entity().getCapacityType(), SearchCriteria.Op.IN);
+        sb.done();
+        SearchCriteria<CapacityVO> sc = sb.create();
+        sc.setParameters("hostId", hostId);
+        sc.setParameters("type", capacityTypes.toArray());
+        return listBy(sc);
+    }
+
     @Override
     public List<Long> listClustersInZoneOrPodByHostCapacities(long id, long vmId, int requiredCpu, long requiredRam, short capacityTypeForOrdering, boolean isZone) {
         TransactionLegacy txn = TransactionLegacy.currentTxn();
diff --git a/server/src/main/java/com/cloud/capacity/CapacityManagerImpl.java b/server/src/main/java/com/cloud/capacity/CapacityManagerImpl.java
index abb67e7e714..486e90b8520 100644
--- a/server/src/main/java/com/cloud/capacity/CapacityManagerImpl.java
+++ b/server/src/main/java/com/cloud/capacity/CapacityManagerImpl.java
@@ -41,6 +41,7 @@ import org.apache.cloudstack.utils.cache.LazyCache;
 import org.apache.cloudstack.utils.cache.SingleCache;
 import org.apache.cloudstack.utils.executor.QueuedExecutor;
 import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.lang3.ObjectUtils;
 import org.apache.log4j.Logger;
 
 import com.cloud.agent.AgentManager;
@@ -801,9 +802,24 @@ public class CapacityManagerImpl extends ManagerBase implements CapacityManager,
             }
         }
 
-        CapacityVO cpuCap = _capacityDao.findByHostIdType(host.getId(), Capacity.CAPACITY_TYPE_CPU);
-        CapacityVO memCap = _capacityDao.findByHostIdType(host.getId(), Capacity.CAPACITY_TYPE_MEMORY);
-        CapacityVO cpuCoreCap = _capacityDao.findByHostIdType(host.getId(), CapacityVO.CAPACITY_TYPE_CPU_CORE);
+        List<CapacityVO> capacities = _capacityDao.listByHostIdTypes(host.getId(), List.of(Capacity.CAPACITY_TYPE_CPU,
+                Capacity.CAPACITY_TYPE_MEMORY,
+                CapacityVO.CAPACITY_TYPE_CPU_CORE));
+        CapacityVO cpuCap = null;
+        CapacityVO memCap = null;
+        CapacityVO cpuCoreCap = null;
+        for (CapacityVO c : capacities) {
+            if (c.getCapacityType() == Capacity.CAPACITY_TYPE_CPU) {
+                cpuCap = c;
+            } else if (c.getCapacityType() == Capacity.CAPACITY_TYPE_MEMORY) {
+                memCap = c;
+            } else if (c.getCapacityType() == Capacity.CAPACITY_TYPE_CPU_CORE) {
+                cpuCoreCap = c;
+            }
+            if (ObjectUtils.allNotNull(cpuCap, memCap, cpuCoreCap)) {
+                break;
+            }
+        }
 
         if (cpuCoreCap != null) {
             long hostTotalCpuCore = host.getCpus().longValue();