resolver: replace AddressMap.Range with Keys (#4891)

Co-authored-by: Menghan Li <menghanl@google.com>

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})
 }
 

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
 }

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)
 	}
 }

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) &&