request-strategy/piece-request-order.go

   1 package requestStrategy
   2
   3 import "github.com/anacrolix/torrent/metainfo"
   4
   5 type Btree interface {
   6         Delete(pieceRequestOrderItem)
   7         Add(pieceRequestOrderItem)
   8         Scan(func(pieceRequestOrderItem) bool)
   9 }
  10
  11 func NewPieceOrder(btree Btree, cap int) *PieceRequestOrder {
  12         return &PieceRequestOrder{
  13                 tree: btree,
  14                 keys: make(map[PieceRequestOrderKey]PieceRequestOrderState, cap),
  15         }
  16 }
  17
  18 type PieceRequestOrder struct {
  19         tree Btree
  20         keys map[PieceRequestOrderKey]PieceRequestOrderState
  21 }
  22
  23 type PieceRequestOrderKey struct {
  24         InfoHash metainfo.Hash
  25         Index    int
  26 }
  27
  28 type PieceRequestOrderState struct {
  29         Priority     piecePriority
  30         Partial      bool
  31         Availability int
  32 }
  33
  34 type pieceRequestOrderItem struct {
  35         key   PieceRequestOrderKey
  36         state PieceRequestOrderState
  37 }
  38
  39 func (me *pieceRequestOrderItem) Less(otherConcrete *pieceRequestOrderItem) bool {
  40         return pieceOrderLess(me, otherConcrete).Less()
  41 }
  42
  43 func (me *PieceRequestOrder) Add(key PieceRequestOrderKey, state PieceRequestOrderState) {
  44         if _, ok := me.keys[key]; ok {
  45                 panic(key)
  46         }
  47         me.tree.Add(pieceRequestOrderItem{key, state})
  48         me.keys[key] = state
  49 }
  50
  51 func (me *PieceRequestOrder) Update(
  52         key PieceRequestOrderKey,
  53         state PieceRequestOrderState,
  54 ) {
  55         oldState, ok := me.keys[key]
  56         if !ok {
  57                 panic("key should have been added already")
  58         }
  59         if state == oldState {
  60                 return
  61         }
  62         me.tree.Delete(pieceRequestOrderItem{key, oldState})
  63         me.tree.Add(pieceRequestOrderItem{key, state})
  64         me.keys[key] = state
  65 }
  66
  67 func (me *PieceRequestOrder) existingItemForKey(key PieceRequestOrderKey) pieceRequestOrderItem {
  68         return pieceRequestOrderItem{
  69                 key:   key,
  70                 state: me.keys[key],
  71         }
  72 }
  73
  74 func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) {
  75         me.tree.Delete(pieceRequestOrderItem{key, me.keys[key]})
  76         delete(me.keys, key)
  77 }
  78
  79 func (me *PieceRequestOrder) Len() int {
  80         return len(me.keys)
  81 }