resolver: Make EndpointMap's Get, Set and Delete operations O(1) (#8179) (#8184)

Author: arjan-bal

resolver/map.go

@@ -18,6 +18,12 @@
 
 package resolver
 
+import (
+	"encoding/base64"
+	"sort"
+	"strings"
+)
+
 type addressMapEntry struct {
 	addr  Address
 	value any
@@ -137,66 +143,61 @@ func (a *AddressMap) Values() []any {
 	return ret
 }
 
-type endpointNode struct {
-	addrs map[string]struct{}
-}
-
-// Equal returns whether the unordered set of addrs are the same between the
-// endpoint nodes.
-func (en *endpointNode) Equal(en2 *endpointNode) bool {
-	if len(en.addrs) != len(en2.addrs) {
-		return false
-	}
-	for addr := range en.addrs {
-		if _, ok := en2.addrs[addr]; !ok {
-			return false
-		}
-	}
-	return true
-}
-
-func toEndpointNode(endpoint Endpoint) endpointNode {
-	en := make(map[string]struct{})
-	for _, addr := range endpoint.Addresses {
-		en[addr.Addr] = struct{}{}
-	}
-	return endpointNode{
-		addrs: en,
-	}
-}
+type endpointMapKey string
 
 // EndpointMap is a map of endpoints to arbitrary values keyed on only the
 // unordered set of address strings within an endpoint. This map is not thread
 // safe, thus it is unsafe to access concurrently. Must be created via
 // NewEndpointMap; do not construct directly.
 type EndpointMap struct {
-	endpoints map[*endpointNode]any
+	endpoints map[endpointMapKey]endpointData
+}
+
+type endpointData struct {
+	// decodedKey stores the original key to avoid decoding when iterating on
+	// EndpointMap keys.
+	decodedKey Endpoint
+	value      any
 }
 
 // NewEndpointMap creates a new EndpointMap.
 func NewEndpointMap() *EndpointMap {
 	return &EndpointMap{
-		endpoints: make(map[*endpointNode]any),
+		endpoints: make(map[endpointMapKey]endpointData),
+	}
+}
+
+// encodeEndpoint returns a string that uniquely identifies the unordered set of
+// addresses within an endpoint.
+func encodeEndpoint(e Endpoint) endpointMapKey {
+	addrs := make([]string, 0, len(e.Addresses))
+	// base64 encoding the address strings restricts the characters present
+	// within the strings. This allows us to use a delimiter without the need of
+	// escape characters.
+	for _, addr := range e.Addresses {
+		addrs = append(addrs, base64.StdEncoding.EncodeToString([]byte(addr.Addr)))
 	}
+	sort.Strings(addrs)
+	// " " should not appear in base64 encoded strings.
+	return endpointMapKey(strings.Join(addrs, " "))
 }
 
 // Get returns the value for the address in the map, if present.
 func (em *EndpointMap) Get(e Endpoint) (value any, ok bool) {
-	en := toEndpointNode(e)
-	if endpoint := em.find(en); endpoint != nil {
-		return em.endpoints[endpoint], true
+	val, found := em.endpoints[encodeEndpoint(e)]
+	if found {
+		return val.value, true
 	}
 	return nil, false
 }
 
 // Set updates or adds the value to the address in the map.
 func (em *EndpointMap) Set(e Endpoint, value any) {
-	en := toEndpointNode(e)
-	if endpoint := em.find(en); endpoint != nil {
-		em.endpoints[endpoint] = value
-		return
+	en := encodeEndpoint(e)
+	em.endpoints[en] = endpointData{
+		decodedKey: Endpoint{Addresses: e.Addresses},
+		value:      value,
 	}
-	em.endpoints[&en] = value
 }
 
 // Len returns the number of entries in the map.
@@ -211,12 +212,8 @@ func (em *EndpointMap) Len() int {
 // used for EndpointMap accesses.
 func (em *EndpointMap) Keys() []Endpoint {
 	ret := make([]Endpoint, 0, len(em.endpoints))
-	for en := range em.endpoints {
-		var endpoint Endpoint
-		for addr := range en.addrs {
-			endpoint.Addresses = append(endpoint.Addresses, Address{Addr: addr})
-		}
-		ret = append(ret, endpoint)
+	for _, en := range em.endpoints {
+		ret = append(ret, en.decodedKey)
 	}
 	return ret
 }
@@ -225,27 +222,13 @@ func (em *EndpointMap) Keys() []Endpoint {
 func (em *EndpointMap) Values() []any {
 	ret := make([]any, 0, len(em.endpoints))
 	for _, val := range em.endpoints {
-		ret = append(ret, val)
+		ret = append(ret, val.value)
 	}
 	return ret
 }
 
-// find returns a pointer to the endpoint node in em if the endpoint node is
-// already present. If not found, nil is returned. The comparisons are done on
-// the unordered set of addresses within an endpoint.
-func (em EndpointMap) find(e endpointNode) *endpointNode {
-	for endpoint := range em.endpoints {
-		if e.Equal(endpoint) {
-			return endpoint
-		}
-	}
-	return nil
-}
-
 // Delete removes the specified endpoint from the map.
 func (em *EndpointMap) Delete(e Endpoint) {
-	en := toEndpointNode(e)
-	if entry := em.find(en); entry != nil {
-		delete(em.endpoints, entry)
-	}
+	en := encodeEndpoint(e)
+	delete(em.endpoints, en)
 }

resolver/map_test.go

@@ -287,3 +287,25 @@ func (s) TestEndpointMap_Values(t *testing.T) {
 		t.Fatalf("em.Values() returned unexpected elements (-want, +got):\n%v", diff)
 	}
 }
+
+// BenchmarkEndpointMap benchmarks map operations that are expected to run
+// faster than O(n). This test doesn't run O(n) operations including listing
+// keys and values.
+func BenchmarkEndpointMap(b *testing.B) {
+	em := NewEndpointMap()
+	for i := range b.N {
+		em.Set(Endpoint{
+			Addresses: []Address{{Addr: fmt.Sprintf("%d.%d.%d.%d", i, i, i, i)}},
+		}, i)
+	}
+	for i := range b.N {
+		em.Get(Endpoint{
+			Addresses: []Address{{Addr: fmt.Sprintf("%d.%d.%d.%d", i, i, i, i)}},
+		})
+	}
+	for i := range b.N {
+		em.Delete(Endpoint{
+			Addresses: []Address{{Addr: fmt.Sprintf("%d.%d.%d.%d", i, i, i, i)}},
+		})
+	}
+}