Pass weights to wrr balancer through attributes. (#3530)

attributes/attributes.go

@@ -50,6 +50,9 @@ func New(kvs ...interface{}) *Attributes {
 // times, the last value overwrites all previous values for that key.  To
 // remove an existing key, use a nil value.
 func (a *Attributes) WithValues(kvs ...interface{}) *Attributes {
+	if a == nil {
+		return New(kvs...)
+	}
 	if len(kvs)%2 != 0 {
 		panic(fmt.Sprintf("attributes.New called with unexpected input: len(kvs) = %v", len(kvs)))
 	}
@@ -66,5 +69,8 @@ func (a *Attributes) WithValues(kvs ...interface{}) *Attributes {
 // Value returns the value associated with these attributes for key, or nil if
 // no value is associated with key.
 func (a *Attributes) Value(key interface{}) interface{} {
+	if a == nil {
+		return nil
+	}
 	return a.m[key]
 }

balancer/weightedroundrobin/weightedroundrobin.go

@@ -19,11 +19,37 @@
 // Package weightedroundrobin defines a weighted roundrobin balancer.
 package weightedroundrobin
 
+import (
+	"google.golang.org/grpc/resolver"
+)
+
 // Name is the name of weighted_round_robin balancer.
 const Name = "weighted_round_robin"
 
-// AddrInfo will be stored inside Address metadata in order to use weighted roundrobin
-// balancer.
+// 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.
 type AddrInfo struct {
 	Weight uint32
 }
+
+// SetAddrInfo returns a copy of addr in which the Attributes field is updated
+// with addrInfo.
+//
+// This is an EXPERIMENTAL API.
+func SetAddrInfo(addr resolver.Address, addrInfo AddrInfo) resolver.Address {
+	addr.Attributes = addr.Attributes.WithValues(attributeKey{}, addrInfo)
+	return addr
+}
+
+// GetAddrInfo returns the AddrInfo stored in the Attributes fields of addr.
+//
+// This is an EXPERIMENTAL API.
+func GetAddrInfo(addr resolver.Address) AddrInfo {
+	v := addr.Attributes.Value(attributeKey{})
+	ai, _ := v.(AddrInfo)
+	return ai
+}

balancer/weightedroundrobin/weightedwoundrobin_test.go

@@ -0,0 +1,82 @@
+/*
+ *
+ * Copyright 2020 gRPC authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package weightedroundrobin
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/resolver"
+)
+
+func TestAddrInfoToAndFromAttributes(t *testing.T) {
+	tests := []struct {
+		desc            string
+		inputAddrInfo   AddrInfo
+		inputAttributes *attributes.Attributes
+		wantAddrInfo    AddrInfo
+	}{
+		{
+			desc:            "empty attributes",
+			inputAddrInfo:   AddrInfo{Weight: 100},
+			inputAttributes: nil,
+			wantAddrInfo:    AddrInfo{Weight: 100},
+		},
+		{
+			desc:            "non-empty attributes",
+			inputAddrInfo:   AddrInfo{Weight: 100},
+			inputAttributes: attributes.New("foo", "bar"),
+			wantAddrInfo:    AddrInfo{Weight: 100},
+		},
+		{
+			desc:            "addrInfo not present in empty attributes",
+			inputAddrInfo:   AddrInfo{},
+			inputAttributes: nil,
+			wantAddrInfo:    AddrInfo{},
+		},
+		{
+			desc:            "addrInfo not present in non-empty attributes",
+			inputAddrInfo:   AddrInfo{},
+			inputAttributes: attributes.New("foo", "bar"),
+			wantAddrInfo:    AddrInfo{},
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			addr := resolver.Address{Attributes: test.inputAttributes}
+			addr = SetAddrInfo(addr, test.inputAddrInfo)
+			gotAddrInfo := GetAddrInfo(addr)
+			if !cmp.Equal(gotAddrInfo, test.wantAddrInfo) {
+				t.Errorf("gotAddrInfo: %v, wantAddrInfo: %v", gotAddrInfo, test.wantAddrInfo)
+			}
+
+		})
+	}
+}
+
+func TestGetAddInfoEmpty(t *testing.T) {
+	addr := resolver.Address{Attributes: attributes.New()}
+	gotAddrInfo := GetAddrInfo(addr)
+	wantAddrInfo := AddrInfo{}
+	if !cmp.Equal(gotAddrInfo, wantAddrInfo) {
+		t.Errorf("gotAddrInfo: %v, wantAddrInfo: %v", gotAddrInfo, wantAddrInfo)
+	}
+}

xds/internal/balancer/edsbalancer/eds_impl.go

@@ -277,9 +277,16 @@ func (edsImpl *edsBalancerImpl) handleEDSResponsePerPriority(bgwc *balancerGroup
 				Addr: lbEndpoint.Address,
 			}
 			if edsImpl.subBalancerBuilder.Name() == weightedroundrobin.Name && lbEndpoint.Weight != 0 {
-				address.Metadata = &weightedroundrobin.AddrInfo{
-					Weight: lbEndpoint.Weight,
-				}
+				ai := weightedroundrobin.AddrInfo{Weight: lbEndpoint.Weight}
+				address = weightedroundrobin.SetAddrInfo(address, ai)
+				// 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
+				// Metadata field here to allow users of this field to migrate
+				// to the new one.
+				// TODO(easwars): Remove this once all users have migrated.
+				// See https://github.com/grpc/grpc-go/issues/3563.
+				address.Metadata = &ai
 			}
 			newAddrs = append(newAddrs, address)
 		}