Review comments #1.

Author: easwars

balancer/weightedroundrobin/weightedroundrobin.go

@@ -19,36 +19,35 @@
 // Package weightedroundrobin defines a weighted roundrobin balancer.
 package weightedroundrobin
 
-import "google.golang.org/grpc/attributes"
+import (
+	"google.golang.org/grpc/resolver"
+)
 
-const (
-	// Name is the name of weighted_round_robin balancer.
-	Name = "weighted_round_robin"
+// Name is the name of weighted_round_robin balancer.
+const Name = "weighted_round_robin"
 
-	// Attribute key used to store AddrInfo in the Attributes field of
-	// resolver.Address.
-	attributeKey = "/balancer/weightedroundrobin/addrInfo"
-)
+// 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 roundrobin
-// balancer.
+// AddrInfo will be stored inside Address metadata in order to use weighted
+// roundrobin balancer.
 type AddrInfo struct {
 	Weight uint32
 }
 
-// AddAddrInfoToAttributes returns a new Attributes containing all key/value
-// pairs in a with ai being added to it.
-func AddAddrInfoToAttributes(ai *AddrInfo, a *attributes.Attributes) *attributes.Attributes {
-	return a.WithValues(attributeKey, ai)
+// SetAddrInfo sets addInfo in the Attributes field of addr.
+func SetAddrInfo(addrInfo *AddrInfo, addr *resolver.Address) {
+	addr.Attributes = addr.Attributes.WithValues(attributeKey{}, addrInfo)
 }
 
-// GetAddrInfoFromAttributes returns the AddrInfo stored in a. Returns nil if no
-// AddrInfo is present in a.
-func GetAddrInfoFromAttributes(a *attributes.Attributes) *AddrInfo {
-	if a == nil {
+// GetAddrInfo returns the AddrInfo stored in the Attributes fields of addr.
+// Returns nil if no AddrInfo is present.
+func GetAddrInfo(addr *resolver.Address) *AddrInfo {
+	if addr == nil || addr.Attributes == nil {
 		return nil
 	}
-	ai := a.Value(attributeKey)
+	ai := addr.Attributes.Value(attributeKey{})
 	if ai == nil {
 		return nil
 	}

balancer/weightedroundrobin/weightedwoundrobin_test.go

@@ -22,8 +22,8 @@ import (
 	"testing"
 
 	"github.com/google/go-cmp/cmp"
-
 	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/resolver"
 )
 
 func TestAddAddrInfoToAndFromAttributes(t *testing.T) {
@@ -61,8 +61,9 @@ func TestAddAddrInfoToAndFromAttributes(t *testing.T) {
 
 	for _, test := range tests {
 		t.Run(test.desc, func(t *testing.T) {
-			outputAttributes := AddAddrInfoToAttributes(test.inputAddrInfo, test.inputAttributes)
-			gotAddrInfo := GetAddrInfoFromAttributes(outputAttributes)
+			addr := &resolver.Address{Attributes: test.inputAttributes}
+			SetAddrInfo(test.inputAddrInfo, addr)
+			gotAddrInfo := GetAddrInfo(addr)
 			if !cmp.Equal(gotAddrInfo, test.wantAddrInfo) {
 				t.Errorf("gotAddrInfo: %v, wantAddrInfo: %v", gotAddrInfo, test.wantAddrInfo)
 			}

xds/internal/balancer/edsbalancer/eds_impl.go

@@ -278,7 +278,7 @@ func (edsImpl *edsBalancerImpl) handleEDSResponsePerPriority(bgwc *balancerGroup
 			}
 			if edsImpl.subBalancerBuilder.Name() == weightedroundrobin.Name && lbEndpoint.Weight != 0 {
 				ai := &weightedroundrobin.AddrInfo{Weight: lbEndpoint.Weight}
-				address.Attributes = weightedroundrobin.AddAddrInfoToAttributes(ai, address.Attributes)
+				weightedroundrobin.SetAddrInfo(ai, &address)
 				// Metadata field in resolver.Address is deprecated. The
 				// attributes field should be used to specify arbitrary
 				// attributes about the address. We still need to populate the