p2p/peer_set.go

   1 package p2p
   2
   3 import (
   4         "sync"
   5 )
   6
   7 // IPeerSet has a (immutable) subset of the methods of PeerSet.
   8 type IPeerSet interface {
   9         Has(key string) bool
  10         Get(key string) *Peer
  11         List() []*Peer
  12         Size() int
  13 }
  14
  15 //-----------------------------------------------------------------------------
  16
  17 // PeerSet is a special structure for keeping a table of peers.
  18 // Iteration over the peers is super fast and thread-safe.
  19 type PeerSet struct {
  20         mtx    sync.Mutex
  21         lookup map[string]*peerSetItem
  22         list   []*Peer
  23 }
  24
  25 type peerSetItem struct {
  26         peer  *Peer
  27         index int
  28 }
  29
  30 // NewPeerSet creates a new peerSet with a list of initial capacity of 256 items.
  31 func NewPeerSet() *PeerSet {
  32         return &PeerSet{
  33                 lookup: make(map[string]*peerSetItem),
  34                 list:   make([]*Peer, 0, 256),
  35         }
  36 }
  37
  38 // Add adds the peer to the PeerSet.
  39 // Returns false if peer with key (PubKeyEd25519) is already set
  40 func (ps *PeerSet) Add(peer *Peer) error {
  41         ps.mtx.Lock()
  42         defer ps.mtx.Unlock()
  43
  44         if ps.lookup[peer.Key] != nil {
  45                 return ErrDuplicatePeer
  46         }
  47
  48         ps.lookup[peer.Key] = &peerSetItem{peer, len(ps.list)}
  49         ps.list = append(ps.list, peer)
  50         return nil
  51 }
  52
  53 func (ps *PeerSet) DoFilter(ip string, pubKey string) error {
  54         if ps.Has(pubKey) {
  55                 return ErrDuplicatePeer
  56         }
  57
  58         return nil
  59 }
  60
  61 // Get looks up a peer by the provided peerKey.
  62 func (ps *PeerSet) Get(peerKey string) *Peer {
  63         ps.mtx.Lock()
  64         defer ps.mtx.Unlock()
  65         item, ok := ps.lookup[peerKey]
  66         if ok {
  67                 return item.peer
  68         }
  69         return nil
  70 }
  71
  72 // Has returns true if the PeerSet contains
  73 // the peer referred to by this peerKey.
  74 func (ps *PeerSet) Has(peerKey string) bool {
  75         ps.mtx.Lock()
  76         defer ps.mtx.Unlock()
  77         _, ok := ps.lookup[peerKey]
  78         return ok
  79 }
  80
  81 // List threadsafe list of peers.
  82 func (ps *PeerSet) List() []*Peer {
  83         ps.mtx.Lock()
  84         defer ps.mtx.Unlock()
  85         return ps.list
  86 }
  87
  88 // Remove discards peer if the peer was previously memoized.
  89 func (ps *PeerSet) Remove(peer *Peer) {
  90         ps.mtx.Lock()
  91         defer ps.mtx.Unlock()
  92         item := ps.lookup[peer.Key]
  93         if item == nil {
  94                 return
  95         }
  96
  97         index := item.index
  98         // Copy the list but without the last element.
  99         // (we must copy because we're mutating the list)
 100         newList := make([]*Peer, len(ps.list)-1)
 101         copy(newList, ps.list)
 102         // If it's the last peer, that's an easy special case.
 103         if index == len(ps.list)-1 {
 104                 ps.list = newList
 105                 delete(ps.lookup, peer.Key)
 106                 return
 107         }
 108
 109         // Move the last item from ps.list to "index" in list.
 110         lastPeer := ps.list[len(ps.list)-1]
 111         lastPeerKey := lastPeer.Key
 112         lastPeerItem := ps.lookup[lastPeerKey]
 113         newList[index] = lastPeer
 114         lastPeerItem.index = index
 115         ps.list = newList
 116         delete(ps.lookup, peer.Key)
 117 }
 118
 119 // Size returns the number of unique items in the peerSet.
 120 func (ps *PeerSet) Size() int {
 121         ps.mtx.Lock()
 122         defer ps.mtx.Unlock()
 123         return len(ps.list)
 124 }