hpds_control_center/mq/index.go

package mq

import (
	"encoding/json"
	"fmt"
	"git.hpds.cc/Component/logging"
	"go.uber.org/zap"
	"hpds_control_center/config"
	"hpds_control_center/internal/balance"
	"hpds_control_center/model"
	"os"
	"time"

	"git.hpds.cc/pavement/hpds_node"
)

var MqList []HpdsMqNode

type HpdsMqNode struct {
	MqType   uint
	Topic    string
	Node     config.HpdsNode
	EndPoint interface{}
	Logger   *logging.Logger
}

func must(logger *logging.Logger, err error) {
	if err != nil {
		if logger != nil {
			logger.With(zap.String("web节点", "错误信息")).Error("启动错误", zap.Error(err))
		} else {
			_, _ = fmt.Fprint(os.Stderr, err)
		}
		os.Exit(1)
	}
}

func NewMqClient(funcs []config.FuncConfig, node config.HpdsNode, logger *logging.Logger) (mqList []HpdsMqNode, err error) {
	mqList = make([]HpdsMqNode, 0)
	for _, v := range funcs {
		switch v.MqType {
		case 2:
			sf := hpds_node.NewStreamFunction(
				v.Name,
				hpds_node.WithMqAddr(fmt.Sprintf("%s:%d", node.Host, node.Port)),
				hpds_node.WithObserveDataTags(v.DataTag),
				hpds_node.WithCredential(node.Token),
			)
			err = sf.Connect()
			must(logger, err)
			nodeInfo := HpdsMqNode{
				MqType:   2,
				Topic:    v.Name,
				Node:     node,
				EndPoint: sf,
			}
			switch v.Name {
			case "task-request":
				_ = sf.SetHandler(TaskRequestHandler)
			default:

			}
			mqList = append(mqList, nodeInfo)
		default:
			ap := hpds_node.NewAccessPoint(
				v.Name,
				hpds_node.WithMqAddr(fmt.Sprintf("%s:%d", node.Host, node.Port)),
				hpds_node.WithCredential(node.Token),
			)
			err = ap.Connect()
			nodeInfo := HpdsMqNode{
				MqType:   1,
				Topic:    v.Name,
				Node:     node,
				EndPoint: ap,
			}
			must(logger, err)
			ap.SetDataTag(v.DataTag)
			mqList = append(mqList, nodeInfo)
		}

	}
	return mqList, err
}

func GetMqClient(topic string, mqType uint) *HpdsMqNode {
	for _, v := range MqList {
		if v.Topic == topic && v.MqType == mqType {
			return &v
		}
	}
	return nil
}

func GenerateAndSendData(stream hpds_node.AccessPoint, data []byte) error {
	_, err := stream.Write(data)
	if err != nil {
		return err
	}
	time.Sleep(1000 * time.Millisecond)
	return nil
}

