realtime re-balance remote agents

engine/orchestration/src/main/java/com/cloud/agent/manager/ClusteredAgentManagerImpl.java

@@ -824,7 +824,7 @@ public void startRebalanceAgents() {
         final List<ManagementServerHostVO> allMS = _mshostDao.listBy(ManagementServerHost.State.Up);
         final QueryBuilder<HostVO> sc = QueryBuilder.create(HostVO.class);
         sc.and(sc.entity().getManagementServerId(), Op.NNULL);
-        sc.and(sc.entity().getType(), Op.EQ, Host.Type.Routing);
+        sc.and(sc.entity().getType(), Op.IN, (Object[]) AgentLoadBalancerPlanner.REBALANCEABLE_HOST_TYPES);
         final List<HostVO> allManagedAgents = sc.list();
 
         int avLoad;
@@ -1037,6 +1037,20 @@ protected boolean rebalanceHost(final long hostId, final long currentOwnerId, fi
     protected boolean rebalanceHost(final long hostId, final long currentOwnerId, final long futureOwnerId, final boolean isConnectionTransfer) throws AgentUnavailableException {
         boolean result = true;
         if (currentOwnerId == _nodeId) {
+            final AgentAttache attache = findAttache(hostId);
+            if (attache != null && !(attache instanceof ClusteredDirectAgentAttache)) {
+                // Indirectly connected agents (KVM hosts, SSVM, CPVM) dial in to a management server rather
+                // than being loaded directly by it, so this management server can't hand the host to a
+                // specific future owner the way it can for direct agents. Disconnect it instead: the agent
+                // reconnects on its own using its indirect agent LB configuration (the "host" global setting
+                // and indirect.agent.lb.algorithm), which is what actually determines its next owner.
+                logger.debug("Host id={} ({}) is an indirectly connected agent; disconnecting it so it reconnects and picks a management server " +
+                        "using its own load balancing configuration", hostId, attache);
+                result = handleDisconnectWithoutInvestigation(attache, Event.AgentDisconnected, true, true);
+                finishRebalance(hostId, futureOwnerId, result ? Event.RebalanceCompleted : Event.RebalanceFailed);
+                return result;
+            }
+
             if (!startRebalance(hostId)) {
                 logger.debug("Failed to start agent rebalancing");
                 finishRebalance(hostId, futureOwnerId, Event.RebalanceFailed);
@@ -1577,11 +1591,11 @@ protected void runInContext() {
                     if (!_agentLbHappened) {
                         QueryBuilder<HostVO> sc = QueryBuilder.create(HostVO.class);
                         sc.and(sc.entity().getManagementServerId(), Op.NNULL);
-                        sc.and(sc.entity().getType(), Op.EQ, Host.Type.Routing);
+                        sc.and(sc.entity().getType(), Op.IN, (Object[]) AgentLoadBalancerPlanner.REBALANCEABLE_HOST_TYPES);
                         final List<HostVO> allManagedRoutingAgents = sc.list();
 
                         sc = QueryBuilder.create(HostVO.class);
-                        sc.and(sc.entity().getType(), Op.EQ, Host.Type.Routing);
+                        sc.and(sc.entity().getType(), Op.IN, (Object[]) AgentLoadBalancerPlanner.REBALANCEABLE_HOST_TYPES);
                         final List<HostVO> allAgents = sc.list();
                         final double allHostsCount = allAgents.size();
                         final double managedHostsCount = allManagedRoutingAgents.size();

engine/orchestration/src/main/java/com/cloud/cluster/agentlb/AgentLoadBalancerPlanner.java

@@ -19,11 +19,19 @@
 import java.util.List;
 
 import com.cloud.cluster.ManagementServerHostVO;
+import com.cloud.host.Host;
 import com.cloud.host.HostVO;
 import com.cloud.utils.component.Adapter;
 
 public interface AgentLoadBalancerPlanner extends Adapter {
 
+    /**
+     * Host types eligible for agent load balancing between management servers: hypervisor hosts of any
+     * hypervisor (KVM, VMware, XenServer, ...) as well as the system VM agents (SSVM, CPVM) that connect
+     * to a management server the same way a KVM host does.
+     */
+    Host.Type[] REBALANCEABLE_HOST_TYPES = {Host.Type.Routing, Host.Type.ConsoleProxy, Host.Type.SecondaryStorageVM};
+
     List<HostVO> getHostsToRebalance(ManagementServerHostVO ms, int avLoad);
 
 }

engine/orchestration/src/main/java/com/cloud/cluster/agentlb/ClusterBasedAgentLoadBalancerPlanner.java

@@ -29,7 +29,6 @@
 import com.cloud.cluster.ManagementServerHostVO;
 import org.springframework.stereotype.Component;
 
-import com.cloud.host.Host;
 import com.cloud.host.HostVO;
 import com.cloud.host.Status;
 import com.cloud.host.dao.HostDao;
@@ -47,7 +46,7 @@ public class ClusterBasedAgentLoadBalancerPlanner extends AdapterBase implements
     public List<HostVO> getHostsToRebalance(ManagementServerHostVO ms, int avLoad) {
         long msId = ms.getMsid();
         QueryBuilder<HostVO> sc = QueryBuilder.create(HostVO.class);
-        sc.and(sc.entity().getType(), Op.EQ, Host.Type.Routing);
+        sc.and(sc.entity().getType(), Op.IN, (Object[]) REBALANCEABLE_HOST_TYPES);
         sc.and(sc.entity().getManagementServerId(), Op.EQ, msId);
         List<HostVO> allHosts = sc.list();
 
@@ -60,7 +59,7 @@ public List<HostVO> getHostsToRebalance(ManagementServerHostVO ms, int avLoad) {
 
         sc = QueryBuilder.create(HostVO.class);
         sc.and(sc.entity().getManagementServerId(), Op.EQ, msId);
-        sc.and(sc.entity().getType(), Op.EQ, Host.Type.Routing);
+        sc.and(sc.entity().getType(), Op.IN, (Object[]) REBALANCEABLE_HOST_TYPES);
         sc.and(sc.entity().getStatus(), Op.EQ, Status.Up);
         List<HostVO> directHosts = sc.list();
 

engine/orchestration/src/test/java/com/cloud/agent/manager/ClusteredAgentManagerImplTest.java

@@ -17,10 +17,13 @@
 
 package com.cloud.agent.manager;
 
+import com.cloud.cluster.agentlb.dao.HostTransferMapDao;
 import com.cloud.configuration.ManagementServiceConfiguration;
+import com.cloud.exception.AgentUnavailableException;
 import com.cloud.ha.HighAvailabilityManagerImpl;
 import com.cloud.host.HostVO;
 import com.cloud.host.Status;
+import com.cloud.host.Status.Event;
 import com.cloud.host.dao.HostDao;
 import com.cloud.resource.ResourceManagerImpl;
 import org.junit.Before;
@@ -33,9 +36,12 @@
 import java.util.ArrayList;
 import java.util.List;
 
+import static org.junit.Assert.assertFalse;
+import static org.junit.Assert.assertTrue;
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.ArgumentMatchers.anyBoolean;
 import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.Mockito.doNothing;
 import static org.mockito.Mockito.doReturn;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.never;
@@ -147,4 +153,45 @@ public void scanDirectAgentToLoadHostWithNonForwardAttacheAndDisconnectedTest()
         verify(clusteredAgentManagerImpl).investigate(agentAttache);
         verify(clusteredAgentManagerImpl).loadDirectlyConnectedHost(hostVO, false);
     }
+
+    // https://github.com/apache/cloudstack/issues/9640
+    // Indirectly connected agents (KVM hosts, SSVM, CPVM) dial in to a management server rather than
+    // being loaded directly by it, so they must be disconnected (and left to reconnect on their own)
+    // instead of going through the direct-agent rebalance dance that expects a ClusteredDirectAgentAttache.
+    @Test
+    public void rebalanceHostDisconnectsIndirectAgentInsteadOfDirectRebalanceTest() throws AgentUnavailableException {
+        ClusteredAgentManagerImpl clusteredAgentManagerImpl = Mockito.spy(new ClusteredAgentManagerImpl());
+        clusteredAgentManagerImpl._nodeId = 1L;
+        clusteredAgentManagerImpl._hostTransferDao = mock(HostTransferMapDao.class);
+
+        long hostId = 10L;
+        AgentAttache indirectAttache = mock(ClusteredAgentAttache.class);
+        when(clusteredAgentManagerImpl.findAttache(hostId)).thenReturn(indirectAttache);
+        doReturn(true).when(clusteredAgentManagerImpl).handleDisconnectWithoutInvestigation(indirectAttache, Event.AgentDisconnected, true, true);
+        doNothing().when(clusteredAgentManagerImpl).finishRebalance(hostId, 2L, Event.RebalanceCompleted);
+
+        boolean result = clusteredAgentManagerImpl.rebalanceHost(hostId, 1L, 2L, false);
+
+        assertTrue(result);
+        verify(clusteredAgentManagerImpl).handleDisconnectWithoutInvestigation(indirectAttache, Event.AgentDisconnected, true, true);
+        verify(clusteredAgentManagerImpl, never()).startRebalance(hostId);
+    }
+
+    @Test
+    public void rebalanceHostStillUsesDirectRebalanceForDirectAgentTest() throws AgentUnavailableException {
+        ClusteredAgentManagerImpl clusteredAgentManagerImpl = Mockito.spy(new ClusteredAgentManagerImpl());
+        clusteredAgentManagerImpl._nodeId = 1L;
+
+        long hostId = 11L;
+        AgentAttache directAttache = mock(ClusteredDirectAgentAttache.class);
+        when(clusteredAgentManagerImpl.findAttache(hostId)).thenReturn(directAttache);
+        doReturn(false).when(clusteredAgentManagerImpl).startRebalance(hostId);
+        doNothing().when(clusteredAgentManagerImpl).finishRebalance(hostId, 2L, Event.RebalanceFailed);
+
+        boolean result = clusteredAgentManagerImpl.rebalanceHost(hostId, 1L, 2L, false);
+
+        assertFalse(result);
+        verify(clusteredAgentManagerImpl).startRebalance(hostId);
+        verify(clusteredAgentManagerImpl, never()).handleDisconnectWithoutInvestigation(any(), any(), anyBoolean(), anyBoolean());
+    }
 }