github.com/project-88388/tendermint-v0.34.14-terra.2@v1.0.0/p2p/peer_set.go (about)

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