github.com/MetalBlockchain/metalgo@v1.11.9/network/peer/set.go (about) 1 // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved. 2 // See the file LICENSE for licensing terms. 3 4 package peer 5 6 import ( 7 "github.com/MetalBlockchain/metalgo/ids" 8 "github.com/MetalBlockchain/metalgo/utils/sampler" 9 ) 10 11 var _ Set = (*peerSet)(nil) 12 13 func NoPrecondition(Peer) bool { 14 return true 15 } 16 17 // Set contains a group of peers. 18 type Set interface { 19 // Add this peer to the set. 20 // 21 // If a peer with the same [peer.ID] is already in the set, then the new 22 // peer instance will replace the old peer instance. 23 // 24 // Add does not change the [peer.ID] returned from calls to [GetByIndex]. 25 Add(peer Peer) 26 27 // GetByID attempts to fetch a [peer] whose [peer.ID] is equal to [nodeID]. 28 // If no such peer exists in the set, then [false] will be returned. 29 GetByID(nodeID ids.NodeID) (Peer, bool) 30 31 // GetByIndex attempts to fetch a peer who has been allocated [index]. If 32 // [index] < 0 or [index] >= [Len], then false will be returned. 33 GetByIndex(index int) (Peer, bool) 34 35 // Remove any [peer] whose [peer.ID] is equal to [nodeID] from the set. 36 Remove(nodeID ids.NodeID) 37 38 // Len returns the number of peers currently in this set. 39 Len() int 40 41 // Sample attempts to return a random slice of peers with length [n]. The 42 // slice will not include any duplicates. Only peers that cause the 43 // [precondition] to return true will be returned in the slice. 44 Sample(n int, precondition func(Peer) bool) []Peer 45 46 // Returns information about all the peers. 47 AllInfo() []Info 48 49 // Info returns information about the requested peers if they are in the 50 // set. 51 Info(nodeIDs []ids.NodeID) []Info 52 } 53 54 type peerSet struct { 55 peersMap map[ids.NodeID]int // nodeID -> peer's index in peersSlice 56 peersSlice []Peer // invariant: len(peersSlice) == len(peersMap) 57 } 58 59 // NewSet returns a set that does not internally manage synchronization. 60 // 61 // Only [Add] and [Remove] require exclusion on the data structure. The 62 // remaining methods are safe for concurrent use. 63 func NewSet() Set { 64 return &peerSet{ 65 peersMap: make(map[ids.NodeID]int), 66 } 67 } 68 69 func (s *peerSet) Add(peer Peer) { 70 nodeID := peer.ID() 71 index, ok := s.peersMap[nodeID] 72 if !ok { 73 s.peersMap[nodeID] = len(s.peersSlice) 74 s.peersSlice = append(s.peersSlice, peer) 75 } else { 76 s.peersSlice[index] = peer 77 } 78 } 79 80 func (s *peerSet) GetByID(nodeID ids.NodeID) (Peer, bool) { 81 index, ok := s.peersMap[nodeID] 82 if !ok { 83 return nil, false 84 } 85 return s.peersSlice[index], true 86 } 87 88 func (s *peerSet) GetByIndex(index int) (Peer, bool) { 89 if index < 0 || index >= len(s.peersSlice) { 90 return nil, false 91 } 92 return s.peersSlice[index], true 93 } 94 95 func (s *peerSet) Remove(nodeID ids.NodeID) { 96 index, ok := s.peersMap[nodeID] 97 if !ok { 98 return 99 } 100 101 lastIndex := len(s.peersSlice) - 1 102 lastPeer := s.peersSlice[lastIndex] 103 lastPeerID := lastPeer.ID() 104 105 s.peersMap[lastPeerID] = index 106 s.peersSlice[index] = lastPeer 107 108 delete(s.peersMap, nodeID) 109 s.peersSlice[lastIndex] = nil 110 s.peersSlice = s.peersSlice[:lastIndex] 111 } 112 113 func (s *peerSet) Len() int { 114 return len(s.peersSlice) 115 } 116 117 func (s *peerSet) Sample(n int, precondition func(Peer) bool) []Peer { 118 if n <= 0 { 119 return nil 120 } 121 122 sampler := sampler.NewUniform() 123 sampler.Initialize(uint64(len(s.peersSlice))) 124 125 peers := make([]Peer, 0, n) 126 for len(peers) < n { 127 index, hasNext := sampler.Next() 128 if !hasNext { 129 // We have run out of peers to attempt to sample. 130 break 131 } 132 peer := s.peersSlice[index] 133 if !precondition(peer) { 134 continue 135 } 136 peers = append(peers, peer) 137 } 138 return peers 139 } 140 141 func (s *peerSet) AllInfo() []Info { 142 peerInfo := make([]Info, len(s.peersSlice)) 143 for i, peer := range s.peersSlice { 144 peerInfo[i] = peer.Info() 145 } 146 return peerInfo 147 } 148 149 func (s *peerSet) Info(nodeIDs []ids.NodeID) []Info { 150 peerInfo := make([]Info, 0, len(nodeIDs)) 151 for _, nodeID := range nodeIDs { 152 if peer, ok := s.GetByID(nodeID); ok { 153 peerInfo = append(peerInfo, peer.Info()) 154 } 155 } 156 return peerInfo 157 }