peerconn.go | 10 ++++++---- peerconn_test.go | 2 +- piece.go | 40 ++++++++++++++++++++-------------------- roaring.go | 16 ++++++++++++++++ torrent.go | 18 +++++++++++++----- torrent_test.go | 4 +++- diff --git a/peerconn.go b/peerconn.go index 8d3f6d819f84fdd72c37c6d7539e8755fe744218..03749e0ef21ccbb782ed55a6b72561774c84874b 100644 --- a/peerconn.go +++ b/peerconn.go @@ -181,10 +181,12 @@ } if !cn.actualRequestState.Interested { return false } - if cn.peerAllowedFast.IterTyped(func(_i int) bool { - i := RequestIndex(_i) - return cn.actualRequestState.Requests.Rank((i+1)*cn.t.chunksPerRegularPiece())- - cn.actualRequestState.Requests.Rank(i*cn.t.chunksPerRegularPiece()) == 0 + if cn.peerAllowedFast.IterTyped(func(i int) bool { + return roaringBitmapRangeCardinality( + &cn.actualRequestState.Requests, + cn.t.pieceRequestIndexOffset(i), + cn.t.pieceRequestIndexOffset(i+1), + ) == 0 }) { return true } diff --git a/peerconn_test.go b/peerconn_test.go index 395dbb62bb1cc8e63f64b22d5eaf11f725ce4354..5e8cc1b56f3317eda9821545621b17a44ecd5631 100644 --- a/peerconn_test.go +++ b/peerconn_test.go @@ -128,7 +128,7 @@ for i := 0; i < b.N; i += 1 { cl.lock() // The chunk must be written to storage everytime, to ensure the // writeSem is unlocked. - t.pieces[0]._dirtyChunks.Clear() + t.pendAllChunkSpecs(0) cn.validReceiveChunks = map[RequestIndex]int{ t.requestIndexFromRequest(newRequestFromMessage(&msg)): 1, } diff --git a/piece.go b/piece.go index 30ac229787eb82fe594317b1510120b5c5929df1..7424bba403476c278ff6c558e18a21f9127b3afa 100644 --- a/piece.go +++ b/piece.go @@ -18,9 +18,6 @@ hash *metainfo.Hash t *Torrent index pieceIndex files []*File - // Chunks we've written to since the last check. The chunk offset and - // length can be determined by the request chunkSize in use. - _dirtyChunks bitmap.Bitmap readerCond chansync.BroadcastCond @@ -56,7 +53,7 @@ return p.t.storage.Piece(p.Info()) } func (p *Piece) pendingChunkIndex(chunkIndex chunkIndexType) bool { - return !p._dirtyChunks.Contains(bitmap.BitIndex(chunkIndex)) + return !p.chunkIndexDirty(chunkIndex) } func (p *Piece) pendingChunk(cs ChunkSpec, chunkSize pp.Integer) bool { @@ -64,24 +61,27 @@ return p.pendingChunkIndex(chunkIndexFromChunkSpec(cs, chunkSize)) } func (p *Piece) hasDirtyChunks() bool { - return p._dirtyChunks.Len() != 0 + return p.numDirtyChunks() != 0 } -func (p *Piece) numDirtyChunks() pp.Integer { - return pp.Integer(p._dirtyChunks.Len()) +func (p *Piece) numDirtyChunks() chunkIndexType { + return chunkIndexType(roaringBitmapRangeCardinality( + &p.t.dirtyChunks, + p.requestIndexOffset(), + p.t.pieceRequestIndexOffset(p.index+1))) } func (p *Piece) unpendChunkIndex(i chunkIndexType) { - p._dirtyChunks.Add(bitmap.BitIndex(i)) + p.t.dirtyChunks.Add(p.requestIndexOffset() + i) p.readerCond.Broadcast() } func (p *Piece) pendChunkIndex(i RequestIndex) { - p._dirtyChunks.Remove(bitmap.BitIndex(i)) + p.t.dirtyChunks.Remove(p.requestIndexOffset() + i) } -func (p *Piece) numChunks() pp.Integer { - return pp.Integer(p.t.pieceNumChunks(p.index)) +func (p *Piece) numChunks() chunkIndexType { + return p.t.pieceNumChunks(p.index) } func (p *Piece) incrementPendingWrites() { @@ -110,12 +110,12 @@ } p.pendingWritesMutex.Unlock() } -func (p *Piece) chunkIndexDirty(chunk pp.Integer) bool { - return p._dirtyChunks.Contains(bitmap.BitIndex(chunk)) +func (p *Piece) chunkIndexDirty(chunk chunkIndexType) bool { + return p.t.dirtyChunks.Contains(p.requestIndexOffset() + chunk) } -func (p *Piece) chunkIndexSpec(chunk pp.Integer) ChunkSpec { - return chunkIndexSpec(chunk, p.length(), p.chunkSize()) +func (p *Piece) chunkIndexSpec(chunk chunkIndexType) ChunkSpec { + return chunkIndexSpec(pp.Integer(chunk), p.length(), p.chunkSize()) } func (p *Piece) numDirtyBytes() (ret pp.Integer) { @@ -141,7 +141,7 @@ func (p *Piece) chunkSize() pp.Integer { return p.t.chunkSize } -func (p *Piece) lastChunkIndex() pp.Integer { +func (p *Piece) lastChunkIndex() chunkIndexType { return p.numChunks() - 1 } @@ -230,7 +230,7 @@ return } func (p *Piece) allChunksDirty() bool { - return p._dirtyChunks.Len() == bitmap.BitRange(p.numChunks()) + return p.numDirtyChunks() == p.numChunks() } func (p *Piece) State() PieceState { @@ -238,8 +238,8 @@ return p.t.PieceState(p.index) } func (p *Piece) iterUndirtiedChunks(f func(cs chunkIndexType)) { - for i := chunkIndexType(0); i < chunkIndexType(p.numChunks()); i++ { - if p.chunkIndexDirty(pp.Integer(i)) { + for i := chunkIndexType(0); i < p.numChunks(); i++ { + if p.chunkIndexDirty(i) { continue } f(i) @@ -247,5 +247,5 @@ } } func (p *Piece) requestIndexOffset() RequestIndex { - return RequestIndex(p.index) * p.t.chunksPerRegularPiece() + return p.t.pieceRequestIndexOffset(p.index) } diff --git a/roaring.go b/roaring.go new file mode 100644 index 0000000000000000000000000000000000000000..6e5409eca106073a85f9a392ebdc6e93db0db267 --- /dev/null +++ b/roaring.go @@ -0,0 +1,16 @@ +package torrent + +import ( + "github.com/RoaringBitmap/roaring" +) + +// Return the number of bits set in the range. To do this we need the rank of the item before the +// first, and the rank of the last item. An off-by-one minefield. Hopefully I haven't missed +// something in roaring's API that provides this. +func roaringBitmapRangeCardinality(bm *roaring.Bitmap, start, end uint32) (card uint64) { + card = bm.Rank(end - 1) + if start != 0 { + card -= bm.Rank(start - 1) + } + return +} diff --git a/torrent.go b/torrent.go index 370865e6bee24c974ec1883e79a23cfffbea8936..4e272048e08546d6d0c5828daf8830004be03e37 100644 --- a/torrent.go +++ b/torrent.go @@ -144,6 +144,8 @@ connPieceInclinationPool sync.Pool // Count of each request across active connections. pendingRequests map[RequestIndex]int + // Chunks we've written to since the corresponding piece was last checked. + dirtyChunks roaring.Bitmap pex pexState } @@ -851,7 +853,9 @@ return uint32((pp.Integer(t.usualPieceSize()) + t.chunkSize - 1) / t.chunkSize) } func (t *Torrent) pendAllChunkSpecs(pieceIndex pieceIndex) { - t.pieces[pieceIndex]._dirtyChunks.Clear() + t.dirtyChunks.RemoveRange( + uint64(t.pieceRequestIndexOffset(pieceIndex)), + uint64(t.pieceRequestIndexOffset(pieceIndex+1))) } func (t *Torrent) pieceLength(piece pieceIndex) pp.Integer { @@ -1037,11 +1041,11 @@ func (t *Torrent) pieceNumPendingChunks(piece pieceIndex) pp.Integer { if t.pieceComplete(piece) { return 0 } - return pp.Integer(t.pieceNumChunks(piece)) - t.pieces[piece].numDirtyChunks() + return pp.Integer(t.pieceNumChunks(piece) - t.pieces[piece].numDirtyChunks()) } func (t *Torrent) pieceAllDirty(piece pieceIndex) bool { - return t.pieces[piece]._dirtyChunks.Len() == bitmap.BitRange(t.pieceNumChunks(piece)) + return t.pieces[piece].allChunksDirty() } func (t *Torrent) readersChanged() { @@ -2267,14 +2271,18 @@ func (t *Torrent) requestIndexToRequest(ri RequestIndex) Request { index := ri / t.chunksPerRegularPiece() return Request{ pp.Integer(index), - t.piece(int(index)).chunkIndexSpec(pp.Integer(ri % t.chunksPerRegularPiece())), + t.piece(int(index)).chunkIndexSpec(ri % t.chunksPerRegularPiece()), } } func (t *Torrent) requestIndexFromRequest(r Request) RequestIndex { - return t.chunksPerRegularPiece()*uint32(r.Index) + uint32(r.Begin/t.chunkSize) + return t.pieceRequestIndexOffset(pieceIndex(r.Index)) + uint32(r.Begin/t.chunkSize) } func (t *Torrent) numChunks() RequestIndex { return RequestIndex((t.Length() + int64(t.chunkSize) - 1) / int64(t.chunkSize)) } + +func (t *Torrent) pieceRequestIndexOffset(piece pieceIndex) RequestIndex { + return RequestIndex(piece) * t.chunksPerRegularPiece() +} diff --git a/torrent_test.go b/torrent_test.go index bd762c9c009bc04c38b752b381fb9178540428c0..bd5db3d42ef1e175515e50f5a9631c14cd14aafd 100644 --- a/torrent_test.go +++ b/torrent_test.go @@ -148,7 +148,9 @@ tt := cl.newTorrent(mi.HashInfoBytes(), badStorage{}) tt.setChunkSize(2) require.NoError(t, tt.setInfoBytesLocked(mi.InfoBytes)) tt.cl.lock() - tt.pieces[1]._dirtyChunks.AddRange(0, 3) + tt.dirtyChunks.AddRange( + uint64(tt.pieceRequestIndexOffset(1)), + uint64(tt.pieceRequestIndexOffset(1)+3)) require.True(t, tt.pieceAllDirty(1)) tt.pieceHashed(1, false, nil) // Dirty chunks should be cleared so we can try again.