Speedup resource count calculation (#8903)

* Speed up resource count calculation

* Refactor resource count calculation

* Start transaction for updateCountByDeltaForIds

engine/schema/src/main/java/com/cloud/configuration/dao/ResourceCountDao.java

@@ -49,10 +49,15 @@ public interface ResourceCountDao extends GenericDao<ResourceCountVO, Long> {
 
     ResourceCountVO findByOwnerAndTypeAndTag(long ownerId, ResourceOwnerType ownerType, ResourceType type, String tag);
 
+    List<ResourceCountVO> findByOwnersAndTypeAndTag(List<Long> ownerIdList, ResourceOwnerType ownerType,
+            ResourceType type, String tag);
+
     List<ResourceCountVO> listResourceCountByOwnerType(ResourceOwnerType ownerType);
 
     Set<Long> listAllRowsToUpdate(long ownerId, ResourceOwnerType ownerType, ResourceType type, String tag);
 
+    boolean updateCountByDeltaForIds(List<Long> ids, boolean increment, long delta);
+
     Set<Long> listRowsToUpdateForDomain(long domainId, ResourceType type, String tag);
 
     long removeEntriesByOwner(long ownerId, ResourceOwnerType ownerType);
@@ -72,4 +77,6 @@ public interface ResourceCountDao extends GenericDao<ResourceCountVO, Long> {
     long countMemoryAllocatedToAccount(long accountId);
 
     void removeResourceCountsForNonMatchingTags(Long ownerId, ResourceOwnerType ownerType, List<ResourceType> types, List<String> tags);
+
+    List<ResourceCountVO> lockRows(Set<Long> ids);
 }

engine/schema/src/main/java/com/cloud/configuration/dao/ResourceCountDaoImpl.java

@@ -24,6 +24,7 @@ import java.util.Arrays;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Set;
+import java.util.stream.Collectors;
 
 import javax.annotation.PostConstruct;
 import javax.inject.Inject;
@@ -56,27 +57,30 @@ public class ResourceCountDaoImpl extends GenericDaoBase<ResourceCountVO, Long>
     private final SearchBuilder<ResourceCountVO> TypeSearch;
     private final SearchBuilder<ResourceCountVO> TypeNullTagSearch;
     private final SearchBuilder<ResourceCountVO> NonMatchingTagsSearch;
-
     private final SearchBuilder<ResourceCountVO> AccountSearch;
     private final SearchBuilder<ResourceCountVO> DomainSearch;
+    private final SearchBuilder<ResourceCountVO> IdsSearch;
 
     @Inject
     private DomainDao _domainDao;
     @Inject
     private AccountDao _accountDao;
 
+    protected static final String INCREMENT_COUNT_BY_IDS_SQL = "UPDATE `cloud`.`resource_count` SET `count` = `count` + ? WHERE `id` IN (?)";
+    protected static final String DECREMENT_COUNT_BY_IDS_SQL = "UPDATE `cloud`.`resource_count` SET `count` = `count` - ? WHERE `id` IN (?)";
+
     public ResourceCountDaoImpl() {
         TypeSearch = createSearchBuilder();
         TypeSearch.and("type", TypeSearch.entity().getType(), SearchCriteria.Op.EQ);
-        TypeSearch.and("accountId", TypeSearch.entity().getAccountId(), SearchCriteria.Op.EQ);
+        TypeSearch.and("accountId", TypeSearch.entity().getAccountId(), SearchCriteria.Op.IN);
-        TypeSearch.and("domainId", TypeSearch.entity().getDomainId(), SearchCriteria.Op.EQ);
+        TypeSearch.and("domainId", TypeSearch.entity().getDomainId(), SearchCriteria.Op.IN);
         TypeSearch.and("tag", TypeSearch.entity().getTag(), SearchCriteria.Op.EQ);
         TypeSearch.done();
 
         TypeNullTagSearch = createSearchBuilder();
         TypeNullTagSearch.and("type", TypeNullTagSearch.entity().getType(), SearchCriteria.Op.EQ);
-        TypeNullTagSearch.and("accountId", TypeNullTagSearch.entity().getAccountId(), SearchCriteria.Op.EQ);
+        TypeNullTagSearch.and("accountId", TypeNullTagSearch.entity().getAccountId(), SearchCriteria.Op.IN);
-        TypeNullTagSearch.and("domainId", TypeNullTagSearch.entity().getDomainId(), SearchCriteria.Op.EQ);
+        TypeNullTagSearch.and("domainId", TypeNullTagSearch.entity().getDomainId(), SearchCriteria.Op.IN);
         TypeNullTagSearch.and("tag", TypeNullTagSearch.entity().getTag(), SearchCriteria.Op.NULL);
         TypeNullTagSearch.done();
 
@@ -90,6 +94,10 @@ public class ResourceCountDaoImpl extends GenericDaoBase<ResourceCountVO, Long>
 
         AccountSearch = createSearchBuilder();
         DomainSearch = createSearchBuilder();
+
+        IdsSearch = createSearchBuilder();
+        IdsSearch.and("id", IdsSearch.entity().getId(), SearchCriteria.Op.IN);
+        IdsSearch.done();
     }
 
     @PostConstruct
@@ -109,6 +117,19 @@ public class ResourceCountDaoImpl extends GenericDaoBase<ResourceCountVO, Long>
 
     @Override
     public ResourceCountVO findByOwnerAndTypeAndTag(long ownerId, ResourceOwnerType ownerType, ResourceType type, String tag) {
+        List<ResourceCountVO> resourceCounts = findByOwnersAndTypeAndTag(List.of(ownerId), ownerType, type, tag);
+        if (CollectionUtils.isNotEmpty(resourceCounts)) {
+            return resourceCounts.get(0);
+        } else {
+            return null;
+        }
+    }
+
+    @Override
+    public List<ResourceCountVO> findByOwnersAndTypeAndTag(List<Long> ownerIdList, ResourceOwnerType ownerType, ResourceType type, String tag) {
+        if (CollectionUtils.isEmpty(ownerIdList)) {
+            return new ArrayList<>();
+        }
         SearchCriteria<ResourceCountVO> sc = tag != null ? TypeSearch.create() : TypeNullTagSearch.create();
         sc.setParameters("type", type);
         if (tag != null) {
@@ -116,13 +137,13 @@ public class ResourceCountDaoImpl extends GenericDaoBase<ResourceCountVO, Long>
         }
 
         if (ownerType == ResourceOwnerType.Account) {
-            sc.setParameters("accountId", ownerId);
+            sc.setParameters("accountId", ownerIdList.toArray());
-            return findOneIncludingRemovedBy(sc);
+            return listIncludingRemovedBy(sc);
         } else if (ownerType == ResourceOwnerType.Domain) {
-            sc.setParameters("domainId", ownerId);
+            sc.setParameters("domainId", ownerIdList.toArray());
-            return findOneIncludingRemovedBy(sc);
+            return listIncludingRemovedBy(sc);
         } else {
-            return null;
+            return new ArrayList<>();
         }
     }
 
@@ -154,6 +175,26 @@ public class ResourceCountDaoImpl extends GenericDaoBase<ResourceCountVO, Long>
         return update(resourceCountVO.getId(), resourceCountVO);
     }
 
+    @Override
+    public boolean updateCountByDeltaForIds(List<Long> ids, boolean increment, long delta) {
+        if (CollectionUtils.isEmpty(ids)) {
+            return false;
+        }
+        String updateSql = increment ? INCREMENT_COUNT_BY_IDS_SQL : DECREMENT_COUNT_BY_IDS_SQL;
+
+        String poolIdsInStr = ids.stream().map(String::valueOf).collect(Collectors.joining(",", "(", ")"));
+        String sql = updateSql.replace("(?)", poolIdsInStr);
+
+        final TransactionLegacy txn = TransactionLegacy.currentTxn();
+        try(PreparedStatement pstmt = txn.prepareStatement(sql);) {
+            pstmt.setLong(1, delta);
+            pstmt.executeUpdate();
+            return true;
+        } catch (SQLException e) {
+            throw new CloudRuntimeException(e);
+        }
+    }
+
     @Override
     public Set<Long> listRowsToUpdateForDomain(long domainId, ResourceType type, String tag) {
         Set<Long> rowIds = new HashSet<Long>();
@@ -345,4 +386,14 @@ public class ResourceCountDaoImpl extends GenericDaoBase<ResourceCountVO, Long>
         }
         remove(sc);
     }
