Pass address weights to wrr balancer through attributes.

Author: easwars

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 {
+		a = New()
+	}
 	if len(kvs)%2 != 0 {
 		panic(fmt.Sprintf("attributes.New called with unexpected input: len(kvs) = %v", len(kvs)))
 	}

balancer/weightedroundrobin/weightedroundrobin.go

@@ -19,11 +19,41 @@
 // Package weightedroundrobin defines a weighted roundrobin balancer.
 package weightedroundrobin
 
-// Name is the name of weighted_round_robin balancer.
-const Name = "weighted_round_robin"
+import "google.golang.org/grpc/attributes"
+
+const (
+	// Name is the name of weighted_round_robin balancer.
+	Name = "weighted_round_robin"
+
+	// Attribute key used to store AddrInfo in the Attributes field of
+	// resolver.Address.
+	attributeKey = "/balancer/weightedroundrobin/addrInfo"
+)
 
 // 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)
+}
+
+// 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 {
+		return nil
+	}
+	ai := a.Value(attributeKey)
+	if ai == nil {
+		return nil
+	}
+	if _, ok := ai.(*AddrInfo); !ok {
+		return nil
+	}
+	return ai.(*AddrInfo)
+}

balancer/weightedroundrobin/weightedwoundrobin_test.go

@@ -0,0 +1,72 @@
+/*
+ *
+ * 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"
+)
+
+func TestAddAddrInfoToAndFromAttributes(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:   nil,
+			inputAttributes: nil,
+			wantAddrInfo:    nil,
+		},
+		{
+			desc:            "addrInfo not present in non-empty attributes",
+			inputAddrInfo:   nil,
+			inputAttributes: attributes.New("foo", "bar"),
+			wantAddrInfo:    nil,
+		},
+	}
+
+	for _, test := range tests {
+		t.Run(test.desc, func(t *testing.T) {
+			outputAttributes := AddAddrInfoToAttributes(test.inputAddrInfo, test.inputAttributes)
+			gotAddrInfo := GetAddrInfoFromAttributes(outputAttributes)
+			if !cmp.Equal(gotAddrInfo, test.wantAddrInfo) {
+				t.Errorf("gotAddrInfo: %v, wantAddrInfo: %v", gotAddrInfo, test.wantAddrInfo)
+			}
+
+		})
+	}
+}

xds/internal/balancer/edsbalancer/eds_impl.go

@@ -277,9 +277,15 @@ 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.Attributes = weightedroundrobin.AddAddrInfoToAttributes(ai, address.Attributes)
+				// 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.
+				address.Metadata = ai
 			}
 			newAddrs = append(newAddrs, address)
 		}