Use roaring.Bitmap directly for completed pieces

Looking at improving the performance around this per https://github.com/anacrolix/torrent/discussions/547#discussion-3522317.
This commit is contained in:
Matt Joiner 2021-08-16 11:07:10 +10:00
parent cc0abf4fd2
commit 22c5a94a6a
5 changed files with 34 additions and 23 deletions

15
file.go
View File

@ -1,6 +1,7 @@
package torrent package torrent
import ( import (
"github.com/RoaringBitmap/roaring"
"github.com/anacrolix/missinggo/v2/bitmap" "github.com/anacrolix/missinggo/v2/bitmap"
"github.com/anacrolix/torrent/metainfo" "github.com/anacrolix/torrent/metainfo"
@ -59,32 +60,32 @@ func fileBytesLeft(
fileEndPieceIndex int, fileEndPieceIndex int,
fileTorrentOffset int64, fileTorrentOffset int64,
fileLength int64, fileLength int64,
torrentCompletedPieces bitmap.Bitmap, torrentCompletedPieces *roaring.Bitmap,
) (left int64) { ) (left int64) {
numPiecesSpanned := fileEndPieceIndex - fileFirstPieceIndex numPiecesSpanned := fileEndPieceIndex - fileFirstPieceIndex
switch numPiecesSpanned { switch numPiecesSpanned {
case 0: case 0:
case 1: case 1:
if !torrentCompletedPieces.Get(bitmap.BitIndex(fileFirstPieceIndex)) { if !torrentCompletedPieces.Contains(bitmap.BitIndex(fileFirstPieceIndex)) {
left += fileLength left += fileLength
} }
default: default:
if !torrentCompletedPieces.Get(bitmap.BitIndex(fileFirstPieceIndex)) { if !torrentCompletedPieces.Contains(bitmap.BitIndex(fileFirstPieceIndex)) {
left += torrentUsualPieceSize - (fileTorrentOffset % torrentUsualPieceSize) left += torrentUsualPieceSize - (fileTorrentOffset % torrentUsualPieceSize)
} }
if !torrentCompletedPieces.Get(bitmap.BitIndex(fileEndPieceIndex - 1)) { if !torrentCompletedPieces.Contains(bitmap.BitIndex(fileEndPieceIndex - 1)) {
left += fileTorrentOffset + fileLength - int64(fileEndPieceIndex-1)*torrentUsualPieceSize left += fileTorrentOffset + fileLength - int64(fileEndPieceIndex-1)*torrentUsualPieceSize
} }
completedMiddlePieces := torrentCompletedPieces.Copy() completedMiddlePieces := torrentCompletedPieces.Clone()
completedMiddlePieces.RemoveRange(0, bitmap.BitRange(fileFirstPieceIndex+1)) completedMiddlePieces.RemoveRange(0, bitmap.BitRange(fileFirstPieceIndex+1))
completedMiddlePieces.RemoveRange(bitmap.BitRange(fileEndPieceIndex-1), bitmap.ToEnd) completedMiddlePieces.RemoveRange(bitmap.BitRange(fileEndPieceIndex-1), bitmap.ToEnd)
left += int64(numPiecesSpanned-2-pieceIndex(completedMiddlePieces.Len())) * torrentUsualPieceSize left += int64(numPiecesSpanned-2-pieceIndex(completedMiddlePieces.GetCardinality())) * torrentUsualPieceSize
} }
return return
} }
func (f *File) bytesLeft() (left int64) { func (f *File) bytesLeft() (left int64) {
return fileBytesLeft(int64(f.t.usualPieceSize()), f.firstPieceIndex(), f.endPieceIndex(), f.offset, f.length, f.t._completedPieces) return fileBytesLeft(int64(f.t.usualPieceSize()), f.firstPieceIndex(), f.endPieceIndex(), f.offset, f.length, &f.t._completedPieces)
} }
// The relative file path for a multi-file torrent, and the torrent name for a // The relative file path for a multi-file torrent, and the torrent name for a

View File

@ -3,7 +3,7 @@ package torrent
import ( import (
"testing" "testing"
"github.com/anacrolix/missinggo/v2/bitmap" "github.com/RoaringBitmap/roaring"
"github.com/stretchr/testify/assert" "github.com/stretchr/testify/assert"
) )
@ -28,14 +28,14 @@ type testFileBytesLeft struct {
endPieceIndex int endPieceIndex int
fileOffset int64 fileOffset int64
fileLength int64 fileLength int64
completedPieces bitmap.Bitmap completedPieces roaring.Bitmap
expected int64 expected int64
name string name string
} }
func (me testFileBytesLeft) Run(t *testing.T) { func (me testFileBytesLeft) Run(t *testing.T) {
t.Run(me.name, func(t *testing.T) { t.Run(me.name, func(t *testing.T) {
assert.EqualValues(t, me.expected, fileBytesLeft(me.usualPieceSize, me.firstPieceIndex, me.endPieceIndex, me.fileOffset, me.fileLength, me.completedPieces)) assert.EqualValues(t, me.expected, fileBytesLeft(me.usualPieceSize, me.firstPieceIndex, me.endPieceIndex, me.fileOffset, me.fileLength, &me.completedPieces))
}) })
} }
@ -108,7 +108,7 @@ func TestFileBytesLeft(t *testing.T) {
fileOffset: 5, fileOffset: 5,
fileLength: 7, fileLength: 7,
expected: 0, expected: 0,
completedPieces: func() (ret bitmap.Bitmap) { completedPieces: func() (ret roaring.Bitmap) {
ret.AddRange(0, 5) ret.AddRange(0, 5)
return return
}(), }(),

View File

@ -646,7 +646,7 @@ func (cn *PeerConn) postBitfield() {
Type: pp.Bitfield, Type: pp.Bitfield,
Bitfield: cn.t.bitfield(), Bitfield: cn.t.bitfield(),
}) })
cn.sentHaves = cn.t._completedPieces.Copy() cn.sentHaves = bitmap.Bitmap{cn.t._completedPieces.Clone()}
} }
func (cn *PeerConn) updateRequests() { func (cn *PeerConn) updateRequests() {

View File

@ -18,6 +18,7 @@ import (
"time" "time"
"unsafe" "unsafe"
"github.com/RoaringBitmap/roaring"
"github.com/anacrolix/dht/v2" "github.com/anacrolix/dht/v2"
"github.com/anacrolix/log" "github.com/anacrolix/log"
"github.com/anacrolix/missinggo/iter" "github.com/anacrolix/missinggo/iter"
@ -131,7 +132,7 @@ type Torrent struct {
// file priorities and completion states elsewhere. // file priorities and completion states elsewhere.
_pendingPieces prioritybitmap.PriorityBitmap _pendingPieces prioritybitmap.PriorityBitmap
// A cache of completed piece indices. // A cache of completed piece indices.
_completedPieces bitmap.Bitmap _completedPieces roaring.Bitmap
// Pieces that need to be hashed. // Pieces that need to be hashed.
piecesQueuedForHash bitmap.Bitmap piecesQueuedForHash bitmap.Bitmap
activePieceHashes int activePieceHashes int
@ -252,7 +253,7 @@ func (t *Torrent) setChunkSize(size pp.Integer) {
} }
func (t *Torrent) pieceComplete(piece pieceIndex) bool { func (t *Torrent) pieceComplete(piece pieceIndex) bool {
return t._completedPieces.Get(bitmap.BitIndex(piece)) return t._completedPieces.Contains(bitmap.BitIndex(piece))
} }
func (t *Torrent) pieceCompleteUncached(piece pieceIndex) storage.Completion { func (t *Torrent) pieceCompleteUncached(piece pieceIndex) storage.Completion {
@ -755,9 +756,13 @@ func (t *Torrent) bytesMissingLocked() int64 {
return t.bytesLeft() return t.bytesLeft()
} }
func iterFlipped(b *roaring.Bitmap, end uint64, cb func(uint32) bool) {
roaring.Flip(b, 0, end).Iterate(cb)
}
func (t *Torrent) bytesLeft() (left int64) { func (t *Torrent) bytesLeft() (left int64) {
bitmap.Flip(t._completedPieces, 0, bitmap.BitRange(t.numPieces())).IterTyped(func(piece int) bool { iterFlipped(&t._completedPieces, uint64(t.numPieces()), func(x uint32) bool {
p := &t.pieces[piece] p := t.piece(pieceIndex(x))
left += int64(p.length() - p.numDirtyBytes()) left += int64(p.length() - p.numDirtyBytes())
return true return true
}) })
@ -792,7 +797,7 @@ func (t *Torrent) numPieces() pieceIndex {
} }
func (t *Torrent) numPiecesCompleted() (num pieceIndex) { func (t *Torrent) numPiecesCompleted() (num pieceIndex) {
return pieceIndex(t._completedPieces.Len()) return pieceIndex(t._completedPieces.GetCardinality())
} }
func (t *Torrent) close() (err error) { func (t *Torrent) close() (err error) {
@ -838,7 +843,7 @@ func (t *Torrent) writeChunk(piece int, begin int64, data []byte) (err error) {
func (t *Torrent) bitfield() (bf []bool) { func (t *Torrent) bitfield() (bf []bool) {
bf = make([]bool, t.numPieces()) bf = make([]bool, t.numPieces())
t._completedPieces.IterTyped(func(piece int) (again bool) { t._completedPieces.Iterate(func(piece uint32) (again bool) {
bf[piece] = true bf[piece] = true
return true return true
}) })
@ -895,14 +900,14 @@ func (t *Torrent) hashPiece(piece pieceIndex) (ret metainfo.Hash, err error) {
} }
func (t *Torrent) haveAnyPieces() bool { func (t *Torrent) haveAnyPieces() bool {
return t._completedPieces.Len() != 0 return t._completedPieces.GetCardinality() != 0
} }
func (t *Torrent) haveAllPieces() bool { func (t *Torrent) haveAllPieces() bool {
if !t.haveInfo() { if !t.haveInfo() {
return false return false
} }
return t._completedPieces.Len() == bitmap.BitRange(t.numPieces()) return t._completedPieces.GetCardinality() == bitmap.BitRange(t.numPieces())
} }
func (t *Torrent) havePiece(index pieceIndex) bool { func (t *Torrent) havePiece(index pieceIndex) bool {
@ -1238,7 +1243,12 @@ func (t *Torrent) updatePieceCompletion(piece pieceIndex) bool {
changed := cached != uncached changed := cached != uncached
complete := uncached.Complete complete := uncached.Complete
p.storageCompletionOk = uncached.Ok p.storageCompletionOk = uncached.Ok
t._completedPieces.Set(bitmap.BitIndex(piece), complete) x := uint32(piece)
if complete {
t._completedPieces.Add(x)
} else {
t._completedPieces.Remove(x)
}
if complete && len(p.dirtiers) != 0 { if complete && len(p.dirtiers) != 0 {
t.logger.Printf("marked piece %v complete but still has dirtiers", piece) t.logger.Printf("marked piece %v complete but still has dirtiers", piece)
} }
@ -1347,7 +1357,7 @@ func (t *Torrent) bytesCompleted() int64 {
if !t.haveInfo() { if !t.haveInfo() {
return 0 return 0
} }
return t.info.TotalLength() - t.bytesLeft() return *t.length - t.bytesLeft()
} }
func (t *Torrent) SetInfoBytes(b []byte) (err error) { func (t *Torrent) SetInfoBytes(b []byte) (err error) {

View File

@ -98,7 +98,7 @@ func BenchmarkUpdatePiecePriorities(b *testing.B) {
} }
assert.Len(b, t.readers, 7) assert.Len(b, t.readers, 7)
for i := 0; i < t.numPieces(); i += 3 { for i := 0; i < t.numPieces(); i += 3 {
t._completedPieces.Set(bitmap.BitIndex(i), true) t._completedPieces.Add(bitmap.BitIndex(i))
} }
t.DownloadPieces(0, t.numPieces()) t.DownloadPieces(0, t.numPieces())
for range iter.N(b.N) { for range iter.N(b.N) {