peerconn.go | 22 ++++++++++++++-------- test/transfer_test.go | 2 +- torrent.go | 4 ++-- diff --git a/peerconn.go b/peerconn.go index 6851d1d602b62b887449121e570529826006f728..2c309bd7bb9bbd3b523a0895a8a464b0ed8ca6f2 100644 --- a/peerconn.go +++ b/peerconn.go @@ -13,6 +13,7 @@ "strconv" "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 @@ peerRequests map[Request]*peerRequestState 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 @@ } 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 @@ return cn.peerMinPieces } 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 @@ cn.raisePeerMinPieces(piece + 1) 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 @@ if pp.Contains(bitmap.BitIndex(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() return nil @@ -1596,15 +1601,16 @@ } // 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 71de58b44b7b53fe0cd1ebe61f62865ebaf06f79..b70cb46c7b1fe246a00a8802278e403467739fd3 100644 --- a/test/transfer_test.go +++ b/test/transfer_test.go @@ -168,7 +168,7 @@ assert.NotEmpty(t, leecherPeerConns) } 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 3981651a614f247cf36756043d481f85482979bf..d184afd3f1f5fe831c3b1bf0b6e4c0632b08bf3b 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 }) }