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 // Get looks up a peer by the provided peerKey.
  54 func (ps *PeerSet) Get(peerKey string) *Peer {
  55         ps.mtx.Lock()
  56         defer ps.mtx.Unlock()
  57         item, ok := ps.lookup[peerKey]
  58         if ok {
  59                 return item.peer
  60         }
  61         return nil
  62 }
  63
  64 // Has returns true if the PeerSet contains
  65 // the peer referred to by this peerKey.
  66 func (ps *PeerSet) Has(peerKey string) bool {
  67         ps.mtx.Lock()
  68         defer ps.mtx.Unlock()
  69         _, ok := ps.lookup[peerKey]
  70         return ok
  71 }
  72
  73 // List threadsafe list of peers.
  74 func (ps *PeerSet) List() []*Peer {
  75         ps.mtx.Lock()
  76         defer ps.mtx.Unlock()
  77         return ps.list
  78 }
  79
  80 // Remove discards peer if the peer was previously memoized.
  81 func (ps *PeerSet) Remove(peer *Peer) {
  82         ps.mtx.Lock()
  83         defer ps.mtx.Unlock()
  84         item := ps.lookup[peer.Key]
  85         if item == nil {
  86                 return
  87         }
  88
  89         index := item.index
  90         // Copy the list but without the last element.
  91         // (we must copy because we're mutating the list)
  92         newList := make([]*Peer, len(ps.list)-1)
  93         copy(newList, ps.list)
  94         // If it's the last peer, that's an easy special case.
  95         if index == len(ps.list)-1 {
  96                 ps.list = newList
  97                 delete(ps.lookup, peer.Key)
  98                 return
  99         }
 100
 101         // Move the last item from ps.list to "index" in list.
 102         lastPeer := ps.list[len(ps.list)-1]
 103         lastPeerKey := lastPeer.Key
 104         lastPeerItem := ps.lookup[lastPeerKey]
 105         newList[index] = lastPeer
 106         lastPeerItem.index = index
 107         ps.list = newList
 108         delete(ps.lookup, peer.Key)
 109 }
 110
 111 // Size returns the number of unique items in the peerSet.
 112 func (ps *PeerSet) Size() int {
 113         ps.mtx.Lock()
 114         defer ps.mtx.Unlock()
 115         return len(ps.list)
 116 }