You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

77 lines
1.8 KiB

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