Change Peer._peerPieces to use raw roaring Bitmap type

The wrapper type was from when roaring didn't support zero-alloc initialization.
This commit is contained in:
Matt Joiner 2021-10-05 17:48:34 +11:00
parent 719d5c6400
commit da1221dd50
3 changed files with 17 additions and 11 deletions

View File

@ -13,6 +13,7 @@ import (
"strings" "strings"
"time" "time"
"github.com/RoaringBitmap/roaring"
"github.com/anacrolix/log" "github.com/anacrolix/log"
"github.com/anacrolix/missinggo/iter" "github.com/anacrolix/missinggo/iter"
"github.com/anacrolix/missinggo/v2/bitmap" "github.com/anacrolix/missinggo/v2/bitmap"
@ -110,7 +111,7 @@ type Peer struct {
PeerPrefersEncryption bool // as indicated by 'e' field in extension handshake PeerPrefersEncryption bool // as indicated by 'e' field in extension handshake
PeerListenPort int PeerListenPort int
// The pieces the peer has claimed to have. // The pieces the peer has claimed to have.
_peerPieces bitmap.Bitmap _peerPieces roaring.Bitmap
// The peer has everything. This can occur due to a special message, when // The peer has everything. This can occur due to a special message, when
// we may not even know the number of pieces in the torrent yet. // we may not even know the number of pieces in the torrent yet.
peerSentHaveAll bool peerSentHaveAll bool
@ -238,7 +239,7 @@ func (cn *Peer) peerHasAllPieces() (all bool, known bool) {
if !cn.t.haveInfo() { if !cn.t.haveInfo() {
return false, false return false, false
} }
return bitmap.Flip(cn._peerPieces, 0, bitmap.BitRange(cn.t.numPieces())).IsEmpty(), true return roaring.Flip(&cn._peerPieces, 0, bitmap.BitRange(cn.t.numPieces())).IsEmpty(), true
} }
func (cn *PeerConn) locker() *lockWithDeferreds { func (cn *PeerConn) locker() *lockWithDeferreds {
@ -259,7 +260,7 @@ func (cn *Peer) bestPeerNumPieces() pieceIndex {
} }
func (cn *Peer) completedString() string { func (cn *Peer) completedString() string {
have := pieceIndex(cn._peerPieces.Len()) have := pieceIndex(cn._peerPieces.GetCardinality())
if cn.peerSentHaveAll { if cn.peerSentHaveAll {
have = cn.bestPeerNumPieces() have = cn.bestPeerNumPieces()
} }
@ -758,7 +759,7 @@ func (cn *PeerConn) peerSentHave(piece pieceIndex) error {
if !cn.peerHasPiece(piece) { if !cn.peerHasPiece(piece) {
cn.t.incPieceAvailability(piece) cn.t.incPieceAvailability(piece)
} }
cn._peerPieces.Set(bitmap.BitIndex(piece), true) cn._peerPieces.Add(uint32(piece))
cn.t.maybeDropMutuallyCompletePeer(&cn.Peer) cn.t.maybeDropMutuallyCompletePeer(&cn.Peer)
if cn.updatePiecePriority(piece) { if cn.updatePiecePriority(piece) {
cn.updateRequests() cn.updateRequests()
@ -789,7 +790,11 @@ func (cn *PeerConn) peerSentBitfield(bf []bool) error {
cn.t.decPieceAvailability(i) cn.t.decPieceAvailability(i)
} }
} }
cn._peerPieces.Set(bitmap.BitIndex(i), have) if have {
cn._peerPieces.Add(uint32(i))
} else {
cn._peerPieces.Remove(uint32(i))
}
} }
cn.peerPiecesChanged() cn.peerPiecesChanged()
return nil return nil
@ -1596,15 +1601,16 @@ func (l connectionTrust) Less(r connectionTrust) bool {
// Returns the pieces the peer could have based on their claims. If we don't know how many pieces // Returns the pieces the peer could have based on their claims. If we don't know how many pieces
// are in the torrent, it could be a very large range the peer has sent HaveAll. // are in the torrent, it could be a very large range the peer has sent HaveAll.
func (cn *PeerConn) PeerPieces() bitmap.Bitmap { func (cn *PeerConn) PeerPieces() *roaring.Bitmap {
cn.locker().RLock() cn.locker().RLock()
defer cn.locker().RUnlock() defer cn.locker().RUnlock()
return cn.newPeerPieces() return cn.newPeerPieces()
} }
// Returns a new Bitmap that includes bits for all pieces the peer could have based on their claims. // Returns a new Bitmap that includes bits for all pieces the peer could have based on their claims.
func (cn *Peer) newPeerPieces() bitmap.Bitmap { func (cn *Peer) newPeerPieces() *roaring.Bitmap {
ret := cn._peerPieces.Copy() // TODO: Can we use copy on write?
ret := cn._peerPieces.Clone()
if cn.peerSentHaveAll { if cn.peerSentHaveAll {
if cn.t.haveInfo() { if cn.t.haveInfo() {
ret.AddRange(0, bitmap.BitRange(cn.t.numPieces())) ret.AddRange(0, bitmap.BitRange(cn.t.numPieces()))

View File

@ -168,7 +168,7 @@ func testClientTransfer(t *testing.T, ps testClientTransferParams) {
} }
foundSeeder := false foundSeeder := false
for _, pc := range leecherPeerConns { for _, pc := range leecherPeerConns {
completed := pc.PeerPieces().Len() completed := pc.PeerPieces().GetCardinality()
t.Logf("peer conn %v has %v completed pieces", pc, completed) t.Logf("peer conn %v has %v completed pieces", pc, completed)
if completed == bitmap.BitRange(leecherTorrent.Info().NumPieces()) { if completed == bitmap.BitRange(leecherTorrent.Info().NumPieces()) {
foundSeeder = true foundSeeder = true

View File

@ -1408,8 +1408,8 @@ func (t *Torrent) decPeerPieceAvailability(p *Peer) {
if !t.haveInfo() { if !t.haveInfo() {
return return
} }
p.newPeerPieces().IterTyped(func(i int) bool { p.newPeerPieces().Iterate(func(i uint32) bool {
p.t.decPieceAvailability(i) p.t.decPieceAvailability(pieceIndex(i))
return true return true
}) })
} }