status-go/vendor/github.com/anacrolix/torrent/request-strategy/piece-request-order.go

97 lines
2.0 KiB
Go

package request_strategy
import (
"fmt"
"github.com/anacrolix/torrent/metainfo"
"github.com/google/btree"
)
func NewPieceOrder() *PieceRequestOrder {
return &PieceRequestOrder{
tree: btree.New(32),
keys: make(map[PieceRequestOrderKey]PieceRequestOrderState),
}
}
type PieceRequestOrder struct {
tree *btree.BTree
keys map[PieceRequestOrderKey]PieceRequestOrderState
}
type PieceRequestOrderKey struct {
InfoHash metainfo.Hash
Index int
}
type PieceRequestOrderState struct {
Priority piecePriority
Partial bool
Availability int
}
type pieceRequestOrderItem struct {
key PieceRequestOrderKey
state PieceRequestOrderState
}
func (me *pieceRequestOrderItem) Less(other btree.Item) bool {
otherConcrete := other.(*pieceRequestOrderItem)
return pieceOrderLess(me, otherConcrete).Less()
}
func (me *PieceRequestOrder) Add(key PieceRequestOrderKey, state PieceRequestOrderState) {
if _, ok := me.keys[key]; ok {
panic(key)
}
if me.tree.ReplaceOrInsert(&pieceRequestOrderItem{
key: key,
state: state,
}) != nil {
panic("shouldn't already have this")
}
me.keys[key] = state
}
func (me *PieceRequestOrder) Update(key PieceRequestOrderKey, state PieceRequestOrderState) {
oldState, ok := me.keys[key]
if !ok {
panic("key should have been added already")
}
if state == oldState {
return
}
item := pieceRequestOrderItem{
key: key,
state: oldState,
}
if me.tree.Delete(&item) == nil {
panic(fmt.Sprintf("%#v", key))
}
item.state = state
if me.tree.ReplaceOrInsert(&item) != nil {
panic(key)
}
me.keys[key] = state
}
func (me *PieceRequestOrder) existingItemForKey(key PieceRequestOrderKey) pieceRequestOrderItem {
return pieceRequestOrderItem{
key: key,
state: me.keys[key],
}
}
func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) {
item := me.existingItemForKey(key)
if me.tree.Delete(&item) == nil {
panic(key)
}
delete(me.keys, key)
// log.Printf("deleting %#v", key)
}
func (me *PieceRequestOrder) Len() int {
return len(me.keys)
}