+
+    @Override
+    public List<ResourceCountVO> lockRows(Set<Long> ids) {
+        if (CollectionUtils.isEmpty(ids)) {
+            return new ArrayList<>();
+        }
+        SearchCriteria<ResourceCountVO> sc = IdsSearch.create();
+        sc.setParameters("id", ids.toArray());
+        return lockRows(sc, null, true);
+    }
 }

engine/schema/src/main/java/org/apache/cloudstack/reservation/dao/ReservationDao.java

@@ -31,4 +31,5 @@ public interface ReservationDao extends GenericDao<ReservationVO, Long> {
     void setResourceId(Resource.ResourceType type, Long resourceId);
     List<Long> getResourceIds(long accountId, Resource.ResourceType type);
     List<ReservationVO> getReservationsForAccount(long accountId, Resource.ResourceType type, String tag);
+    void removeByIds(List<Long> reservationIds);
 }

engine/schema/src/main/java/org/apache/cloudstack/reservation/dao/ReservationDaoImpl.java

@@ -29,6 +29,7 @@ import com.cloud.utils.db.GenericDaoBase;
 import com.cloud.utils.db.SearchBuilder;
 import com.cloud.utils.db.SearchCriteria;
 import org.apache.cloudstack.user.ResourceReservation;
+import org.apache.commons.collections.CollectionUtils;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 
@@ -40,6 +41,7 @@ public class ReservationDaoImpl extends GenericDaoBase<ReservationVO, Long> impl
     private static final String RESOURCE_ID = "resourceId";
     private static final String ACCOUNT_ID = "accountId";
     private static final String DOMAIN_ID = "domainId";
