github.com/MetalBlockchain/metalgo@v1.11.9/network/peer/set.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package peer
     5  
     6  import (
     7  	"github.com/MetalBlockchain/metalgo/ids"
     8  	"github.com/MetalBlockchain/metalgo/utils/sampler"
     9  )
    10  
    11  var _ Set = (*peerSet)(nil)
    12  
    13  func NoPrecondition(Peer) bool {
    14  	return true
    15  }
    16  
    17  // Set contains a group of peers.
    18  type Set interface {
    19  	// Add this peer to the set.
    20  	//
    21  	// If a peer with the same [peer.ID] is already in the set, then the new
    22  	// peer instance will replace the old peer instance.
    23  	//
    24  	// Add does not change the [peer.ID] returned from calls to [GetByIndex].
    25  	Add(peer Peer)
    26  
    27  	// GetByID attempts to fetch a [peer] whose [peer.ID] is equal to [nodeID].
    28  	// If no such peer exists in the set, then [false] will be returned.
    29  	GetByID(nodeID ids.NodeID) (Peer, bool)
    30  
    31  	// GetByIndex attempts to fetch a peer who has been allocated [index]. If
    32  	// [index] < 0 or [index] >= [Len], then false will be returned.
    33  	GetByIndex(index int) (Peer, bool)
    34  
    35  	// Remove any [peer] whose [peer.ID] is equal to [nodeID] from the set.
    36  	Remove(nodeID ids.NodeID)
    37  
    38  	// Len returns the number of peers currently in this set.
    39  	Len() int
    40  
    41  	// Sample attempts to return a random slice of peers with length [n]. The
    42  	// slice will not include any duplicates. Only peers that cause the
    43  	// [precondition] to return true will be returned in the slice.
    44  	Sample(n int, precondition func(Peer) bool) []Peer
    45  
    46  	// Returns information about all the peers.
    47  	AllInfo() []Info
    48  
    49  	// Info returns information about the requested peers if they are in the
    50  	// set.
    51  	Info(nodeIDs []ids.NodeID) []Info
    52  }
    53  
    54  type peerSet struct {
    55  	peersMap   map[ids.NodeID]int // nodeID -> peer's index in peersSlice
    56  	peersSlice []Peer             // invariant: len(peersSlice) == len(peersMap)
    57  }
    58  
    59  // NewSet returns a set that does not internally manage synchronization.
    60  //
    61  // Only [Add] and [Remove] require exclusion on the data structure. The
    62  // remaining methods are safe for concurrent use.
    63  func NewSet() Set {
    64  	return &peerSet{
    65  		peersMap: make(map[ids.NodeID]int),
    66  	}
    67  }
    68  
    69  func (s *peerSet) Add(peer Peer) {
    70  	nodeID := peer.ID()
    71  	index, ok := s.peersMap[nodeID]
    72  	if !ok {
    73  		s.peersMap[nodeID] = len(s.peersSlice)
    74  		s.peersSlice = append(s.peersSlice, peer)
    75  	} else {
    76  		s.peersSlice[index] = peer
    77  	}
    78  }
    79  
    80  func (s *peerSet) GetByID(nodeID ids.NodeID) (Peer, bool) {
    81  	index, ok := s.peersMap[nodeID]
    82  	if !ok {
    83  		return nil, false
    84  	}
    85  	return s.peersSlice[index], true
    86  }
    87  
    88  func (s *peerSet) GetByIndex(index int) (Peer, bool) {
    89  	if index < 0 || index >= len(s.peersSlice) {
    90  		return nil, false
    91  	}
    92  	return s.peersSlice[index], true
    93  }
    94  
    95  func (s *peerSet) Remove(nodeID ids.NodeID) {
    96  	index, ok := s.peersMap[nodeID]
    97  	if !ok {
    98  		return
    99  	}
   100  
   101  	lastIndex := len(s.peersSlice) - 1
   102  	lastPeer := s.peersSlice[lastIndex]
   103  	lastPeerID := lastPeer.ID()
   104  
   105  	s.peersMap[lastPeerID] = index
   106  	s.peersSlice[index] = lastPeer
   107  
   108  	delete(s.peersMap, nodeID)
   109  	s.peersSlice[lastIndex] = nil
   110  	s.peersSlice = s.peersSlice[:lastIndex]
   111  }
   112  
   113  func (s *peerSet) Len() int {
   114  	return len(s.peersSlice)
   115  }
   116  
   117  func (s *peerSet) Sample(n int, precondition func(Peer) bool) []Peer {
   118  	if n <= 0 {
   119  		return nil
   120  	}
   121  
   122  	sampler := sampler.NewUniform()
   123  	sampler.Initialize(uint64(len(s.peersSlice)))
   124  
   125  	peers := make([]Peer, 0, n)
   126  	for len(peers) < n {
   127  		index, hasNext := sampler.Next()
   128  		if !hasNext {
   129  			// We have run out of peers to attempt to sample.
   130  			break
   131  		}
   132  		peer := s.peersSlice[index]
   133  		if !precondition(peer) {
   134  			continue
   135  		}
   136  		peers = append(peers, peer)
   137  	}
   138  	return peers
   139  }
   140  
   141  func (s *peerSet) AllInfo() []Info {
   142  	peerInfo := make([]Info, len(s.peersSlice))
   143  	for i, peer := range s.peersSlice {
   144  		peerInfo[i] = peer.Info()
   145  	}
   146  	return peerInfo
   147  }
   148  
   149  func (s *peerSet) Info(nodeIDs []ids.NodeID) []Info {
   150  	peerInfo := make([]Info, 0, len(nodeIDs))
   151  	for _, nodeID := range nodeIDs {
   152  		if peer, ok := s.GetByID(nodeID); ok {
   153  			peerInfo = append(peerInfo, peer.Info())
   154  		}
   155  	}
   156  	return peerInfo
   157  }