[Quota] Improve Quota balance calculation flow (#8581)

framework/quota/src/main/java/org/apache/cloudstack/quota/QuotaManagerImpl.java

@@ -22,6 +22,7 @@ import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Date;
 import java.util.HashMap;
+import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.TimeZone;
@@ -54,6 +55,7 @@ import org.apache.commons.collections.CollectionUtils;
 import org.apache.commons.lang3.BooleanUtils;
 import org.apache.commons.lang3.StringUtils;
 import org.apache.commons.lang3.math.NumberUtils;
+import org.apache.commons.lang3.time.DateUtils;
 import org.apache.log4j.Logger;
 import org.springframework.stereotype.Component;
 
@@ -156,78 +158,81 @@ public class QuotaManagerImpl extends ManagerBase implements QuotaManager {
             return;
         }
 
-        QuotaUsageVO firstQuotaUsage = accountQuotaUsages.get(0);
+        Date startDate = accountQuotaUsages.get(0).getStartDate();
-        Date startDate = firstQuotaUsage.getStartDate();
+        Date endDate = accountQuotaUsages.get(0).getEndDate();
-        Date endDate = firstQuotaUsage.getStartDate();
+        Date lastQuotaUsageEndDate = accountQuotaUsages.get(accountQuotaUsages.size() - 1).getEndDate();
 
-        s_logger.info(String.format("Processing quota balance for account [%s] between [%s] and [%s].", accountToString, startDate,
+        LinkedHashSet<Pair<Date, Date>> periods = accountQuotaUsages.stream()
-                accountQuotaUsages.get(accountQuotaUsages.size() - 1).getEndDate()));
+                .map(quotaUsageVO -> new Pair<>(quotaUsageVO.getStartDate(), quotaUsageVO.getEndDate()))
+                .collect(Collectors.toCollection(LinkedHashSet::new));
+
+        s_logger.info(String.format("Processing quota balance for account[%s] between [%s] and [%s].", accountToString, startDate, lastQuotaUsageEndDate));
 
-        BigDecimal aggregatedUsage = BigDecimal.ZERO;
         long accountId = accountVo.getAccountId();
         long domainId = accountVo.getDomainId();
+        BigDecimal accountBalance = retrieveBalanceForUsageCalculation(accountId, domainId, startDate, accountToString);
 
-        aggregatedUsage = getUsageValueAccordingToLastQuotaUsageEntryAndLastQuotaBalance(accountId, domainId, startDate, endDate, aggregatedUsage, accountToString);
+        for (Pair<Date, Date> period : periods) {
+            startDate = period.first();
+            endDate = period.second();
 
-        for (QuotaUsageVO quotaUsage : accountQuotaUsages) {
+            accountBalance = calculateBalanceConsideringCreditsAddedAndQuotaUsed(accountBalance, accountQuotaUsages, accountId, domainId, startDate, endDate, accountToString);
-            Date quotaUsageStartDate = quotaUsage.getStartDate();
+            _quotaBalanceDao.saveQuotaBalance(new QuotaBalanceVO(accountId, domainId, accountBalance, endDate));
-            Date quotaUsageEndDate = quotaUsage.getEndDate();
-            BigDecimal quotaUsed = quotaUsage.getQuotaUsed();
-
-            if (quotaUsed.equals(BigDecimal.ZERO)) {
-                aggregatedUsage = aggregatedUsage.add(aggregateCreditBetweenDates(accountId, domainId, quotaUsageStartDate, quotaUsageEndDate, accountToString));
-                continue;
-            }
-
-            if (startDate.compareTo(quotaUsageStartDate) == 0) {
-                aggregatedUsage = aggregatedUsage.subtract(quotaUsed);
-                continue;
-            }
-
-            _quotaBalanceDao.saveQuotaBalance(new QuotaBalanceVO(accountId, domainId, aggregatedUsage, endDate));
-
-            aggregatedUsage = BigDecimal.ZERO;
-            startDate = quotaUsageStartDate;
-            endDate = quotaUsageEndDate;
-
-            QuotaBalanceVO lastRealBalanceEntry = _quotaBalanceDao.findLastBalanceEntry(accountId, domainId, endDate);
-            Date lastBalanceDate = new Date(0);
-
-            if (lastRealBalanceEntry != null) {
-                lastBalanceDate = lastRealBalanceEntry.getUpdatedOn();
-                aggregatedUsage = aggregatedUsage.add(lastRealBalanceEntry.getCreditBalance());
-            }
-
-            aggregatedUsage = aggregatedUsage.add(aggregateCreditBetweenDates(accountId, domainId, lastBalanceDate, endDate, accountToString));
-            aggregatedUsage = aggregatedUsage.subtract(quotaUsed);
         }
-
+        saveQuotaAccount(accountId, accountBalance, endDate);
-        _quotaBalanceDao.saveQuotaBalance(new QuotaBalanceVO(accountId, domainId, aggregatedUsage, endDate));
-        saveQuotaAccount(accountId, aggregatedUsage, endDate);
     }
 
-    protected BigDecimal getUsageValueAccordingToLastQuotaUsageEntryAndLastQuotaBalance(long accountId, long domainId, Date startDate, Date endDate, BigDecimal aggregatedUsage,
+    /**
-            String accountToString) {
+     * Calculates the balance for the given account considering the specified period. The balance is calculated as follows:
+     * <ol>
+     *     <li>The credits added in this period are added to the balance.</li>
+     *     <li>All quota consumed in this period are subtracted from the account balance.</li>
+     * </ol>
+     */
+    protected BigDecimal calculateBalanceConsideringCreditsAddedAndQuotaUsed(BigDecimal accountBalance, List<QuotaUsageVO> accountQuotaUsages, long accountId, long domainId,
+                                                                             Date startDate, Date endDate, String accountToString) {
+        accountBalance = accountBalance.add(aggregateCreditBetweenDates(accountId, domainId, startDate, endDate, accountToString));
+
+        for (QuotaUsageVO quotaUsageVO : accountQuotaUsages) {
+            if (DateUtils.isSameInstant(quotaUsageVO.getStartDate(), startDate)) {
+                accountBalance = accountBalance.subtract(quotaUsageVO.getQuotaUsed());
+            }
+        }
+        return accountBalance;
+    }
+
+    /**
+     * Retrieves the initial balance prior to the period of the quota processing.
+     * <ul>
+     *     <li>
+     *         If it is the first time of processing for the account, the credits prior to the quota processing are added, and the first balance is persisted in the DB.
+     *     </li>
+     *     <li>
+     *         Otherwise, the last real balance of the account is retrieved.
+     *     </li>
+     * </ul>
+     */
+    protected BigDecimal retrieveBalanceForUsageCalculation(long accountId, long domainId, Date startDate, String accountToString) {
+        BigDecimal accountBalance = BigDecimal.ZERO;
         QuotaUsageVO lastQuotaUsage = _quotaUsageDao.findLastQuotaUsageEntry(accountId, domainId, startDate);
 
         if (lastQuotaUsage == null) {
-            aggregatedUsage = aggregatedUsage.add(aggregateCreditBetweenDates(accountId, domainId, new Date(0), startDate, accountToString));
+            accountBalance = accountBalance.add(aggregateCreditBetweenDates(accountId, domainId, new Date(0), startDate, accountToString));
-            QuotaBalanceVO firstBalance = new QuotaBalanceVO(accountId, domainId, aggregatedUsage, startDate);
+            QuotaBalanceVO firstBalance = new QuotaBalanceVO(accountId, domainId, accountBalance, startDate);
 
             s_logger.debug(String.format("Persisting the first quota balance [%s] for account [%s].", firstBalance, accountToString));
             _quotaBalanceDao.saveQuotaBalance(firstBalance);
         } else {
-            QuotaBalanceVO lastRealBalance = _quotaBalanceDao.findLastBalanceEntry(accountId, domainId, endDate);
+            QuotaBalanceVO lastRealBalance = _quotaBalanceDao.findLastBalanceEntry(accountId, domainId, startDate);
 
-            if (lastRealBalance != null) {
+            if (lastRealBalance == null) {
-                aggregatedUsage = aggregatedUsage.add(lastRealBalance.getCreditBalance());
-                aggregatedUsage = aggregatedUsage.add(aggregateCreditBetweenDates(accountId, domainId, lastRealBalance.getUpdatedOn(), endDate, accountToString));
-            } else {
                 s_logger.warn(String.format("Account [%s] has quota usage entries, however it does not have a quota balance.", accountToString));
+            } else {
+                accountBalance = accountBalance.add(lastRealBalance.getCreditBalance());
             }
         }
 
-        return aggregatedUsage;
+        return accountBalance;
     }
 
     protected void saveQuotaAccount(long accountId, BigDecimal aggregatedUsage, Date endDate) {

test/integration/plugins/quota/test_quota_balance.py

@@ -0,0 +1,191 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""
+Test cases for validating the Quota balance of accounts
+"""
+
+from marvin.cloudstackTestCase import *
+from marvin.lib.utils import *
+from marvin.lib.base import *
+from marvin.lib.common import *
+from nose.plugins.attrib import attr
+
+
+class TestQuotaBalance(cloudstackTestCase):
+
+    @classmethod
+    def setUpClass(cls):
+        testClient = super(TestQuotaBalance, cls).getClsTestClient()
+        cls.apiclient = testClient.getApiClient()
+        cls.services = testClient.getParsedTestDataConfig()
+        cls.mgtSvrDetails = cls.config.__dict__["mgtSvr"][0].__dict__
+
+        # Get Zone, Domain and templates
+        cls.domain = get_domain(cls.apiclient)
+        cls.zone = get_zone(cls.apiclient, cls.testClient.getZoneForTests())
+        cls.zone
+
+        # Create Account
+        cls.account = Account.create(
+            cls.apiclient,
+            cls.services["account"],
+            domainid=cls.domain.id
+        )
+        cls._cleanup = [
+            cls.account,
+        ]
+
+        cls.services["account"] = cls.account.name
+
+        if not is_config_suitable(apiclient=cls.apiclient, name='quota.enable.service', value='true'):
+            cls.debug("Quota service is not enabled, therefore the configuration `quota.enable.service` will be set to `true` and the management server will be restarted.")
+            Configurations.update(cls.apiclient, "quota.enable.service", "true")
+            cls.restartServer()
+
+        return
+
+    @classmethod
+    def restartServer(cls):
+        """Restart management server"""
+
+        cls.debug("Restarting management server")
+        sshClient = SshClient(
+            cls.mgtSvrDetails["mgtSvrIp"],
+            22,
+            cls.mgtSvrDetails["user"],
+            cls.mgtSvrDetails["passwd"]
+        )
+
+        command = "service cloudstack-management restart"
+        sshClient.execute(command)
+
+        # Waits for management to come up in 5 mins, when it's up it will continue
+        timeout = time.time() + 300
+        while time.time() < timeout:
+            if cls.isManagementUp() is True:
+                time.sleep(30)
+                return
+            time.sleep(5)
+        return cls.fail("Management server did not come up, failing")
+
+    @classmethod
+    def isManagementUp(cls):
+        try:
+            cls.apiclient.listInfrastructure(listInfrastructure.listInfrastructureCmd())
+            return True
+        except Exception:
+            return False
+
+    @classmethod
+    def tearDownClass(cls):
+        try:
+            # Cleanup resources used
+            cleanup_resources(cls.apiclient, cls._cleanup)
+        except Exception as e:
+            raise Exception("Warning: Exception during cleanup : %s" % e)
+        return
+
+    def setUp(self):
+        self.apiclient = self.testClient.getApiClient()
+        self.dbclient = self.testClient.getDbConnection()
+        self.cleanup = []
+        self.tariffs = []
+        return
+
+    def tearDown(self):
+        try:
+            cleanup_resources(self.apiclient, self.cleanup)
+            self.delete_tariffs()
+        except Exception as e:
+            raise Exception("Warning: Exception during cleanup : %s" % e)
+        return
+
+    def delete_tariffs(self):
+        for tariff in self.tariffs:
+            cmd = quotaTariffDelete.quotaTariffDeleteCmd()
+            cmd.id = tariff.uuid
+            self.apiclient.quotaTariffDelete(cmd)
+
+    @attr(tags=["advanced", "smoke", "quota"], required_hardware="false")
+    def test_quota_balance(self):
+        """
+        Test Quota balance
+
+        Validate the following
+        1. Add credits to an account
+        2. Create Quota tariff for the usage type 21 (VM_DISK_IO_READ)
+        3. Simulate quota usage by inserting a row in the `cloud_usage` table
+        4. Update the balance of the account by calling the API quotaUpdate
+        5. Verify the balance of the account according to the tariff created
+        """
+
+        # Create quota tariff for the usage type 21 (VM_DISK_IO_READ)
+        cmd = quotaTariffCreate.quotaTariffCreateCmd()
+        cmd.name = 'Tariff'
+        cmd.value = '10'
+        cmd.usagetype = '21'
+        self.tariffs.append(self.apiclient.quotaTariffCreate(cmd))
+
+        # Add credits to the account
+        cmd = quotaCredits.quotaCreditsCmd()
+        cmd.account = self.account.name
+        cmd.domainid = self.domain.id
+        cmd.value = 100
+        self.apiclient.quotaCredits(cmd)
+
+        # Fetch account ID from account_uuid
+        account_id_select = f"SELECT id FROM account WHERE uuid = '{self.account.id}';"
+        self.debug(account_id_select)
+        qresultset = self.dbclient.execute(account_id_select)
+        account_id = qresultset[0][0]
+
+        # Fetch domain ID from domain_uuid
+        domain_id_select = f"SELECT id FROM `domain` d WHERE uuid = '{self.domain.id}';"
+        self.debug(domain_id_select)
+        qresultset = self.dbclient.execute(domain_id_select)
+        domain_id = qresultset[0][0]
+
+        # Fetch zone ID from zone_uuid
+        zone_id_select = f"SELECT id from data_center dc where dc.uuid = '{self.zone.id}';"
+        self.debug(zone_id_select)
+        qresultset = self.dbclient.execute(zone_id_select)
+        zone_id = qresultset[0][0]
+
+        start_date = datetime.datetime.now() + datetime.timedelta(seconds=1)
+        end_date = datetime.datetime.now() + datetime.timedelta(hours=1)
+
+        # Manually insert a usage regarding the usage type 21 (VM_DISK_IO_READ)
+        sql_query = (f"INSERT INTO cloud_usage.cloud_usage (zone_id,account_id,domain_id,description,usage_display,usage_type,raw_usage,vm_instance_id,vm_name,offering_id,template_id,"
+                     f"usage_id,`type`,`size`,network_id,start_date,end_date,virtual_size,cpu_speed,cpu_cores,memory,quota_calculated,is_hidden,state)"
+                     f" VALUES ('{zone_id}','{account_id}','{domain_id}','Test','1 Hrs',21,1,NULL,NULL,NULL,NULL,NULL,'VirtualMachine',NULL,NULL,'{start_date}','{end_date}',NULL,NULL,NULL,NULL,0,0,NULL);")
+        self.debug(sql_query)
+        self.dbclient.execute(sql_query)
+
+        # Update quota to calculate the balance of the account
+        cmd = quotaUpdate.quotaUpdateCmd()
+        self.apiclient.quotaUpdate(cmd)
+
+        # Retrieve the quota balance of the account
+        cmd = quotaBalance.quotaBalanceCmd()
+        cmd.domainid = self.account.domainid
+        cmd.account = self.account.name
+        response = self.apiclient.quotaBalance(cmd)
+
+        self.debug(f"The quota balance for the account {self.account.name} is {response.balance}.")
+        self.assertEqual(response.balance.startquota, 90, f"The `startQuota` response field is supposed to be 90 but was {response.balance.startquota}.")
+
+        return

usage/src/main/java/com/cloud/usage/UsageManagerImpl.java

@@ -425,9 +425,14 @@ public class UsageManagerImpl extends ManagerBase implements UsageManager, Runna
                 cal.add(Calendar.MILLISECOND, -1);
                 endDate = cal.getTime().getTime();
             } else {
-                endDate = cal.getTime().getTime(); // current time
                 cal.add(Calendar.MINUTE, -1 * _aggregationDuration);
+                cal.set(Calendar.SECOND, 0);
+                cal.set(Calendar.MILLISECOND, 0);
                 startDate = cal.getTime().getTime();
+
+                cal.add(Calendar.MINUTE, _aggregationDuration);
+                cal.add(Calendar.MILLISECOND, -1);
+                endDate = cal.getTime().getTime();
             }
 
             parse(job, startDate, endDate);