go_dreamfactory/lego/sys/rpcx/selector.go

package rpcx

import (
	"context"
	"go_dreamfactory/lego/sys/log"
	"regexp"
	"strings"

	"github.com/smallnest/rpcx/share"
	"github.com/valyala/fastrand"
)

var rex_nogather = regexp.MustCompile(`\!\[([^)]+)\]`)
var rex_noid = regexp.MustCompile(`\!([^)]+)`)
var rex_gather = regexp.MustCompile(`\[([^)]+)\]`)

func newSelector(fn func(map[string]*ServiceNode)) *Selector {
	return &Selector{
		updateServerEvent: fn,
		servers:           make(map[string]*ServiceNode),
		serversType:       make(map[string][]*ServiceNode),
		i:                 make(map[string]int),
	}
}

type ServiceNode struct {
	ServiceTag  string `json:"stag"`    //服务集群标签
	ServiceId   string `json:"sid"`     //服务id
	ServiceType string `json:"stype"`   //服务类型
	Version     string `json:"version"` //服务版本
	ServiceAddr string `json:"addr"`    //服务地址
}

type Selector struct {
	updateServerEvent func(map[string]*ServiceNode)
	servers           map[string]*ServiceNode
	serversType       map[string][]*ServiceNode
	i                 map[string]int
}

///servicePath = (worker)/(worker/worker_1)/(worker/!worker_1)/(worker/[worker_1,worker_2])/(worker/![worker_1,worker_2])
func (this *Selector) Select(ctx context.Context, servicePath, serviceMethod string, args interface{}) string {
	routrules := ctx.Value(share.ReqMetaDataKey).(map[string]string)[CallRoutRulesKey]
	service := strings.Split(routrules, "/")
	leng := len(service)
	if leng == 1 {
		if nodes, ok := this.serversType[service[0]]; ok {
			i, ok := this.i[service[0]]
			if !ok {
				i = 0
			}
			i = i % len(nodes)
			this.i[service[0]] = i + 1
			return nodes[i].ServiceAddr
		}
	} else if leng == 2 {
		result := this.ParseRoutRules(service[1])
		if len(result) == 0 {
			log.Errorf("Select no found any node servicePath:%v serviceMethod:%v RoutRules:%v \n", servicePath, serviceMethod, routrules)
			return ""
		}
		i := fastrand.Uint32n(uint32(len(result)))
		if node, ok := this.servers[result[i]]; ok {
			return node.ServiceAddr
		}
	}
	log.Errorf("Select no found any node servicePath:%v serviceMethod:%v RoutRules:%v \n", servicePath, serviceMethod, routrules)
	return ""
}

//更新服务列表
func (this *Selector) UpdateServer(servers map[string]string) {
	ss := make(map[string]*ServiceNode)
	sst := make(map[string][]*ServiceNode)
	for _, v := range servers {
		if node, err := smetaToServiceNode(v); err != nil {
			log.Errorf("smetaToServiceNode:%s err:%v", v, err)
			continue
		} else {
			ss[node.ServiceId] = node
			if ssts, ok := sst[node.ServiceType]; !ok {
				sst[node.ServiceType] = make([]*ServiceNode, 0)
				sst[node.ServiceType] = append(sst[node.ServiceType], node)
			} else {
				ssts = append(ssts, node)
			}
		}

	}
	this.servers = ss
	this.serversType = sst
	if this.updateServerEvent != nil {
		go this.updateServerEvent(ss)
	}
}

//路由规则解析
func (this *Selector) ParseRoutRules(rules string) (result []string) {
	result = make([]string, 0)

	//解析 ![sid,sid] 格式规则
	if out := rex_nogather.FindAllStringSubmatch(rules, -1); len(out) == 1 && len(out[0]) == 2 {
		if nogather := strings.Split(out[0][1], ","); len(nogather) > 0 {
			for k, _ := range this.servers {
				iskeep := false
				for _, v := range nogather {
					if k == v {
						iskeep = true
						break
					}
				}
				if !iskeep {
					result = append(result, k)
				}
			}
			return
		}
	}
	//解析 !sid 格式规则
	if out := rex_noid.FindAllStringSubmatch(rules, -1); len(out) == 1 && len(out[0]) == 2 {
		for k, _ := range this.servers {
			iskeep := false
			if k == out[0][1] {
				iskeep = true
				break
			}
			if !iskeep {
				result = append(result, k)
			}
		}
		return
	}
	//解析 [sid,sid] 格式规则
	if out := rex_gather.FindAllStringSubmatch(rules, -1); len(out) == 1 && len(out[0]) == 2 {
		if nogather := strings.Split(out[0][1], ","); len(nogather) > 0 {
			for k, _ := range this.servers {
				iskeep := false
				for _, v := range nogather {
					if k == v {
						iskeep = true
						break
					}
				}
				if iskeep {
					result = append(result, k)
				}
			}
			return
		}
	}
	if _, ok := this.servers[rules]; ok {
		result = append(result, rules)
	}
	return
}