github.com/badrootd/celestia-core@v0.0.0-20240305091328-aa4207a4b25d/mempool/cat/peers.go (about)

     1  package cat
     2  
     3  import (
     4  	"fmt"
     5  
     6  	tmsync "github.com/badrootd/celestia-core/libs/sync"
     7  	"github.com/badrootd/celestia-core/mempool"
     8  	"github.com/badrootd/celestia-core/p2p"
     9  )
    10  
    11  const firstPeerID = mempool.UnknownPeerID + 1
    12  
    13  // mempoolIDs is a thread-safe map of peer IDs to shorter uint16 IDs used by the Reactor for tracking peer
    14  // messages and peer state such as what transactions peers have seen
    15  type mempoolIDs struct {
    16  	mtx       tmsync.RWMutex
    17  	peerMap   map[p2p.ID]uint16   // quick lookup table for peer ID to short ID
    18  	nextID    uint16              // assumes that a node will never have over 65536 active peers
    19  	activeIDs map[uint16]p2p.Peer // used to check if a given peerID key is used, the value doesn't matter
    20  }
    21  
    22  func newMempoolIDs() *mempoolIDs {
    23  	return &mempoolIDs{
    24  		peerMap:   make(map[p2p.ID]uint16),
    25  		activeIDs: make(map[uint16]p2p.Peer),
    26  		nextID:    firstPeerID, // reserve unknownPeerID(0) for mempoolReactor.BroadcastTx
    27  	}
    28  }
    29  
    30  // ReserveForPeer searches for the next unused ID and assigns it to the
    31  // peer.
    32  func (ids *mempoolIDs) ReserveForPeer(peer p2p.Peer) {
    33  	ids.mtx.Lock()
    34  	defer ids.mtx.Unlock()
    35  
    36  	if _, ok := ids.peerMap[peer.ID()]; ok {
    37  		panic("duplicate peer added to mempool")
    38  	}
    39  
    40  	curID := ids.nextPeerID()
    41  	ids.peerMap[peer.ID()] = curID
    42  	ids.activeIDs[curID] = peer
    43  }
    44  
    45  // nextPeerID returns the next unused peer ID to use.
    46  // This assumes that ids's mutex is already locked.
    47  func (ids *mempoolIDs) nextPeerID() uint16 {
    48  	if len(ids.activeIDs) == mempool.MaxActiveIDs {
    49  		panic(fmt.Sprintf("node has maximum %d active IDs and wanted to get one more", mempool.MaxActiveIDs))
    50  	}
    51  
    52  	_, idExists := ids.activeIDs[ids.nextID]
    53  	for idExists {
    54  		ids.nextID++
    55  		_, idExists = ids.activeIDs[ids.nextID]
    56  	}
    57  	curID := ids.nextID
    58  	ids.nextID++
    59  	return curID
    60  }
    61  
    62  // Reclaim returns the ID reserved for the peer back to unused pool.
    63  func (ids *mempoolIDs) Reclaim(peerID p2p.ID) uint16 {
    64  	ids.mtx.Lock()
    65  	defer ids.mtx.Unlock()
    66  
    67  	removedID, ok := ids.peerMap[peerID]
    68  	if ok {
    69  		delete(ids.activeIDs, removedID)
    70  		delete(ids.peerMap, peerID)
    71  		return removedID
    72  	}
    73  	return 0
    74  }
    75  
    76  // GetIDForPeer returns the shorthand ID reserved for the peer.
    77  func (ids *mempoolIDs) GetIDForPeer(peerID p2p.ID) uint16 {
    78  	ids.mtx.RLock()
    79  	defer ids.mtx.RUnlock()
    80  
    81  	id, exists := ids.peerMap[peerID]
    82  	if !exists {
    83  		return 0
    84  	}
    85  	return id
    86  }
    87  
    88  // GetPeer returns the peer for the given shorthand ID.
    89  func (ids *mempoolIDs) GetPeer(id uint16) p2p.Peer {
    90  	ids.mtx.RLock()
    91  	defer ids.mtx.RUnlock()
    92  
    93  	return ids.activeIDs[id]
    94  }
    95  
    96  // GetAll returns all active peers.
    97  func (ids *mempoolIDs) GetAll() map[uint16]p2p.Peer {
    98  	ids.mtx.RLock()
    99  	defer ids.mtx.RUnlock()
   100  
   101  	// make a copy of the map.
   102  	peers := make(map[uint16]p2p.Peer, len(ids.activeIDs))
   103  	for id, peer := range ids.activeIDs {
   104  		peers[id] = peer
   105  	}
   106  	return peers
   107  }
   108  
   109  // Len returns the number of active peers.
   110  func (ids *mempoolIDs) Len() int {
   111  	ids.mtx.RLock()
   112  	defer ids.mtx.RUnlock()
   113  
   114  	return len(ids.activeIDs)
   115  }