From 735668d20c66e562e6bd1b61c9486f709aefc8da Mon Sep 17 00:00:00 2001
From: wanyaoqi <18528551+wanyaoqi@users.noreply.github.com>
Date: Thu, 7 Nov 2024 23:16:07 +0800
Subject: [PATCH] fix(scheduler,region,host): guest add alloc extra cpu
 (#21557)

---
 pkg/apis/compute/api.go                       |  4 +++
 pkg/apis/compute/guests.go                    | 26 +++++++++++---
 pkg/apis/scheduler/api.go                     |  9 ++---
 pkg/compute/guestdrivers/base.go              |  7 +++-
 pkg/compute/guestdrivers/kvm.go               |  4 +++
 pkg/compute/models/guest_actions.go           |  6 +++-
 pkg/compute/models/guests.go                  |  7 ++--
 pkg/compute/tasks/guest_change_config_task.go | 18 ++++++++++
 pkg/hostman/guestman/desc/desc.go             |  3 +-
 pkg/hostman/guestman/pod.go                   |  7 ++--
 pkg/mcclient/options/compute/servers.go       | 10 +++---
 .../predicates/guest/cpu_predicate.go         |  2 +-
 pkg/scheduler/cache/candidate/hosts.go        | 10 ++++++
 pkg/scheduler/core/result.go                  | 34 +++++++++++++++----
 14 files changed, 119 insertions(+), 28 deletions(-)

diff --git a/pkg/apis/compute/api.go b/pkg/apis/compute/api.go
index 1bb133a481..82f35a7b8d 100644
--- a/pkg/apis/compute/api.go
+++ b/pkg/apis/compute/api.go
@@ -474,6 +474,10 @@ type ServerCreateInput struct {
 	// default: 1
 	CpuSockets int `json:"cpu_sockets"`
 
+	// 额外分配 cpu 数量
+	// required: false
+	ExtraCpuCount int `json:"extra_cpu_count"`
+
 	// 用户自定义启动脚本
 	// 支持 #cloud-config yaml 格式及shell脚本
 	// 支持特殊user data平台: Aliyun, Qcloud, Azure, Apsara, Ucloud
diff --git a/pkg/apis/compute/guests.go b/pkg/apis/compute/guests.go
index 08a3948fe6..094687a5be 100644
--- a/pkg/apis/compute/guests.go
+++ b/pkg/apis/compute/guests.go
@@ -818,6 +818,8 @@ type ServerChangeConfigInput struct {
 
 	// cpu大小
 	VcpuCount *int `json:"vcpu_count"`
+	// 任务分配CPU大小
+	ExtraCpuCount *int `json:"extra_cpu_count"`
 	// 内存大小, 1024M, 1G
 	VmemSize string `json:"vmem_size"`
 
@@ -953,7 +955,8 @@ type SCpuNumaPin struct {
 	SizeMB *int `json:"size_mb"`
 	NodeId int  `json:"node_id"`
 
-	VcpuPin []SVCpuPin `json:"vcpu_pin"`
+	VcpuPin       []SVCpuPin `json:"vcpu_pin"`
+	ExtraCpuCount int        `json:"extra_cpu_count"`
 }
 
 type ServerSetBootIndexInput struct {
@@ -1309,10 +1312,11 @@ type ServerChangeBandwidthInput struct {
 }
 
 type ServerChangeConfigSpecs struct {
-	CpuSockets   int    `json:"cpu_sockets"`
-	VcpuCount    int    `json:"vcpu_count"`
-	VmemSize     int    `json:"vmem_size"`
-	InstanceType string `json:"instance_type"`
+	CpuSockets    int    `json:"cpu_sockets"`
+	VcpuCount     int    `json:"vcpu_count"`
+	ExtraCpuCount int    `json:"extra_cpu_count"`
+	VmemSize      int    `json:"vmem_size"`
+	InstanceType  string `json:"instance_type"`
 }
 
 type DiskResizeSpec struct {
@@ -1353,6 +1357,18 @@ func (conf ServerChangeConfigSettings) AddedCpu() int {
 	return addCpu
 }
 
+func (conf ServerChangeConfigSettings) ExtraCpuChanged() bool {
+	return conf.ExtraCpuCount != conf.Old.ExtraCpuCount
+}
+
+func (conf ServerChangeConfigSettings) AddedExtraCpu() int {
+	addCpu := conf.ExtraCpuCount - conf.Old.ExtraCpuCount
+	if addCpu < 0 {
+		addCpu = 0
+	}
+	return addCpu
+}
+
 func (conf ServerChangeConfigSettings) MemChanged() bool {
 	return conf.VmemSize != conf.Old.VmemSize
 }
diff --git a/pkg/apis/scheduler/api.go b/pkg/apis/scheduler/api.go
index 715fc36b0d..fc10b49290 100644
--- a/pkg/apis/scheduler/api.go
+++ b/pkg/apis/scheduler/api.go
@@ -87,7 +87,7 @@ type ScheduleInput struct {
 	OsArch          string `json:"os_arch"`
 	ResetCpuNumaPin bool   `json:"reset_cpu_numa_pin"`
 
-	// For Migrate
+	ExtraCpuCount   int           `json:"extra_cpu_count"`
 	CpuNumaPin      []SCpuNumaPin `json:"cpu_numa_pin"`
 	PreferNumaNodes []int         `json:"prefer_numa_nodes"`
 
@@ -199,9 +199,10 @@ type SCpuPin struct {
 }
 
 type SCpuNumaPin struct {
-	CpuPin    []int
-	NodeId    int
-	MemSizeMB *int
+	CpuPin        []int
+	NodeId        int
+	MemSizeMB     *int
+	ExtraCpuCount int
 }
 
 type CandidateResource struct {
diff --git a/pkg/compute/guestdrivers/base.go b/pkg/compute/guestdrivers/base.go
index 7068c89450..8481f25c3d 100644
--- a/pkg/compute/guestdrivers/base.go
+++ b/pkg/compute/guestdrivers/base.go
@@ -564,6 +564,7 @@ func (base *SBaseGuestDriver) ValidateGuestChangeConfigInput(ctx context.Context
 	confs.Old.VcpuCount = guest.VcpuCount
 	confs.Old.CpuSockets = guest.CpuSockets
 	confs.Old.VmemSize = guest.VmemSize
+	confs.Old.ExtraCpuCount = guest.ExtraCpuCount
 
 	region, err := guest.GetRegion()
 	if err != nil {
@@ -586,6 +587,10 @@ func (base *SBaseGuestDriver) ValidateGuestChangeConfigInput(ctx context.Context
 		} else {
 			confs.VcpuCount = guest.VcpuCount
 		}
+		if input.ExtraCpuCount != nil {
+			confs.ExtraCpuCount = *input.ExtraCpuCount
+		}
+
 		if len(input.VmemSize) > 0 {
 			if !regutils.MatchSize(input.VmemSize) {
 				return nil, httperrors.NewBadRequestError("Memory size %q must be number[+unit], like 256M, 1G or 256", input.VmemSize)
@@ -678,7 +683,7 @@ func (base *SBaseGuestDriver) ValidateGuestChangeConfigInput(ctx context.Context
 	}
 
 	// schedulr forecast
-	schedDesc := guest.ChangeConfToSchedDesc(confs.AddedCpu(), confs.AddedMem(), schedInputDisks)
+	schedDesc := guest.ChangeConfToSchedDesc(confs.AddedCpu(), confs.AddedExtraCpu(), confs.AddedMem(), schedInputDisks)
 	s := auth.GetAdminSession(ctx, options.Options.Region)
 	canChangeConf, res, err := scheduler.SchedManager.DoScheduleForecast(s, schedDesc, 1)
 	if err != nil {
diff --git a/pkg/compute/guestdrivers/kvm.go b/pkg/compute/guestdrivers/kvm.go
index 77b3642a50..0e57c3af5f 100644
--- a/pkg/compute/guestdrivers/kvm.go
+++ b/pkg/compute/guestdrivers/kvm.go
@@ -1252,6 +1252,10 @@ func (kvm *SKVMGuestDriver) ValidateGuestChangeConfigInput(ctx context.Context,
 		return nil, errors.Wrap(err, "SBaseGuestDriver.ValidateGuestChangeConfigInput")
 	}
 
+	if confs.ExtraCpuChanged() && guest.Status != api.VM_READY {
+		return nil, httperrors.NewInvalidStatusError("Can't change extra cpus on vm status %s", guest.Status)
+	}
+
 	for i := range input.ResetTrafficLimits {
 		input.ResetTrafficLimits[i].Mac = netutils.FormatMacAddr(input.ResetTrafficLimits[i].Mac)
 		_, err := guest.GetGuestnetworkByMac(input.ResetTrafficLimits[i].Mac)
diff --git a/pkg/compute/models/guest_actions.go b/pkg/compute/models/guest_actions.go
index dd3668c113..e0d4b9d848 100644
--- a/pkg/compute/models/guest_actions.go
+++ b/pkg/compute/models/guest_actions.go
@@ -3218,6 +3218,9 @@ func (self *SGuest) PerformChangeConfig(ctx context.Context, userCred mcclient.T
 	if added := confs.AddedCpu(); added > 0 {
 		pendingUsage.Cpu = added
 	}
+	if added := confs.AddedExtraCpu(); added > 0 {
+		pendingUsage.Cpu += added
+	}
 	if added := confs.AddedMem(); added > 0 {
 		pendingUsage.Memory = added
 	}
@@ -3241,7 +3244,7 @@ func (self *SGuest) PerformChangeConfig(ctx context.Context, userCred mcclient.T
 	return nil, nil
 }
 
-func (self *SGuest) ChangeConfToSchedDesc(addCpu, addMem int, schedInputDisks []*api.DiskConfig) *schedapi.ScheduleInput {
+func (self *SGuest) ChangeConfToSchedDesc(addCpu, addExtraCpu, addMem int, schedInputDisks []*api.DiskConfig) *schedapi.ScheduleInput {
 	region, _ := self.GetRegion()
 	devs, _ := self.GetIsolatedDevices()
 	desc := &schedapi.ScheduleInput{
@@ -3260,6 +3263,7 @@ func (self *SGuest) ChangeConfToSchedDesc(addCpu, addMem int, schedInputDisks []
 		OsArch:            self.OsArch,
 		ChangeConfig:      true,
 		HasIsolatedDevice: len(devs) > 0,
+		ExtraCpuCount:     addExtraCpu,
 	}
 	return desc
 }
diff --git a/pkg/compute/models/guests.go b/pkg/compute/models/guests.go
index f0fe37ac33..8540151330 100644
--- a/pkg/compute/models/guests.go
+++ b/pkg/compute/models/guests.go
@@ -129,6 +129,8 @@ type SGuest struct {
 	VmemSize int `nullable:"false" list:"user" create:"required"`
 	// CPU 内存绑定信息
 	CpuNumaPin jsonutils.JSONObject `nullable:"true" get:"user" update:"user" create:"optional"`
+	// 额外分配的 CPU 数量
+	ExtraCpuCount int `nullable:"false" default:"0" list:"user" create:"optional"`
 
 	// 启动顺序
 	BootOrder string `width:"8" charset:"ascii" nullable:"true" default:"cdn" list:"user" update:"user" create:"optional"`
@@ -1324,8 +1326,9 @@ func (guest *SGuest) SetCpuNumaPin(
 		vcpuId := 0
 		for i := range schedCpuNumaPin {
 			cpuNumaPin[i] = api.SCpuNumaPin{
-				SizeMB: schedCpuNumaPin[i].MemSizeMB,
-				NodeId: schedCpuNumaPin[i].NodeId,
+				SizeMB:        schedCpuNumaPin[i].MemSizeMB,
+				NodeId:        schedCpuNumaPin[i].NodeId,
+				ExtraCpuCount: schedCpuNumaPin[i].ExtraCpuCount,
 			}
 
 			if len(schedCpuNumaPin[i].CpuPin) > 0 {
diff --git a/pkg/compute/tasks/guest_change_config_task.go b/pkg/compute/tasks/guest_change_config_task.go
index 993f366b75..c9b580f26e 100644
--- a/pkg/compute/tasks/guest_change_config_task.go
+++ b/pkg/compute/tasks/guest_change_config_task.go
@@ -86,6 +86,24 @@ func (task *GuestChangeConfigTask) SaveScheduleResult(ctx context.Context, obj I
 	guest := task.GetObject().(*models.SGuest)
 	task.Params.Set("sched_session_id", jsonutils.NewString(target.SessionId))
 
+	confs, err := task.getChangeConfigSetting()
+	if err != nil {
+		task.markStageFailed(ctx, guest, jsonutils.NewString(err.Error()))
+		return
+	}
+	if confs.ExtraCpuChanged() {
+		_, err = db.Update(guest, func() error {
+			if confs.ExtraCpuCount > 0 {
+				guest.ExtraCpuCount = confs.ExtraCpuCount
+			}
+			return nil
+		})
+		if err != nil {
+			task.markStageFailed(ctx, guest, jsonutils.NewString(err.Error()))
+			return
+		}
+	}
+
 	if len(target.CpuNumaPin) > 0 {
 		task.Params.Set("cpu_numa_pin", jsonutils.Marshal(target.CpuNumaPin))
 	}
diff --git a/pkg/hostman/guestman/desc/desc.go b/pkg/hostman/guestman/desc/desc.go
index ee550509dd..35a7533a01 100644
--- a/pkg/hostman/guestman/desc/desc.go
+++ b/pkg/hostman/guestman/desc/desc.go
@@ -60,7 +60,8 @@ type SCpuNumaPin struct {
 	Unregular bool
 	NodeId    *uint16 `json:",omitempty"`
 
-	VcpuPin []SVCpuPin `json:",omitempty"`
+	VcpuPin       []SVCpuPin `json:",omitempty"`
+	ExtraCpuCount int        `json:"extra_cpu_count"`
 }
 
 type SVCpuPin struct {
diff --git a/pkg/hostman/guestman/pod.go b/pkg/hostman/guestman/pod.go
index 983374dc08..b8384792f5 100644
--- a/pkg/hostman/guestman/pod.go
+++ b/pkg/hostman/guestman/pod.go
@@ -1459,6 +1459,7 @@ func (s *sPodGuestInstance) createContainer(ctx context.Context, userCred mcclie
 
 	var cpuSetCpus string
 	var cpuSetMems string
+	var extraCpuCount int
 	{
 		cpuSets := sets.NewString()
 		cpuMemSets := sets.NewString()
@@ -1470,8 +1471,10 @@ func (s *sPodGuestInstance) createContainer(ctx context.Context, userCred mcclie
 				if cpuNumaPin.NodeId != nil && cpuNumaPin.SizeMB > 0 {
 					cpuMemSets.Insert(fmt.Sprintf("%d", int(*cpuNumaPin.NodeId)))
 				}
+				if cpuNumaPin.ExtraCpuCount > 0 {
+					extraCpuCount += cpuNumaPin.ExtraCpuCount
+				}
 			}
-
 			cpuSetCpus = strings.Join(cpuSets.List(), ",")
 			cpuSetMems = strings.Join(cpuMemSets.List(), ",")
 		} else if len(s.Desc.VcpuPin) > 0 {
@@ -1507,7 +1510,7 @@ func (s *sPodGuestInstance) createContainer(ctx context.Context, userCred mcclie
 			Resources: &runtimeapi.LinuxContainerResources{
 				// REF: https://docs.docker.com/config/containers/resource_constraints/#configure-the-default-cfs-scheduler
 				CpuPeriod: s.getDefaultCPUPeriod(),
-				CpuQuota:  s.GetDesc().Cpu * s.getDefaultCPUPeriod(),
+				CpuQuota:  (s.GetDesc().Cpu + int64(extraCpuCount)) * s.getDefaultCPUPeriod(),
 				//CpuShares:              defaultCPUPeriod,
 				MemoryLimitInBytes:     s.GetDesc().Mem * 1024 * 1024,
 				OomScoreAdj:            0,
diff --git a/pkg/mcclient/options/compute/servers.go b/pkg/mcclient/options/compute/servers.go
index b3118412f4..b6c583afa5 100644
--- a/pkg/mcclient/options/compute/servers.go
+++ b/pkg/mcclient/options/compute/servers.go
@@ -454,6 +454,7 @@ type ServerCreateOptionalOptions struct {
 	Iso              string   `help:"ISO image ID" metavar:"IMAGE_ID" json:"cdrom"`
 	IsoBootIndex     *int8    `help:"Iso bootindex" metavar:"IMAGE_BOOT_INDEX" json:"cdrom_boot_index"`
 	VcpuCount        int      `help:"#CPU cores of VM server, default 1" default:"1" metavar:"<SERVER_CPU_COUNT>" json:"vcpu_count" token:"ncpu"`
+	ExtraCpuCount    int      `help:"Extra allocate cpu count" json:"extra_cpu_count"`
 	InstanceType     string   `help:"instance flavor"`
 	Vga              string   `help:"VGA driver" choices:"std|vmware|cirrus|qxl|virtio"`
 	Vdi              string   `help:"VDI protocool" choices:"vnc|spice"`
@@ -1039,10 +1040,11 @@ func (o *ServerRebuildRootOptions) Description() string {
 
 type ServerChangeConfigOptions struct {
 	ServerIdOptions
-	VcpuCount  *int     `help:"New number of Virtual CPU cores" json:"vcpu_count" token:"ncpu"`
-	CpuSockets *int     `help:"Cpu sockets"`
-	VmemSize   string   `help:"New memory size" json:"vmem_size" token:"vmem"`
-	Disk       []string `help:"Data disk description, from the 1st data disk to the last one, empty string if no change for this data disk"`
+	VcpuCount     *int     `help:"New number of Virtual CPU cores" json:"vcpu_count" token:"ncpu"`
+	ExtraCpuCount *int     `help:"Extra allocate cpu count" json:"extra_cpu_count"`
+	CpuSockets    *int     `help:"Cpu sockets"`
+	VmemSize      string   `help:"New memory size" json:"vmem_size" token:"vmem"`
+	Disk          []string `help:"Data disk description, from the 1st data disk to the last one, empty string if no change for this data disk"`
 
 	InstanceType string `help:"Instance Type, e.g. S2.SMALL2 for qcloud"`
 
diff --git a/pkg/scheduler/algorithm/predicates/guest/cpu_predicate.go b/pkg/scheduler/algorithm/predicates/guest/cpu_predicate.go
index c0d210172f..da7f780a11 100644
--- a/pkg/scheduler/algorithm/predicates/guest/cpu_predicate.go
+++ b/pkg/scheduler/algorithm/predicates/guest/cpu_predicate.go
@@ -81,7 +81,7 @@ func (f *CPUPredicate) Execute(ctx context.Context, u *core.Unit, c core.Candida
 	}
 
 	freeCPUCount := getter.FreeCPUCount(useRsvd)
-	reqCPUCount := int64(d.Ncpu)
+	reqCPUCount := int64(d.Ncpu + d.ExtraCpuCount)
 	if freeCPUCount < reqCPUCount {
 		totalCPUCount := getter.TotalCPUCount(useRsvd)
 		h.AppendInsufficientResourceError(reqCPUCount, totalCPUCount, freeCPUCount)
diff --git a/pkg/scheduler/cache/candidate/hosts.go b/pkg/scheduler/cache/candidate/hosts.go
index 688e428d92..c9a8d81839 100644
--- a/pkg/scheduler/cache/candidate/hosts.go
+++ b/pkg/scheduler/cache/candidate/hosts.go
@@ -1526,6 +1526,11 @@ func (b *HostBuilder) fillGuestsResourceInfo(desc *HostDesc, host *computemodels
 				if err := guest.CpuNumaPin.Unmarshal(&cpuNumaPin); err != nil {
 					return errors.Wrap(err, "unmarshal cpu numa pin")
 				}
+				for i := range cpuNumaPin {
+					if cpuNumaPin[i].ExtraCpuCount > 0 {
+						creatingCPUCount += int64(cpuNumaPin[i].ExtraCpuCount)
+					}
+				}
 				guestsCpuNumaPin = append(guestsCpuNumaPin, cpuNumaPin...)
 			}
 		} else if !IsGuestStoppedStatus(guest) {
@@ -1538,6 +1543,11 @@ func (b *HostBuilder) fillGuestsResourceInfo(desc *HostDesc, host *computemodels
 				if err := guest.CpuNumaPin.Unmarshal(&cpuNumaPin); err != nil {
 					return errors.Wrap(err, "unmarshal cpu numa pin")
 				}
+				for i := range cpuNumaPin {
+					if cpuNumaPin[i].ExtraCpuCount > 0 {
+						creatingCPUCount += int64(cpuNumaPin[i].ExtraCpuCount)
+					}
+				}
 				guestsCpuNumaPin = append(guestsCpuNumaPin, cpuNumaPin...)
 			}
 		}
diff --git a/pkg/scheduler/core/result.go b/pkg/scheduler/core/result.go
index bda5bfa011..bcd14fd31f 100644
--- a/pkg/scheduler/core/result.go
+++ b/pkg/scheduler/core/result.go
@@ -86,14 +86,34 @@ func (item *SchedResultItem) ToCandidateResource(storageUsed *StorageUsed) *sche
 }
 
 func (item *SchedResultItem) selectCpuNumaPin() []schedapi.SCpuNumaPin {
-	// if !item.Candidater.Getter().Host().EnableNumaAllocate {
-	// 	return nil
-	// }
-
-	if item.SchedData.LiveMigrate && len(item.SchedData.CpuNumaPin) > 0 {
-		return item.Candidater.AllocCpuNumaPinWithNodeCount(item.SchedData.Ncpu, item.SchedData.Memory, len(item.SchedData.CpuNumaPin))
+	vcpuCount := item.SchedData.Ncpu
+	if item.SchedData.ExtraCpuCount > 0 {
+		vcpuCount += item.SchedData.ExtraCpuCount
 	}
-	return item.Candidater.AllocCpuNumaPin(item.SchedData.Ncpu, item.SchedData.Memory*1024, item.SchedData.PreferNumaNodes)
+
+	var res []schedapi.SCpuNumaPin
+	if item.SchedData.LiveMigrate && len(item.SchedData.CpuNumaPin) > 0 {
+		res = item.Candidater.AllocCpuNumaPinWithNodeCount(vcpuCount, item.SchedData.Memory*1024, len(item.SchedData.CpuNumaPin))
+	} else {
+		res = item.Candidater.AllocCpuNumaPin(vcpuCount, item.SchedData.Memory*1024, item.SchedData.PreferNumaNodes)
+	}
+
+	if item.SchedData.ExtraCpuCount > 0 {
+		extraCpuCnt := item.SchedData.ExtraCpuCount
+		for extraCpuCnt > 0 {
+			cpuMaxIdx := 0
+			cpuMax := -1
+			for i := range res {
+				if len(res[i].CpuPin)-res[i].ExtraCpuCount > cpuMax {
+					cpuMax = len(res[i].CpuPin) - res[i].ExtraCpuCount
+					cpuMaxIdx = i
+				}
+			}
+			res[cpuMaxIdx].ExtraCpuCount += 1
+			extraCpuCnt -= 1
+		}
+	}
+	return res
 }
 
 func (item *SchedResultItem) getDisks(used *StorageUsed) []*schedapi.CandidateDisk {