diff --git a/peerconn.go b/peerconn.go index 6851d1d6..2c309bd7 100644 --- a/peerconn.go +++ b/peerconn.go @@ -13,6 +13,7 @@ import ( "strings" "time" + "github.com/RoaringBitmap/roaring" "github.com/anacrolix/log" "github.com/anacrolix/missinggo/iter" "github.com/anacrolix/missinggo/v2/bitmap" @@ -110,7 +111,7 @@ type Peer struct { PeerPrefersEncryption bool // as indicated by 'e' field in extension handshake PeerListenPort int // 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 // we may not even know the number of pieces in the torrent yet. peerSentHaveAll bool @@ -238,7 +239,7 @@ func (cn *Peer) peerHasAllPieces() (all bool, known bool) { if !cn.t.haveInfo() { 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 { @@ -259,7 +260,7 @@ func (cn *Peer) bestPeerNumPieces() pieceIndex { } func (cn *Peer) completedString() string { - have := pieceIndex(cn._peerPieces.Len()) + have := pieceIndex(cn._peerPieces.GetCardinality()) if cn.peerSentHaveAll { have = cn.bestPeerNumPieces() } @@ -758,7 +759,7 @@ func (cn *PeerConn) peerSentHave(piece pieceIndex) error { if !cn.peerHasPiece(piece) { cn.t.incPieceAvailability(piece) } - cn._peerPieces.Set(bitmap.BitIndex(piece), true) + cn._peerPieces.Add(uint32(piece)) cn.t.maybeDropMutuallyCompletePeer(&cn.Peer) if cn.updatePiecePriority(piece) { cn.updateRequests() @@ -789,7 +790,11 @@ func (cn *PeerConn) peerSentBitfield(bf []bool) error { 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() 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 // 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() defer cn.locker().RUnlock() return cn.newPeerPieces() } // Returns a new Bitmap that includes bits for all pieces the peer could have based on their claims. -func (cn *Peer) newPeerPieces() bitmap.Bitmap { - ret := cn._peerPieces.Copy() +func (cn *Peer) newPeerPieces() *roaring.Bitmap { + // TODO: Can we use copy on write? + ret := cn._peerPieces.Clone() if cn.peerSentHaveAll { if cn.t.haveInfo() { ret.AddRange(0, bitmap.BitRange(cn.t.numPieces())) diff --git a/test/transfer_test.go b/test/transfer_test.go index 71de58b4..b70cb46c 100644 --- a/test/transfer_test.go +++ b/test/transfer_test.go @@ -168,7 +168,7 @@ func testClientTransfer(t *testing.T, ps testClientTransferParams) { } foundSeeder := false for _, pc := range leecherPeerConns { - completed := pc.PeerPieces().Len() + completed := pc.PeerPieces().GetCardinality() t.Logf("peer conn %v has %v completed pieces", pc, completed) if completed == bitmap.BitRange(leecherTorrent.Info().NumPieces()) { foundSeeder = true diff --git a/torrent.go b/torrent.go index 3981651a..d184afd3 100644 --- a/torrent.go +++ b/torrent.go @@ -1408,8 +1408,8 @@ func (t *Torrent) decPeerPieceAvailability(p *Peer) { if !t.haveInfo() { return } - p.newPeerPieces().IterTyped(func(i int) bool { - p.t.decPieceAvailability(i) + p.newPeerPieces().Iterate(func(i uint32) bool { + p.t.decPieceAvailability(pieceIndex(i)) return true }) }