Batch network lookups per NIC in NetworkOrchestrator to eliminate N+1 queries during migration

engine/orchestration/src/main/java/org/apache/cloudstack/engine/orchestration/NetworkOrchestrator.java

@@ -2085,13 +2085,26 @@ public void finalizeUpdateInSequence(Network network, boolean success) {
         }
     }
 
+    private Map<Long, NetworkVO> batchLoadNetworksForNics(List<NicVO> nics) {
+        List<Long> networkIds = nics.stream().map(NicVO::getNetworkId).distinct().collect(Collectors.toList());
+        Map<Long, NetworkVO> result = new HashMap<>(networkIds.size());
+        for (NetworkVO network : _networksDao.listByIds(networkIds)) {
+            result.put(network.getId(), network);
+        }
+        return result;
+    }
+
     @Override
     public void setHypervisorHostname(VirtualMachineProfile vm, DeployDestination dest, boolean migrationSuccessful) throws ResourceUnavailableException {
         String hypervisorHostName = VirtualMachineManager.getHypervisorHostname(dest.getHost().getName());
         if (StringUtils.isNotEmpty(hypervisorHostName)) {
             final List<NicVO> nics = _nicDao.listByVmId(vm.getId());
+            final Map<Long, NetworkVO> networksById = batchLoadNetworksForNics(nics);
             for (final NicVO nic : nics) {
-                final NetworkVO network = _networksDao.findById(nic.getNetworkId());
+                final NetworkVO network = networksById.get(nic.getNetworkId());
+                if (network == null) {
+                    continue;
+                }
                 final Integer networkRate = _networkModel.getNetworkRate(network.getId(), vm.getId());
                 final NicProfile profile = new NicProfile(nic, network, nic.getBroadcastUri(), nic.getIsolationUri(), networkRate, _networkModel.isSecurityGroupSupportedInNetwork(network),
                         _networkModel.getNetworkTag(vm.getHypervisorType(), network));
@@ -2266,8 +2279,13 @@ public void prepareNicForMigration(final VirtualMachineProfile vm, final DeployD
         }
         final List<NicVO> nics = _nicDao.listByVmId(vm.getId());
         final ReservationContext context = new ReservationContextImpl(UUID.randomUUID().toString(), null, null);
+        final Map<Long, NetworkVO> networksById = batchLoadNetworksForNics(nics);
         for (final NicVO nic : nics) {
-            final NetworkVO network = _networksDao.findById(nic.getNetworkId());
+            final NetworkVO network = networksById.get(nic.getNetworkId());
+            if (network == null) {
+                logger.warn("Network {} not found for nic {} during migration prep, skipping", nic.getNetworkId(), nic);
+                continue;
+            }
             final Integer networkRate = _networkModel.getNetworkRate(network.getId(), vm.getId());
 
             final NetworkGuru guru = AdapterBase.getAdapterByName(networkGurus, network.getGuruName());
@@ -2311,9 +2329,14 @@ public void prepareNicForMigration(final VirtualMachineProfile vm, final DeployD
     public void prepareAllNicsForMigration(final VirtualMachineProfile vm, final DeployDestination dest) {
         final List<NicVO> nics = _nicDao.listByVmId(vm.getId());
         final ReservationContext context = new ReservationContextImpl(UUID.randomUUID().toString(), null, null);
+        final Map<Long, NetworkVO> networksById = batchLoadNetworksForNics(nics);
         Long guestNetworkId = null;
         for (final NicVO nic : nics) {
-            final NetworkVO network = _networksDao.findById(nic.getNetworkId());
+            final NetworkVO network = networksById.get(nic.getNetworkId());
+            if (network == null) {
+                logger.warn("Network {} not found for nic {} during migration prep, skipping", nic.getNetworkId(), nic);
+                continue;
+            }
             if (network.getTrafficType().equals(TrafficType.Guest) && network.getGuestType().equals(GuestType.Isolated)) {
                 guestNetworkId = network.getId();
             }

engine/schema/src/main/java/com/cloud/network/dao/NetworkDao.java

@@ -32,6 +32,8 @@ public interface NetworkDao extends GenericDao<NetworkVO, Long>, StateDao<State,
 
     List<NetworkVO> listByOwner(long ownerId);
 
+    List<NetworkVO> listByIds(List<Long> ids);
+
     List<NetworkVO> listByGuestType(GuestType type);
 
     List<NetworkVO> listBy(long accountId, long offeringId, long dataCenterId);

engine/schema/src/main/java/com/cloud/network/dao/NetworkDaoImpl.java

@@ -563,6 +563,19 @@ public List<NetworkVO> listByOwner(final long ownerId) {
         return listBy(sc);
     }
 
+    @Override
+    public List<NetworkVO> listByIds(List<Long> ids) {
+        if (ids == null || ids.isEmpty()) {
+            return new ArrayList<>();
+        }
+        SearchBuilder<NetworkVO> sb = createSearchBuilder();
+        sb.and("id", sb.entity().getId(), Op.IN);
+        sb.done();
+        SearchCriteria<NetworkVO> sc = sb.create();
+        sc.setParameters("id", ids.toArray());
+        return listBy(sc);
+    }
+
     @Override
     public void addDomainToNetwork(final long networkId, final long domainId, final Boolean subdomainAccess) {
         addDomainToNetworknetwork(networkId, domainId, subdomainAccess);

engine/schema/src/test/java/com/cloud/network/dao/NetworkDaoImplTest.java

@@ -26,11 +26,13 @@
 import org.junit.Assert;
 import org.junit.Test;
 import org.junit.runner.RunWith;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mock;
 import org.mockito.Mockito;
 import org.mockito.junit.MockitoJUnitRunner;
 
 
+import java.util.Collections;
 import java.util.List;
 
 @RunWith(MockitoJUnitRunner.class)
@@ -68,4 +70,26 @@ public void listByPhysicalNetworkTrafficTypeTestSetParametersValidation() throws
             txn.close();
         }
     }
+
+    @Test
+    public void listByIdsTestNullOrEmptyIdsReturnsEmptyList() {
+        NetworkDaoImpl networkDaoImplSpy = Mockito.spy(NetworkDaoImpl.class);
+
+        Assert.assertTrue(networkDaoImplSpy.listByIds(null).isEmpty());
+        Assert.assertTrue(networkDaoImplSpy.listByIds(Collections.emptyList()).isEmpty());
+        Mockito.verify(networkDaoImplSpy, Mockito.never()).listBy(Mockito.any(SearchCriteria.class));
+    }
+
+    @Test
+    public void listByIdsTestBuildsInClause() {
+        NetworkDaoImpl networkDaoImplSpy = Mockito.spy(NetworkDaoImpl.class);
+        Mockito.doReturn(listNetworkVoMock).when(networkDaoImplSpy).listBy(Mockito.any(SearchCriteria.class));
+
+        List<NetworkVO> result = networkDaoImplSpy.listByIds(List.of(1L, 2L, 3L));
+
+        Assert.assertEquals(listNetworkVoMock, result);
+        ArgumentCaptor<SearchCriteria> scCaptor = ArgumentCaptor.forClass(SearchCriteria.class);
+        Mockito.verify(networkDaoImplSpy).listBy(scCaptor.capture());
+        Assert.assertEquals("networks.id IN (?,?,?)", scCaptor.getValue().getWhereClause().trim());
+    }
 }

server/src/test/java/com/cloud/vpc/dao/MockNetworkDaoImpl.java

@@ -280,4 +280,9 @@ public List<NetworkVO> listByNetworkDomainsAndAccountIds(Set<String> uniqueNtwkD
     public List<NetworkVO> listByNetworkDomainsAndDomainIds(Set<String> uniqueNtwkDomains, Set<Long> domainIds) {
         return List.of();
     }
+
+    @Override
+    public List<NetworkVO> listByIds(List<Long> ids) {
+        return List.of();
+    }
 }