balancer: support hierarchical paths in addresses

balancer/base/hierarchical.go

@@ -0,0 +1,96 @@
+/*
+ *
+ * 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 base
+
+import (
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/resolver"
+)
+
+type hierarchicalPathKeyType string
+
+const hierarchicalPathKey = hierarchicalPathKeyType("grpc.internal.address.hierarchical_path")
+
+// RetrieveHierarchicalPath returns the hierarchical path of addr.
+func RetrieveHierarchicalPath(addr resolver.Address) []string {
+	attrs := addr.Attributes
+	if attrs == nil {
+		return nil
+	}
+	path, ok := attrs.Value(hierarchicalPathKey).([]string)
+	if !ok {
+		return nil
+	}
+	return path
+}
+
+// OverrideHierarchicalPath overrides the hierarchical path in addr with path.
+func OverrideHierarchicalPath(addr resolver.Address, path []string) resolver.Address {
+	if addr.Attributes == nil {
+		addr.Attributes = attributes.New(hierarchicalPathKey, path)
+		return addr
+	}
+	addr.Attributes = addr.Attributes.WithValues(hierarchicalPathKey, path)
+	return addr
+}
+
+// SplitHierarchicalAddresses splits a slice of addresses into groups based on
+// the first hierarchy path. The first hierarchy path will be removed from the
+// result.
+//
+// Input:
+// [
+//   {addr0, path: [p0, wt0]}
+//   {addr1, path: [p0, wt1]}
+//   {addr2, path: [p1, wt2]}
+//   {addr3, path: [p1, wt3]}
+// ]
+//
+// Addresses will be split into p0/p1, and the p0/p1 will be removed from the
+// path.
+//
+// Output:
+// {
+//   p0: [
+//     {addr0, path: [wt0]},
+//     {addr1, path: [wt1]},
+//   ],
+//   p1: [
+//     {addr2, path: [wt2]},
+//     {addr3, path: [wt3]},
+//   ],
+// }
+func SplitHierarchicalAddresses(addrs []resolver.Address) map[string][]resolver.Address {
+	ret := make(map[string][]resolver.Address)
+	for _, addr := range addrs {
+		oldPath := RetrieveHierarchicalPath(addr)
+		if len(oldPath) == 0 {
+			// When hierarchical path is not set, or has no path in it, skip the
+			// address. Another option is to return this address with path "",
+			// this shouldn't conflict with anything because "" isn't a valid
+			// path.
+			continue
+		}
+		curPath := oldPath[0]
+		newPath := oldPath[1:]
+		newAddr := OverrideHierarchicalPath(addr, newPath)
+		ret[curPath] = append(ret[curPath], newAddr)
+	}
+	return ret
+}

balancer/base/hierarchical_test.go

@@ -0,0 +1,197 @@
+/*
+ *
+ * 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 base
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/resolver"
+)
+
+func TestRetrieveHierarchicalPath(t *testing.T) {
+	tests := []struct {
+		name string
+		addr resolver.Address
+		want []string
+	}{
+		{
+			name: "not set",
+			addr: resolver.Address{},
+			want: nil,
+		},
+		{
+			name: "set",
+			addr: resolver.Address{
+				Attributes: attributes.New(hierarchicalPathKey, []string{"a", "b"}),
+			},
+			want: []string{"a", "b"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := RetrieveHierarchicalPath(tt.addr); !cmp.Equal(got, tt.want) {
+				t.Errorf("RetrieveHierarchicalPath() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestOverrideHierarchicalPath(t *testing.T) {
+	tests := []struct {
+		name string
+		addr resolver.Address
+		path []string
+	}{
+		{
+			name: "before is not set",
+			addr: resolver.Address{},
+			path: []string{"a", "b"},
+		},
+		{
+			name: "before is set",
+			addr: resolver.Address{
+				Attributes: attributes.New(hierarchicalPathKey, []string{"before", "a", "b"}),
+			},
+			path: []string{"a", "b"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			newAddr := OverrideHierarchicalPath(tt.addr, tt.path)
+			newPath := RetrieveHierarchicalPath(newAddr)
+			if !cmp.Equal(newPath, tt.path) {
+				t.Errorf("path after OverrideHierarchicalPath() = %v, want %v", newPath, tt.path)
+			}
+		})
+	}
+}
+
+func TestSplitHierarchicalAddresses(t *testing.T) {
+	tests := []struct {
+		name  string
+		addrs []resolver.Address
+		want  map[string][]resolver.Address
+	}{
+		{
+			name: "all with hierarchy",
+			addrs: []resolver.Address{
+				{Addr: "a0", Attributes: attributes.New(hierarchicalPathKey, []string{"a"})},
+				{Addr: "a1", Attributes: attributes.New(hierarchicalPathKey, []string{"a"})},
+				{Addr: "b0", Attributes: attributes.New(hierarchicalPathKey, []string{"b"})},
+				{Addr: "b1", Attributes: attributes.New(hierarchicalPathKey, []string{"b"})},
+			},
+			want: map[string][]resolver.Address{
+				"a": []resolver.Address{
+					{Addr: "a0", Attributes: attributes.New(hierarchicalPathKey, []string{})},
+					{Addr: "a1", Attributes: attributes.New(hierarchicalPathKey, []string{})},
+				},
+				"b": []resolver.Address{
+					{Addr: "b0", Attributes: attributes.New(hierarchicalPathKey, []string{})},
+					{Addr: "b1", Attributes: attributes.New(hierarchicalPathKey, []string{})},
+				},
+			},
+		},
+		{
+			// Addresses without hierarchy are ignored.
+			name: "without hierarchy",
+			addrs: []resolver.Address{
+				{Addr: "a0", Attributes: attributes.New(hierarchicalPathKey, []string{"a"})},
+				{Addr: "a1", Attributes: attributes.New(hierarchicalPathKey, []string{"a"})},
+				{Addr: "b0", Attributes: nil},
+				{Addr: "b1", Attributes: nil},
+			},
+			want: map[string][]resolver.Address{
+				"a": []resolver.Address{
+					{Addr: "a0", Attributes: attributes.New(hierarchicalPathKey, []string{})},
+					{Addr: "a1", Attributes: attributes.New(hierarchicalPathKey, []string{})},
+				},
+			},
+		},
+		{
+			// If hierarchy is set to a wrong type (which should never happen),
+			// the address is ignored.
+			name: "wrong type",
+			addrs: []resolver.Address{
+				{Addr: "a0", Attributes: attributes.New(hierarchicalPathKey, []string{"a"})},
+				{Addr: "a1", Attributes: attributes.New(hierarchicalPathKey, []string{"a"})},
+				{Addr: "b0", Attributes: attributes.New(hierarchicalPathKey, "b")},
+				{Addr: "b1", Attributes: attributes.New(hierarchicalPathKey, 314)},
+			},
+			want: map[string][]resolver.Address{
+				"a": []resolver.Address{
+					{Addr: "a0", Attributes: attributes.New(hierarchicalPathKey, []string{})},
+					{Addr: "a1", Attributes: attributes.New(hierarchicalPathKey, []string{})},
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := SplitHierarchicalAddresses(tt.addrs); !cmp.Equal(got, tt.want, cmp.AllowUnexported(attributes.Attributes{})) {
+				t.Errorf("SplitHierarchicalAddresses() = %v, want %v", got, tt.want)
+				t.Errorf("diff: %v", cmp.Diff(got, tt.want, cmp.AllowUnexported(attributes.Attributes{})))
+			}
+		})
+	}
+}
+
+func TestSplitHierarchicalAddressesE2E(t *testing.T) {
+	hierarchy := map[string]map[string][]string{
+		"p0": map[string][]string{
+			"wt0": []string{"addr0", "addr1"},
+			"wt1": []string{"addr2", "addr3"},
+		},
+		"p1": map[string][]string{
+			"wt10": []string{"addr10", "addr11"},
+			"wt11": []string{"addr12", "addr13"},
+		},
+	}
+
+	var addrsWithHierarchy []resolver.Address
+	for p, wts := range hierarchy {
+		path1 := []string{p}
+		for wt, addrs := range wts {
+			path2 := append([]string(nil), path1...)
+			path2 = append(path2, wt)
+			for _, addr := range addrs {
+				a := resolver.Address{
+					Addr:       addr,
+					Attributes: attributes.New(hierarchicalPathKey, path2),
+				}
+				addrsWithHierarchy = append(addrsWithHierarchy, a)
+			}
+		}
+	}
+
+	gotHierarchy := make(map[string]map[string][]string)
+	for p1, wts := range SplitHierarchicalAddresses(addrsWithHierarchy) {
+		gotHierarchy[p1] = make(map[string][]string)
+		for p2, addrs := range SplitHierarchicalAddresses(wts) {
+			for _, addr := range addrs {
+				gotHierarchy[p1][p2] = append(gotHierarchy[p1][p2], addr.Addr)
+			}
+		}
+	}
+
+	if !cmp.Equal(gotHierarchy, hierarchy) {
+		t.Errorf("diff: %v", cmp.Diff(gotHierarchy, hierarchy))
+	}
+}