prioritized-peers.go

   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 }