mirror of
https://github.com/fluencelabs/tendermint
synced 2025-06-24 18:21:38 +00:00
Refactor to move common libraries out of project
This commit is contained in:
Jae Kwon
227
p2p/peer_set.go
227
p2p/peer_set.go
@ -1,227 +0,0 @@
|
||||
package p2p
|
||||
|
||||
import (
|
||||
"net"
|
||||
"strings"
|
||||
"sync"
|
||||
)
|
||||
|
||||
// IPeerSet has a (immutable) subset of the methods of PeerSet.
|
||||
type IPeerSet interface {
|
||||
Has(key string) bool
|
||||
Get(key string) *Peer
|
||||
List() []*Peer
|
||||
Size() int
|
||||
}
|
||||
|
||||
//-----------------------------------------------------------------------------
|
||||
|
||||
var (
|
||||
maxPeersPerIPRange = [4]int{11, 7, 5, 3} // ...
|
||||
)
|
||||
|
||||
// PeerSet is a special structure for keeping a table of peers.
|
||||
// Iteration over the peers is super fast and thread-safe.
|
||||
// We also track how many peers per IP range and avoid too many
|
||||
type PeerSet struct {
|
||||
mtx sync.Mutex
|
||||
lookup map[string]*peerSetItem
|
||||
list []*Peer
|
||||
connectedIPs *nestedCounter
|
||||
}
|
||||
|
||||
type peerSetItem struct {
|
||||
peer *Peer
|
||||
index int
|
||||
}
|
||||
|
||||
func NewPeerSet() *PeerSet {
|
||||
return &PeerSet{
|
||||
lookup: make(map[string]*peerSetItem),
|
||||
list: make([]*Peer, 0, 256),
|
||||
connectedIPs: NewNestedCounter(),
|
||||
}
|
||||
}
|
||||
|
||||
// Returns false if peer with key (PubKeyEd25519) is already in set
|
||||
// or if we have too many peers from the peer's IP range
|
||||
func (ps *PeerSet) Add(peer *Peer) error {
|
||||
ps.mtx.Lock()
|
||||
defer ps.mtx.Unlock()
|
||||
if ps.lookup[peer.Key] != nil {
|
||||
return ErrSwitchDuplicatePeer
|
||||
}
|
||||
|
||||
// ensure we havent maxed out connections for the peer's IP range yet
|
||||
// and update the IP range counters
|
||||
if !ps.incrIPRangeCounts(peer.Host) {
|
||||
return ErrSwitchMaxPeersPerIPRange
|
||||
}
|
||||
|
||||
index := len(ps.list)
|
||||
// Appending is safe even with other goroutines
|
||||
// iterating over the ps.list slice.
|
||||
ps.list = append(ps.list, peer)
|
||||
ps.lookup[peer.Key] = &peerSetItem{peer, index}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (ps *PeerSet) Has(peerKey string) bool {
|
||||
ps.mtx.Lock()
|
||||
defer ps.mtx.Unlock()
|
||||
_, ok := ps.lookup[peerKey]
|
||||
return ok
|
||||
}
|
||||
|
||||
func (ps *PeerSet) Get(peerKey string) *Peer {
|
||||
ps.mtx.Lock()
|
||||
defer ps.mtx.Unlock()
|
||||
item, ok := ps.lookup[peerKey]
|
||||
if ok {
|
||||
return item.peer
|
||||
} else {
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
func (ps *PeerSet) Remove(peer *Peer) {
|
||||
ps.mtx.Lock()
|
||||
defer ps.mtx.Unlock()
|
||||
item := ps.lookup[peer.Key]
|
||||
if item == nil {
|
||||
return
|
||||
}
|
||||
|
||||
// update the IP range counters
|
||||
ps.decrIPRangeCounts(peer.Host)
|
||||
|
||||
index := item.index
|
||||
// Copy the list but without the last element.
|
||||
// (we must copy because we're mutating the list)
|
||||
newList := make([]*Peer, len(ps.list)-1)
|
||||
copy(newList, ps.list)
|
||||
// If it's the last peer, that's an easy special case.
|
||||
if index == len(ps.list)-1 {
|
||||
ps.list = newList
|
||||
delete(ps.lookup, peer.Key)
|
||||
return
|
||||
}
|
||||
|
||||
// Move the last item from ps.list to "index" in list.
|
||||
lastPeer := ps.list[len(ps.list)-1]
|
||||
lastPeerKey := lastPeer.Key
|
||||
lastPeerItem := ps.lookup[lastPeerKey]
|
||||
newList[index] = lastPeer
|
||||
lastPeerItem.index = index
|
||||
ps.list = newList
|
||||
delete(ps.lookup, peer.Key)
|
||||
|
||||
}
|
||||
|
||||
func (ps *PeerSet) Size() int {
|
||||
ps.mtx.Lock()
|
||||
defer ps.mtx.Unlock()
|
||||
return len(ps.list)
|
||||
}
|
||||
|
||||
// threadsafe list of peers.
|
||||
func (ps *PeerSet) List() []*Peer {
|
||||
ps.mtx.Lock()
|
||||
defer ps.mtx.Unlock()
|
||||
return ps.list
|
||||
}
|
||||
|
||||
//-----------------------------------------------------------------------------
|
||||
// track the number of IPs we're connected to for each IP address range
|
||||
|
||||
// forms an IP address hierarchy tree with counts
|
||||
// the struct itself is not thread safe and should always only be accessed with the ps.mtx locked
|
||||
type nestedCounter struct {
|
||||
count int
|
||||
children map[string]*nestedCounter
|
||||
}
|
||||
|
||||
func NewNestedCounter() *nestedCounter {
|
||||
nc := new(nestedCounter)
|
||||
nc.children = make(map[string]*nestedCounter)
|
||||
return nc
|
||||
}
|
||||
|
||||
// Check if we have too many IPs in the IP range of the incoming connection
|
||||
// Thread safe
|
||||
func (ps *PeerSet) HasMaxForIPRange(conn net.Conn) (ok bool) {
|
||||
ps.mtx.Lock()
|
||||
defer ps.mtx.Unlock()
|
||||
ip, _, _ := net.SplitHostPort(conn.RemoteAddr().String())
|
||||
ipBytes := strings.Split(ip, ".")
|
||||
|
||||
c := ps.connectedIPs
|
||||
for i, ipByte := range ipBytes {
|
||||
if c, ok = c.children[ipByte]; !ok {
|
||||
return false
|
||||
}
|
||||
if maxPeersPerIPRange[i] <= c.count {
|
||||
return true
|
||||
}
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// Increments counts for this address' IP range
|
||||
// Returns false if we already have enough connections
|
||||
// Not thread safe (only called by ps.Add())
|
||||
func (ps *PeerSet) incrIPRangeCounts(address string) bool {
|
||||
addrParts := strings.Split(address, ".")
|
||||
|
||||
c := ps.connectedIPs
|
||||
return incrNestedCounters(c, addrParts, 0)
|
||||
}
|
||||
|
||||
// Recursively descend the IP hierarchy, checking if we have
|
||||
// max peers for each range and incrementing if not.
|
||||
// Returns false if incr failed because max peers reached for some range counter.
|
||||
func incrNestedCounters(c *nestedCounter, ipBytes []string, index int) bool {
|
||||
ipByte := ipBytes[index]
|
||||
child := c.children[ipByte]
|
||||
if child == nil {
|
||||
child = NewNestedCounter()
|
||||
c.children[ipByte] = child
|
||||
}
|
||||
if index+1 < len(ipBytes) {
|
||||
if !incrNestedCounters(child, ipBytes, index+1) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
if maxPeersPerIPRange[index] <= child.count {
|
||||
return false
|
||||
} else {
|
||||
child.count += 1
|
||||
return true
|
||||
}
|
||||
}
|
||||
|
||||
// Decrement counts for this address' IP range
|
||||
func (ps *PeerSet) decrIPRangeCounts(address string) {
|
||||
addrParts := strings.Split(address, ".")
|
||||
|
||||
c := ps.connectedIPs
|
||||
decrNestedCounters(c, addrParts, 0)
|
||||
}
|
||||
|
||||
// Recursively descend the IP hierarchy, decrementing by one.
|
||||
// If the counter is zero, deletes the child.
|
||||
func decrNestedCounters(c *nestedCounter, ipBytes []string, index int) {
|
||||
ipByte := ipBytes[index]
|
||||
child := c.children[ipByte]
|
||||
if child == nil {
|
||||
log.Error("p2p/peer_set decrNestedCounters encountered a missing child counter")
|
||||
return
|
||||
}
|
||||
if index+1 < len(ipBytes) {
|
||||
decrNestedCounters(child, ipBytes, index+1)
|
||||
}
|
||||
child.count -= 1
|
||||
if child.count <= 0 {
|
||||
delete(c.children, ipByte)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user