Use roaring.Bitmap directly for completed pieces

Looking at improving the performance around this per https://github.com/anacrolix/torrent/discussions/547#discussion-3522317.

diff --git a/file.go b/file.go
index 60317a23ce2adf9931a1608fefb7380578c982d2..45defe6e87cbdc3714990da10f6d8383a95b6a6f 100644 (file)
--- a/file.go
+++ b/file.go
@@ -1,6 +1,7 @@
 package torrent
 
 import (
+       "github.com/RoaringBitmap/roaring"
        "github.com/anacrolix/missinggo/v2/bitmap"
 
        "github.com/anacrolix/torrent/metainfo"
@@ -59,32 +60,32 @@ func fileBytesLeft(
        fileEndPieceIndex int,
        fileTorrentOffset int64,
        fileLength int64,
-       torrentCompletedPieces bitmap.Bitmap,
+       torrentCompletedPieces *roaring.Bitmap,
 ) (left int64) {
        numPiecesSpanned := fileEndPieceIndex - fileFirstPieceIndex
        switch numPiecesSpanned {
        case 0:
        case 1:
-               if !torrentCompletedPieces.Get(bitmap.BitIndex(fileFirstPieceIndex)) {
+               if !torrentCompletedPieces.Contains(bitmap.BitIndex(fileFirstPieceIndex)) {
                        left += fileLength
                }
        default:
-               if !torrentCompletedPieces.Get(bitmap.BitIndex(fileFirstPieceIndex)) {
+               if !torrentCompletedPieces.Contains(bitmap.BitIndex(fileFirstPieceIndex)) {
                        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
                }
-               completedMiddlePieces := torrentCompletedPieces.Copy()
+               completedMiddlePieces := torrentCompletedPieces.Clone()
                completedMiddlePieces.RemoveRange(0, bitmap.BitRange(fileFirstPieceIndex+1))
                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
 }
 
 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

diff --git a/file_test.go b/file_test.go
index 0443cf3bd5cc63a9e9766f45a9d4a489f709c43f..3ad9aca7b0426d1102d5571a5cc91238ff58662a 100644 (file)
--- a/file_test.go
+++ b/file_test.go
@@ -3,7 +3,7 @@ package torrent
 import (
        "testing"
 
-       "github.com/anacrolix/missinggo/v2/bitmap"
+       "github.com/RoaringBitmap/roaring"
        "github.com/stretchr/testify/assert"
 )
 
@@ -28,14 +28,14 @@ type testFileBytesLeft struct {
        endPieceIndex   int
        fileOffset      int64
        fileLength      int64
-       completedPieces bitmap.Bitmap
+       completedPieces roaring.Bitmap
        expected        int64
        name            string
 }
 
 func (me testFileBytesLeft) Run(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,
                fileLength:      7,
                expected:        0,
-               completedPieces: func() (ret bitmap.Bitmap) {
+               completedPieces: func() (ret roaring.Bitmap) {
                        ret.AddRange(0, 5)
                        return
                }(),

diff --git a/peerconn.go b/peerconn.go
index dc9f9efbe3a60a7e6d13bad8fe7e549bc77ca4ed..00131ea897e94b024073ca8f23b6e2d12aab65c0 100644 (file)
--- a/peerconn.go
+++ b/peerconn.go
@@ -646,7 +646,7 @@ func (cn *PeerConn) postBitfield() {
                Type:     pp.Bitfield,
                Bitfield: cn.t.bitfield(),
        })
-       cn.sentHaves = cn.t._completedPieces.Copy()
+       cn.sentHaves = bitmap.Bitmap{cn.t._completedPieces.Clone()}
 }
 
 func (cn *PeerConn) updateRequests() {

diff --git a/torrent.go b/torrent.go
index 5c77aa395045425a9fbb9cd31c3d18c17b93e99a..e7184c2a19b567b949aa7ee896f109ab00d34dd2 100644 (file)
--- a/torrent.go
+++ b/torrent.go
@@ -18,6 +18,7 @@ import (
        "time"
        "unsafe"
 
+       "github.com/RoaringBitmap/roaring"
        "github.com/anacrolix/dht/v2"
        "github.com/anacrolix/log"
        "github.com/anacrolix/missinggo/iter"
@@ -131,7 +132,7 @@ type Torrent struct {
        // file priorities and completion states elsewhere.
        _pendingPieces prioritybitmap.PriorityBitmap
        // A cache of completed piece indices.
-       _completedPieces bitmap.Bitmap
+       _completedPieces roaring.Bitmap
        // Pieces that need to be hashed.
        piecesQueuedForHash bitmap.Bitmap
        activePieceHashes   int
@@ -252,7 +253,7 @@ func (t *Torrent) setChunkSize(size pp.Integer) {
 }
 
 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 {
@@ -755,9 +756,13 @@ func (t *Torrent) bytesMissingLocked() int64 {
        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) {
-       bitmap.Flip(t._completedPieces, 0, bitmap.BitRange(t.numPieces())).IterTyped(func(piece int) bool {
-               p := &t.pieces[piece]
+       iterFlipped(&t._completedPieces, uint64(t.numPieces()), func(x uint32) bool {
+               p := t.piece(pieceIndex(x))
                left += int64(p.length() - p.numDirtyBytes())
                return true
        })
@@ -792,7 +797,7 @@ func (t *Torrent) numPieces() pieceIndex {
 }
 
 func (t *Torrent) numPiecesCompleted() (num pieceIndex) {
-       return pieceIndex(t._completedPieces.Len())
+       return pieceIndex(t._completedPieces.GetCardinality())
 }
 
 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) {
        bf = make([]bool, t.numPieces())
-       t._completedPieces.IterTyped(func(piece int) (again bool) {
+       t._completedPieces.Iterate(func(piece uint32) (again bool) {
                bf[piece] = true
                return true
        })
@@ -895,14 +900,14 @@ func (t *Torrent) hashPiece(piece pieceIndex) (ret metainfo.Hash, err error) {
 }
 
 func (t *Torrent) haveAnyPieces() bool {
-       return t._completedPieces.Len() != 0
+       return t._completedPieces.GetCardinality() != 0
 }
 
 func (t *Torrent) haveAllPieces() bool {
        if !t.haveInfo() {
                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 {
@@ -1238,7 +1243,12 @@ func (t *Torrent) updatePieceCompletion(piece pieceIndex) bool {
        changed := cached != uncached
        complete := uncached.Complete
        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 {
                t.logger.Printf("marked piece %v complete but still has dirtiers", piece)
        }
@@ -1347,7 +1357,7 @@ func (t *Torrent) bytesCompleted() int64 {
        if !t.haveInfo() {
                return 0
        }
-       return t.info.TotalLength() - t.bytesLeft()
+       return *t.length - t.bytesLeft()
 }
 
 func (t *Torrent) SetInfoBytes(b []byte) (err error) {

diff --git a/torrent_test.go b/torrent_test.go
index f47f36c4d54715fb5caacab016150060e880e4bc..b413a39e8c928926cb6c0dfa60a6872540c9d25c 100644 (file)
--- a/torrent_test.go
+++ b/torrent_test.go
@@ -98,7 +98,7 @@ func BenchmarkUpdatePiecePriorities(b *testing.B) {
        }
        assert.Len(b, t.readers, 7)
        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())
        for range iter.N(b.N) {