github.com/bytom/bytom@v1.1.2-0.20221014091027-bbcba3df6075/p2p/peer_set.go (about)

     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  }