76 lines
1.9 KiB
Go
76 lines
1.9 KiB
Go
package rpcx
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"strings"
|
|
|
|
"github.com/smallnest/rpcx/share"
|
|
)
|
|
|
|
func newSelector() *Selector {
|
|
return &Selector{
|
|
servers: make(map[string]*ServiceNode),
|
|
serversType: make(map[string][]*ServiceNode),
|
|
i: make(map[string]int),
|
|
}
|
|
}
|
|
|
|
type ServiceNode struct {
|
|
ServiceId string `json:"sid"` //服务id
|
|
ServiceType string `json:"stype"` //服务类型
|
|
Version string `json:"version"` //服务版本
|
|
ServiceAddr string `json:"addr"` //服务地址
|
|
}
|
|
|
|
type Selector struct {
|
|
servers map[string]*ServiceNode
|
|
serversType map[string][]*ServiceNode
|
|
i map[string]int
|
|
}
|
|
|
|
///servicePath = [stype] or [stype/sid]
|
|
func (this *Selector) Select(ctx context.Context, servicePath, serviceMethod string, args interface{}) string {
|
|
fmt.Printf("Select servicePath:%v serviceMethod:%v RoutRules:%v \n", servicePath, serviceMethod, ctx.Value("RoutRules"))
|
|
routrules := ctx.Value(share.ReqMetaDataKey).(map[string]string)["RoutRules"]
|
|
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 {
|
|
if node, ok := this.servers[service[1]]; ok {
|
|
return node.ServiceAddr
|
|
}
|
|
}
|
|
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 {
|
|
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
|
|
}
|