From a8820548896fd17d5c309c447a10fa2baf8f9c7c Mon Sep 17 00:00:00 2001
From: Zexi Li <zexi.li@icloud.com>
Date: Mon, 21 Mar 2022 12:52:51 +0800
Subject: [PATCH] fix(region): guest and host cpuset one-to-one validate

---
 pkg/compute/models/guest_actions.go | 44 +++++++++++++++++++++++------
 pkg/compute/models/hosts.go         | 16 +++++++++++
 2 files changed, 52 insertions(+), 8 deletions(-)

diff --git a/pkg/compute/models/guest_actions.go b/pkg/compute/models/guest_actions.go
index 7a60480610..d884c6ae5e 100644
--- a/pkg/compute/models/guest_actions.go
+++ b/pkg/compute/models/guest_actions.go
@@ -5407,6 +5407,25 @@ func (self *SGuest) PerformCpuset(ctx context.Context, userCred mcclient.TokenCr
 		return nil, httperrors.NewInputParameterError("Host cores %v not contains input %v", allCores, data.CPUS)
 	}
 
+	host, err := self.GetHost()
+	if err != nil {
+		return nil, err
+	}
+
+	pinnedMap, err := host.GetPinnedCpusetCores(ctx, userCred)
+	if err != nil {
+		return nil, errors.Wrap(err, "Get host pinned cpu cores")
+	}
+
+	pinnedSets := sets.NewInt()
+	for _, pinned := range pinnedMap {
+		pinnedSets.Insert(pinned...)
+	}
+
+	if pinnedSets.HasAny(data.CPUS...) {
+		return nil, httperrors.NewInputParameterError("More than one of input cores %v already setted in host %v", data.CPUS, pinnedSets.List())
+	}
+
 	if err := self.SetMetadata(ctx, api.VM_METADATA_CGROUP_CPUSET, data, userCred); err != nil {
 		return nil, errors.Wrap(err, "set metadata")
 	}
@@ -5446,6 +5465,18 @@ func (self *SGuest) PerformCpusetRemove(ctx context.Context, userCred mcclient.T
 	return resp, nil
 }
 
+func (self *SGuest) getPinnedCpusetCores(ctx context.Context, userCred mcclient.TokenCredential) ([]int, error) {
+	obj := self.GetMetadataJson(ctx, api.VM_METADATA_CGROUP_CPUSET, userCred)
+	if obj == nil {
+		return nil, nil
+	}
+	pinnedInput := new(api.ServerCPUSetInput)
+	if err := obj.Unmarshal(pinnedInput); err != nil {
+		return nil, errors.Wrap(err, "Unmarshal to ServerCPUSetInput")
+	}
+	return pinnedInput.CPUS, nil
+}
+
 func (self *SGuest) GetDetailsCpusetCores(ctx context.Context, userCred mcclient.TokenCredential, input *api.ServerGetCPUSetCoresInput) (*api.ServerGetCPUSetCoresResp, error) {
 	allCores, err := self.getHostLogicalCores()
 	if err != nil {
@@ -5457,14 +5488,11 @@ func (self *SGuest) GetDetailsCpusetCores(ctx context.Context, userCred mcclient
 	}
 
 	// fetch cpuset pinned
-	obj := self.GetMetadataJson(ctx, api.VM_METADATA_CGROUP_CPUSET, userCred)
-	if obj != nil {
-		pinnedInput := new(api.ServerCPUSetInput)
-		if err := obj.Unmarshal(pinnedInput); err != nil {
-			log.Errorf("Unmarshal %q to ServerCPUSetInput: %v", obj, err)
-		} else {
-			resp.PinnedCores = pinnedInput.CPUS
-		}
+	pinned, err := self.getPinnedCpusetCores(ctx, userCred)
+	if err != nil {
+		log.Errorf("getPinnedCpusetCores error: %v", err)
+	} else {
+		resp.PinnedCores = pinned
 	}
 
 	return resp, nil
diff --git a/pkg/compute/models/hosts.go b/pkg/compute/models/hosts.go
index 45b052dabd..c420b16f78 100644
--- a/pkg/compute/models/hosts.go
+++ b/pkg/compute/models/hosts.go
@@ -5877,3 +5877,19 @@ func (manager *SHostManager) InitializeData() error {
 func (self *SHost) PerformProbeIsolatedDevices(ctx context.Context, userCred mcclient.TokenCredential, query jsonutils.JSONObject, data jsonutils.JSONObject) (jsonutils.JSONObject, error) {
 	return self.GetHostDriver().RequestProbeIsolatedDevices(ctx, userCred, self, data)
 }
+
+func (self *SHost) GetPinnedCpusetCores(ctx context.Context, userCred mcclient.TokenCredential) (map[string][]int, error) {
+	gsts, err := self.GetGuests()
+	if err != nil {
+		return nil, errors.Wrap(err, "Get all guests")
+	}
+	ret := make(map[string][]int, 0)
+	for _, gst := range gsts {
+		pinned, err := gst.getPinnedCpusetCores(ctx, userCred)
+		if err != nil {
+			return nil, errors.Wrapf(err, "get guest %s pinned cpuset cores", gst.GetName())
+		}
+		ret[gst.GetId()] = pinned
+	}
+	return ret, nil
+}