Add AjwernerBtree as an alternative btree backing for piece request order

Also add some scans to benchmarks. Make a few changes to reduce allocations using piece request order.

go.mod

@@ -4,6 +4,7 @@ go 1.18
 
 require (
 	github.com/RoaringBitmap/roaring v0.9.4
+	github.com/ajwerner/btree v0.0.0-20211201061316-91c8b66ad617
 	github.com/alexflint/go-arg v1.4.2
 	github.com/anacrolix/args v0.4.1-0.20211104085705-59f0fe94eb8f
 	github.com/anacrolix/chansync v0.3.0
@@ -15,7 +16,7 @@ require (
 	github.com/anacrolix/missinggo v1.3.0
 	github.com/anacrolix/missinggo/perf v1.0.0
 	github.com/anacrolix/missinggo/v2 v2.5.2
-	github.com/anacrolix/multiless v0.2.0
+	github.com/anacrolix/multiless v0.2.1-0.20211218050420-533661eef5dc
 	github.com/anacrolix/squirrel v0.4.0
 	github.com/anacrolix/sync v0.4.0
 	github.com/anacrolix/tagflag v1.3.0

go.sum

@@ -12,6 +12,8 @@ github.com/RoaringBitmap/roaring v0.9.4 h1:ckvZSX5gwCRaJYBNe7syNawCU5oruY9gQmjXl
 github.com/RoaringBitmap/roaring v0.9.4/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-20211201061316-91c8b66ad617 h1:sxP5D87Mq99SZMHhYBmq1yY4AAVkfNY5Wn02B9crYs0=
+github.com/ajwerner/btree v0.0.0-20211201061316-91c8b66ad617/go.mod h1:q37NoqncT41qKc048STsifIt69LfUJ8SrWWcz/yam5k=
 github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
@@ -59,8 +61,8 @@ github.com/anacrolix/missinggo/v2 v2.5.2/go.mod h1:yNvsLrtZYRYCOI+KRH/JM8TodHjtI
 github.com/anacrolix/mmsg v0.0.0-20180515031531-a4a3ba1fc8bb/go.mod h1:x2/ErsYUmT77kezS63+wzZp8E3byYB0gzirM/WMBLfw=
 github.com/anacrolix/mmsg v1.0.0 h1:btC7YLjOn29aTUAExJiVUhQOuf/8rhm+/nWCMAnL3Hg=
 github.com/anacrolix/mmsg v1.0.0/go.mod h1:x8kRaJY/dCrY9Al0PEcj1mb/uFHwP6GCJ9fLl4thEPc=
-github.com/anacrolix/multiless v0.2.0 h1:HtGBBOQcHaJM59RP3ysITId7AMIgiNF4xJucaFh14Ms=
-github.com/anacrolix/multiless v0.2.0/go.mod h1:TrCLEZfIDbMVfLoQt5tOoiBS/uq4y8+ojuEVVvTNPX4=
+github.com/anacrolix/multiless v0.2.1-0.20211218050420-533661eef5dc h1:K047jUtd0Xv4SEpv/5DoBgDvj4ZNpT1SOVtMlFpRrh0=
+github.com/anacrolix/multiless v0.2.1-0.20211218050420-533661eef5dc/go.mod h1:TrCLEZfIDbMVfLoQt5tOoiBS/uq4y8+ojuEVVvTNPX4=
 github.com/anacrolix/squirrel v0.4.0 h1:MOxYwh0mD2rcEJT+N2tFj2Z3dBDrWn6kRxP4qsBxfyk=
 github.com/anacrolix/squirrel v0.4.0/go.mod h1:dJyE7VefQvX0KAKMkOQDGOVEs91a+LvXQ2BjEKl/DIw=
 github.com/anacrolix/stm v0.2.0/go.mod h1:zoVQRvSiGjGoTmbM0vSLIiaKjWtNPeTvXUSdJQA4hsg=

request-strategy/ajwerner-btree.go

@@ -0,0 +1,44 @@
+package request_strategy
+
+import (
+	"github.com/ajwerner/btree"
+)
+
+type ajwernerBtree struct {
+	btree btree.Set[pieceRequestOrderItem]
+}
+
+var _ Btree = (*ajwernerBtree)(nil)
+
+func NewAjwernerBtree() *ajwernerBtree {
+	return &ajwernerBtree{
+		btree: btree.MakeSet(func(t, t2 pieceRequestOrderItem) int {
+			return pieceOrderLess(&t, &t2).OrderingInt()
+		}),
+	}
+}
+
+func mustValue[V any](b bool, panicValue V) {
+	if !b {
+		panic(panicValue)
+	}
+}
+
+func (a *ajwernerBtree) Delete(item pieceRequestOrderItem) {
+	mustValue(a.btree.Delete(item), item)
+}
+
+func (a *ajwernerBtree) Add(item pieceRequestOrderItem) {
+	_, overwrote := a.btree.Upsert(item)
+	mustValue(!overwrote, item)
+}
+
+func (a *ajwernerBtree) Scan(f func(pieceRequestOrderItem) bool) {
+	it := a.btree.Iterator()
+	it.First()
+	for it.First(); it.Valid(); it.Next() {
+		if !f(it.Cur()) {
+			break
+		}
+	}
+}

request-strategy/order.go

@@ -51,11 +51,7 @@ func GetRequestablePieces(input Input, pro *PieceRequestOrder, f func(ih metainf
 		storageLeft = &cap
 	}
 	var allTorrentsUnverifiedBytes int64
-	min, ok := pro.tree.Min()
-	if !ok {
-		return
-	}
-	pro.tree.Ascend(min, func(_i pieceRequestOrderItem) bool {
+	pro.tree.Scan(func(_i pieceRequestOrderItem) bool {
 		ih := _i.key.InfoHash
 		var t Torrent = input.Torrent(ih)
 		var piece Piece = t.Piece(_i.key.Index)

request-strategy/piece-request-order.go

@@ -1,27 +1,23 @@
 package request_strategy
 
-import (
-	"fmt"
+import "github.com/anacrolix/torrent/metainfo"
 
-	"github.com/anacrolix/torrent/metainfo"
-	"github.com/tidwall/btree"
-)
+type Btree interface {
+	Delete(pieceRequestOrderItem)
+	Add(pieceRequestOrderItem)
+	Scan(func(pieceRequestOrderItem) bool)
+}
 
-func NewPieceOrder() *PieceRequestOrder {
+func NewPieceOrder(btree Btree, cap int) *PieceRequestOrder {
 	return &PieceRequestOrder{
-		tree: btree.NewOptions(
-			func(a, b pieceRequestOrderItem) bool {
-				return a.Less(&b)
-			},
-			btree.Options{NoLocks: true}),
-		keys: make(map[PieceRequestOrderKey]PieceRequestOrderState),
+		tree: btree,
+		keys: make(map[PieceRequestOrderKey]PieceRequestOrderState, cap),
 	}
 }
 
 type PieceRequestOrder struct {
-	tree     *btree.BTree[pieceRequestOrderItem]
-	keys     map[PieceRequestOrderKey]PieceRequestOrderState
-	PathHint *btree.PathHint
+	tree Btree
+	keys map[PieceRequestOrderKey]PieceRequestOrderState
 }
 
 type PieceRequestOrderKey struct {
@@ -48,17 +44,10 @@ func (me *PieceRequestOrder) Add(key PieceRequestOrderKey, state PieceRequestOrd
 	if _, ok := me.keys[key]; ok {
 		panic(key)
 	}
-	if _, ok := me.tree.SetHint(pieceRequestOrderItem{
-		key:   key,
-		state: state,
-	}, me.PathHint); ok {
-		panic("shouldn't already have this")
-	}
+	me.tree.Add(pieceRequestOrderItem{key, state})
 	me.keys[key] = state
 }
 
-type PieceRequestOrderPathHint = btree.PathHint
-
 func (me *PieceRequestOrder) Update(
 	key PieceRequestOrderKey,
 	state PieceRequestOrderState,
@@ -70,17 +59,8 @@ func (me *PieceRequestOrder) Update(
 	if state == oldState {
 		return
 	}
-	item := pieceRequestOrderItem{
-		key:   key,
-		state: oldState,
-	}
-	if _, ok := me.tree.DeleteHint(item, me.PathHint); !ok {
-		panic(fmt.Sprintf("%#v", key))
-	}
-	item.state = state
-	if _, ok := me.tree.SetHint(item, me.PathHint); ok {
-		panic(key)
-	}
+	me.tree.Delete(pieceRequestOrderItem{key, oldState})
+	me.tree.Add(pieceRequestOrderItem{key, state})
 	me.keys[key] = state
 }
 
@@ -92,10 +72,7 @@ func (me *PieceRequestOrder) existingItemForKey(key PieceRequestOrderKey) pieceR
 }
 
 func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) {
-	item := me.existingItemForKey(key)
-	if _, ok := me.tree.DeleteHint(item, me.PathHint); !ok {
-		panic(key)
-	}
+	me.tree.Delete(pieceRequestOrderItem{key, me.keys[key]})
 	delete(me.keys, key)
 }
 

request-strategy/piece-request-order_test.go

@@ -6,22 +6,25 @@ import (
 	"github.com/bradfitz/iter"
 )
 
-func benchmarkPieceRequestOrder(
+func benchmarkPieceRequestOrder[B Btree](
 	b *testing.B,
-	hintForPiece func(index int) *PieceRequestOrderPathHint,
+	// Initialize the next run, and return a Btree
+	newBtree func() B,
+	// Set any path hinting for the specified piece
+	hintForPiece func(index int),
 	numPieces int,
 ) {
 	b.ResetTimer()
 	b.ReportAllocs()
 	for range iter.N(b.N) {
-		pro := NewPieceOrder()
+		pro := NewPieceOrder(newBtree(), numPieces)
 		state := PieceRequestOrderState{}
 		doPieces := func(m func(PieceRequestOrderKey)) {
 			for i := range iter.N(numPieces) {
 				key := PieceRequestOrderKey{
 					Index: i,
 				}
-				pro.PathHint = hintForPiece(i)
+				hintForPiece(i)
 				m(key)
 			}
 		}
@@ -32,15 +35,24 @@ func benchmarkPieceRequestOrder(
 		doPieces(func(key PieceRequestOrderKey) {
 			pro.Update(key, state)
 		})
+		pro.tree.Scan(func(item pieceRequestOrderItem) bool {
+			return true
+		})
 		doPieces(func(key PieceRequestOrderKey) {
 			state.Priority = piecePriority(key.Index / 4)
 			pro.Update(key, state)
 		})
-		// state.Priority = 0
+		pro.tree.Scan(func(item pieceRequestOrderItem) bool {
+			return item.key.Index < 1000
+		})
+		state.Priority = 0
 		state.Availability++
 		doPieces(func(key PieceRequestOrderKey) {
 			pro.Update(key, state)
 		})
+		pro.tree.Scan(func(item pieceRequestOrderItem) bool {
+			return item.key.Index < 1000
+		})
 		state.Availability--
 		doPieces(func(key PieceRequestOrderKey) {
 			pro.Update(key, state)
@@ -52,23 +64,43 @@ func benchmarkPieceRequestOrder(
 	}
 }
 
+func zero[T any](t *T) {
+	var zt T
+	*t = zt
+}
+
 func BenchmarkPieceRequestOrder(b *testing.B) {
 	const numPieces = 2000
-	b.Run("NoPathHints", func(b *testing.B) {
-		benchmarkPieceRequestOrder(b, func(int) *PieceRequestOrderPathHint {
-			return nil
-		}, numPieces)
+	b.Run("TidwallBtree", func(b *testing.B) {
+		b.Run("NoPathHints", func(b *testing.B) {
+			benchmarkPieceRequestOrder(b, NewTidwallBtree, func(int) {}, numPieces)
+		})
+		b.Run("SharedPathHint", func(b *testing.B) {
+			var pathHint PieceRequestOrderPathHint
+			var btree *tidwallBtree
+			benchmarkPieceRequestOrder(
+				b, func() *tidwallBtree {
+					zero(&pathHint)
+					btree = NewTidwallBtree()
+					btree.PathHint = &pathHint
+					return btree
+				}, func(int) {}, numPieces,
+			)
+		})
+		b.Run("PathHintPerPiece", func(b *testing.B) {
+			pathHints := make([]PieceRequestOrderPathHint, numPieces)
+			var btree *tidwallBtree
+			benchmarkPieceRequestOrder(
+				b, func() *tidwallBtree {
+					btree = NewTidwallBtree()
+					return btree
+				}, func(index int) {
+					btree.PathHint = &pathHints[index]
+				}, numPieces,
+			)
+		})
 	})
-	b.Run("SharedPathHint", func(b *testing.B) {
-		var pathHint PieceRequestOrderPathHint
-		benchmarkPieceRequestOrder(b, func(int) *PieceRequestOrderPathHint {
-			return &pathHint
-		}, numPieces)
-	})
-	b.Run("PathHintPerPiece", func(b *testing.B) {
-		pathHints := make([]PieceRequestOrderPathHint, numPieces)
-		benchmarkPieceRequestOrder(b, func(index int) *PieceRequestOrderPathHint {
-			return &pathHints[index]
-		}, numPieces)
+	b.Run("AjwernerBtree", func(b *testing.B) {
+		benchmarkPieceRequestOrder(b, NewAjwernerBtree, func(index int) {}, numPieces)
 	})
 }

request-strategy/tidwall-btree.go

@@ -0,0 +1,37 @@
+package request_strategy
+
+import (
+	"github.com/tidwall/btree"
+)
+
+type tidwallBtree struct {
+	tree     *btree.BTree[pieceRequestOrderItem]
+	PathHint *btree.PathHint
+}
+
+func (me *tidwallBtree) Scan(f func(pieceRequestOrderItem) bool) {
+	me.tree.Scan(f)
+}
+
+func NewTidwallBtree() *tidwallBtree {
+	return &tidwallBtree{
+		tree: btree.NewOptions(
+			func(a, b pieceRequestOrderItem) bool {
+				return a.Less(&b)
+			},
+			btree.Options{NoLocks: true}),
+	}
+}
+
+func (me *tidwallBtree) Add(item pieceRequestOrderItem) {
+	if _, ok := me.tree.SetHint(item, me.PathHint); ok {
+		panic("shouldn't already have this")
+	}
+}
+
+type PieceRequestOrderPathHint = btree.PathHint
+
+func (me *tidwallBtree) Delete(item pieceRequestOrderItem) {
+	_, deleted := me.tree.DeleteHint(item, me.PathHint)
+	mustValue(deleted, item)
+}

torrent-piece-request-order.go

@@ -45,7 +45,7 @@ func (t *Torrent) initPieceRequestOrder() {
 	key := t.clientPieceRequestOrderKey()
 	cpro := t.cl.pieceRequestOrder
 	if cpro[key] == nil {
-		cpro[key] = request_strategy.NewPieceOrder()
+		cpro[key] = request_strategy.NewPieceOrder(request_strategy.NewTidwallBtree(), t.numPieces())
 	}
 }