func TaskRequestHandler(data []byte) (byte, []byte) {
	fmt.Println("接收数据", string(data))
	cmd := new(InstructionReq)
	err := json.Unmarshal(data, cmd)
	if err != nil {
		return 0x0B, []byte(err.Error())
	}
	switch cmd.Command {
	case TaskAdd:
		payload := cmd.Payload.(map[string]interface{})
		if payload["nodeId"].(float64) == 0 {
			//根据业务属性进行分配节点
			m := model.GetModelById(int64(payload["modelId"].(float64)))
			var nodeList []model.Node
			//todo 需要增加模型下发记录
			if m.IsLightWeight {
				nodeList = model.GetLightWeight(m.ModelId)
			} else {
				nodeList = model.GetAllNode(m.ModelId)
			}
			if nodeList != nil {
				if len(nodeList) > 1 {
					//这里采用加权算法，权重采用CPU占用+mem使用+任务执行状态
					list := model.GetNodeState(nodeList)
					lb := balance.LoadBalanceFactory(balance.LbWeightRoundRobin)
					for _, v := range list {
						_ = lb.Add(v)
					}
					nodeId, _ := lb.Get(0)

					payload["nodeId"] = nodeId
					cmd := &InstructionReq{
						Command: TaskExecute,
						Payload: payload,
					}

					pData, _ := json.Marshal(cmd)
					cli := GetMqClient("task-execute", 1)
					if cli != nil {
						_ = GenerateAndSendData(cli.EndPoint.(hpds_node.AccessPoint), pData)
					}
					model.UpdateTaskExecuteNode(int64(payload["taskId"].(float64)), nodeId)
				} else {
					payload["nodeId"] = nodeList[0].NodeId
					issue := new(model.IssueModel)
					h, _ := model.DB.Where("model_id=? and node_id =?", int64(payload["modelId"].(float64)), nodeList[0].NodeId).Get(issue)
					if !h {

					}
					payload["issueResult"] = issue.IssueResult
					cmd := &InstructionReq{
						Command: TaskExecute,
						Payload: payload,
					}
					pData, _ := json.Marshal(cmd)
					cli := GetMqClient("task-execute", 1)
					if cli != nil {
						_ = GenerateAndSendData(cli.EndPoint.(hpds_node.AccessPoint), pData)
					}
					model.UpdateTaskExecuteNode(int64(payload["taskId"].(float64)), nodeList[0].NodeId)
				}
			} else {

			}

		} else {
			cmd := &InstructionReq{
				Command: TaskExecute,
				Payload: payload,
			}
			pData, _ := json.Marshal(cmd)
			cli := GetMqClient("task-execute", 1)
			if cli != nil {
				_ = GenerateAndSendData(cli.EndPoint.(hpds_node.AccessPoint), pData)
			}
		}

	case ModelIssue:
		payload := cmd.Payload.(map[string]interface{})
		cmd := &InstructionReq{
			Command: ModelIssueRepeater,
			Payload: payload,
		}
		pData, _ := json.Marshal(cmd)
		cli := GetMqClient("task-execute", 1)
		if cli != nil {
			_ = GenerateAndSendData(cli.EndPoint.(hpds_node.AccessPoint), pData)
		}
	case ModelIssueResponse:
		payload := cmd.Payload.(map[string]interface{})
		//查找下发记录
		item := new(model.IssueModel)
		h, _ := model.DB.Where("model_id = ? and node_id = ?", payload["modelId"].(int64), payload["nodeId"].(int64)).Get(item)
		pData, _ := json.Marshal(payload)
		if h {
			item.Status = 1
			item.IssueResult = string(pData)
			item.UpdateAt = time.Now().Unix()
			_, _ = model.DB.ID(item.Id).AllCols().Update(item)
		} else {
			item.ModelId = payload["modelId"].(int64)
			item.NodeId = payload["nodeId"].(int64)
			item.Status = 1
			item.IssueResult = string(pData)
			item.CreateAt = time.Now().Unix()
			item.UpdateAt = time.Now().Unix()
			_, _ = model.DB.Insert(item)
		}
	case TaskResponse:
		payload := cmd.Payload.(map[string]interface{})
		item := new(model.TaskResult)
		item.TaskId = int64(payload["taskId"].(float64))
		item.TaskCode = payload["taskCode"].(string)
		item.NodeId = int64(payload["nodeId"].(float64))
		item.ModelId = int64(payload["modelId"].(float64))
		item.StartTime = int64(payload["startTime"].(float64))
		item.FinishTime = int64(payload["finishTime"].(float64))
		item.SubDataset = payload["subDataset"].(string)
		item.DatasetId = int64(payload["datasetArr"].(float64))
		item.SrcPath = payload["srcPath"].(string)
		item.Result = payload["body"].(string)
		_, _ = model.DB.Insert(item)
		//fn := payload["fileName"].(string)
		//dec := base64.NewDecoder(base64.StdEncoding, strings.NewReader(payload["file"].(string)))

	default:

	}
	return byte(cmd.Command), nil
}