init

naming/balancer/key/key.go

@@ -0,0 +1,88 @@
+package key
+
+import (
+	"github.com/rs/zerolog/log"
+	"google.golang.org/grpc/balancer"
+	"google.golang.org/grpc/balancer/base"
+	"google.golang.org/grpc/grpclog"
+	"sync"
+)
+
+const (
+	Name       = "picker_key"
+	DefaultKey = "X-rand-string-balabala"
+)
+
+func init() {
+	balancer.Register(newBuilder(DefaultKey))
+}
+
+func newBuilder(key string) balancer.Builder {
+	return base.NewBalancerBuilder(
+		Name,
+		&keyPickerBuilder{key: key},
+		base.Config{HealthCheck: false},
+	)
+}
+
+type keyPickerBuilder struct {
+	key string
+}
+
+func (b *keyPickerBuilder) Build(info base.PickerBuildInfo) balancer.Picker {
+	grpclog.Infof("keyPickerBuilder: newPicker called with info: %v", info)
+	if len(info.ReadySCs) == 0 {
+		return base.NewErrPicker(balancer.ErrNoSubConnAvailable)
+	}
+
+	picker := &keyPicker{
+		addr2sc: make(map[string]balancer.SubConn),
+		key:     b.key,
+	}
+
+	for sc, conInfo := range info.ReadySCs {
+		addr := conInfo.Address.Addr
+		picker.setAddr2sc(addr, sc)
+	}
+
+	return picker
+}
+
+type keyPicker struct {
+	addr2sc map[string]balancer.SubConn
+	sync.RWMutex
+	key string
+}
+
+func (p *keyPicker) setAddr2sc(addr string, sc balancer.SubConn) {
+	p.Lock()
+	defer p.Unlock()
+	p.addr2sc[addr] = sc
+	log.Debug().Str("addr", addr).Msg("set addr 2 sc")
+}
+
+func (p *keyPicker) getAddr2sc(addr string) (balancer.SubConn, error) {
+	p.RLock()
+	defer p.RUnlock()
+	sc, ok := p.addr2sc[addr]
+	if ok {
+		return sc, nil
+	}
+	log.Warn().Str("addr", addr).Msg("server warn")
+	// TODO 如果找不到就随便给一个
+	for _, tsc := range p.addr2sc {
+
+		return tsc, nil
+	}
+	return nil, nil
+}
+
+func (p *keyPicker) Pick(info balancer.PickInfo) (balancer.PickResult, error) {
+	var res balancer.PickResult
+
+	addr, _ := info.Ctx.Value(p.key).(string)
+
+	sc, _ := p.getAddr2sc(addr)
+	res.SubConn = sc
+	return res, nil
+}