From 7692b745e8a5604979fed4d1ec5db1dd0c668c09 Mon Sep 17 00:00:00 2001
From: Fabricio Duarte <fabricio.duarte.jr@gmail.com>
Date: Wed, 28 Aug 2024 07:06:16 -0300
Subject: [PATCH] Fix NPE when sending copy command to least busy SSVM (#9125)

---
 .../main/java/com/cloud/host/dao/HostDao.java |  2 ++
 .../java/com/cloud/host/dao/HostDaoImpl.java  | 25 ++++++++++++++++++
 .../image/BaseImageStoreDriverImpl.java       | 26 ++-----------------
 3 files changed, 29 insertions(+), 24 deletions(-)

diff --git a/engine/schema/src/main/java/com/cloud/host/dao/HostDao.java b/engine/schema/src/main/java/com/cloud/host/dao/HostDao.java
index 08380ed8b40..a2df6db44e5 100644
--- a/engine/schema/src/main/java/com/cloud/host/dao/HostDao.java
+++ b/engine/schema/src/main/java/com/cloud/host/dao/HostDao.java
@@ -169,4 +169,6 @@ public interface HostDao extends GenericDao<HostVO, Long>, StateDao<Status, Stat
     List<HostVO> findHostsWithTagRuleThatMatchComputeOferringTags(String computeOfferingTags);
 
     List<Long> findClustersThatMatchHostTagRule(String computeOfferingTags);
+
+    List<Long> listSsvmHostsWithPendingMigrateJobsOrderedByJobCount();
 }
diff --git a/engine/schema/src/main/java/com/cloud/host/dao/HostDaoImpl.java b/engine/schema/src/main/java/com/cloud/host/dao/HostDaoImpl.java
index 170c6a45fc3..63950294654 100644
--- a/engine/schema/src/main/java/com/cloud/host/dao/HostDaoImpl.java
+++ b/engine/schema/src/main/java/com/cloud/host/dao/HostDaoImpl.java
@@ -1371,6 +1371,31 @@ public class HostDaoImpl extends GenericDaoBase<HostVO, Long> implements HostDao
         return new ArrayList<>(result);
     }
 
+    @Override
+    public List<Long> listSsvmHostsWithPendingMigrateJobsOrderedByJobCount() {
+        String query = "SELECT cel.host_id, COUNT(*) " +
+                "FROM cmd_exec_log cel " +
+                "JOIN host h ON cel.host_id = h.id " +
+                "WHERE h.removed IS NULL " +
+                "GROUP BY cel.host_id " +
+                "ORDER BY 2";
+
+        TransactionLegacy txn = TransactionLegacy.currentTxn();
+        List<Long> result = new ArrayList<>();
+
+        PreparedStatement pstmt;
+        try {
+            pstmt = txn.prepareAutoCloseStatement(query);
+            ResultSet rs = pstmt.executeQuery();
+            while (rs.next()) {
+                result.add((long) rs.getInt(1));
+            }
+        } catch (SQLException e) {
+            logger.warn("SQLException caught while listing SSVMs with least migrate jobs.", e);
+        }
+        return result;
+    }
+
     private String getHostIdsByComputeTags(List<String> offeringTags){
         List<String> questionMarks = new ArrayList();
         offeringTags.forEach((tag) -> { questionMarks.add("?"); });
diff --git a/engine/storage/src/main/java/org/apache/cloudstack/storage/image/BaseImageStoreDriverImpl.java b/engine/storage/src/main/java/org/apache/cloudstack/storage/image/BaseImageStoreDriverImpl.java
index 9606da1e337..db3f798a68a 100644
--- a/engine/storage/src/main/java/org/apache/cloudstack/storage/image/BaseImageStoreDriverImpl.java
+++ b/engine/storage/src/main/java/org/apache/cloudstack/storage/image/BaseImageStoreDriverImpl.java
@@ -20,9 +20,6 @@ package org.apache.cloudstack.storage.image;
 
 import java.net.URI;
 import java.net.URISyntaxException;
-import java.sql.PreparedStatement;
-import java.sql.ResultSet;
-import java.sql.SQLException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Date;
@@ -87,7 +84,6 @@ import com.cloud.storage.dao.VMTemplateZoneDao;
 import com.cloud.storage.dao.VolumeDao;
 import com.cloud.storage.download.DownloadMonitor;
 import com.cloud.utils.NumbersUtil;
-import com.cloud.utils.db.TransactionLegacy;
 import com.cloud.utils.exception.CloudRuntimeException;
 import com.cloud.utils.net.Proxy;
 import com.cloud.vm.VirtualMachineManager;
@@ -425,8 +421,9 @@ public abstract class BaseImageStoreDriverImpl implements ImageStoreDriver {
     private Answer sendToLeastBusyEndpoint(List<EndPoint> eps, CopyCommand cmd) {
         Answer answer = null;
         EndPoint endPoint = null;
-        List<Long> epIds = ssvmWithLeastMigrateJobs();
 
+        logger.debug("Picking SSVM from the pool with least commands running on it.");
+        List<Long> epIds = hostDao.listSsvmHostsWithPendingMigrateJobsOrderedByJobCount();
         if (epIds.isEmpty()) {
             Collections.shuffle(eps);
             endPoint = eps.get(0);
@@ -533,23 +530,4 @@ public abstract class BaseImageStoreDriverImpl implements ImageStoreDriver {
     private Integer getCopyCmdsCountToSpecificSSVM(Long ssvmId) {
         return _cmdExecLogDao.getCopyCmdCountForSSVM(ssvmId);
     }
-
-    private List<Long> ssvmWithLeastMigrateJobs() {
-        logger.debug("Picking ssvm from the pool with least commands running on it");
-        String query = "select host_id, count(*) from cmd_exec_log group by host_id order by 2;";
-        TransactionLegacy txn = TransactionLegacy.currentTxn();
-
-        List<Long> result = new ArrayList<Long>();
-        PreparedStatement pstmt = null;
-        try {
-            pstmt = txn.prepareAutoCloseStatement(query);
-            ResultSet rs = pstmt.executeQuery();
-            while (rs.next()) {
-                result.add((long) rs.getInt(1));
-            }
-        } catch (SQLException e) {
-            logger.debug("SQLException caught", e);
-        }
-        return result;
-    }
 }