+    private static final String IDS = "ids";
     private final SearchBuilder<ReservationVO> listResourceByAccountAndTypeSearch;
     private final SearchBuilder<ReservationVO> listAccountAndTypeSearch;
     private final SearchBuilder<ReservationVO> listAccountAndTypeAndNoTagSearch;
@@ -47,6 +49,7 @@ public class ReservationDaoImpl extends GenericDaoBase<ReservationVO, Long> impl
     private final SearchBuilder<ReservationVO> listDomainAndTypeSearch;
     private final SearchBuilder<ReservationVO> listDomainAndTypeAndNoTagSearch;
     private final SearchBuilder<ReservationVO> listResourceByAccountAndTypeAndNoTagSearch;
+    private final SearchBuilder<ReservationVO> listIdsSearch;
 
     public ReservationDaoImpl() {
 
@@ -87,6 +90,10 @@ public class ReservationDaoImpl extends GenericDaoBase<ReservationVO, Long> impl
         listDomainAndTypeAndNoTagSearch.and(RESOURCE_TYPE, listDomainAndTypeAndNoTagSearch.entity().getResourceType(), SearchCriteria.Op.EQ);
         listDomainAndTypeAndNoTagSearch.and(RESOURCE_TAG, listDomainAndTypeAndNoTagSearch.entity().getTag(), SearchCriteria.Op.NULL);
         listDomainAndTypeAndNoTagSearch.done();
+
+        listIdsSearch = createSearchBuilder();
+        listIdsSearch.and(IDS, listIdsSearch.entity().getId(), SearchCriteria.Op.IN);
+        listIdsSearch.done();
     }
 
     @Override
@@ -161,4 +168,13 @@ public class ReservationDaoImpl extends GenericDaoBase<ReservationVO, Long> impl
         }
         return listBy(sc);
     }
+
+    @Override
+    public void removeByIds(List<Long> reservationIds) {
+        if (CollectionUtils.isNotEmpty(reservationIds)) {
+            SearchCriteria<ReservationVO> sc = listIdsSearch.create();
+            sc.setParameters(IDS, reservationIds.toArray());
+            remove(sc);
+        }
+    }
 }

server/src/main/java/com/cloud/resourcelimit/ResourceLimitManagerImpl.java

