github.com/devwanda/aphelion-staking@v0.33.9/p2p/peer_set.go (about) 1 package p2p 2 3 import ( 4 "net" 5 "sync" 6 ) 7 8 // IPeerSet has a (immutable) subset of the methods of PeerSet. 9 type IPeerSet interface { 10 Has(key ID) bool 11 HasIP(ip net.IP) bool 12 Get(key ID) Peer 13 List() []Peer 14 Size() int 15 } 16 17 //----------------------------------------------------------------------------- 18 19 // PeerSet is a special structure for keeping a table of peers. 20 // Iteration over the peers is super fast and thread-safe. 21 type PeerSet struct { 22 mtx sync.Mutex 23 lookup map[ID]*peerSetItem 24 list []Peer 25 } 26 27 type peerSetItem struct { 28 peer Peer 29 index int 30 } 31 32 // NewPeerSet creates a new peerSet with a list of initial capacity of 256 items. 33 func NewPeerSet() *PeerSet { 34 return &PeerSet{ 35 lookup: make(map[ID]*peerSetItem), 36 list: make([]Peer, 0, 256), 37 } 38 } 39 40 // Add adds the peer to the PeerSet. 41 // It returns an error carrying the reason, if the peer is already present. 42 func (ps *PeerSet) Add(peer Peer) error { 43 ps.mtx.Lock() 44 defer ps.mtx.Unlock() 45 46 if ps.lookup[peer.ID()] != nil { 47 return ErrSwitchDuplicatePeerID{peer.ID()} 48 } 49 50 index := len(ps.list) 51 // Appending is safe even with other goroutines 52 // iterating over the ps.list slice. 53 ps.list = append(ps.list, peer) 54 ps.lookup[peer.ID()] = &peerSetItem{peer, index} 55 return nil 56 } 57 58 // Has returns true if the set contains the peer referred to by this 59 // peerKey, otherwise false. 60 func (ps *PeerSet) Has(peerKey ID) bool { 61 ps.mtx.Lock() 62 _, ok := ps.lookup[peerKey] 63 ps.mtx.Unlock() 64 return ok 65 } 66 67 // HasIP returns true if the set contains the peer referred to by this IP 68 // address, otherwise false. 69 func (ps *PeerSet) HasIP(peerIP net.IP) bool { 70 ps.mtx.Lock() 71 defer ps.mtx.Unlock() 72 73 return ps.hasIP(peerIP) 74 } 75 76 // hasIP does not acquire a lock so it can be used in public methods which 77 // already lock. 78 func (ps *PeerSet) hasIP(peerIP net.IP) bool { 79 for _, item := range ps.lookup { 80 if item.peer.RemoteIP().Equal(peerIP) { 81 return true 82 } 83 } 84 85 return false 86 } 87 88 // Get looks up a peer by the provided peerKey. Returns nil if peer is not 89 // found. 90 func (ps *PeerSet) Get(peerKey ID) Peer { 91 ps.mtx.Lock() 92 defer ps.mtx.Unlock() 93 item, ok := ps.lookup[peerKey] 94 if ok { 95 return item.peer 96 } 97 return nil 98 } 99 100 // Remove discards peer by its Key, if the peer was previously memoized. 101 // Returns true if the peer was removed, and false if it was not found. 102 // in the set. 103 func (ps *PeerSet) Remove(peer Peer) bool { 104 ps.mtx.Lock() 105 defer ps.mtx.Unlock() 106 107 item := ps.lookup[peer.ID()] 108 if item == nil { 109 return false 110 } 111 112 index := item.index 113 // Create a new copy of the list but with one less item. 114 // (we must copy because we'll be mutating the list). 115 newList := make([]Peer, len(ps.list)-1) 116 copy(newList, ps.list) 117 // If it's the last peer, that's an easy special case. 118 if index == len(ps.list)-1 { 119 ps.list = newList 120 delete(ps.lookup, peer.ID()) 121 return true 122 } 123 124 // Replace the popped item with the last item in the old list. 125 lastPeer := ps.list[len(ps.list)-1] 126 lastPeerKey := lastPeer.ID() 127 lastPeerItem := ps.lookup[lastPeerKey] 128 newList[index] = lastPeer 129 lastPeerItem.index = index 130 ps.list = newList 131 delete(ps.lookup, peer.ID()) 132 return true 133 } 134 135 // Size returns the number of unique items in the peerSet. 136 func (ps *PeerSet) Size() int { 137 ps.mtx.Lock() 138 defer ps.mtx.Unlock() 139 return len(ps.list) 140 } 141 142 // List returns the threadsafe list of peers. 143 func (ps *PeerSet) List() []Peer { 144 ps.mtx.Lock() 145 defer ps.mtx.Unlock() 146 return ps.list 147 }