github.com/devwanda/aphelion-staking@v0.33.9/p2p/peer_set.go (about)

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