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.

148 lines
3.4 KiB

9 years ago
7 years ago
9 years ago
7 years ago
7 years ago
9 years ago
7 years ago
9 years ago
7 years ago
9 years ago
9 years ago
7 years ago
9 years ago
7 years ago
7 years ago
9 years ago
7 years ago
9 years ago
9 years ago
9 years ago
9 years ago
7 years ago
7 years ago
7 years ago
9 years ago
9 years ago
9 years ago
7 years ago
9 years ago
9 years ago
7 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
9 years ago
7 years ago
9 years ago
  1. package p2p
  2. import (
  3. "net"
  4. tmsync "github.com/tendermint/tendermint/libs/sync"
  5. )
  6. // IPeerSet has a (immutable) subset of the methods of PeerSet.
  7. type IPeerSet interface {
  8. Has(key NodeID) bool
  9. HasIP(ip net.IP) bool
  10. Get(key NodeID) Peer
  11. List() []Peer
  12. Size() int
  13. }
  14. //-----------------------------------------------------------------------------
  15. // PeerSet is a special structure for keeping a table of peers.
  16. // Iteration over the peers is super fast and thread-safe.
  17. type PeerSet struct {
  18. mtx tmsync.Mutex
  19. lookup map[NodeID]*peerSetItem
  20. list []Peer
  21. }
  22. type peerSetItem struct {
  23. peer Peer
  24. index int
  25. }
  26. // NewPeerSet creates a new peerSet with a list of initial capacity of 256 items.
  27. func NewPeerSet() *PeerSet {
  28. return &PeerSet{
  29. lookup: make(map[NodeID]*peerSetItem),
  30. list: make([]Peer, 0, 256),
  31. }
  32. }
  33. // Add adds the peer to the PeerSet.
  34. // It returns an error carrying the reason, if the peer is already present.
  35. func (ps *PeerSet) Add(peer Peer) error {
  36. ps.mtx.Lock()
  37. defer ps.mtx.Unlock()
  38. if ps.lookup[peer.ID()] != nil {
  39. return ErrSwitchDuplicatePeerID{peer.ID()}
  40. }
  41. index := len(ps.list)
  42. // Appending is safe even with other goroutines
  43. // iterating over the ps.list slice.
  44. ps.list = append(ps.list, peer)
  45. ps.lookup[peer.ID()] = &peerSetItem{peer, index}
  46. return nil
  47. }
  48. // Has returns true if the set contains the peer referred to by this
  49. // peerKey, otherwise false.
  50. func (ps *PeerSet) Has(peerKey NodeID) bool {
  51. ps.mtx.Lock()
  52. _, ok := ps.lookup[peerKey]
  53. ps.mtx.Unlock()
  54. return ok
  55. }
  56. // HasIP returns true if the set contains the peer referred to by this IP
  57. // address, otherwise false.
  58. func (ps *PeerSet) HasIP(peerIP net.IP) bool {
  59. ps.mtx.Lock()
  60. defer ps.mtx.Unlock()
  61. return ps.hasIP(peerIP)
  62. }
  63. // hasIP does not acquire a lock so it can be used in public methods which
  64. // already lock.
  65. func (ps *PeerSet) hasIP(peerIP net.IP) bool {
  66. for _, item := range ps.lookup {
  67. if item.peer.RemoteIP().Equal(peerIP) {
  68. return true
  69. }
  70. }
  71. return false
  72. }
  73. // Get looks up a peer by the provided peerKey. Returns nil if peer is not
  74. // found.
  75. func (ps *PeerSet) Get(peerKey NodeID) Peer {
  76. ps.mtx.Lock()
  77. defer ps.mtx.Unlock()
  78. item, ok := ps.lookup[peerKey]
  79. if ok {
  80. return item.peer
  81. }
  82. return nil
  83. }
  84. // Remove discards peer by its Key, if the peer was previously memoized.
  85. // Returns true if the peer was removed, and false if it was not found.
  86. // in the set.
  87. func (ps *PeerSet) Remove(peer Peer) bool {
  88. ps.mtx.Lock()
  89. defer ps.mtx.Unlock()
  90. item := ps.lookup[peer.ID()]
  91. if item == nil {
  92. return false
  93. }
  94. index := item.index
  95. // Create a new copy of the list but with one less item.
  96. // (we must copy because we'll be mutating the list).
  97. newList := make([]Peer, len(ps.list)-1)
  98. copy(newList, ps.list)
  99. // If it's the last peer, that's an easy special case.
  100. if index == len(ps.list)-1 {
  101. ps.list = newList
  102. delete(ps.lookup, peer.ID())
  103. return true
  104. }
  105. // Replace the popped item with the last item in the old list.
  106. lastPeer := ps.list[len(ps.list)-1]
  107. lastPeerKey := lastPeer.ID()
  108. lastPeerItem := ps.lookup[lastPeerKey]
  109. newList[index] = lastPeer
  110. lastPeerItem.index = index
  111. ps.list = newList
  112. delete(ps.lookup, peer.ID())
  113. return true
  114. }
  115. // Size returns the number of unique items in the peerSet.
  116. func (ps *PeerSet) Size() int {
  117. ps.mtx.Lock()
  118. defer ps.mtx.Unlock()
  119. return len(ps.list)
  120. }
  121. // List returns the threadsafe list of peers.
  122. func (ps *PeerSet) List() []Peer {
  123. ps.mtx.Lock()
  124. defer ps.mtx.Unlock()
  125. return ps.list
  126. }