waf-platform/EdgeAdmin/internal/web/actions/default/ns/nsnodeutils/utils.go

// Copyright 2022 Liuxiangchao iwind.liu@gmail.com. All rights reserved. Official site: https://goedge.cn .
//go:build plus

package nsnodeutils

import (
	"context"
	"encoding/json"
	"github.com/TeaOSLab/EdgeAdmin/internal/configs"
	"github.com/TeaOSLab/EdgeAdmin/internal/rpc"
	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
	"sort"
	"strconv"
	"sync"
)

// MessageResult 和节点消息通讯结果定义
type MessageResult struct {
	NodeId   int64  `json:"nodeId"`
	NodeName string `json:"nodeName"`
	IsOK     bool   `json:"isOk"`
	Message  string `json:"message"`
}

// SendMessageToNSCluster 向NS集群发送命令消息
func SendMessageToNSCluster(ctx context.Context, clusterId int64, code string, msg interface{}, timeoutSeconds int32) (results []*MessageResult, err error) {
	results = []*MessageResult{}

	msgJSON, err := json.Marshal(msg)
	if err != nil {
		return results, err
	}

	defaultRPCClient, err := rpc.SharedRPC()
	if err != nil {
		return results, err
	}

	// 获取所有节点
	nodesResp, err := defaultRPCClient.NSNodeRPC().FindAllNSNodesWithNSClusterId(ctx, &pb.FindAllNSNodesWithNSClusterIdRequest{
		NsClusterId: clusterId,
	})
	if err != nil {
		return results, err
	}
	var nodes = nodesResp.NsNodes
	if len(nodes) == 0 {
		return results, nil
	}

	var rpcMap = map[int64]*rpc.RPCClient{} // apiNodeId => RPCClient
	var locker = &sync.Mutex{}

	var wg = &sync.WaitGroup{}
	wg.Add(len(nodes))

	for _, node := range nodes {
		// TODO 检查是否在线

		if len(node.ConnectedAPINodeIds) == 0 {
			locker.Lock()
			results = append(results, &MessageResult{
				NodeId:   node.Id,
				NodeName: node.Name,
				IsOK:     false,
				Message:  "节点尚未连接到API",
			})
			locker.Unlock()
			wg.Done()
			continue
		}

		// 获取API节点信息
		apiNodeId := node.ConnectedAPINodeIds[0]
		rpcClient, ok := rpcMap[apiNodeId]
		if !ok {
			apiNodeResp, err := defaultRPCClient.APINodeRPC().FindEnabledAPINode(ctx, &pb.FindEnabledAPINodeRequest{ApiNodeId: apiNodeId})
			if err != nil {
				locker.Lock()
				results = append(results, &MessageResult{
					NodeId:   node.Id,
					NodeName: node.Name,
					IsOK:     false,
					Message:  "无法读取对应的API节点信息：" + err.Error(),
				})
				locker.Unlock()
				wg.Done()
				continue
			}

			if apiNodeResp.ApiNode == nil {
				locker.Lock()
				results = append(results, &MessageResult{
					NodeId:   node.Id,
					NodeName: node.Name,
					IsOK:     false,
					Message:  "无法读取对应的API节点信息：API节点ID：" + strconv.FormatInt(apiNodeId, 10),
				})
				locker.Unlock()
				wg.Done()
				continue
			}
			apiNode := apiNodeResp.ApiNode

			apiRPCClient, err := rpc.NewRPCClient(&configs.APIConfig{
				RPCEndpoints: apiNode.AccessAddrs,
				NodeId:       apiNode.UniqueId,
				Secret:       apiNode.Secret,
			}, false)
			if err != nil {
				locker.Lock()
				results = append(results, &MessageResult{
					NodeId:   node.Id,
					NodeName: node.Name,
					IsOK:     false,
					Message:  "初始化API节点错误：API节点ID：" + strconv.FormatInt(apiNodeId, 10) + "：" + err.Error(),
				})
				locker.Unlock()
				wg.Done()
				continue
			}
			rpcMap[apiNodeId] = apiRPCClient
			rpcClient = apiRPCClient
		}

		// 发送消息
		go func(node *pb.NSNode) {
			defer wg.Done()

			result, err := rpcClient.NSNodeRPC().SendCommandToNSNode(ctx, &pb.NSNodeStreamMessage{
				NsNodeId:       node.Id,
				TimeoutSeconds: timeoutSeconds,
				Code:           code,
				DataJSON:       msgJSON,
			})
			if err != nil {
				locker.Lock()
				results = append(results, &MessageResult{
					NodeId:   node.Id,
					NodeName: node.Name,
					IsOK:     false,
					Message:  "API返回错误：" + err.Error(),
				})
				locker.Unlock()
				return
			}

			locker.Lock()
			results = append(results, &MessageResult{
				NodeId:   node.Id,
				NodeName: node.Name,
				IsOK:     result.IsOk,
				Message:  result.Message,
			})
			locker.Unlock()
		}(node)
	}
	wg.Wait()

	// 对结果进行排序
	if len(results) > 0 {
		sort.Slice(results, func(i, j int) bool {
			return results[i].NodeId < results[j].NodeId
		})
	}

	// 关闭RPC
	for _, rpcClient := range rpcMap {
		_ = rpcClient.Close()
	}

	return
}