@@ -21,6 +21,7 @@ import static com.cloud.utils.NumbersUtil.toHumanReadableSize;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
@@ -203,17 +204,12 @@ public class ResourceLimitManagerImpl extends ManagerBase implements ResourceLim
 
     @SuppressWarnings("unchecked")
     protected void removeResourceReservationIfNeededAndIncrementResourceCount(final long accountId, final ResourceType type, String tag, final long numToIncrement) {
+        Object obj = CallContext.current().getContextParameter(CheckedReservation.getResourceReservationContextParameterKey(type));
+        List<Long> reservationIds = (List<Long>)obj; // This complains an unchecked casting warning
         Transaction.execute(new TransactionCallbackWithExceptionNoReturn<CloudRuntimeException>() {
             @Override
             public void doInTransactionWithoutResult(TransactionStatus status) throws CloudRuntimeException {
-
+                reservationDao.removeByIds(reservationIds);
-                Object obj = CallContext.current().getContextParameter(CheckedReservation.getResourceReservationContextParameterKey(type));
-                if (obj instanceof List) {
-                    List<Long> reservationIds = (List<Long>)obj; // This complains an unchecked casting warning
-                    for (Long reservationId : reservationIds) {
-                        reservationDao.remove(reservationId);
-                    }
-                }
                 if (!updateResourceCountForAccount(accountId, type, tag, true, numToIncrement)) {
                     // we should fail the operation (resource creation) when failed to update the resource count
                     throw new CloudRuntimeException("Failed to increment resource count of type " + type + " for account id=" + accountId);
@@ -579,13 +575,6 @@ public class ResourceLimitManagerImpl extends ManagerBase implements ResourceLim
         return _resourceCountDao.lockRows(sc, null, true);
     }
 
-    private List<ResourceCountVO> lockDomainRows(long domainId, final ResourceType type, String tag) {
-        Set<Long> rowIdsToLock = _resourceCountDao.listAllRowsToUpdate(domainId, ResourceOwnerType.Domain, type, tag);
-        SearchCriteria<ResourceCountVO> sc = ResourceCountSearch.create();
-        sc.setParameters("id", rowIdsToLock.toArray());
-        return _resourceCountDao.lockRows(sc, null, true);
-    }
-
     @Override
     public long findDefaultResourceLimitForDomain(ResourceType resourceType) {
         Long resourceLimit = null;
@@ -613,13 +602,11 @@ public class ResourceLimitManagerImpl extends ManagerBase implements ResourceLim
     }
 
     @Override
-    @DB
     public void checkResourceLimit(final Account account, final ResourceType type, long... count) throws ResourceAllocationException {
         checkResourceLimitWithTag(account, type, null, count);
     }
 
     @Override
-    @DB
     public void checkResourceLimitWithTag(final Account account, final ResourceType type, String tag, long... count) throws ResourceAllocationException {
         final long numResources = ((count.length == 0) ? 1 : count[0]);
         Project project = null;
@@ -1124,7 +1111,6 @@ public class ResourceLimitManagerImpl extends ManagerBase implements ResourceLim
         return recalculateResourceCount(accountId, domainId, typeId, null);
     }
 
-    @DB
     protected boolean updateResourceCountForAccount(final long accountId, final ResourceType type, String tag, final boolean increment, final long delta) {
         if (logger.isDebugEnabled()) {
             String convertedDelta = String.valueOf(delta);
@@ -1134,25 +1120,8 @@ public class ResourceLimitManagerImpl extends ManagerBase implements ResourceLim
             String typeStr = StringUtils.isNotEmpty(tag) ? String.format("%s (tag: %s)", type, tag) : type.getName();
             logger.debug("Updating resource Type = " + typeStr + " count for Account = " + accountId + " Operation = " + (increment ? "increasing" : "decreasing") + " Amount = " + convertedDelta);
         }
-        try {
+        Set<Long> rowIdsToUpdate = _resourceCountDao.listAllRowsToUpdate(accountId, ResourceOwnerType.Account, type, tag);
-            return Transaction.execute(new TransactionCallback<Boolean>() {
+        return _resourceCountDao.updateCountByDeltaForIds(new ArrayList<>(rowIdsToUpdate), increment, delta);
-                @Override
-                public Boolean doInTransaction(TransactionStatus status) {
-                    boolean result = true;
-                    List<ResourceCountVO> rowsToUpdate = lockAccountAndOwnerDomainRows(accountId, type, tag);
-                    for (ResourceCountVO rowToUpdate : rowsToUpdate) {
-                        if (!_resourceCountDao.updateById(rowToUpdate.getId(), increment, delta)) {
-                            logger.trace("Unable to update resource count for the row " + rowToUpdate);
-                            result = false;
-                        }
-                    }
-                    return result;
-                }
-            });
-        } catch (Exception ex) {
-            logger.error("Failed to update resource count for account id=" + accountId);
-            return false;
-        }
     }
 
     /**
@@ -1163,43 +1132,63 @@ public class ResourceLimitManagerImpl extends ManagerBase implements ResourceLim
      * @param type the resource type to do the recalculation for
      * @return the resulting new resource count
      */
-    @DB
     protected long recalculateDomainResourceCount(final long domainId, final ResourceType type, String tag) {
-        return Transaction.execute(new TransactionCallback<Long>() {
+        List<AccountVO> accounts = _accountDao.findActiveAccountsForDomain(domainId);
-            @Override
+        List<DomainVO> childDomains = _domainDao.findImmediateChildrenForParent(domainId);
-            public Long doInTransaction(TransactionStatus status) {
-                long newResourceCount = 0;
-                lockDomainRows(domainId, type, tag);
-                ResourceCountVO domainRC = _resourceCountDao.findByOwnerAndTypeAndTag(domainId, ResourceOwnerType.Domain, type, tag);
-                long oldResourceCount = domainRC.getCount();
 
-                List<DomainVO> domainChildren = _domainDao.findImmediateChildrenForParent(domainId);
+        if (CollectionUtils.isNotEmpty(childDomains)) {
-                // for each child domain update the resource count
+            for (DomainVO childDomain : childDomains) {
+                recalculateDomainResourceCount(childDomain.getId(), type, tag);
+            }
+        }
+        if (CollectionUtils.isNotEmpty(accounts)) {
+            for (AccountVO account : accounts) {
+                recalculateAccountResourceCount(account.getId(), type, tag);
+            }
+        }
+
+        return Transaction.execute((TransactionCallback<Long>) status -> {
+            long newResourceCount = 0L;
+            List<Long> domainIdList = childDomains.stream().map(DomainVO::getId).collect(Collectors.toList());
+            domainIdList.add(domainId);
+            List<Long> accountIdList = accounts.stream().map(AccountVO::getId).collect(Collectors.toList());
+            List<ResourceCountVO> domainRCList = _resourceCountDao.findByOwnersAndTypeAndTag(domainIdList, ResourceOwnerType.Domain, type, tag);
+            List<ResourceCountVO> accountRCList = _resourceCountDao.findByOwnersAndTypeAndTag(accountIdList, ResourceOwnerType.Account, type, tag);
+
+            Set<Long> rowIdsToLock = new HashSet<>();
+            if (domainRCList != null) {
+                rowIdsToLock.addAll(domainRCList.stream().map(ResourceCountVO::getId).collect(Collectors.toList()));
+            }
+            if (accountRCList != null) {
+                rowIdsToLock.addAll(accountRCList.stream().map(ResourceCountVO::getId).collect(Collectors.toList()));
+            }
+            // lock the resource count rows for current domain, immediate child domain & accounts
+            List<ResourceCountVO> resourceCounts = _resourceCountDao.lockRows(rowIdsToLock);
+
+            long oldResourceCount = 0L;
+            ResourceCountVO domainRC = null;
 
             // calculate project count here
             if (type == ResourceType.project) {
                 newResourceCount += _projectDao.countProjectsForDomain(domainId);
             }
 
-                for (DomainVO childDomain : domainChildren) {
+            for (ResourceCountVO resourceCount : resourceCounts) {
-                    long childDomainResourceCount = recalculateDomainResourceCount(childDomain.getId(), type, tag);
+                if (resourceCount.getResourceOwnerType() == ResourceOwnerType.Domain && resourceCount.getDomainId() == domainId) {
-                    newResourceCount += childDomainResourceCount; // add the child domain count to parent domain count
+                    oldResourceCount = resourceCount.getCount();
+                    domainRC = resourceCount;
+                } else {
+                    newResourceCount += resourceCount.getCount();
                 }
-
-                List<AccountVO> accounts = _accountDao.findActiveAccountsForDomain(domainId);
-                for (AccountVO account : accounts) {
-                    long accountResourceCount = recalculateAccountResourceCount(account.getId(), type, tag);
-                    newResourceCount += accountResourceCount; // add account's resource count to parent domain count
             }
-                _resourceCountDao.setResourceCount(domainId, ResourceOwnerType.Domain, type, tag, newResourceCount);
 
             if (oldResourceCount != newResourceCount) {
+                domainRC.setCount(newResourceCount);
+                _resourceCountDao.update(domainRC.getId(), domainRC);
                 logger.warn("Discrepency in the resource count has been detected " + "(original count = " + oldResourceCount + " correct count = " + newResourceCount + ") for Type = " + type
                         + " for Domain ID = " + domainId + " is fixed during resource count recalculation.");
             }
-
             return newResourceCount;
-            }
         });
     }
 
@@ -1238,16 +1227,10 @@ public class ResourceLimitManagerImpl extends ManagerBase implements ResourceLim
         final ResourceCountVO accountRC = _resourceCountDao.findByOwnerAndTypeAndTag(accountId, ResourceOwnerType.Account, type, tag);
         if (accountRC != null) {
             oldCount = accountRC.getCount();
-        }
-
             if (newCount == null || !newCount.equals(oldCount)) {
-            Transaction.execute(new TransactionCallbackNoReturn() {
+                accountRC.setCount((newCount == null) ? 0 : newCount);
-                @Override
+                _resourceCountDao.update(accountRC.getId(), accountRC);
-                public void doInTransactionWithoutResult(TransactionStatus status) {
-                    lockAccountAndOwnerDomainRows(accountId, type, tag);
-                    _resourceCountDao.setResourceCount(accountId, ResourceOwnerType.Account, type, tag, (newCount == null) ? 0 : newCount);
             }
-            });
         }
 
         // No need to log message for primary and secondary storage because both are recalculating the