balancer: support hierarchical paths in addresses

internal/hierarchy/hierarchy.go

@@ -0,0 +1,99 @@
+/*
+ *
+ * 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 hierarchy contains functions to set and get hierarchy string from
+// addresses.
+//
+// This package is experimental.
+package hierarchy
+
+import (
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/resolver"
+)
+
+type pathKeyType string
+
+const pathKey = pathKeyType("grpc.internal.address.hierarchical_path")
+
+// Get returns the hierarchical path of addr.
+func Get(addr resolver.Address) []string {
+	attrs := addr.Attributes
+	if attrs == nil {
+		return nil
+	}
+	path, ok := attrs.Value(pathKey).([]string)
+	if !ok {
+		return nil
+	}
+	return path
+}
+
+// Set overrides the hierarchical path in addr with path.
+func Set(addr resolver.Address, path []string) resolver.Address {
+	if addr.Attributes == nil {
+		addr.Attributes = attributes.New(pathKey, path)
+		return addr
+	}
+	addr.Attributes = addr.Attributes.WithValues(pathKey, path)
+	return addr
+}
+
+// Group 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]},
+//   ],
+// }
+//
+// If hierarchical path is not set, or has no path in it, the address is
+// dropped.
+func Group(addrs []resolver.Address) map[string][]resolver.Address {
+	ret := make(map[string][]resolver.Address)
+	for _, addr := range addrs {
+		oldPath := Get(addr)
+		if len(oldPath) == 0 {
+			continue
+		}
+		curPath := oldPath[0]
+		newPath := oldPath[1:]
+		newAddr := Set(addr, newPath)
+		ret[curPath] = append(ret[curPath], newAddr)
+	}
+	return ret
+}

internal/hierarchy/hierarchy_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 hierarchy
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/attributes"
+	"google.golang.org/grpc/resolver"
+)
+
+func TestGet(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(pathKey, []string{"a", "b"}),
+			},
+			want: []string{"a", "b"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := Get(tt.addr); !cmp.Equal(got, tt.want) {
+				t.Errorf("Get() = %v, want %v", got, tt.want)
+			}
+		})
+	}
+}
+
+func TestSet(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(pathKey, []string{"before", "a", "b"}),
+			},
+			path: []string{"a", "b"},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			newAddr := Set(tt.addr, tt.path)
+			newPath := Get(newAddr)
+			if !cmp.Equal(newPath, tt.path) {
+				t.Errorf("path after Set() = %v, want %v", newPath, tt.path)
+			}
+		})
+	}
+}
+
+func TestGroup(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(pathKey, []string{"a"})},
+				{Addr: "a1", Attributes: attributes.New(pathKey, []string{"a"})},
+				{Addr: "b0", Attributes: attributes.New(pathKey, []string{"b"})},
+				{Addr: "b1", Attributes: attributes.New(pathKey, []string{"b"})},
+			},
+			want: map[string][]resolver.Address{
+				"a": {
+					{Addr: "a0", Attributes: attributes.New(pathKey, []string{})},
+					{Addr: "a1", Attributes: attributes.New(pathKey, []string{})},
+				},
+				"b": {
+					{Addr: "b0", Attributes: attributes.New(pathKey, []string{})},
+					{Addr: "b1", Attributes: attributes.New(pathKey, []string{})},
+				},
+			},
+		},
+		{
+			// Addresses without hierarchy are ignored.
+			name: "without hierarchy",
+			addrs: []resolver.Address{
+				{Addr: "a0", Attributes: attributes.New(pathKey, []string{"a"})},
+				{Addr: "a1", Attributes: attributes.New(pathKey, []string{"a"})},
+				{Addr: "b0", Attributes: nil},
+				{Addr: "b1", Attributes: nil},
+			},
+			want: map[string][]resolver.Address{
+				"a": {
+					{Addr: "a0", Attributes: attributes.New(pathKey, []string{})},
+					{Addr: "a1", Attributes: attributes.New(pathKey, []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(pathKey, []string{"a"})},
+				{Addr: "a1", Attributes: attributes.New(pathKey, []string{"a"})},
+				{Addr: "b0", Attributes: attributes.New(pathKey, "b")},
+				{Addr: "b1", Attributes: attributes.New(pathKey, 314)},
+			},
+			want: map[string][]resolver.Address{
+				"a": {
+					{Addr: "a0", Attributes: attributes.New(pathKey, []string{})},
+					{Addr: "a1", Attributes: attributes.New(pathKey, []string{})},
+				},
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			if got := Group(tt.addrs); !cmp.Equal(got, tt.want, cmp.AllowUnexported(attributes.Attributes{})) {
+				t.Errorf("Group() = %v, want %v", got, tt.want)
+				t.Errorf("diff: %v", cmp.Diff(got, tt.want, cmp.AllowUnexported(attributes.Attributes{})))
+			}
+		})
+	}
+}
+
+func TestGroupE2E(t *testing.T) {
+	hierarchy := map[string]map[string][]string{
+		"p0": {
+			"wt0": {"addr0", "addr1"},
+			"wt1": {"addr2", "addr3"},
+		},
+		"p1": {
+			"wt10": {"addr10", "addr11"},
+			"wt11": {"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(pathKey, path2),
+				}
+				addrsWithHierarchy = append(addrsWithHierarchy, a)
+			}
+		}
+	}
+
+	gotHierarchy := make(map[string]map[string][]string)
+	for p1, wts := range Group(addrsWithHierarchy) {
+		gotHierarchy[p1] = make(map[string][]string)
+		for p2, addrs := range Group(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))
+	}
+}