From d5a8f1d875667dcab6130b214029f2ec74603db0 Mon Sep 17 00:00:00 2001
From: Min Chen <min.chen@citrix.com>
Date: Mon, 15 Sep 2014 17:37:51 -0700
Subject: [PATCH] CLOUDSTACK-7553: Clean up cached agentMap and pingMap in case
 of agents connecting back to a different MS.

---
 .../cloud/agent/manager/AgentManagerImpl.java    | 16 ++++++++++++++++
 .../agent/manager/ClusteredAgentManagerImpl.java | 12 +++++++++---
 2 files changed, 25 insertions(+), 3 deletions(-)

diff --git a/engine/orchestration/src/com/cloud/agent/manager/AgentManagerImpl.java b/engine/orchestration/src/com/cloud/agent/manager/AgentManagerImpl.java
index f1f6eb9e589..2d5eb2b129e 100755
--- a/engine/orchestration/src/com/cloud/agent/manager/AgentManagerImpl.java
+++ b/engine/orchestration/src/com/cloud/agent/manager/AgentManagerImpl.java
@@ -1387,6 +1387,22 @@ public class AgentManagerImpl extends ManagerBase implements AgentManager, Handl
         _executor.submit(new DisconnectTask(attache, event, true));
     }
 
+    protected boolean isHostOwnerSwitched(final long hostId) {
+        HostVO host = _hostDao.findById(hostId);
+        if (host == null) {
+            s_logger.warn("Can't find the host " + hostId);
+            return false;
+        }
+        return isHostOwnerSwitched(host);
+    }
+
+    protected boolean isHostOwnerSwitched(HostVO host) {
+        if (host.getStatus() == Status.Up && host.getManagementServerId() != null && host.getManagementServerId() != _nodeId) {
+            return true;
+        }
+        return false;
+    }
+
     private void disconnectInternal(final long hostId, final Status.Event event, boolean invstigate) {
         AgentAttache attache = findAttache(hostId);
 
diff --git a/engine/orchestration/src/com/cloud/agent/manager/ClusteredAgentManagerImpl.java b/engine/orchestration/src/com/cloud/agent/manager/ClusteredAgentManagerImpl.java
index 5c6319d1a79..bf28e2bf698 100755
--- a/engine/orchestration/src/com/cloud/agent/manager/ClusteredAgentManagerImpl.java
+++ b/engine/orchestration/src/com/cloud/agent/manager/ClusteredAgentManagerImpl.java
@@ -254,7 +254,10 @@ public class ClusteredAgentManagerImpl extends AgentManagerImpl implements Clust
             _agents.put(id, attache);
         }
         if (old != null) {
-            old.disconnect(Status.Removed);
+            if (s_logger.isDebugEnabled()) {
+                s_logger.debug("Remove stale agent attache from current management server");
+            }
+            removeAgent(old, Status.Removed);
         }
         return attache;
     }
@@ -547,8 +550,11 @@ public class ClusteredAgentManagerImpl extends AgentManagerImpl implements Clust
         }
 
         AgentAttache agent = findAttache(hostId);
-        if (agent == null) {
-            if (host.getStatus() == Status.Up && (host.getManagementServerId() != null && host.getManagementServerId() != _nodeId)) {
+        if (agent == null || !agent.forForward()) {
+            if (isHostOwnerSwitched(host)) {
+                if (s_logger.isDebugEnabled()) {
+                    s_logger.debug("Host " + hostId + " has switched to another management server, need to update agent map with a forwarding agent attache");
+                }
                 agent = createAttache(hostId);
             }
         }