Use reusable roaring iterators

client.go

@@ -19,25 +19,24 @@ import (
 	"strings"
 	"time"
 
+	"github.com/davecgh/go-spew/spew"
+	"github.com/dustin/go-humanize"
+	gbtree "github.com/google/btree"
+	"github.com/pion/datachannel"
+	"golang.org/x/time/rate"
+
+	"github.com/anacrolix/chansync"
 	"github.com/anacrolix/chansync/events"
 	"github.com/anacrolix/dht/v2"
 	"github.com/anacrolix/dht/v2/krpc"
 	"github.com/anacrolix/generics"
+	. "github.com/anacrolix/generics"
 	"github.com/anacrolix/log"
 	"github.com/anacrolix/missinggo/perf"
 	"github.com/anacrolix/missinggo/v2"
 	"github.com/anacrolix/missinggo/v2/bitmap"
 	"github.com/anacrolix/missinggo/v2/pproffd"
 	"github.com/anacrolix/sync"
-	request_strategy "github.com/anacrolix/torrent/request-strategy"
-	"github.com/davecgh/go-spew/spew"
-	"github.com/dustin/go-humanize"
-	"github.com/google/btree"
-	"github.com/pion/datachannel"
-	"golang.org/x/time/rate"
-
-	"github.com/anacrolix/chansync"
-	. "github.com/anacrolix/generics"
 
 	"github.com/anacrolix/torrent/bencode"
 	"github.com/anacrolix/torrent/internal/limiter"
@@ -45,6 +44,7 @@ import (
 	"github.com/anacrolix/torrent/metainfo"
 	"github.com/anacrolix/torrent/mse"
 	pp "github.com/anacrolix/torrent/peer_protocol"
+	request_strategy "github.com/anacrolix/torrent/request-strategy"
 	"github.com/anacrolix/torrent/storage"
 	"github.com/anacrolix/torrent/tracker"
 	"github.com/anacrolix/torrent/webtorrent"
@@ -1179,7 +1179,7 @@ func (cl *Client) newTorrentOpt(opts AddTorrentOpts) (t *Torrent) {
 		cl:       cl,
 		infoHash: opts.InfoHash,
 		peers: prioritizedPeers{
-			om: btree.New(32),
+			om: gbtree.New(32),
 			getPrio: func(p PeerInfo) peerPriority {
 				ipPort := p.addr()
 				return bep40PriorityIgnoreError(cl.publicAddr(ipPort.IP), ipPort)

go.mod

@@ -4,7 +4,7 @@ go 1.18
 
 require (
 	crawshaw.io/sqlite v0.3.3-0.20210127221821-98b1f83c5508
-	github.com/RoaringBitmap/roaring v0.9.4
+	github.com/RoaringBitmap/roaring v1.0.1-0.20220510143707-3f418c4f42a4
 	github.com/ajwerner/btree v0.0.0-20211221152037-f427b3e689c0
 	github.com/alexflint/go-arg v1.4.2
 	github.com/anacrolix/args v0.5.1-0.20220509024600-c3b77d0b61ac

go.sum

@@ -11,8 +11,8 @@ github.com/Julusian/godocdown v0.0.0-20170816220326-6d19f8ff2df8/go.mod h1:INZr5
 github.com/RoaringBitmap/roaring v0.4.7/go.mod h1:8khRDP4HmeXns4xIj9oGrKSz7XTQiJx2zgh7AcNke4w=
 github.com/RoaringBitmap/roaring v0.4.17/go.mod h1:D3qVegWTmfCaX4Bl5CrBE9hfrSrrXIr8KVNvRsDi1NI=
 github.com/RoaringBitmap/roaring v0.4.23/go.mod h1:D0gp8kJQgE1A4LQ5wFLggQEyvDi06Mq5mKs52e1TwOo=
-github.com/RoaringBitmap/roaring v0.9.4 h1:ckvZSX5gwCRaJYBNe7syNawCU5oruY9gQmjXlp4riwo=
+github.com/RoaringBitmap/roaring v1.0.1-0.20220510143707-3f418c4f42a4 h1:LxR70Si8X8IlYTeKiPReOBt9pla6wEJDlkawfwDdrB0=
-github.com/RoaringBitmap/roaring v0.9.4/go.mod h1:icnadbWcNyfEHlYdr+tDlOTih1Bf/h+rzPpv4sbomAA=
+github.com/RoaringBitmap/roaring v1.0.1-0.20220510143707-3f418c4f42a4/go.mod h1:icnadbWcNyfEHlYdr+tDlOTih1Bf/h+rzPpv4sbomAA=
 github.com/Shopify/sarama v1.19.0/go.mod h1:FVkBWblsNy7DGZRfXLU0O9RCGt5g3g3yEuWXgklEdEo=
 github.com/Shopify/toxiproxy v2.1.4+incompatible/go.mod h1:OXgGpZ6Cli1/URJOF1DMxUHB2q5Ap20/P/eIdh4G0pI=
 github.com/ajwerner/btree v0.0.0-20211221152037-f427b3e689c0 h1:byYvvbfSo3+9efR4IeReh77gVs4PnNDR3AMOE9NJ7a0=

piece.go

@@ -1,7 +1,6 @@
 package torrent
 
 import (
-	"encoding/gob"
 	"fmt"
 	"sync"
 
@@ -41,8 +40,6 @@ type Piece struct {
 	// Connections that have written data to this piece since its last check.
 	// This can include connections that have closed.
 	dirtiers map[*Peer]struct{}
-
-	undirtiedChunksIter undirtiedChunksIter
 }
 
 func (p *Piece) String() string {
@@ -244,10 +241,6 @@ func (p *Piece) State() PieceState {
 	return p.t.PieceState(p.index)
 }
 
-func init() {
-	gob.Register(undirtiedChunksIter{})
-}
-
 func (p *Piece) requestIndexOffset() RequestIndex {
 	return p.t.pieceRequestIndexOffset(p.index)
 }

requesting.go

@@ -11,9 +11,10 @@ import (
 
 	"github.com/anacrolix/log"
 	"github.com/anacrolix/multiless"
+	"github.com/anacrolix/torrent/typed-roaring"
 	"github.com/lispad/go-generics-tools/binheap"
 
-	request_strategy "github.com/anacrolix/torrent/request-strategy"
+	"github.com/anacrolix/torrent/request-strategy"
 )
 
 func (t *Torrent) requestStrategyPieceOrderState(i int) request_strategy.PieceRequestOrderState {
@@ -195,6 +196,8 @@ func (p *Peer) getDesiredRequestState() (desired desiredRequestState) {
 		pieceStates:    t.requestPieceStates,
 		requestIndexes: t.requestIndexes,
 	}
+	// Caller-provided allocation for roaring bitmap iteration.
+	var it typedRoaring.Iterator[RequestIndex]
 	request_strategy.GetRequestablePieces(
 		input,
 		t.getPieceRequestOrder(),
@@ -207,8 +210,7 @@ func (p *Peer) getDesiredRequestState() (desired desiredRequestState) {
 			}
 			requestHeap.pieceStates[pieceIndex] = pieceExtra
 			allowedFast := p.peerAllowedFast.Contains(pieceIndex)
-			p.t.piece(pieceIndex).undirtiedChunksIter.Iter(func(ci request_strategy.ChunkIndex) {
+			t.iterUndirtiedRequestIndexesInPiece(&it, pieceIndex, func(r request_strategy.RequestIndex) {
-				r := p.t.pieceRequestIndexOffset(pieceIndex) + ci
 				if !allowedFast {
 					// We must signal interest to request this. TODO: We could set interested if the
 					// peers pieces (minus the allowed fast set) overlap with our missing pieces if

torrent.go

@@ -390,11 +390,6 @@ func (t *Torrent) makePieces() {
 		beginFile := pieceFirstFileIndex(piece.torrentBeginOffset(), files)
 		endFile := pieceEndFileIndex(piece.torrentEndOffset(), files)
 		piece.files = files[beginFile:endFile]
-		piece.undirtiedChunksIter = undirtiedChunksIter{
-			TorrentDirtyChunks: &t.dirtyChunks,
-			StartRequestIndex:  piece.requestIndexOffset(),
-			EndRequestIndex:    piece.requestIndexOffset() + piece.numChunks(),
-		}
 	}
 }
 
@@ -2513,6 +2508,20 @@ func (t *Torrent) pieceIndexOfRequestIndex(ri RequestIndex) pieceIndex {
 	return pieceIndex(ri / t.chunksPerRegularPiece())
 }
 
+func (t *Torrent) iterUndirtiedRequestIndexesInPiece(
+	reuseIter *typedRoaring.Iterator[RequestIndex],
+	piece pieceIndex,
+	f func(RequestIndex),
+) {
+	reuseIter.Initialize(&t.dirtyChunks)
+	pieceRequestIndexOffset := t.pieceRequestIndexOffset(piece)
+	iterBitmapUnsetInRange(
+		reuseIter,
+		pieceRequestIndexOffset, pieceRequestIndexOffset+t.pieceNumChunks(piece),
+		f,
+	)
+}
+
 type requestState struct {
 	peer *Peer
 	when time.Time

typed-roaring/bitmap.go

@@ -42,6 +42,7 @@ func (me *Bitmap[T]) Remove(x T) {
 	me.Bitmap.Remove(uint32(x))
 }
 
-func (me *Bitmap[T]) Iterator() Iterator[T] {
+// Returns an uninitialized iterator for the type of the receiver.
-	return Iterator[T]{me.Bitmap.Iterator()}
+func (Bitmap[T]) IteratorType() Iterator[T] {
+	return Iterator[T]{}
 }

typed-roaring/iterator.go

@@ -5,13 +5,17 @@ import (
 )
 
 type Iterator[T BitConstraint] struct {
-	roaring.IntPeekable
+	roaring.IntIterator
 }
 
-func (t Iterator[T]) Next() T {
+func (t *Iterator[T]) Next() T {
-	return T(t.IntPeekable.Next())
+	return T(t.IntIterator.Next())
 }
 
-func (t Iterator[T]) AdvanceIfNeeded(minVal T) {
+func (t *Iterator[T]) AdvanceIfNeeded(minVal T) {
-	t.IntPeekable.AdvanceIfNeeded(uint32(minVal))
+	t.IntIterator.AdvanceIfNeeded(uint32(minVal))
+}
+
+func (t *Iterator[T]) Initialize(a *Bitmap[T]) {
+	t.IntIterator.Initialize(&a.Bitmap)
 }

undirtied-chunks-iter.go

@@ -4,31 +4,20 @@ import (
 	"github.com/anacrolix/torrent/typed-roaring"
 )
 
-// Use an iterator to jump between dirty bits.
+func iterBitmapUnsetInRange[T typedRoaring.BitConstraint](it *typedRoaring.Iterator[T], start, end T, f func(T)) {
-type undirtiedChunksIter struct {
+	it.AdvanceIfNeeded(start)
-	TorrentDirtyChunks *typedRoaring.Bitmap[RequestIndex]
+	lastDirty := start - 1
-	StartRequestIndex  RequestIndex
-	EndRequestIndex    RequestIndex
-}
-
-func (me *undirtiedChunksIter) Iter(f func(chunkIndexType)) {
-	it := me.TorrentDirtyChunks.Iterator()
-	startIndex := me.StartRequestIndex
-	endIndex := me.EndRequestIndex
-	it.AdvanceIfNeeded(startIndex)
-	lastDirty := startIndex - 1
 	for it.HasNext() {
 		next := it.Next()
-		if next >= endIndex {
+		if next >= end {
 			break
 		}
 		for index := lastDirty + 1; index < next; index++ {
-			f(index - startIndex)
+			f(index)
 		}
 		lastDirty = next
 	}
-	for index := lastDirty + 1; index < endIndex; index++ {
+	for index := lastDirty + 1; index < end; index++ {
-		f(index - startIndex)
+		f(index)
 	}
-	return
 }

undirtied-chunks-iter_test.go

@@ -6,17 +6,14 @@ import (
 	typedRoaring "github.com/anacrolix/torrent/typed-roaring"
 )
 
-func BenchmarkUndirtiedChunksIter(b *testing.B) {
+func BenchmarkIterUndirtiedRequestIndexesInPiece(b *testing.B) {
 	var bitmap typedRoaring.Bitmap[RequestIndex]
-	a := undirtiedChunksIter{
+	it := bitmap.IteratorType()
-		TorrentDirtyChunks: &bitmap,
-		StartRequestIndex:  69,
-		EndRequestIndex:    420,
-	}
 	b.ReportAllocs()
 	for i := 0; i < b.N; i++ {
-		a.Iter(func(chunkIndex chunkIndexType) {
+		// This is the worst case, when Torrent.iterUndirtiedRequestIndexesInPiece can't find a
-
+		// usable cached iterator. This should be the only allocation.
-		})
+		it.Initialize(&bitmap)
+		iterBitmapUnsetInRange(&it, 69, 420, func(RequestIndex) {})
 	}
 }