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
19package resolver
20
21type addressMapEntry struct {
22 addr Address
23 value any
24}
25
26// AddressMap is a map of addresses to arbitrary values taking into account
27// Attributes. BalancerAttributes are ignored, as are Metadata and Type.
28// Multiple accesses may not be performed concurrently. Must be created via
29// NewAddressMap; do not construct directly.
30type AddressMap struct {
31 // The underlying map is keyed by an Address with fields that we don't care
32 // about being set to their zero values. The only fields that we care about
33 // are `Addr`, `ServerName` and `Attributes`. Since we need to be able to
34 // distinguish between addresses with same `Addr` and `ServerName`, but
35 // different `Attributes`, we cannot store the `Attributes` in the map key.
36 //
37 // The comparison operation for structs work as follows:
38 // Struct values are comparable if all their fields are comparable. Two
39 // struct values are equal if their corresponding non-blank fields are equal.
40 //
41 // The value type of the map contains a slice of addresses which match the key
42 // in their `Addr` and `ServerName` fields and contain the corresponding value
43 // associated with them.
44 m map[Address]addressMapEntryList
45}
46
47func toMapKey(addr *Address) Address {
48 return Address{Addr: addr.Addr, ServerName: addr.ServerName}
49}
50
51type addressMapEntryList []*addressMapEntry
52
53// NewAddressMap creates a new AddressMap.
54func NewAddressMap() *AddressMap {
55 return &AddressMap{m: make(map[Address]addressMapEntryList)}
56}
57
58// find returns the index of addr in the addressMapEntry slice, or -1 if not
59// present.
60func (l addressMapEntryList) find(addr Address) int {
61 for i, entry := range l {
62 // Attributes are the only thing to match on here, since `Addr` and
63 // `ServerName` are already equal.
64 if entry.addr.Attributes.Equal(addr.Attributes) {
65 return i
66 }
67 }
68 return -1
69}
70
71// Get returns the value for the address in the map, if present.
72func (a *AddressMap) Get(addr Address) (value any, ok bool) {
73 addrKey := toMapKey(&addr)
74 entryList := a.m[addrKey]
75 if entry := entryList.find(addr); entry != -1 {
76 return entryList[entry].value, true
77 }
78 return nil, false
79}
80
81// Set updates or adds the value to the address in the map.
82func (a *AddressMap) Set(addr Address, value any) {
83 addrKey := toMapKey(&addr)
84 entryList := a.m[addrKey]
85 if entry := entryList.find(addr); entry != -1 {
86 entryList[entry].value = value
87 return
88 }
89 a.m[addrKey] = append(entryList, &addressMapEntry{addr: addr, value: value})
90}
91
92// Delete removes addr from the map.
93func (a *AddressMap) Delete(addr Address) {
94 addrKey := toMapKey(&addr)
95 entryList := a.m[addrKey]
96 entry := entryList.find(addr)
97 if entry == -1 {
98 return
99 }
100 if len(entryList) == 1 {
101 entryList = nil
102 } else {
103 copy(entryList[entry:], entryList[entry+1:])
104 entryList = entryList[:len(entryList)-1]
105 }
106 a.m[addrKey] = entryList
107}
108
109// Len returns the number of entries in the map.
110func (a *AddressMap) Len() int {
111 ret := 0
112 for _, entryList := range a.m {
113 ret += len(entryList)
114 }
115 return ret
116}
117
118// Keys returns a slice of all current map keys.
119func (a *AddressMap) Keys() []Address {
120 ret := make([]Address, 0, a.Len())
121 for _, entryList := range a.m {
122 for _, entry := range entryList {
123 ret = append(ret, entry.addr)
124 }
125 }
126 return ret
127}
128
129// Values returns a slice of all current map values.
130func (a *AddressMap) Values() []any {
131 ret := make([]any, 0, a.Len())
132 for _, entryList := range a.m {
133 for _, entry := range entryList {
134 ret = append(ret, entry.value)
135 }
136 }
137 return ret
138}
139
140type endpointNode struct {
141 addrs map[string]struct{}
142}
143
144// Equal returns whether the unordered set of addrs are the same between the
145// endpoint nodes.
146func (en *endpointNode) Equal(en2 *endpointNode) bool {
147 if len(en.addrs) != len(en2.addrs) {
148 return false
149 }
150 for addr := range en.addrs {
151 if _, ok := en2.addrs[addr]; !ok {
152 return false
153 }
154 }
155 return true
156}
157
158func toEndpointNode(endpoint Endpoint) endpointNode {
159 en := make(map[string]struct{})
160 for _, addr := range endpoint.Addresses {
161 en[addr.Addr] = struct{}{}
162 }
163 return endpointNode{
164 addrs: en,
165 }
166}
167
168// EndpointMap is a map of endpoints to arbitrary values keyed on only the
169// unordered set of address strings within an endpoint. This map is not thread
170// safe, thus it is unsafe to access concurrently. Must be created via
171// NewEndpointMap; do not construct directly.
172type EndpointMap struct {
173 endpoints map[*endpointNode]any
174}
175
176// NewEndpointMap creates a new EndpointMap.
177func NewEndpointMap() *EndpointMap {
178 return &EndpointMap{
179 endpoints: make(map[*endpointNode]any),
180 }
181}
182
183// Get returns the value for the address in the map, if present.
184func (em *EndpointMap) Get(e Endpoint) (value any, ok bool) {
185 en := toEndpointNode(e)
186 if endpoint := em.find(en); endpoint != nil {
187 return em.endpoints[endpoint], true
188 }
189 return nil, false
190}
191
192// Set updates or adds the value to the address in the map.
193func (em *EndpointMap) Set(e Endpoint, value any) {
194 en := toEndpointNode(e)
195 if endpoint := em.find(en); endpoint != nil {
196 em.endpoints[endpoint] = value
197 return
198 }
199 em.endpoints[&en] = value
200}
201
202// Len returns the number of entries in the map.
203func (em *EndpointMap) Len() int {
204 return len(em.endpoints)
205}
206
207// Keys returns a slice of all current map keys, as endpoints specifying the
208// addresses present in the endpoint keys, in which uniqueness is determined by
209// the unordered set of addresses. Thus, endpoint information returned is not
210// the full endpoint data (drops duplicated addresses and attributes) but can be
211// used for EndpointMap accesses.
212func (em *EndpointMap) Keys() []Endpoint {
213 ret := make([]Endpoint, 0, len(em.endpoints))
214 for en := range em.endpoints {
215 var endpoint Endpoint
216 for addr := range en.addrs {
217 endpoint.Addresses = append(endpoint.Addresses, Address{Addr: addr})
218 }
219 ret = append(ret, endpoint)
220 }
221 return ret
222}
223
224// Values returns a slice of all current map values.
225func (em *EndpointMap) Values() []any {
226 ret := make([]any, 0, len(em.endpoints))
227 for _, val := range em.endpoints {
228 ret = append(ret, val)
229 }
230 return ret
231}
232
233// find returns a pointer to the endpoint node in em if the endpoint node is
234// already present. If not found, nil is returned. The comparisons are done on
235// the unordered set of addresses within an endpoint.
236func (em EndpointMap) find(e endpointNode) *endpointNode {
237 for endpoint := range em.endpoints {
238 if e.Equal(endpoint) {
239 return endpoint
240 }
241 }
242 return nil
243}
244
245// Delete removes the specified endpoint from the map.
246func (em *EndpointMap) Delete(e Endpoint) {
247 en := toEndpointNode(e)
248 if entry := em.find(en); entry != nil {
249 delete(em.endpoints, entry)
250 }
251}