From a51b977953d244e020dcffaa29df84690d11b8da Mon Sep 17 00:00:00 2001
From: Qu Xuan <quxuan@yunionyun.com>
Date: Wed, 15 Sep 2021 11:58:12 +0800
Subject: [PATCH] fix(region): disk choice master host with os arch

---
 pkg/apis/const.go                     | 16 +++++++++++++---
 pkg/compute/models/disks.go           | 26 ++++++++++++++++++++++++++
 pkg/compute/tasks/disk_create_task.go |  7 ++++++-
 3 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/pkg/apis/const.go b/pkg/apis/const.go
index a3923fc5fc..6c3ed14ff4 100644
--- a/pkg/apis/const.go
+++ b/pkg/apis/const.go
@@ -74,10 +74,20 @@ const (
 	OS_ARCH_AARCH64 = "aarch64"
 )
 
-func IsARM(osArch string) bool {
-	return utils.IsInStringArray(osArch, []string{
+var (
+	ARCH_X86 = []string{
+		OS_ARCH_X86,
+		OS_ARCH_I386,
+		OS_ARCH_X86_32,
+		OS_ARCH_X86_64,
+	}
+	ARCH_ARM = []string{
 		OS_ARCH_ARM,
 		OS_ARCH_AARCH32,
 		OS_ARCH_AARCH64,
-	})
+	}
+)
+
+func IsARM(osArch string) bool {
+	return utils.IsInStringArray(osArch, ARCH_ARM)
 }
diff --git a/pkg/compute/models/disks.go b/pkg/compute/models/disks.go
index ade9d77788..38b17f1401 100644
--- a/pkg/compute/models/disks.go
+++ b/pkg/compute/models/disks.go
@@ -1223,6 +1223,32 @@ func (self *SDisk) GetPathAtHost(host *SHost) string {
 	return ""
 }
 
+func (self *SDisk) GetMasterHost() (*SHost, error) {
+	hosts := HostManager.Query().SubQuery()
+	hoststorages := HoststorageManager.Query().SubQuery()
+
+	q := hosts.Query().Join(hoststorages, sqlchemy.Equals(hoststorages.Field("host_id"), hosts.Field("id")))
+	q = q.Filter(sqlchemy.Equals(hoststorages.Field("storage_id"), self.StorageId))
+	q = q.IsTrue("enabled")
+	q = q.Equals("host_status", api.HOST_ONLINE).Asc("id")
+	guest := self.GetGuest()
+	if guest != nil && len(guest.OsArch) > 0 {
+		switch guest.OsArch {
+		case apis.OS_ARCH_X86:
+			q = q.In("cpu_architecture", apis.ARCH_X86)
+		case apis.OS_ARCH_ARM:
+			q = q.In("cpu_architecture", apis.ARCH_ARM)
+		}
+	}
+	host := SHost{}
+	host.SetModelManager(HostManager, &host)
+	err := q.First(&host)
+	if err != nil {
+		return nil, errors.Wrapf(err, "q.First")
+	}
+	return &host, nil
+}
+
 func (self *SDisk) GetFetchUrl() string {
 	storage, _ := self.GetStorage()
 	if storage == nil {
diff --git a/pkg/compute/tasks/disk_create_task.go b/pkg/compute/tasks/disk_create_task.go
index 1b2834cff1..31a7bfb2c0 100644
--- a/pkg/compute/tasks/disk_create_task.go
+++ b/pkg/compute/tasks/disk_create_task.go
@@ -20,6 +20,7 @@ import (
 
 	"yunion.io/x/jsonutils"
 	"yunion.io/x/log"
+	"yunion.io/x/pkg/errors"
 	"yunion.io/x/pkg/util/stringutils"
 
 	api "yunion.io/x/onecloud/pkg/apis/compute"
@@ -54,7 +55,11 @@ func (self *DiskCreateTask) OnStorageCacheImageComplete(ctx context.Context, dis
 		db.OpsLog.LogEvent(disk, db.ACT_DELOCATE, disk.GetShortDesc(ctx), self.GetUserCred())
 	}
 	storage, _ := disk.GetStorage()
-	host := storage.GetMasterHost()
+	host, err := disk.GetMasterHost()
+	if err != nil {
+		self.OnStartAllocateFailed(ctx, disk, jsonutils.NewString(errors.Wrapf(err, "GetMasterHost").Error()))
+		return
+	}
 	db.OpsLog.LogEvent(disk, db.ACT_ALLOCATING, disk.GetShortDesc(ctx), self.GetUserCred())
 	disk.SetStatus(self.GetUserCred(), api.DISK_STARTALLOC, fmt.Sprintf("Disk start alloc use host %s(%s)", host.Name, host.Id))
 	if rebuild && storage.StorageType == api.STORAGE_RBD {