1 /*
2 *
3 * Copyright 2021 gRPC authors.
4 *
5 * Licensed under the Apache License, Version 2.0 (the "License");
6 * you may not use this file except in compliance with the License.
7 * You may obtain a copy of the License at
8 *
9 * http://www.apache.org/licenses/LICENSE-2.0
10 *
11 * Unless required by applicable law or agreed to in writing, software
12 * distributed under the License is distributed on an "AS IS" BASIS,
13 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 * See the License for the specific language governing permissions and
15 * limitations under the License.
16 *
17 */
18 19 package resolver
20 21 import (
22 "encoding/base64"
23 "sort"
24 "strings"
25 )
26 27 type addressMapEntry[T any] struct {
28 addr Address
29 value T
30 }
31 32 // AddressMap is an AddressMapV2[any]. It will be deleted in an upcoming
33 // release of grpc-go.
34 //
35 // Deprecated: use the generic AddressMapV2 type instead.
36 type AddressMap = AddressMapV2[any]
37 38 // AddressMapV2 is a map of addresses to arbitrary values taking into account
39 // Attributes. BalancerAttributes are ignored, as are Metadata and Type.
40 // Multiple accesses may not be performed concurrently. Must be created via
41 // NewAddressMap; do not construct directly.
42 type AddressMapV2[T any] struct {
43 // The underlying map is keyed by an Address with fields that we don't care
44 // about being set to their zero values. The only fields that we care about
45 // are `Addr`, `ServerName` and `Attributes`. Since we need to be able to
46 // distinguish between addresses with same `Addr` and `ServerName`, but
47 // different `Attributes`, we cannot store the `Attributes` in the map key.
48 //
49 // The comparison operation for structs work as follows:
50 // Struct values are comparable if all their fields are comparable. Two
51 // struct values are equal if their corresponding non-blank fields are equal.
52 //
53 // The value type of the map contains a slice of addresses which match the key
54 // in their `Addr` and `ServerName` fields and contain the corresponding value
55 // associated with them.
56 m map[Address]addressMapEntryList[T]
57 }
58 59 func toMapKey(addr *Address) Address {
60 return Address{Addr: addr.Addr, ServerName: addr.ServerName}
61 }
62 63 type addressMapEntryList[T any] []*addressMapEntry[T]
64 65 // NewAddressMap creates a new AddressMapV2[any].
66 //
67 // Deprecated: use the generic NewAddressMapV2 constructor instead.
68 func NewAddressMap() *AddressMap {
69 return NewAddressMapV2[any]()
70 }
71 72 // NewAddressMapV2 creates a new AddressMapV2.
73 func NewAddressMapV2[T any]() *AddressMapV2[T] {
74 return &AddressMapV2[T]{m: make(map[Address]addressMapEntryList[T])}
75 }
76 77 // find returns the index of addr in the addressMapEntry slice, or -1 if not
78 // present.
79 func (l addressMapEntryList[T]) find(addr Address) int {
80 for i, entry := range l {
81 // Attributes are the only thing to match on here, since `Addr` and
82 // `ServerName` are already equal.
83 if entry.addr.Attributes.Equal(addr.Attributes) {
84 return i
85 }
86 }
87 return -1
88 }
89 90 // Get returns the value for the address in the map, if present.
91 func (a *AddressMapV2[T]) Get(addr Address) (value T, ok bool) {
92 addrKey := toMapKey(&addr)
93 entryList := a.m[addrKey]
94 if entry := entryList.find(addr); entry != -1 {
95 return entryList[entry].value, true
96 }
97 return value, false
98 }
99 100 // Set updates or adds the value to the address in the map.
101 func (a *AddressMapV2[T]) Set(addr Address, value T) {
102 addrKey := toMapKey(&addr)
103 entryList := a.m[addrKey]
104 if entry := entryList.find(addr); entry != -1 {
105 entryList[entry].value = value
106 return
107 }
108 a.m[addrKey] = append(entryList, &addressMapEntry[T]{addr: addr, value: value})
109 }
110 111 // Delete removes addr from the map.
112 func (a *AddressMapV2[T]) Delete(addr Address) {
113 addrKey := toMapKey(&addr)
114 entryList := a.m[addrKey]
115 entry := entryList.find(addr)
116 if entry == -1 {
117 return
118 }
119 if len(entryList) == 1 {
120 entryList = nil
121 } else {
122 copy(entryList[entry:], entryList[entry+1:])
123 entryList = entryList[:len(entryList)-1]
124 }
125 a.m[addrKey] = entryList
126 }
127 128 // Len returns the number of entries in the map.
129 func (a *AddressMapV2[T]) Len() int {
130 ret := 0
131 for _, entryList := range a.m {
132 ret += len(entryList)
133 }
134 return ret
135 }
136 137 // Keys returns a slice of all current map keys.
138 func (a *AddressMapV2[T]) Keys() []Address {
139 ret := make([]Address, 0, a.Len())
140 for _, entryList := range a.m {
141 for _, entry := range entryList {
142 ret = append(ret, entry.addr)
143 }
144 }
145 return ret
146 }
147 148 // Values returns a slice of all current map values.
149 func (a *AddressMapV2[T]) Values() []T {
150 ret := make([]T, 0, a.Len())
151 for _, entryList := range a.m {
152 for _, entry := range entryList {
153 ret = append(ret, entry.value)
154 }
155 }
156 return ret
157 }
158 159 type endpointMapKey string
160 161 // EndpointMap is a map of endpoints to arbitrary values keyed on only the
162 // unordered set of address strings within an endpoint. This map is not thread
163 // safe, thus it is unsafe to access concurrently. Must be created via
164 // NewEndpointMap; do not construct directly.
165 type EndpointMap[T any] struct {
166 endpoints map[endpointMapKey]endpointData[T]
167 }
168 169 type endpointData[T any] struct {
170 // decodedKey stores the original key to avoid decoding when iterating on
171 // EndpointMap keys.
172 decodedKey Endpoint
173 value T
174 }
175 176 // NewEndpointMap creates a new EndpointMap.
177 func NewEndpointMap[T any]() *EndpointMap[T] {
178 return &EndpointMap[T]{
179 endpoints: make(map[endpointMapKey]endpointData[T]),
180 }
181 }
182 183 // encodeEndpoint returns a string that uniquely identifies the unordered set of
184 // addresses within an endpoint.
185 func encodeEndpoint(e Endpoint) endpointMapKey {
186 addrs := make([]string, 0, len(e.Addresses))
187 // base64 encoding the address strings restricts the characters present
188 // within the strings. This allows us to use a delimiter without the need of
189 // escape characters.
190 for _, addr := range e.Addresses {
191 addrs = append(addrs, base64.StdEncoding.EncodeToString([]byte(addr.Addr)))
192 }
193 sort.Strings(addrs)
194 // " " should not appear in base64 encoded strings.
195 return endpointMapKey(strings.Join(addrs, " "))
196 }
197 198 // Get returns the value for the address in the map, if present.
199 func (em *EndpointMap[T]) Get(e Endpoint) (value T, ok bool) {
200 val, found := em.endpoints[encodeEndpoint(e)]
201 if found {
202 return val.value, true
203 }
204 return value, false
205 }
206 207 // Set updates or adds the value to the address in the map.
208 func (em *EndpointMap[T]) Set(e Endpoint, value T) {
209 en := encodeEndpoint(e)
210 em.endpoints[en] = endpointData[T]{
211 decodedKey: Endpoint{Addresses: e.Addresses},
212 value: value,
213 }
214 }
215 216 // Len returns the number of entries in the map.
217 func (em *EndpointMap[T]) Len() int {
218 return len(em.endpoints)
219 }
220 221 // Keys returns a slice of all current map keys, as endpoints specifying the
222 // addresses present in the endpoint keys, in which uniqueness is determined by
223 // the unordered set of addresses. Thus, endpoint information returned is not
224 // the full endpoint data (drops duplicated addresses and attributes) but can be
225 // used for EndpointMap accesses.
226 func (em *EndpointMap[T]) Keys() []Endpoint {
227 ret := make([]Endpoint, 0, len(em.endpoints))
228 for _, en := range em.endpoints {
229 ret = append(ret, en.decodedKey)
230 }
231 return ret
232 }
233 234 // Values returns a slice of all current map values.
235 func (em *EndpointMap[T]) Values() []T {
236 ret := make([]T, 0, len(em.endpoints))
237 for _, val := range em.endpoints {
238 ret = append(ret, val.value)
239 }
240 return ret
241 }
242 243 // Delete removes the specified endpoint from the map.
244 func (em *EndpointMap[T]) Delete(e Endpoint) {
245 en := encodeEndpoint(e)
246 delete(em.endpoints, en)
247 }
248