diff --git a/internal/hierarchy/hierarchy.go b/internal/hierarchy/hierarchy.go new file mode 100644 index 000000000000..17185d95d38e --- /dev/null +++ b/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 +} diff --git a/internal/hierarchy/hierarchy_test.go b/internal/hierarchy/hierarchy_test.go new file mode 100644 index 000000000000..fc62f82b0850 --- /dev/null +++ b/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)) + } +}