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.

75 lines
1.7 KiB

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