From 78f906ec65690bf69a7a9ca414c573f3dcd584e0 Mon Sep 17 00:00:00 2001
From: mouuii <49775493+mouuii@users.noreply.github.com>
Date: Mon, 16 Aug 2021 09:50:18 +0800
Subject: [PATCH] pm hostid bug (#177)

Signed-off-by: mouuii <zhouchengbin@koderover.com>

Co-authored-by: mouuii <zhouchengbin@koderover.com>
---
 .../common/repository/mongodb/private_key.go  | 23 ++++++++++++++-----
 .../aslan/core/service/service/service.go     |  6 ++++-
 .../aslan/core/system/service/private_key.go  |  4 +++-
 .../service/workflow/workflow_task.go         |  2 +-
 .../core/service/scheduler/schedule_env.go    |  4 ++++
 5 files changed, 30 insertions(+), 9 deletions(-)

diff --git a/pkg/microservice/aslan/core/common/repository/mongodb/private_key.go b/pkg/microservice/aslan/core/common/repository/mongodb/private_key.go
index 238c6e5d..26b899b7 100644
--- a/pkg/microservice/aslan/core/common/repository/mongodb/private_key.go
+++ b/pkg/microservice/aslan/core/common/repository/mongodb/private_key.go
@@ -63,15 +63,26 @@ func (c *PrivateKeyColl) EnsureIndex(ctx context.Context) error {
 	return err
 }
 
-func (c *PrivateKeyColl) Find(id string) (*models.PrivateKey, error) {
+type FindPrivateKeyOption struct {
+	ID      string
+	Address string
+}
+
+func (c *PrivateKeyColl) Find(option FindPrivateKeyOption) (*models.PrivateKey, error) {
 	privateKey := new(models.PrivateKey)
-	oid, err := primitive.ObjectIDFromHex(id)
-	if err != nil {
-		return nil, err
+	query := bson.M{}
+	if option.ID != "" {
+		oid, err := primitive.ObjectIDFromHex(option.ID)
+		if err != nil {
+			return nil, err
+		}
+		query["_id"] = oid
+	}
+	if option.Address != "" {
+		query["ip"] = option.Address
 	}
 
-	query := bson.M{"_id": oid}
-	err = c.FindOne(context.TODO(), query).Decode(privateKey)
+	err := c.FindOne(context.TODO(), query).Decode(privateKey)
 	return privateKey, err
 }
 
diff --git a/pkg/microservice/aslan/core/service/service/service.go b/pkg/microservice/aslan/core/service/service/service.go
index c20d41db..44ab00ea 100644
--- a/pkg/microservice/aslan/core/service/service/service.go
+++ b/pkg/microservice/aslan/core/service/service/service.go
@@ -395,7 +395,11 @@ func UpdateServiceTemplate(args *commonservice.ServiceTmplObject) error {
 			existEnv = true
 		}
 
-		if _, err := commonrepo.NewPrivateKeyColl().Find(envStatus.HostID); err == nil {
+		op := commonrepo.FindPrivateKeyOption{
+			Address: envStatus.Address,
+			ID:      envStatus.HostID,
+		}
+		if _, err := commonrepo.NewPrivateKeyColl().Find(op); err == nil {
 			existHost = true
 		}
 
diff --git a/pkg/microservice/aslan/core/system/service/private_key.go b/pkg/microservice/aslan/core/system/service/private_key.go
index 8e32a0d8..d34ca4a0 100644
--- a/pkg/microservice/aslan/core/system/service/private_key.go
+++ b/pkg/microservice/aslan/core/system/service/private_key.go
@@ -35,7 +35,9 @@ func ListPrivateKeys(log *zap.SugaredLogger) ([]*commonmodels.PrivateKey, error)
 }
 
 func GetPrivateKey(id string, log *zap.SugaredLogger) (*commonmodels.PrivateKey, error) {
-	resp, err := commonrepo.NewPrivateKeyColl().Find(id)
+	resp, err := commonrepo.NewPrivateKeyColl().Find(commonrepo.FindPrivateKeyOption{
+		ID: id,
+	})
 	if err != nil {
 		log.Errorf("PrivateKey.Find %s error: %v", id, err)
 		return resp, e.ErrGetPrivateKey
diff --git a/pkg/microservice/aslan/core/workflow/service/workflow/workflow_task.go b/pkg/microservice/aslan/core/workflow/service/workflow/workflow_task.go
index 54514be7..551515ea 100644
--- a/pkg/microservice/aslan/core/workflow/service/workflow/workflow_task.go
+++ b/pkg/microservice/aslan/core/workflow/service/workflow/workflow_task.go
@@ -1639,7 +1639,7 @@ func BuildModuleToSubTasks(moduleName, version, target, serviceName, productName
 			privateKeys := make([]*task.SSH, 0)
 			for _, sshID := range module.SSHs {
 				//私钥信息可能被更新，而构建中存储的信息是旧的，需要根据id获取最新的私钥信息
-				latestKeyInfo, err := commonrepo.NewPrivateKeyColl().Find(sshID)
+				latestKeyInfo, err := commonrepo.NewPrivateKeyColl().Find(commonrepo.FindPrivateKeyOption{ID: sshID})
 				if err != nil || latestKeyInfo == nil {
 					log.Errorf("PrivateKey.Find failed, id:%s, err:%v", sshID, err)
 					continue
diff --git a/pkg/microservice/cron/core/service/scheduler/schedule_env.go b/pkg/microservice/cron/core/service/scheduler/schedule_env.go
index 335ff891..8116a5e9 100644
--- a/pkg/microservice/cron/core/service/scheduler/schedule_env.go
+++ b/pkg/microservice/cron/core/service/scheduler/schedule_env.go
@@ -166,6 +166,10 @@ func (c *CronClient) RunScheduledService(svc *service.Service, healthCheck *serv
 			}
 			key := fmt.Sprintf("%s-%s-%d-%s-%s", healthCheck.Protocol, tmpEnvStatus.Address, healthCheck.Port, healthCheck.Path, tmpEnvStatus.EnvName)
 			envStatusKeys.Insert(key)
+			if tmpEnvStatus.Address == "" {
+				tmpEnvStatus.Address = envStatus.Address
+			}
+
 			tmpEnvStatus.Status = envStatus.Status
 		}
 		currentEnvStatusKey := fmt.Sprintf("%s-%s-%d-%s-%s", healthCheck.Protocol, envStatus.Address, healthCheck.Port, healthCheck.Path, envStatus.EnvName)
-- 
GitLab