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/google/btree"
   8 )
   9
  10 func NewPieceOrder() *PieceRequestOrder {
  11         return &PieceRequestOrder{
  12                 tree: btree.New(32),
  13                 keys: make(map[PieceRequestOrderKey]PieceRequestOrderState),
  14         }
  15 }
  16
  17 type PieceRequestOrder struct {
  18         tree *btree.BTree
  19         keys map[PieceRequestOrderKey]PieceRequestOrderState
  20 }
  21
  22 type PieceRequestOrderKey struct {
  23         InfoHash metainfo.Hash
  24         Index    int
  25 }
  26
  27 type PieceRequestOrderState struct {
  28         Priority     piecePriority
  29         Partial      bool
  30         Availability int64
  31 }
  32
  33 type pieceRequestOrderItem struct {
  34         key   PieceRequestOrderKey
  35         state PieceRequestOrderState
  36 }
  37
  38 func (me pieceRequestOrderItem) Less(other btree.Item) bool {
  39         otherConcrete := other.(pieceRequestOrderItem)
  40         return pieceOrderLess(
  41                 pieceOrderInput{
  42                         PieceRequestOrderState: me.state,
  43                         PieceRequestOrderKey:   me.key,
  44                 },
  45                 pieceOrderInput{
  46                         PieceRequestOrderState: otherConcrete.state,
  47                         PieceRequestOrderKey:   otherConcrete.key,
  48                 },
  49         ).Less()
  50 }
  51
  52 func (me *PieceRequestOrder) Add(key PieceRequestOrderKey, state PieceRequestOrderState) {
  53         if _, ok := me.keys[key]; ok {
  54                 panic(key)
  55         }
  56         if me.tree.ReplaceOrInsert(pieceRequestOrderItem{
  57                 key:   key,
  58                 state: state,
  59         }) != nil {
  60                 panic("shouldn't already have this")
  61         }
  62         me.keys[key] = state
  63 }
  64
  65 func (me *PieceRequestOrder) Update(key PieceRequestOrderKey, state PieceRequestOrderState) {
  66         if me.tree.Delete(me.existingItemForKey(key)) == nil {
  67                 panic(fmt.Sprintf("%#v", key))
  68         }
  69         if me.tree.ReplaceOrInsert(pieceRequestOrderItem{
  70                 key:   key,
  71                 state: state,
  72         }) != nil {
  73                 panic(key)
  74         }
  75         me.keys[key] = state
  76 }
  77
  78 func (me *PieceRequestOrder) existingItemForKey(key PieceRequestOrderKey) pieceRequestOrderItem {
  79         return pieceRequestOrderItem{
  80                 key:   key,
  81                 state: me.keys[key],
  82         }
  83 }
  84
  85 func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) {
  86         if me.tree.Delete(me.existingItemForKey(key)) == nil {
  87                 panic(key)
  88         }
  89         delete(me.keys, key)
  90         // log.Printf("deleting %#v", key)
  91 }