github.com/aakash4dev/cometbft@v0.38.2/mempool/ids.go (about)

     1  package mempool
     2  
     3  import (
     4  	"fmt"
     5  
     6  	cmtsync "github.com/aakash4dev/cometbft/libs/sync"
     7  	"github.com/aakash4dev/cometbft/p2p"
     8  )
     9  
    10  type mempoolIDs struct {
    11  	mtx       cmtsync.RWMutex
    12  	peerMap   map[p2p.ID]uint16
    13  	nextID    uint16              // assumes that a node will never have over 65536 active peers
    14  	activeIDs map[uint16]struct{} // used to check if a given peerID key is used, the value doesn't matter
    15  }
    16  
    17  // Reserve searches for the next unused ID and assigns it to the
    18  // peer.
    19  func (ids *mempoolIDs) ReserveForPeer(peer p2p.Peer) {
    20  	ids.mtx.Lock()
    21  	defer ids.mtx.Unlock()
    22  
    23  	curID := ids.nextPeerID()
    24  	ids.peerMap[peer.ID()] = curID
    25  	ids.activeIDs[curID] = struct{}{}
    26  }
    27  
    28  // nextPeerID returns the next unused peer ID to use.
    29  // This assumes that ids's mutex is already locked.
    30  func (ids *mempoolIDs) nextPeerID() uint16 {
    31  	if len(ids.activeIDs) == MaxActiveIDs {
    32  		panic(fmt.Sprintf("node has maximum %d active IDs and wanted to get one more", MaxActiveIDs))
    33  	}
    34  
    35  	_, idExists := ids.activeIDs[ids.nextID]
    36  	for idExists {
    37  		ids.nextID++
    38  		_, idExists = ids.activeIDs[ids.nextID]
    39  	}
    40  	curID := ids.nextID
    41  	ids.nextID++
    42  	return curID
    43  }
    44  
    45  // Reclaim returns the ID reserved for the peer back to unused pool.
    46  func (ids *mempoolIDs) Reclaim(peer p2p.Peer) {
    47  	ids.mtx.Lock()
    48  	defer ids.mtx.Unlock()
    49  
    50  	removedID, ok := ids.peerMap[peer.ID()]
    51  	if ok {
    52  		delete(ids.activeIDs, removedID)
    53  		delete(ids.peerMap, peer.ID())
    54  	}
    55  }
    56  
    57  // GetForPeer returns an ID reserved for the peer.
    58  func (ids *mempoolIDs) GetForPeer(peer p2p.Peer) uint16 {
    59  	ids.mtx.RLock()
    60  	defer ids.mtx.RUnlock()
    61  
    62  	return ids.peerMap[peer.ID()]
    63  }
    64  
    65  func newMempoolIDs() *mempoolIDs {
    66  	return &mempoolIDs{
    67  		peerMap:   make(map[p2p.ID]uint16),
    68  		activeIDs: map[uint16]struct{}{0: {}},
    69  		nextID:    1, // reserve unknownPeerID(0) for mempoolReactor.BroadcastTx
    70  	}
    71  }