request-strategy/piece-request-order.go

   1 package request_strategy
   2
   3 import (
   4         "fmt"
   5
   6         "github.com/anacrolix/torrent/metainfo"
   7         "github.com/tidwall/btree"
   8 )
   9
  10 func NewPieceOrder() *PieceRequestOrder {
  11         return &PieceRequestOrder{
  12                 tree: btree.NewOptions(
  13                         func(a, b pieceRequestOrderItem) bool {
  14                                 return a.Less(&b)
  15                         },
  16                         btree.Options{NoLocks: true}),
  17                 keys: make(map[PieceRequestOrderKey]PieceRequestOrderState),
  18         }
  19 }
  20
  21 type PieceRequestOrder struct {
  22         tree     *btree.BTree[pieceRequestOrderItem]
  23         keys     map[PieceRequestOrderKey]PieceRequestOrderState
  24         PathHint *btree.PathHint
  25 }
  26
  27 type PieceRequestOrderKey struct {
  28         InfoHash metainfo.Hash
  29         Index    int
  30 }
  31
  32 type PieceRequestOrderState struct {
  33         Priority     piecePriority
  34         Partial      bool
  35         Availability int
  36 }
  37
  38 type pieceRequestOrderItem struct {
  39         key   PieceRequestOrderKey
  40         state PieceRequestOrderState
  41 }
  42
  43 func (me *pieceRequestOrderItem) Less(otherConcrete *pieceRequestOrderItem) bool {
  44         return pieceOrderLess(me, otherConcrete).Less()
  45 }
  46
  47 func (me *PieceRequestOrder) Add(key PieceRequestOrderKey, state PieceRequestOrderState) {
  48         if _, ok := me.keys[key]; ok {
  49                 panic(key)
  50         }
  51         if _, ok := me.tree.SetHint(pieceRequestOrderItem{
  52                 key:   key,
  53                 state: state,
  54         }, me.PathHint); ok {
  55                 panic("shouldn't already have this")
  56         }
  57         me.keys[key] = state
  58 }
  59
  60 type PieceRequestOrderPathHint = btree.PathHint
  61
  62 func (me *PieceRequestOrder) Update(
  63         key PieceRequestOrderKey,
  64         state PieceRequestOrderState,
  65 ) {
  66         oldState, ok := me.keys[key]
  67         if !ok {
  68                 panic("key should have been added already")
  69         }
  70         if state == oldState {
  71                 return
  72         }
  73         item := pieceRequestOrderItem{
  74                 key:   key,
  75                 state: oldState,
  76         }
  77         if _, ok := me.tree.DeleteHint(item, me.PathHint); !ok {
  78                 panic(fmt.Sprintf("%#v", key))
  79         }
  80         item.state = state
  81         if _, ok := me.tree.SetHint(item, me.PathHint); ok {
  82                 panic(key)
  83         }
  84         me.keys[key] = state
  85 }
  86
  87 func (me *PieceRequestOrder) existingItemForKey(key PieceRequestOrderKey) pieceRequestOrderItem {
  88         return pieceRequestOrderItem{
  89                 key:   key,
  90                 state: me.keys[key],
  91         }
  92 }
  93
  94 func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) {
  95         item := me.existingItemForKey(key)
  96         if _, ok := me.tree.DeleteHint(item, me.PathHint); !ok {
  97                 panic(key)
  98         }
  99         delete(me.keys, key)
 100 }
 101
 102 func (me *PieceRequestOrder) Len() int {
 103         return len(me.keys)
 104 }