gRPC負載均衡（自定義負載均衡策略）

摘要：最後，我們只需要在服務端註冊服務時候附帶權重，然後客戶端在服務發現時把權重 Set 到 resolver.Address 中，最後客戶端把負載論衡策略改成 weight 就完成了。= nil { //非數字字符默認權重爲1 nodeWeight = 1 } //把服務地址權重存儲到resolver.Address的元數據中 addr = weight.SetAddrInfo(addr, weight.AddrInfo{Weight: nodeWeight}) s.serverList[key] = addr s.cc.UpdateState(resolver.State{Addresses: s.getServices()}) log.Println("put key :", key, "wieght:", val) }。

前言

上篇文章介紹瞭如何實現gRPC負載均衡，但目前官方只提供了 pick_first 和 round_robin 兩種負載均衡策略，輪詢法 round_robin 不能滿足因服務器配置不同而承擔不同負載量，這篇文章將介紹如何實現自定義負載均衡策略-- 加權隨機法 。

加權隨機法 可以根據服務器的處理能力而分配不同的權重，從而實現處理能力高的服務器可承擔更多的請求，處理能力低的服務器少承擔請求。

自定義負載均衡策略

gRPC提供了 V2PickerBuilder 和 V2Picker 接口讓我們實現自己的負載均衡策略。

type V2PickerBuilder interface {
	Build(info PickerBuildInfo) balancer.V2Picker
}

V2PickerBuilder 接口：創建V2版本的子連接選擇器。

Build 方法：返回一個V2選擇器，將用於gRPC選擇子連接。

type V2Picker interface {
	Pick(info PickInfo) (PickResult, error)
}

V2Picker 接口：用於gRPC選擇子連接去發送請求。 Pick 方法：子連接選擇

問題來了，我們需要把服務器地址的權重添加進去，但是地址 resolver.Address 並沒有提供權重的屬性。官方給的答覆是：把權重存儲到地址的元數據 metadata 中。

// attributeKey is the type used as the key to store AddrInfo in the Attributes
// field of resolver.Address.
type attributeKey struct{}

// AddrInfo will be stored inside Address metadata in order to use weighted balancer.
type AddrInfo struct {
	Weight int
}

// SetAddrInfo returns a copy of addr in which the Attributes field is updated
// with addrInfo.
func SetAddrInfo(addr resolver.Address, addrInfo AddrInfo) resolver.Address {
	addr.Attributes = attributes.New()
	addr.Attributes = addr.Attributes.WithValues(attributeKey{}, addrInfo)
	return addr
}

// GetAddrInfo returns the AddrInfo stored in the Attributes fields of addr.
func GetAddrInfo(addr resolver.Address) AddrInfo {
	v := addr.Attributes.Value(attributeKey{})
	ai, _ := v.(AddrInfo)
	return ai
}

定義 AddrInfo 結構體並添加權重 Weight 屬性， Set 方法把 Weight 存儲到 resolver.Address 中， Get 方法從 resolver.Address 獲取 Weight 。

解決權重存儲問題後，接下來我們實現加權隨機法負載均衡策略。

首先實現 V2PickerBuilder 接口，返回子連接選擇器。

func (*rrPickerBuilder) Build(info base.PickerBuildInfo) balancer.V2Picker {
	grpclog.Infof("weightPicker: newPicker called with info: %v", info)
	if len(info.ReadySCs) == 0 {
		return base.NewErrPickerV2(balancer.ErrNoSubConnAvailable)
	}
	var scs []balancer.SubConn
	for subConn, addr := range info.ReadySCs {
		node := GetAddrInfo(addr.Address)
		if node.Weight <= 0 {
			node.Weight = minWeight
		} else if node.Weight > 5 {
			node.Weight = maxWeight
		}
		for i := 0; i < node.Weight; i++ {
			scs = append(scs, subConn)
		}
	}
	return &rrPicker{
		subConns: scs,
	}
}

加權隨機法 中，我使用空間換時間的方式，把權重轉成地址個數（例如 addr1 的權重是 3 ，那麼添加 3 個子連接到切片中； addr2 權重爲 1 ，則添加 1 個子連接；選擇子連接時候，按子連接切片長度生成隨機數，以隨機數作爲下標就是選中的子連接），避免重複計算權重。考慮到內存佔用，權重定義從 1 到 5 權重。

接下來實現子連接的選擇，獲取隨機數，選擇子連接

type rrPicker struct {
	subConns []balancer.SubConn
	mu sync.Mutex
}

func (p *rrPicker) Pick(balancer.PickInfo) (balancer.PickResult, error) {
	p.mu.Lock()
	index := rand.Intn(len(p.subConns))
	sc := p.subConns[index]
	p.mu.Unlock()
	return balancer.PickResult{SubConn: sc}, nil
}

關鍵代碼完成後，我們把加權隨機法負載均衡策略命名爲 weight ，並註冊到gRPC的負載均衡策略中。

// Name is the name of weight balancer.
const Name = "weight"
// NewBuilder creates a new weight balancer builder.
func newBuilder() balancer.Builder {
	return base.NewBalancerBuilderV2(Name, &rrPickerBuilder{}, base.Config{HealthCheck: false})
}

func init() {
	balancer.Register(newBuilder())
}

完整代碼 weight.go

最後，我們只需要在服務端註冊服務時候附帶權重，然後客戶端在服務發現時把權重 Set 到 resolver.Address 中，最後客戶端把負載論衡策略改成 weight 就完成了。

//SetServiceList 設置服務地址
func (s *ServiceDiscovery) SetServiceList(key, val string) {
	s.lock.Lock()
	defer s.lock.Unlock()
	//獲取服務地址
	addr := resolver.Address{Addr: strings.TrimPrefix(key, s.prefix)}
	//獲取服務地址權重
	nodeWeight, err := strconv.Atoi(val)
	if err != nil {
		//非數字字符默認權重爲1
		nodeWeight = 1
	}
	//把服務地址權重存儲到resolver.Address的元數據中
	addr = weight.SetAddrInfo(addr, weight.AddrInfo{Weight: nodeWeight})
	s.serverList[key] = addr
	s.cc.UpdateState(resolver.State{Addresses: s.getServices()})
	log.Println("put key :", key, "wieght:", val)
}

客戶端使用 weight 負載均衡策略

func main() {
	r := etcdv3.NewServiceDiscovery(EtcdEndpoints)
	resolver.Register(r)
	// 連接服務器
	conn, err := grpc.Dial(
		fmt.Sprintf("%s:///%s", r.Scheme(), SerName),
		grpc.WithBalancerName("weight"),
		grpc.WithInsecure(),
	)
	if err != nil {
		log.Fatalf("net.Connect err: %v", err)
	}
	defer conn.Close()

運行效果：

運行 服務1 ，權重爲 1