From 390cc11e7ad99826637efe90f94f9b409ff09364 Mon Sep 17 00:00:00 2001
From: rainzm <mjoycarry@gmail.com>
Date: Wed, 23 Sep 2020 15:32:51 +0800
Subject: [PATCH] feat(region): sync capacityUsed for storage after deploying
 or deleting guest

---
 pkg/compute/models/guest_actions.go    | 47 ++++++++++++++++++++++++++
 pkg/compute/tasks/guest_create_task.go |  5 +++
 pkg/compute/tasks/guest_delete_task.go | 25 ++++++++++++++
 3 files changed, 77 insertions(+)

diff --git a/pkg/compute/models/guest_actions.go b/pkg/compute/models/guest_actions.go
index daf64d4207..c7377fa805 100644
--- a/pkg/compute/models/guest_actions.go
+++ b/pkg/compute/models/guest_actions.go
@@ -3660,6 +3660,53 @@ func (self *SGuest) PerformRenew(ctx context.Context, userCred mcclient.TokenCre
 	return nil, nil
 }
 
+func (self *SGuest) GetStorages() []*SStorage {
+	disks := self.GetDisks()
+	storageMap := make(map[string]*SStorage)
+	for i := range disks {
+		storage := disks[i].GetStorage()
+		if _, ok := storageMap[storage.GetId()]; !ok {
+			storageMap[storage.GetId()] = storage
+		}
+	}
+	ret := make([]*SStorage, 0, len(storageMap))
+	for _, s := range storageMap {
+		ret = append(ret, s)
+	}
+	return ret
+}
+
+func (self *SGuest) SyncCapacityUsedForStorage(ctx context.Context, storageIds []string) error {
+	if self.Hypervisor != api.HYPERVISOR_ESXI {
+		return nil
+	}
+	var storages []*SStorage
+	if len(storageIds) == 0 {
+		storages = self.GetStorages()
+	} else {
+		q := StorageManager.Query()
+		if len(storageIds) == 1 {
+			q = q.Equals("id", storageIds[0])
+		} else {
+			q = q.In("id", storageIds[0])
+		}
+		ss := make([]SStorage, 0, len(storageIds))
+		err := db.FetchModelObjects(StorageManager, q, &ss)
+		if err != nil {
+			return errors.Wrap(err, "FetchModelObjects")
+		}
+		storages = make([]*SStorage, len(ss))
+		for i := range ss {
+			storages[i] = &ss[i]
+		}
+	}
+	for _, s := range storages {
+		err := s.SyncCapacityUsed(ctx)
+		return errors.Wrapf(err, "unable to SyncCapacityUsed for storage %q", s.GetId())
+	}
+	return nil
+}
+
 func (self *SGuest) startGuestRenewTask(ctx context.Context, userCred mcclient.TokenCredential, duration string, parentTaskId string) error {
 	self.SetStatus(userCred, api.VM_RENEWING, "")
 	data := jsonutils.NewDict()
diff --git a/pkg/compute/tasks/guest_create_task.go b/pkg/compute/tasks/guest_create_task.go
index 7ded0ea8ce..a359df397c 100644
--- a/pkg/compute/tasks/guest_create_task.go
+++ b/pkg/compute/tasks/guest_create_task.go
@@ -109,6 +109,11 @@ func (self *GuestCreateTask) StartDeployGuest(ctx context.Context, guest *models
 
 func (self *GuestCreateTask) OnDeployGuestDescComplete(ctx context.Context, obj db.IStandaloneModel, data jsonutils.JSONObject) {
 	guest := obj.(*models.SGuest)
+	// sync capacityUsed for storage
+	err := guest.SyncCapacityUsedForStorage(ctx, nil)
+	if err != nil {
+		log.Errorf("unable to SyncCapacityUsedForStorage: %v", err)
+	}
 
 	// bind eip
 	{
diff --git a/pkg/compute/tasks/guest_delete_task.go b/pkg/compute/tasks/guest_delete_task.go
index 4071ff8eba..d4aee3efcb 100644
--- a/pkg/compute/tasks/guest_delete_task.go
+++ b/pkg/compute/tasks/guest_delete_task.go
@@ -36,6 +36,10 @@ type GuestDeleteTask struct {
 	SGuestBaseTask
 }
 
+var (
+	STORAGEIDS = "storage_ids"
+)
+
 func init() {
 	taskman.RegisterTask(GuestDeleteTask{})
 }
@@ -251,6 +255,15 @@ func (self *GuestDeleteTask) OnPendingDeleteCompleteFailed(ctx context.Context,
 }
 
 func (self *GuestDeleteTask) StartDeleteGuest(ctx context.Context, guest *models.SGuest) {
+	// Temporary storageids to sync capacityUsed after delete
+	{
+		storages := guest.GetStorages()
+		storageIds := make([]string, len(storages))
+		for i := range storages {
+			storageIds[i] = storages[i].GetId()
+		}
+		self.Params.Set(STORAGEIDS, jsonutils.NewStringArray(storageIds))
+	}
 	// No snapshot
 	self.SetStage("OnGuestDetachDisksComplete", nil)
 	guest.GetDriver().RequestDetachDisksFromGuestForDelete(ctx, guest, self)
@@ -303,6 +316,18 @@ func (self *GuestDeleteTask) OnGuestDeleteComplete(ctx context.Context, obj db.I
 	guest.EjectIso(self.UserCred)
 	guest.DeleteEip(ctx, self.UserCred)
 	guest.GetDriver().OnDeleteGuestFinalCleanup(ctx, guest, self.UserCred)
+	// sync capacity used for storage
+	ja, err := self.Params.GetArray(STORAGEIDS)
+	if err == nil {
+		storageIds := make([]string, len(ja))
+		for i := range ja {
+			storageIds[i], _ = ja[i].GetString()
+		}
+		err = guest.SyncCapacityUsedForStorage(ctx, storageIds)
+		if err != nil {
+			log.Errorf("unable to SyncCapacityUsedForStoarage: %v", err)
+		}
+	}
 	self.DeleteGuest(ctx, guest)
 }