github.com/number571/tendermint@v0.34.11-gost/internal/mempool/ids.go (about)

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