request-strategy/piece-request-order.go

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