diff --git a/client.go b/client.go index 8a686767..9d0777e4 100644 --- a/client.go +++ b/client.go @@ -75,7 +75,7 @@ type Client struct { dopplegangerAddrs map[string]struct{} badPeerIPs map[string]struct{} torrents map[InfoHash]*Torrent - pieceRequestOrder map[storage.TorrentCapacity]*request_strategy.PieceRequestOrder + pieceRequestOrder map[interface{}]*request_strategy.PieceRequestOrder acceptLimiter map[ipStr]int dialRateLimiter *rate.Limiter diff --git a/request-strategy/piece-request-order.go b/request-strategy/piece-request-order.go index 0b1d6b4e..d906b835 100644 --- a/request-strategy/piece-request-order.go +++ b/request-strategy/piece-request-order.go @@ -92,3 +92,7 @@ func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) { delete(me.keys, key) // log.Printf("deleting %#v", key) } + +func (me *PieceRequestOrder) Len() int { + return len(me.keys) +} diff --git a/requesting.go b/requesting.go index aaaa4f9c..c1ffa697 100644 --- a/requesting.go +++ b/requesting.go @@ -160,7 +160,7 @@ func (p *Peer) getDesiredRequestState() (desired desiredRequestState) { } request_strategy.GetRequestablePieces( input, - p.t.cl.pieceRequestOrder[p.t.storage.Capacity], + p.t.getPieceRequestOrder(), func(ih InfoHash, pieceIndex int) { if ih != p.t.infoHash { return diff --git a/torrent-piece-request-order.go b/torrent-piece-request-order.go index e1ab9670..7b1ef978 100644 --- a/torrent-piece-request-order.go +++ b/torrent-piece-request-order.go @@ -1,7 +1,51 @@ package torrent +import ( + request_strategy "github.com/anacrolix/torrent/request-strategy" +) + func (t *Torrent) updatePieceRequestOrder(pieceIndex int) { - t.cl.pieceRequestOrder[t.storage.Capacity].Update( + t.cl.pieceRequestOrder[t.clientPieceRequestOrderKey()].Update( t.pieceRequestOrderKey(pieceIndex), t.requestStrategyPieceOrderState(pieceIndex)) } + +func (t *Torrent) clientPieceRequestOrderKey() interface{} { + if t.storage.Capacity == nil { + return t + } + return t.storage.Capacity +} + +func (t *Torrent) deletePieceRequestOrder() { + cpro := t.cl.pieceRequestOrder + key := t.clientPieceRequestOrderKey() + pro := cpro[key] + for i := 0; i < t.numPieces(); i++ { + pro.Delete(t.pieceRequestOrderKey(i)) + } + if pro.Len() == 0 { + delete(cpro, key) + } +} + +func (t *Torrent) initPieceRequestOrder() { + if t.cl.pieceRequestOrder == nil { + t.cl.pieceRequestOrder = make(map[interface{}]*request_strategy.PieceRequestOrder) + } + key := t.clientPieceRequestOrderKey() + cpro := t.cl.pieceRequestOrder + if cpro[key] == nil { + cpro[key] = request_strategy.NewPieceOrder() + } +} + +func (t *Torrent) addRequestOrderPiece(i int) { + t.cl.pieceRequestOrder[t.clientPieceRequestOrderKey()].Add( + t.pieceRequestOrderKey(i), + t.requestStrategyPieceOrderState(i)) +} + +func (t *Torrent) getPieceRequestOrder() *request_strategy.PieceRequestOrder { + return t.cl.pieceRequestOrder[t.clientPieceRequestOrderKey()] +} diff --git a/torrent.go b/torrent.go index db24603b..ade7b933 100644 --- a/torrent.go +++ b/torrent.go @@ -437,12 +437,7 @@ func (t *Torrent) pieceRequestOrderKey(i int) request_strategy.PieceRequestOrder // This seems to be all the follow-up tasks after info is set, that can't fail. func (t *Torrent) onSetInfo() { - if t.cl.pieceRequestOrder == nil { - t.cl.pieceRequestOrder = make(map[storage.TorrentCapacity]*request_strategy.PieceRequestOrder) - } - if t.cl.pieceRequestOrder[t.storage.Capacity] == nil { - t.cl.pieceRequestOrder[t.storage.Capacity] = request_strategy.NewPieceOrder() - } + t.initPieceRequestOrder() for i := range t.pieces { p := &t.pieces[i] // Need to add availability before updating piece completion, as that may result in conns @@ -451,9 +446,7 @@ func (t *Torrent) onSetInfo() { panic(p.availability) } p.availability = int64(t.pieceAvailabilityFromPeers(i)) - t.cl.pieceRequestOrder[t.storage.Capacity].Add( - t.pieceRequestOrderKey(i), - t.requestStrategyPieceOrderState(i)) + t.addRequestOrderPiece(i) t.updatePieceCompletion(pieceIndex(i)) if !t.initialPieceCheckDisabled && !p.storageCompletionOk { // t.logger.Printf("piece %s completion unknown, queueing check", p) @@ -818,12 +811,6 @@ func (t *Torrent) numPiecesCompleted() (num pieceIndex) { return pieceIndex(t._completedPieces.GetCardinality()) } -func (t *Torrent) deletePieceRequestOrder() { - for i := 0; i < t.numPieces(); i++ { - t.cl.pieceRequestOrder[t.storage.Capacity].Delete(t.pieceRequestOrderKey(i)) - } -} - func (t *Torrent) close(wg *sync.WaitGroup) (err error) { t.closed.Set() if t.storage != nil {