kratos/pkg/naming/opt.go

package naming

import (
	"encoding/json"
	"fmt"
	"math/rand"
	"net/url"
	"os"
	"sort"

	"github.com/bilibili/kratos/pkg/conf/env"
	"github.com/bilibili/kratos/pkg/log"

	"github.com/dgryski/go-farm"
)

// BuildOptions build options.
type BuildOptions struct {
	Filter     func(map[string][]*Instance) map[string][]*Instance
	Subset     func([]*Instance, int) []*Instance
	SubsetSize int
	ClientZone string
	Scheduler  func(*InstancesInfo) []*Instance
}

// BuildOpt build option interface.
type BuildOpt interface {
	Apply(*BuildOptions)
}

type funcOpt struct {
	f func(*BuildOptions)
}

func (f *funcOpt) Apply(opt *BuildOptions) {
	f.f(opt)
}

// Filter filter option.
func Filter(schema string, clusters map[string]struct{}) BuildOpt {
	return &funcOpt{f: func(opt *BuildOptions) {
		opt.Filter = func(inss map[string][]*Instance) map[string][]*Instance {
			newInss := make(map[string][]*Instance)
			for zone := range inss {
				var instances []*Instance
				for _, ins := range inss[zone] {
					//如果r.clusters的长度大于0说明需要进行集群选择
					if len(clusters) > 0 {
						if _, ok := clusters[ins.Metadata[MetaCluster]]; !ok {
							continue
						}
					}
					var addr string
					for _, a := range ins.Addrs {
						u, err := url.Parse(a)
						if err == nil && u.Scheme == schema {
							addr = u.Host
						}
					}
					if addr == "" {
						fmt.Fprintf(os.Stderr, "resolver: app(%s,%s) no valid grpc address(%v) found!", ins.AppID, ins.Hostname, ins.Addrs)
						log.Warn("resolver: invalid rpc address(%s,%s,%v) found!", ins.AppID, ins.Hostname, ins.Addrs)
						continue
					}
					instances = append(instances, ins)
				}
				newInss[zone] = instances
			}
			return newInss
		}
	}}
}

func defulatSubset(inss []*Instance, size int) []*Instance {
	backends := inss
	if len(backends) <= int(size) {
		return backends
	}
	clientID := env.Hostname
	sort.Slice(backends, func(i, j int) bool {
		return backends[i].Hostname < backends[j].Hostname
	})
	count := len(backends) / size
	// hash得到ID
	id := farm.Fingerprint64([]byte(clientID))
	// 获得rand轮数
	round := int64(id / uint64(count))

	s := rand.NewSource(round)
	ra := rand.New(s)
	//  根据source洗牌
	ra.Shuffle(len(backends), func(i, j int) {
		backends[i], backends[j] = backends[j], backends[i]
	})
	start := (id % uint64(count)) * uint64(size)
	return backends[int(start) : int(start)+int(size)]
}

// Subset Subset option.
func Subset(defaultSize int) BuildOpt {
	return &funcOpt{f: func(opt *BuildOptions) {
		opt.SubsetSize = defaultSize
		opt.Subset = defulatSubset
	}}
}

// ScheduleNode ScheduleNode option.
func ScheduleNode(clientZone string) BuildOpt {
	return &funcOpt{f: func(opt *BuildOptions) {
		opt.ClientZone = clientZone
		opt.Scheduler = func(app *InstancesInfo) (instances []*Instance) {
			type Zone struct {
				inss   []*Instance
				weight int64
				name   string
				score  float64
			}
			var zones []*Zone

			if app.Scheduler != nil {
				si, err := json.Marshal(app.Scheduler)
				if err == nil {
					log.Info("schedule info: %s", string(si))
				}
				if strategy, ok := app.Scheduler.Clients[clientZone]; ok {
					var min *Zone
					for name, zone := range strategy.Zones {
						inss := app.Instances[name]
						if len(inss) == 0 {
							continue
						}
						z := &Zone{
							inss:   inss,
							weight: zone.Weight,
							name:   name,
							score:  float64(len(inss)) / float64(zone.Weight),
						}
						if min == nil || z.score < min.score {
							min = z
						}
						zones = append(zones, z)
					}
					if opt.SubsetSize != 0 && len(min.inss) > opt.SubsetSize {
						min.score = float64(opt.SubsetSize) / float64(min.weight)
					}
					for _, z := range zones {
						nums := int(min.score * float64(z.weight))
						if nums == 0 {
							nums = 1
						}
						if nums < len(z.inss) {
							if opt.Subset != nil {
								z.inss = opt.Subset(z.inss, nums)
							} else {
								z.inss = defulatSubset(z.inss, nums)
							}
						}
					}
				}
			}
			for _, zone := range zones {
				for _, ins := range zone.inss {
					instances = append(instances, ins)
				}
			}
			//如果没有拿到节点，则选择直接获取
			if len(instances) == 0 {
				instances = app.Instances[clientZone]
				if len(instances) == 0 {
					for _, value := range app.Instances {
						instances = append(instances, value...)
					}
				}
			}
			return
		}
	}}
}