balancer: start populating weight by edsbalancer for weighted_round_robin

balancer/base/balancer.go

@@ -20,6 +20,7 @@ package base
 
 import (
 	"context"
+	"google.golang.org/grpc/balancer/wrr"
 
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/connectivity"
@@ -38,7 +39,7 @@ func (bb *baseBuilder) Build(cc balancer.ClientConn, opt balancer.BuildOptions)
 		cc:            cc,
 		pickerBuilder: bb.pickerBuilder,
 
-		subConns: make(map[resolver.Address]balancer.SubConn),
+		subConns: make(map[resolver.Address]AddrInfo),
 		scStates: make(map[balancer.SubConn]connectivity.State),
 		csEvltr:  &balancer.ConnectivityStateEvaluator{},
 		// Initialize picker to a picker that always return
@@ -60,7 +61,7 @@ type baseBalancer struct {
 	csEvltr *balancer.ConnectivityStateEvaluator
 	state   connectivity.State
 
-	subConns map[resolver.Address]balancer.SubConn
+	subConns map[resolver.Address]AddrInfo
 	scStates map[balancer.SubConn]connectivity.State
 	picker   balancer.Picker
 	config   Config
@@ -85,15 +86,19 @@ func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) {
 				grpclog.Warningf("base.baseBalancer: failed to create new SubConn: %v", err)
 				continue
 			}
-			b.subConns[a] = sc
+			addrInfo := AddrInfo{SubConn: sc}
+			if wrrInfo, ok := a.Metadata.(*wrr.Info); ok {
+				addrInfo.Weight = wrrInfo.Weight
+			}
+			b.subConns[a] = addrInfo
 			b.scStates[sc] = connectivity.Idle
 			sc.Connect()
 		}
 	}
 	for a, sc := range b.subConns {
 		// a was removed by resolver.
 		if _, ok := addrsSet[a]; !ok {
-			b.cc.RemoveSubConn(sc)
+			b.cc.RemoveSubConn(sc.SubConn)
 			delete(b.subConns, a)
 			// Keep the state of this sc in b.scStates until sc's state becomes Shutdown.
 			// The entry will be deleted in HandleSubConnStateChange.
@@ -110,15 +115,28 @@ func (b *baseBalancer) regeneratePicker() {
 		b.picker = NewErrPicker(balancer.ErrTransientFailure)
 		return
 	}
-	readySCs := make(map[resolver.Address]balancer.SubConn)
-
-	// Filter out all ready SCs from full subConn map.
-	for addr, sc := range b.subConns {
-		if st, ok := b.scStates[sc]; ok && st == connectivity.Ready {
-			readySCs[addr] = sc
+	pickerBuilderV2, ok := b.pickerBuilder.(PickerBuilderV2)
+	if ok {
+		readySCs := make(map[resolver.Address]AddrInfo)
+
+		// Filter out all ready SCs from full subConn map.
+		for addr, sc := range b.subConns {
+			if st, ok := b.scStates[sc.SubConn]; ok && st == connectivity.Ready {
+				readySCs[addr] = sc
+			}
+		}
+		b.picker = pickerBuilderV2.BuildV2(readySCs)
+	} else {
+		readySCs := make(map[resolver.Address]balancer.SubConn)
+
+		// Filter out all ready SCs from full subConn map.
+		for addr, sc := range b.subConns {
+			if st, ok := b.scStates[sc.SubConn]; ok && st == connectivity.Ready {
+				readySCs[addr] = sc.SubConn
+			}
 		}
+		b.picker = b.pickerBuilder.Build(readySCs)
 	}
-	b.picker = b.pickerBuilder.Build(readySCs)
 }
 
 func (b *baseBalancer) HandleSubConnStateChange(sc balancer.SubConn, s connectivity.State) {

balancer/base/base.go

@@ -42,6 +42,16 @@ type PickerBuilder interface {
 	Build(readySCs map[resolver.Address]balancer.SubConn) balancer.Picker
 }
 
+type AddrInfo struct {
+	SubConn balancer.SubConn
+	Weight uint32
+}
+
+type PickerBuilderV2 interface {
+	PickerBuilder
+	BuildV2(readySCs map[resolver.Address]AddrInfo) balancer.Picker
+}
+
 // NewBalancerBuilder returns a balancer builder. The balancers
 // built by this builder will use the picker builder to build pickers.
 func NewBalancerBuilder(name string, pb PickerBuilder) balancer.Builder {

balancer/wrr/wrr.go

@@ -0,0 +1,9 @@
+package wrr
+
+// Name is the name of weighted_round_robin balancer.
+const Name = "weighted_round_robin"
+
+// Information that should be stored inside Address metadata in order to use wrr.
+type Info struct{
+	Weight uint32
+}

balancer/xds/edsbalancer/edsbalancer.go

@@ -20,6 +20,7 @@ package edsbalancer
 import (
 	"context"
 	"encoding/json"
+	"google.golang.org/grpc/balancer/wrr"
 	"net"
 	"reflect"
 	"strconv"
@@ -227,9 +228,13 @@ func (xdsB *EDSBalancer) HandleEDSResponse(edsResp *edspb.ClusterLoadAssignment)
 		var newAddrs []resolver.Address
 		for _, lbEndpoint := range locality.GetLbEndpoints() {
 			socketAddress := lbEndpoint.GetEndpoint().GetAddress().GetSocketAddress()
-			newAddrs = append(newAddrs, resolver.Address{
+			address := resolver.Address{
 				Addr: net.JoinHostPort(socketAddress.GetAddress(), strconv.Itoa(int(socketAddress.GetPortValue()))),
-			})
+			}
+			if lbEndpoint.GetLoadBalancingWeight().GetValue() != 0 {
+				address.Metadata = &wrr.Info{Weight:lbEndpoint.GetLoadBalancingWeight().GetValue()}
+			}
+			newAddrs = append(newAddrs, address)
 		}
 		var weightChanged, addrsChanged bool
 		config, ok := xdsB.lidToConfig[lid]