github.com/anacrolix/torrent@v1.61.0/prioritized-peers.go (about)

     1  package torrent
     2  
     3  import (
     4  	"hash/maphash"
     5  
     6  	"github.com/anacrolix/multiless"
     7  	"github.com/google/btree"
     8  )
     9  
    10  // Peers are stored with their priority at insertion. Their priority may
    11  // change if our apparent IP changes, we don't currently handle that.
    12  type prioritizedPeersItem struct {
    13  	prio peerPriority
    14  	p    PeerInfo
    15  }
    16  
    17  var hashSeed = maphash.MakeSeed()
    18  
    19  func (me prioritizedPeersItem) addrHash() int64 {
    20  	var h maphash.Hash
    21  	h.SetSeed(hashSeed)
    22  	h.WriteString(me.p.Addr.String())
    23  	return int64(h.Sum64())
    24  }
    25  
    26  func (me prioritizedPeersItem) Less(than btree.Item) bool {
    27  	other := than.(prioritizedPeersItem)
    28  	return multiless.New().Bool(
    29  		me.p.Trusted, other.p.Trusted).Uint32(
    30  		me.prio, other.prio).Int64(
    31  		me.addrHash(), other.addrHash(),
    32  	).Less()
    33  }
    34  
    35  type prioritizedPeers struct {
    36  	om      *btree.BTree
    37  	getPrio func(PeerInfo) peerPriority
    38  }
    39  
    40  func (me *prioritizedPeers) Each(f func(PeerInfo)) {
    41  	me.om.Ascend(func(i btree.Item) bool {
    42  		f(i.(prioritizedPeersItem).p)
    43  		return true
    44  	})
    45  }
    46  
    47  func (me *prioritizedPeers) Len() int {
    48  	if me == nil || me.om == nil {
    49  		return 0
    50  	}
    51  	return me.om.Len()
    52  }
    53  
    54  // Returns true if a peer is replaced.
    55  func (me *prioritizedPeers) Add(p PeerInfo) bool {
    56  	return me.om.ReplaceOrInsert(prioritizedPeersItem{me.getPrio(p), p}) != nil
    57  }
    58  
    59  // Returns true if a peer is replaced.
    60  func (me *prioritizedPeers) AddReturningReplacedPeer(p PeerInfo) (ret PeerInfo, ok bool) {
    61  	item := me.om.ReplaceOrInsert(prioritizedPeersItem{me.getPrio(p), p})
    62  	if item == nil {
    63  		return
    64  	}
    65  	ret = item.(prioritizedPeersItem).p
    66  	ok = true
    67  	return
    68  }
    69  
    70  func (me *prioritizedPeers) DeleteMin() (ret prioritizedPeersItem, ok bool) {
    71  	i := me.om.DeleteMin()
    72  	if i == nil {
    73  		return
    74  	}
    75  	ret = i.(prioritizedPeersItem)
    76  	ok = true
    77  	return
    78  }
    79  
    80  func (me *prioritizedPeers) PopMax() PeerInfo {
    81  	return me.om.DeleteMax().(prioritizedPeersItem).p
    82  }