From 0b3850c14f9f3420d92bf6eb3477959b8f28e4ad Mon Sep 17 00:00:00 2001
From: Prachi Damle <prachi@cloud.com>
Date: Wed, 3 Apr 2013 16:53:57 -0700
Subject: [PATCH] Added cleanup of affinitygroups when a VM is expunging and
 when the account is deleted.

---
 server/src/com/cloud/user/AccountManagerImpl.java        | 9 ++++++++-
 .../cloudstack/affinity/AffinityGroupServiceImpl.java    | 5 ++---
 2 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/server/src/com/cloud/user/AccountManagerImpl.java b/server/src/com/cloud/user/AccountManagerImpl.java
index bc93df8e756..fe714c530b9 100755
--- a/server/src/com/cloud/user/AccountManagerImpl.java
+++ b/server/src/com/cloud/user/AccountManagerImpl.java
@@ -42,6 +42,7 @@ import org.apache.cloudstack.acl.ControlledEntity;
 import org.apache.cloudstack.acl.RoleType;
 import org.apache.cloudstack.acl.SecurityChecker;
 import org.apache.cloudstack.acl.SecurityChecker.AccessType;
+import org.apache.cloudstack.affinity.dao.AffinityGroupDao;
 import org.apache.cloudstack.api.command.admin.account.UpdateAccountCmd;
 import org.apache.cloudstack.api.command.admin.user.DeleteUserCmd;
 import org.apache.cloudstack.api.command.admin.user.RegisterCmd;
@@ -220,6 +221,8 @@ public class AccountManagerImpl extends ManagerBase implements AccountManager, M
     @Inject
     private AutoScaleManager _autoscaleMgr;
     @Inject VolumeManager volumeMgr;
+    @Inject
+    private AffinityGroupDao _affinityGroupDao;
 
     private List<UserAuthenticator> _userAuthenticators;
     List<UserAuthenticator> _userPasswordEncoders;
@@ -238,7 +241,7 @@ public class AccountManagerImpl extends ManagerBase implements AccountManager, M
     public List<UserAuthenticator> getUserAuthenticators() {
     	return _userAuthenticators;
     }
-    
+
     public void setUserAuthenticators(List<UserAuthenticator> authenticators) {
     	_userAuthenticators = authenticators;
     }
@@ -623,6 +626,10 @@ public class AccountManagerImpl extends ManagerBase implements AccountManager, M
             int numRemoved = _securityGroupDao.removeByAccountId(accountId);
             s_logger.info("deleteAccount: Deleted " + numRemoved + " network groups for account " + accountId);
 
+            // Cleanup affinity groups
+            int numAGRemoved = _affinityGroupDao.removeByAccountId(accountId);
+            s_logger.info("deleteAccount: Deleted " + numAGRemoved + " affinity groups for account " + accountId);
+
             // Delete all the networks
             boolean networksDeleted = true;
             s_logger.debug("Deleting networks for account " + account.getId());
diff --git a/server/src/org/apache/cloudstack/affinity/AffinityGroupServiceImpl.java b/server/src/org/apache/cloudstack/affinity/AffinityGroupServiceImpl.java
index e143db414e3..17f740118f1 100644
--- a/server/src/org/apache/cloudstack/affinity/AffinityGroupServiceImpl.java
+++ b/server/src/org/apache/cloudstack/affinity/AffinityGroupServiceImpl.java
@@ -240,6 +240,7 @@ public class AffinityGroupServiceImpl extends ManagerBase implements AffinityGro
     @Override
     public boolean configure(final String name, final Map<String, Object> params) throws ConfigurationException {
         _name = name;
+        VirtualMachine.State.getStateMachine().registerListener(this);
         return true;
     }
 
@@ -308,9 +309,7 @@ public class AffinityGroupServiceImpl extends ManagerBase implements AffinityGro
                 throw new InvalidParameterValueException("Unable to find affinity group by id " + affinityGroupId);
             }
         }
-        if (affinityGroupIds != null && !affinityGroupIds.isEmpty()) {
-            _affinityGroupVMMapDao.updateMap(vmId, affinityGroupIds);
-        }
+        _affinityGroupVMMapDao.updateMap(vmId, affinityGroupIds);
         // APIResponseHelper will pull out the updated affinitygroups.
         return vmInstance;