map.go 7.1 KB

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