forked from cometbft/cometbft
-
Notifications
You must be signed in to change notification settings - Fork 1
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Move mempoolIDs back to mempool/ids.go (cometbft#467)
- Loading branch information
Showing
4 changed files
with
104 additions
and
115 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,3 +1,71 @@ | ||
| package mempool | ||
|
|
||
| // These functions were moved into v0/reactor.go and v1/reactor.go | ||
| import ( | ||
| "fmt" | ||
|
|
||
| cmtsync "github.com/cometbft/cometbft/libs/sync" | ||
| "github.com/cometbft/cometbft/p2p" | ||
| ) | ||
|
|
||
| type mempoolIDs struct { | ||
| mtx cmtsync.RWMutex | ||
| peerMap map[p2p.ID]uint16 | ||
| nextID uint16 // assumes that a node will never have over 65536 active peers | ||
| activeIDs map[uint16]struct{} // used to check if a given peerID key is used, the value doesn't matter | ||
| } | ||
|
|
||
| // Reserve searches for the next unused ID and assigns it to the | ||
| // peer. | ||
| func (ids *mempoolIDs) ReserveForPeer(peer p2p.Peer) { | ||
| ids.mtx.Lock() | ||
| defer ids.mtx.Unlock() | ||
|
|
||
| curID := ids.nextPeerID() | ||
| ids.peerMap[peer.ID()] = curID | ||
| ids.activeIDs[curID] = struct{}{} | ||
| } | ||
|
|
||
| // nextPeerID returns the next unused peer ID to use. | ||
| // This assumes that ids's mutex is already locked. | ||
| func (ids *mempoolIDs) nextPeerID() uint16 { | ||
| if len(ids.activeIDs) == MaxActiveIDs { | ||
| panic(fmt.Sprintf("node has maximum %d active IDs and wanted to get one more", MaxActiveIDs)) | ||
| } | ||
|
|
||
| _, idExists := ids.activeIDs[ids.nextID] | ||
| for idExists { | ||
| ids.nextID++ | ||
| _, idExists = ids.activeIDs[ids.nextID] | ||
| } | ||
| curID := ids.nextID | ||
| ids.nextID++ | ||
| return curID | ||
| } | ||
|
|
||
| // Reclaim returns the ID reserved for the peer back to unused pool. | ||
| func (ids *mempoolIDs) Reclaim(peer p2p.Peer) { | ||
| ids.mtx.Lock() | ||
| defer ids.mtx.Unlock() | ||
|
|
||
| removedID, ok := ids.peerMap[peer.ID()] | ||
| if ok { | ||
| delete(ids.activeIDs, removedID) | ||
| delete(ids.peerMap, peer.ID()) | ||
| } | ||
| } | ||
|
|
||
| // GetForPeer returns an ID reserved for the peer. | ||
| func (ids *mempoolIDs) GetForPeer(peer p2p.Peer) uint16 { | ||
| ids.mtx.RLock() | ||
| defer ids.mtx.RUnlock() | ||
|
|
||
| return ids.peerMap[peer.ID()] | ||
| } | ||
|
|
||
| func newMempoolIDs() *mempoolIDs { | ||
| return &mempoolIDs{ | ||
| peerMap: make(map[p2p.ID]uint16), | ||
| activeIDs: map[uint16]struct{}{0: {}}, | ||
| nextID: 1, // reserve unknownPeerID(0) for mempoolReactor.BroadcastTx | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,23 +1,42 @@ | ||
| package mempool | ||
|
|
||
| // import ( | ||
| // "testing" | ||
| import ( | ||
| "net" | ||
| "testing" | ||
|
|
||
| // "github.com/stretchr/testify/require" | ||
| // "github.com/cometbft/cometbft/types" | ||
| // ) | ||
| "github.com/cometbft/cometbft/p2p/mock" | ||
| "github.com/stretchr/testify/assert" | ||
| ) | ||
|
|
||
| // func TestMempoolIDsBasic(t *testing.T) { | ||
| // ids := NewMempoolIDs() | ||
| func TestMempoolIDsBasic(t *testing.T) { | ||
| ids := newMempoolIDs() | ||
|
|
||
| // peerID, err := types.NewNodeID("0011223344556677889900112233445566778899") | ||
| // require.NoError(t, err) | ||
| peer := mock.NewPeer(net.IP{127, 0, 0, 1}) | ||
|
|
||
| // ids.ReserveForPeer(peerID) | ||
| // require.EqualValues(t, 1, ids.GetForPeer(peerID)) | ||
| // ids.Reclaim(peerID) | ||
| ids.ReserveForPeer(peer) | ||
| assert.EqualValues(t, 1, ids.GetForPeer(peer)) | ||
| ids.Reclaim(peer) | ||
|
|
||
| // ids.ReserveForPeer(peerID) | ||
| // require.EqualValues(t, 2, ids.GetForPeer(peerID)) | ||
| // ids.Reclaim(peerID) | ||
| // } | ||
| ids.ReserveForPeer(peer) | ||
| assert.EqualValues(t, 2, ids.GetForPeer(peer)) | ||
| ids.Reclaim(peer) | ||
| } | ||
|
|
||
| func TestMempoolIDsPanicsIfNodeRequestsOvermaxActiveIDs(t *testing.T) { | ||
| if testing.Short() { | ||
| return | ||
| } | ||
|
|
||
| // 0 is already reserved for UnknownPeerID | ||
| ids := newMempoolIDs() | ||
|
|
||
| for i := 0; i < MaxActiveIDs-1; i++ { | ||
| peer := mock.NewPeer(net.IP{127, 0, 0, 1}) | ||
| ids.ReserveForPeer(peer) | ||
| } | ||
|
|
||
| assert.Panics(t, func() { | ||
| peer := mock.NewPeer(net.IP{127, 0, 0, 1}) | ||
| ids.ReserveForPeer(peer) | ||
| }) | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters