internal/pieceordering/pieceordering.go

   1 // Implements ordering of torrent piece indices for such purposes as download
   2 // prioritization.
   3 package pieceordering
   4
   5 import (
   6         "math/rand"
   7
   8         "github.com/ryszard/goskiplist/skiplist"
   9 )
  10
  11 // Maintains piece integers by their ascending assigned keys.
  12 type Instance struct {
  13         // Contains the ascending priority keys. The keys contain a slice of piece
  14         // indices.
  15         sl *skiplist.SkipList
  16         // Maps from piece index back to its key, so that it can be remove
  17         // efficiently from the skip list.
  18         pieceKeys map[int]int
  19 }
  20
  21 func New() *Instance {
  22         return &Instance{
  23                 sl: skiplist.NewIntMap(),
  24         }
  25 }
  26
  27 // Add the piece with the given key. If the piece is already present, change
  28 // its key.
  29 func (me *Instance) SetPiece(piece, key int) {
  30         if existingKey, ok := me.pieceKeys[piece]; ok {
  31                 if existingKey == key {
  32                         return
  33                 }
  34                 me.removeKeyPiece(existingKey, piece)
  35         }
  36         var itemSl []int
  37         if exItem, ok := me.sl.Get(key); ok {
  38                 itemSl = exItem.([]int)
  39         }
  40         me.sl.Set(key, append(itemSl, piece))
  41         if me.pieceKeys == nil {
  42                 me.pieceKeys = make(map[int]int)
  43         }
  44         me.pieceKeys[piece] = key
  45         me.shuffleItem(key)
  46 }
  47
  48 // Shuffle the piece indices that share a given key.
  49 func (me *Instance) shuffleItem(key int) {
  50         _item, ok := me.sl.Get(key)
  51         if !ok {
  52                 return
  53         }
  54         item := _item.([]int)
  55         for i := range item {
  56                 j := i + rand.Intn(len(item)-i)
  57                 item[i], item[j] = item[j], item[i]
  58         }
  59         me.sl.Set(key, item)
  60 }
  61
  62 func (me *Instance) removeKeyPiece(key, piece int) {
  63         item, ok := me.sl.Get(key)
  64         if !ok {
  65                 panic("no item for key")
  66         }
  67         itemSl := item.([]int)
  68         for i, piece1 := range itemSl {
  69                 if piece1 == piece {
  70                         itemSl[i] = itemSl[len(itemSl)-1]
  71                         itemSl = itemSl[:len(itemSl)-1]
  72                         break
  73                 }
  74         }
  75         if len(itemSl) == 0 {
  76                 me.sl.Delete(key)
  77         } else {
  78                 me.sl.Set(key, itemSl)
  79         }
  80 }
  81
  82 func (me *Instance) DeletePiece(piece int) {
  83         key, ok := me.pieceKeys[piece]
  84         if !ok {
  85                 return
  86         }
  87         me.removeKeyPiece(key, piece)
  88         delete(me.pieceKeys, piece)
  89 }
  90
  91 // Returns the piece with the lowest key.
  92 func (me *Instance) First() Element {
  93         i := me.sl.SeekToFirst()
  94         if i == nil {
  95                 return nil
  96         }
  97         return &element{i, i.Value().([]int)}
  98 }
  99
 100 func (me *Instance) Empty() bool {
 101         return me.sl.Len() == 0
 102 }
 103
 104 type Element interface {
 105         Piece() int
 106         Next() Element
 107 }
 108
 109 type element struct {
 110         i  skiplist.Iterator
 111         sl []int
 112 }
 113
 114 func (e *element) Next() Element {
 115         e.sl = e.sl[1:]
 116         if len(e.sl) > 0 {
 117                 return e
 118         }
 119         ok := e.i.Next()
 120         if !ok {
 121                 return nil
 122         }
 123         e.sl = e.i.Value().([]int)
 124         return e
 125 }
 126
 127 func (e *element) Piece() int {
 128         return e.sl[0]
 129 }