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         item := me.existingItemForKey(key)
  67         if item.state == state {
  68                 return
  69         }
  70         if me.tree.Delete(&item) == nil {
  71                 panic(fmt.Sprintf("%#v", key))
  72         }
  73         item.state = state
  74         if me.tree.ReplaceOrInsert(&item) != nil {
  75                 panic(key)
  76         }
  77         me.keys[key] = state
  78 }
  79
  80 func (me *PieceRequestOrder) existingItemForKey(key PieceRequestOrderKey) pieceRequestOrderItem {
  81         return pieceRequestOrderItem{
  82                 key:   key,
  83                 state: me.keys[key],
  84         }
  85 }
  86
  87 func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) {
  88         item := me.existingItemForKey(key)
  89         if me.tree.Delete(&item) == nil {
  90                 panic(key)
  91         }
  92         delete(me.keys, key)
  93         // log.Printf("deleting %#v", key)
  94 }