From f00baa6c3c8455ef1db2ac64f7b89a63ec7d2776 Mon Sep 17 00:00:00 2001 From: Doug Fawley Date: Wed, 20 Oct 2021 15:07:37 -0700 Subject: [PATCH] resolver: replace AddressMap.Range with Keys (#4891) Co-authored-by: Menghan Li --- balancer/base/balancer.go | 10 ++++++---- resolver/map.go | 14 ++++++++++---- resolver/map_test.go | 35 ++++++++++++----------------------- resolver/resolver.go | 2 +- 4 files changed, 29 insertions(+), 32 deletions(-) diff --git a/balancer/base/balancer.go b/balancer/base/balancer.go index 908c6e3376e0..a67074a3ad06 100644 --- a/balancer/base/balancer.go +++ b/balancer/base/balancer.go @@ -115,7 +115,8 @@ func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) error { sc.Connect() } } - b.subConns.Range(func(a resolver.Address, sci interface{}) { + for _, a := range b.subConns.Keys() { + sci, _ := b.subConns.Get(a) sc := sci.(balancer.SubConn) // a was removed by resolver. if _, ok := addrsSet.Get(a); !ok { @@ -124,7 +125,7 @@ func (b *baseBalancer) UpdateClientConnState(s balancer.ClientConnState) error { // Keep the state of this sc in b.scStates until sc's state becomes Shutdown. // The entry will be deleted in UpdateSubConnState. } - }) + } // If resolver state contains no addresses, return an error so ClientConn // will trigger re-resolve. Also records this as an resolver error, so when // the overall state turns transient failure, the error message will have @@ -162,12 +163,13 @@ func (b *baseBalancer) regeneratePicker() { readySCs := make(map[balancer.SubConn]SubConnInfo) // Filter out all ready SCs from full subConn map. - b.subConns.Range(func(addr resolver.Address, sci interface{}) { + for _, addr := range b.subConns.Keys() { + sci, _ := b.subConns.Get(addr) sc := sci.(balancer.SubConn) if st, ok := b.scStates[sc]; ok && st == connectivity.Ready { readySCs[sc] = SubConnInfo{Address: addr} } - }) + } b.picker = b.pickerBuilder.Build(PickerBuildInfo{ReadySCs: readySCs}) } diff --git a/resolver/map.go b/resolver/map.go index bfde61b331c5..e87ecd0eeb38 100644 --- a/resolver/map.go +++ b/resolver/map.go @@ -90,14 +90,20 @@ func (a *AddressMap) Delete(addr Address) { // Len returns the number of entries in the map. func (a *AddressMap) Len() int { - return len(a.m) + ret := 0 + for _, entryList := range a.m { + ret += len(entryList) + } + return ret } -// Range invokes f for each entry in the map. -func (a *AddressMap) Range(f func(addr Address, value interface{})) { +// Keys returns a slice of all current map keys. +func (a *AddressMap) Keys() []Address { + ret := make([]Address, 0, a.Len()) for _, entryList := range a.m { for _, entry := range entryList { - f(entry.addr, entry.value) + ret = append(ret, entry.addr) } } + return ret } diff --git a/resolver/map_test.go b/resolver/map_test.go index 86191d82bbb3..26d539bcd6d3 100644 --- a/resolver/map_test.go +++ b/resolver/map_test.go @@ -19,8 +19,11 @@ package resolver import ( + "fmt" + "sort" "testing" + "github.com/google/go-cmp/cmp" "google.golang.org/grpc/attributes" ) @@ -117,7 +120,7 @@ func (s) TestAddressMap_Delete(t *testing.T) { } } -func (s) TestAddressMap_Range(t *testing.T) { +func (s) TestAddressMap_Keys(t *testing.T) { addrMap := NewAddressMap() addrMap.Set(addr1, 1) addrMap.Set(addr2, 2) @@ -127,27 +130,13 @@ func (s) TestAddressMap_Range(t *testing.T) { addrMap.Set(addr6, 6) addrMap.Set(addr7, 7) // aliases addr1 - want := map[int]bool{2: true, 3: true, 4: true, 5: true, 6: true, 7: true} - test := func(a1, a2 Address, n int, v interface{}) { - if a1.Addr == a2.Addr && a1.Attributes == a2.Attributes && a1.ServerName == a2.ServerName { - if ok := want[n]; !ok { - t.Fatal("matched address multiple times:", a1, n, want) - } - if n != v.(int) { - t.Fatalf("%v read value %v; want %v:", a1, v, n) - } - delete(want, n) - } - } - addrMap.Range(func(a Address, v interface{}) { - test(a, addr1, 7, v) - test(a, addr2, 2, v) - test(a, addr3, 3, v) - test(a, addr4, 4, v) - test(a, addr5, 5, v) - test(a, addr6, 6, v) - }) - if len(want) != 0 { - t.Fatalf("did not find expected addresses; remaining: %v", want) + want := []Address{addr1, addr2, addr3, addr4, addr5, addr6} + got := addrMap.Keys() + if d := cmp.Diff(want, got, cmp.Transformer("sort", func(in []Address) []Address { + out := append([]Address(nil), in...) + sort.Slice(out, func(i, j int) bool { return fmt.Sprint(out[i]) < fmt.Sprint(out[j]) }) + return out + })); d != "" { + t.Fatalf("addrMap.Keys returned unexpected elements (-want, +got):\n%v", d) } } diff --git a/resolver/resolver.go b/resolver/resolver.go index 873b932b20d6..e28b68026062 100644 --- a/resolver/resolver.go +++ b/resolver/resolver.go @@ -139,7 +139,7 @@ type Address struct { // Equal returns whether a and o are identical. Metadata is compared directly, // not with any recursive introspection. -func (a *Address) Equal(o *Address) bool { +func (a *Address) Equal(o Address) bool { return a.Addr == o.Addr && a.ServerName == o.ServerName && a.Attributes.Equal(o.Attributes) && a.BalancerAttributes.Equal(o.BalancerAttributes) &&