Fix resource count discrepancy while associating IP address to a network or vpc (#9563)

Author: vishesh92

server/src/main/java/com/cloud/network/IpAddressManagerImpl.java

@@ -35,6 +35,7 @@
 
 import com.cloud.network.dao.PublicIpQuarantineDao;
 import com.cloud.network.vo.PublicIpQuarantineVO;
+import com.cloud.resourcelimit.CheckedReservation;
 import org.apache.cloudstack.acl.ControlledEntity.ACLType;
 import org.apache.cloudstack.acl.SecurityChecker.AccessType;
 import org.apache.cloudstack.annotation.AnnotationService;
@@ -53,6 +54,7 @@
 import org.apache.cloudstack.region.PortableIpDao;
 import org.apache.cloudstack.region.PortableIpVO;
 import org.apache.cloudstack.region.Region;
+import org.apache.cloudstack.reservation.dao.ReservationDao;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.log4j.Logger;
 
@@ -261,6 +263,8 @@ public class IpAddressManagerImpl extends ManagerBase implements IpAddressManage
     @Inject
     ResourceLimitService _resourceLimitMgr;
 
+    @Inject
+    ReservationDao reservationDao;
     @Inject
     NetworkOfferingServiceMapDao _ntwkOfferingSrvcDao;
     @Inject
@@ -1556,21 +1560,25 @@ public IPAddressVO associateIPToGuestNetwork(long ipId, long networkId, boolean
 
         s_logger.debug("Associating ip " + ipToAssoc + " to network " + network);
 
-        IPAddressVO ip = _ipAddressDao.findById(ipId);
-        //update ip address with networkId
-        ip.setAssociatedWithNetworkId(networkId);
-        ip.setSourceNat(isSourceNat);
-        _ipAddressDao.update(ipId, ip);
-
         boolean success = false;
-        try {
+        IPAddressVO ip = null;
+        try (CheckedReservation publicIpReservation = new CheckedReservation(owner, ResourceType.public_ip, 1l, reservationDao, _resourceLimitMgr)) {
+            ip = _ipAddressDao.findById(ipId);
+            //update ip address with networkId
+            ip.setAssociatedWithNetworkId(networkId);
+            ip.setSourceNat(isSourceNat);
+            _ipAddressDao.update(ipId, ip);
+
             success = applyIpAssociations(network, false);
             if (success) {
                 s_logger.debug("Successfully associated ip address " + ip.getAddress().addr() + " to network " + network);
             } else {
                 s_logger.warn("Failed to associate ip address " + ip.getAddress().addr() + " to network " + network);
             }
             return _ipAddressDao.findById(ipId);
+        } catch (Exception e) {
+            s_logger.error(String.format("Failed to associate ip address %s to network %s", ipToAssoc, network), e);
+            throw new CloudRuntimeException(String.format("Failed to associate ip address %s to network %s", ipToAssoc, network), e);
         } finally {
             if (!success && releaseOnFailure) {
                 if (ip != null) {

server/src/main/java/com/cloud/network/vpc/VpcManagerImpl.java

@@ -42,6 +42,7 @@
 import javax.inject.Inject;
 import javax.naming.ConfigurationException;
 
+import com.cloud.resourcelimit.CheckedReservation;
 import org.apache.cloudstack.acl.ControlledEntity.ACLType;
 import org.apache.cloudstack.alert.AlertService;
 import org.apache.cloudstack.annotation.AnnotationService;
@@ -63,6 +64,7 @@
 import org.apache.cloudstack.framework.config.dao.ConfigurationDao;
 import org.apache.cloudstack.managed.context.ManagedContextRunnable;
 import org.apache.cloudstack.query.QueryService;
+import org.apache.cloudstack.reservation.dao.ReservationDao;
 import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang3.ObjectUtils;
 import org.apache.log4j.Logger;
@@ -237,6 +239,8 @@ public class VpcManagerImpl extends ManagerBase implements VpcManager, VpcProvis
     @Inject
     ResourceLimitService _resourceLimitMgr;
     @Inject
+    ReservationDao reservationDao;
+    @Inject
     VpcServiceMapDao _vpcSrvcDao;
     @Inject
     DataCenterDao _dcDao;
@@ -2927,9 +2931,10 @@ public IpAddress associateIPToVpc(final long ipId, final long vpcId) throws Reso
         s_logger.debug("Associating ip " + ipToAssoc + " to vpc " + vpc);
 
         final boolean isSourceNatFinal = isSrcNatIpRequired(vpc.getVpcOfferingId()) && getExistingSourceNatInVpc(vpc.getAccountId(), vpcId) == null;
-        Transaction.execute(new TransactionCallbackNoReturn() {
-            @Override
-            public void doInTransactionWithoutResult(final TransactionStatus status) {
+        try (CheckedReservation publicIpReservation = new CheckedReservation(owner, ResourceType.public_ip, 1l, reservationDao, _resourceLimitMgr)) {
+            Transaction.execute(new TransactionCallbackNoReturn() {
+                @Override
+                public void doInTransactionWithoutResult(final TransactionStatus status) {
                 final IPAddressVO ip = _ipAddressDao.findById(ipId);
                 // update ip address with networkId
                 ip.setVpcId(vpcId);
@@ -2939,8 +2944,12 @@ public void doInTransactionWithoutResult(final TransactionStatus status) {
 
                 // mark ip as allocated
                 _ipAddrMgr.markPublicIpAsAllocated(ip);
-            }
-        });
+                }
+            });
+        } catch (Exception e) {
+            s_logger.error("Failed to associate ip " + ipToAssoc + " to vpc " + vpc, e);
+            throw new CloudRuntimeException("Failed to associate ip " + ipToAssoc + " to vpc " + vpc, e);
+        }
 
         s_logger.debug("Successfully assigned ip " + ipToAssoc + " to vpc " + vpc);
         CallContext.current().putContextParameter(IpAddress.class, ipToAssoc.getUuid());