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 func NewPeerSet() *PeerSet {
  31         return &PeerSet{
  32                 lookup: make(map[string]*peerSetItem),
  33                 list:   make([]*Peer, 0, 256),
  34         }
  35 }
  36
  37 // Returns false if peer with key (PubKeyEd25519) is already set
  38 func (ps *PeerSet) Add(peer *Peer) error {
  39         ps.mtx.Lock()
  40         defer ps.mtx.Unlock()
  41         if ps.lookup[peer.Key] != nil {
  42                 return ErrSwitchDuplicatePeer
  43         }
  44
  45         index := len(ps.list)
  46         // Appending is safe even with other goroutines
  47         // iterating over the ps.list slice.
  48         ps.list = append(ps.list, peer)
  49         ps.lookup[peer.Key] = &peerSetItem{peer, index}
  50         return nil
  51 }
  52
  53 func (ps *PeerSet) Has(peerKey string) bool {
  54         ps.mtx.Lock()
  55         defer ps.mtx.Unlock()
  56         _, ok := ps.lookup[peerKey]
  57         return ok
  58 }
  59
  60 func (ps *PeerSet) Get(peerKey string) *Peer {
  61         ps.mtx.Lock()
  62         defer ps.mtx.Unlock()
  63         item, ok := ps.lookup[peerKey]
  64         if ok {
  65                 return item.peer
  66         } else {
  67                 return nil
  68         }
  69 }
  70
  71 func (ps *PeerSet) Remove(peer *Peer) {
  72         ps.mtx.Lock()
  73         defer ps.mtx.Unlock()
  74         item := ps.lookup[peer.Key]
  75         if item == nil {
  76                 return
  77         }
  78
  79         index := item.index
  80         // Copy the list but without the last element.
  81         // (we must copy because we're mutating the list)
  82         newList := make([]*Peer, len(ps.list)-1)
  83         copy(newList, ps.list)
  84         // If it's the last peer, that's an easy special case.
  85         if index == len(ps.list)-1 {
  86                 ps.list = newList
  87                 delete(ps.lookup, peer.Key)
  88                 return
  89         }
  90
  91         // Move the last item from ps.list to "index" in list.
  92         lastPeer := ps.list[len(ps.list)-1]
  93         lastPeerKey := lastPeer.Key
  94         lastPeerItem := ps.lookup[lastPeerKey]
  95         newList[index] = lastPeer
  96         lastPeerItem.index = index
  97         ps.list = newList
  98         delete(ps.lookup, peer.Key)
  99
 100 }
 101
 102 func (ps *PeerSet) Size() int {
 103         ps.mtx.Lock()
 104         defer ps.mtx.Unlock()
 105         return len(ps.list)
 106 }
 107
 108 // threadsafe list of peers.
 109 func (ps *PeerSet) List() []*Peer {
 110         ps.mtx.Lock()
 111         defer ps.mtx.Unlock()
 112         return ps.list
 113 }