From f2e4a9af638489e2aada4ae7c81f738540b6ee05 Mon Sep 17 00:00:00 2001
From: Qiu Jian <qiujian@yunionyun.com>
Date: Sun, 19 Jul 2020 13:16:07 +0800
Subject: [PATCH] fix: quota usage not updated after disk resizing

---
 pkg/cloudcommon/db/quotas/interface.go |  1 +
 pkg/cloudcommon/db/quotas/quotas.go    | 11 +++++++++++
 pkg/cloudcommon/db/quotas/register.go  | 14 ++++++++++++++
 pkg/compute/tasks/disk_resize_task.go  | 17 +++++++++++++++++
 4 files changed, 43 insertions(+)

diff --git a/pkg/cloudcommon/db/quotas/interface.go b/pkg/cloudcommon/db/quotas/interface.go
index 7cfca841bf..101e5c6fde 100644
--- a/pkg/cloudcommon/db/quotas/interface.go
+++ b/pkg/cloudcommon/db/quotas/interface.go
@@ -71,6 +71,7 @@ type IQuotaManager interface {
 	checkSetPendingQuota(ctx context.Context, userCred mcclient.TokenCredential, quota IQuota) error
 	cancelPendingUsage(ctx context.Context, userCred mcclient.TokenCredential, localUsage IQuota, cancelUsage IQuota, save bool) error
 	cancelUsage(ctx context.Context, userCred mcclient.TokenCredential, usage IQuota) error
+	addUsage(ctx context.Context, userCred mcclient.TokenCredential, usage IQuota) error
 	getQuotaCount(ctx context.Context, request IQuota, pendingKey IQuotaKeys) (int, error)
 
 	FetchIdNames(ctx context.Context, idMap map[string]map[string]string) (map[string]map[string]string, error)
diff --git a/pkg/cloudcommon/db/quotas/quotas.go b/pkg/cloudcommon/db/quotas/quotas.go
index 3f15c112db..3b9e2f9c85 100644
--- a/pkg/cloudcommon/db/quotas/quotas.go
+++ b/pkg/cloudcommon/db/quotas/quotas.go
@@ -113,6 +113,17 @@ func (manager *SQuotaBaseManager) _cancelUsage(ctx context.Context, userCred mcc
 	return manager.changeUsage(ctx, userCred, usage, false)
 }
 
+func (manager *SQuotaBaseManager) addUsage(ctx context.Context, userCred mcclient.TokenCredential, usage IQuota) error {
+	LockQuota(ctx, manager, usage)
+	defer ReleaseQuota(ctx, manager, usage)
+
+	return manager._addUsage(ctx, userCred, usage)
+}
+
+func (manager *SQuotaBaseManager) _addUsage(ctx context.Context, userCred mcclient.TokenCredential, usage IQuota) error {
+	return manager.changeUsage(ctx, userCred, usage, true)
+}
+
 func (manager *SQuotaBaseManager) changeUsage(ctx context.Context, userCred mcclient.TokenCredential, usage IQuota, isAdd bool) error {
 	usages, err := manager.usageStore.GetParentQuotas(ctx, usage.GetKeys())
 	if err != nil {
diff --git a/pkg/cloudcommon/db/quotas/register.go b/pkg/cloudcommon/db/quotas/register.go
index 7b8e607537..5a7eac1877 100644
--- a/pkg/cloudcommon/db/quotas/register.go
+++ b/pkg/cloudcommon/db/quotas/register.go
@@ -87,6 +87,20 @@ func cancelUsage(ctx context.Context, userCred mcclient.TokenCredential, usage I
 	}
 }
 
+func AddUsages(ctx context.Context, userCred mcclient.TokenCredential, usages []db.IUsage) {
+	for _, usage := range usages {
+		addUsage(ctx, userCred, usage.(IQuota))
+	}
+}
+
+func addUsage(ctx context.Context, userCred mcclient.TokenCredential, usage IQuota) {
+	manager := getQuotaManager(usage)
+	err := manager.addUsage(ctx, userCred, usage)
+	if err != nil {
+		log.Errorf("cancelUsage %s fail: %s", jsonutils.Marshal(usage), err)
+	}
+}
+
 func GetQuotaCount(ctx context.Context, request IQuota, pendingKeys IQuotaKeys) (int, error) {
 	manager := getQuotaManager(request)
 	return manager.getQuotaCount(ctx, request, pendingKeys)
diff --git a/pkg/compute/tasks/disk_resize_task.go b/pkg/compute/tasks/disk_resize_task.go
index b93b803794..051681a712 100644
--- a/pkg/compute/tasks/disk_resize_task.go
+++ b/pkg/compute/tasks/disk_resize_task.go
@@ -23,6 +23,7 @@ import (
 
 	api "yunion.io/x/onecloud/pkg/apis/compute"
 	"yunion.io/x/onecloud/pkg/cloudcommon/db"
+	"yunion.io/x/onecloud/pkg/cloudcommon/db/quotas"
 	"yunion.io/x/onecloud/pkg/cloudcommon/db/taskman"
 	"yunion.io/x/onecloud/pkg/compute/models"
 	"yunion.io/x/onecloud/pkg/mcclient"
@@ -105,6 +106,7 @@ func (self *DiskResizeTask) OnDiskResizeComplete(ctx context.Context, disk *mode
 		self.OnStartResizeDiskFailed(ctx, disk, err)
 		return
 	}
+	diff := int(sizeMb) - disk.DiskSize
 	oldStatus := disk.Status
 	_, err = db.Update(disk, func() error {
 		disk.Status = api.DISK_READY
@@ -119,6 +121,21 @@ func (self *DiskResizeTask) OnDiskResizeComplete(ctx context.Context, disk *mode
 	self.SetDiskReady(ctx, disk, self.GetUserCred(), "")
 	notes := fmt.Sprintf("%s=>%s", oldStatus, disk.Status)
 	db.OpsLog.LogEvent(disk, db.ACT_UPDATE_STATUS, notes, self.UserCred)
+
+	if diff > 0 {
+		var addUsage models.SQuota
+		keys, err := disk.GetQuotaKeys()
+		if err != nil { // just log and ignore
+			log.Errorf("disk.GetQuotaKeys fail %s", err)
+		} else {
+			addUsage.SetKeys(keys)
+			addUsage.Storage = diff
+			quotas.AddUsages(ctx, self.UserCred, []db.IUsage{&addUsage})
+		}
+	} else if diff < 0 {
+		// unlikely
+	}
+
 	self.CleanHostSchedCache(disk)
 	db.OpsLog.LogEvent(disk, db.ACT_RESIZE, disk.GetShortDesc(ctx), self.UserCred)
 	logclient.AddActionLogWithStartable(self, disk, logclient.ACT_RESIZE, nil, self.UserCred, true)