connection.go | 14 +++++++++++++- connection_test.go | 2 +- piece.go | 11 +++++++++-- torrent.go | 10 ++++++---- diff --git a/connection.go b/connection.go index f4bef805698e119c04e68ee23ab5986df9b139d8..794f59860ba79aa544d10fe29845a57dc46cc545 100644 --- a/connection.go +++ b/connection.go @@ -541,7 +541,19 @@ // And pieces that we already have. skip.Union(cn.t.completedPieces) // Return an iterator over the different priority classes, minus the skip // pieces. - return iterBitmapsDistinct(skip, now, readahead, cn.t.pendingPieces) + return iter.Chain( + iterBitmapsDistinct(skip, now, readahead), + func(cb iter.Callback) { + cn.t.pendingPieces.IterTyped(func(piece int) bool { + if skip.Contains(piece) { + return true + } + more := cb(piece) + skip.Add(piece) + return more + }) + }, + ) } // The connection should download highest priority pieces first, without any diff --git a/connection_test.go b/connection_test.go index 3f6dfd93b17b282c023adbdaff18750b977a8fef..f22d26b9ba5727a16c23361a7b60aff068bb70ab 100644 --- a/connection_test.go +++ b/connection_test.go @@ -97,7 +97,7 @@ pieceStateChanges: pubsub.NewPubSub(), } t.setChunkSize(defaultChunkSize) t.makePieces() - t.pendingPieces.Add(0) + t.pendingPieces.Set(0, PiecePriorityNormal.BitmapPriority()) r, w := io.Pipe() cn := &connection{ t: t, diff --git a/piece.go b/piece.go index d5130d7b734d20bb45b075fc3b1eb15de5ba9a93..71c6041787289e217d4db4a39057dc3a7a85adfb 100644 --- a/piece.go +++ b/piece.go @@ -21,13 +21,20 @@ *pp = maybe } } +// Priority for use in PriorityBitmap +func (me piecePriority) BitmapPriority() int { + return -int(me) +} + const ( PiecePriorityNone piecePriority = iota // Not wanted. PiecePriorityNormal // Wanted. + PiecePriorityHigh // Wanted a lot. PiecePriorityReadahead // May be required soon. - // Succeeds a piece where a read occurred. Currently the same as Now, apparently due to issues with caching. + // Succeeds a piece where a read occurred. Currently the same as Now, + // apparently due to issues with caching. PiecePriorityNext - PiecePriorityNow // A Reader is reading in this piece. + PiecePriorityNow // A Reader is reading in this piece. Highest urgency. ) type Piece struct { diff --git a/torrent.go b/torrent.go index 301e2d6d67c4c8406f099add71465e47a48832dd..32f66b4e7c53ea667694d260a18cbbbd30e9a5ae 100644 --- a/torrent.go +++ b/torrent.go @@ -16,6 +16,8 @@ "sync" "text/tabwriter" "time" + "github.com/anacrolix/missinggo/prioritybitmap" + "github.com/anacrolix/dht" "github.com/anacrolix/missinggo" "github.com/anacrolix/missinggo/bitmap" @@ -113,7 +115,7 @@ readerReadaheadPieces bitmap.Bitmap // The indexes of pieces we want with normal priority, that aren't // currently available. - pendingPieces bitmap.Bitmap + pendingPieces prioritybitmap.PriorityBitmap // A cache of completed piece indices. completedPieces bitmap.Bitmap // Pieces that need to be hashed. @@ -965,13 +967,13 @@ } if t.havePiece(piece) { return } - t.pendingPieces.Add(piece) + t.pendingPieces.Set(piece, PiecePriorityNormal.BitmapPriority()) t.updatePiecePriority(piece) } func (t *Torrent) unpendPieces(unpend bitmap.Bitmap) { - t.pendingPieces.Sub(unpend) - unpend.IterTyped(func(piece int) (again bool) { + unpend.IterTyped(func(piece int) (more bool) { + t.pendingPieces.Remove(piece) t.updatePiecePriority(piece) return true })