Fix base balancer set nil attributes into readySCs

balancer/base/balancer.go

@@ -22,6 +22,7 @@ import (
 	"errors"
 	"fmt"
 
+	"google.golang.org/grpc/attributes"
 	"google.golang.org/grpc/balancer"
 	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/grpclog"
@@ -41,7 +42,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]subConnInfo),
 		scStates: make(map[balancer.SubConn]connectivity.State),
 		csEvltr:  &balancer.ConnectivityStateEvaluator{},
 		config:   bb.config,
@@ -57,14 +58,19 @@ func (bb *baseBuilder) Name() string {
 	return bb.name
 }
 
+type subConnInfo struct {
+	subConn balancer.SubConn
+	attrs   *attributes.Attributes
+}
+
 type baseBalancer struct {
 	cc            balancer.ClientConn
 	pickerBuilder PickerBuilder
 
 	csEvltr *balancer.ConnectivityStateEvaluator
 	state   connectivity.State
 
-	subConns map[resolver.Address]balancer.SubConn // `attributes` is stripped from the keys of this map (the addresses)
+	subConns map[resolver.Address]subConnInfo // `attributes` is stripped from the keys of this map (the addresses)
 	scStates map[balancer.SubConn]connectivity.State
 	picker   balancer.Picker
 	config   Config
@@ -114,7 +120,7 @@ func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
 		aNoAttrs := a
 		aNoAttrs.Attributes = nil
 		addrsSet[aNoAttrs] = struct{}{}
-		if sc, ok := b.subConns[aNoAttrs]; !ok {
+		if scInfo, ok := b.subConns[aNoAttrs]; !ok {
 			// a is a new address (not existing in b.subConns).
 			//
 			// When creating SubConn, the original address with attributes is
@@ -125,7 +131,7 @@ func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
 				logger.Warningf("base.baseBalancer: failed to create new SubConn: %v", err)
 				continue
 			}
-			b.subConns[aNoAttrs] = sc
+			b.subConns[aNoAttrs] = subConnInfo{subConn: sc, attrs: a.Attributes}
 			b.scStates[sc] = connectivity.Idle
 			sc.Connect()
 		} else {
@@ -135,13 +141,15 @@ func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) error {
 			// The SubConn does a reflect.DeepEqual of the new and old
 			// addresses. So this is a noop if the current address is the same
 			// as the old one (including attributes).
-			b.cc.UpdateAddresses(sc, []resolver.Address{a})
+			scInfo.attrs = a.Attributes
+			b.subConns[aNoAttrs] = scInfo
+			b.cc.UpdateAddresses(scInfo.subConn, []resolver.Address{a})
 		}
 	}
-	for a, sc := range b.subConns {
+	for a, scInfo := range b.subConns {
 		// a was removed by resolver.
 		if _, ok := addrsSet[a]; !ok {
-			b.cc.RemoveSubConn(sc)
+			b.cc.RemoveSubConn(scInfo.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 UpdateSubConnState.
@@ -184,9 +192,10 @@ func (b *baseBalancer) regeneratePicker() {
 	readySCs := make(map[balancer.SubConn]SubConnInfo)
 
 	// 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[sc] = SubConnInfo{Address: addr}
+	for addr, scInfo := range b.subConns {
+		if st, ok := b.scStates[scInfo.subConn]; ok && st == connectivity.Ready {
+			addr.Attributes = scInfo.attrs
+			readySCs[scInfo.subConn] = SubConnInfo{Address: addr}
 		}
 	}
 	b.picker = b.pickerBuilder.Build(PickerBuildInfo{ReadySCs: readySCs})