From 1bc84e316f539ecce518d1f7e5b810fa4f45e7e3 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Sat, 11 Dec 2021 13:46:29 +1100
Subject: [PATCH] Don't group Torrent piece request orders without a storage
 capacity together

---
 client.go                               |  2 +-
 request-strategy/piece-request-order.go |  4 +++
 requesting.go                           |  2 +-
 torrent-piece-request-order.go          | 46 ++++++++++++++++++++++++-
 torrent.go                              | 17 ++-------
 5 files changed, 53 insertions(+), 18 deletions(-)

diff --git a/client.go b/client.go
index 8a686767..9d0777e4 100644
--- a/client.go
+++ b/client.go
@@ -75,7 +75,7 @@ type Client struct {
 	dopplegangerAddrs map[string]struct{}
 	badPeerIPs        map[string]struct{}
 	torrents          map[InfoHash]*Torrent
-	pieceRequestOrder map[storage.TorrentCapacity]*request_strategy.PieceRequestOrder
+	pieceRequestOrder map[interface{}]*request_strategy.PieceRequestOrder
 
 	acceptLimiter   map[ipStr]int
 	dialRateLimiter *rate.Limiter
diff --git a/request-strategy/piece-request-order.go b/request-strategy/piece-request-order.go
index 0b1d6b4e..d906b835 100644
--- a/request-strategy/piece-request-order.go
+++ b/request-strategy/piece-request-order.go
@@ -92,3 +92,7 @@ func (me *PieceRequestOrder) Delete(key PieceRequestOrderKey) {
 	delete(me.keys, key)
 	// log.Printf("deleting %#v", key)
 }
+
+func (me *PieceRequestOrder) Len() int {
+	return len(me.keys)
+}
diff --git a/requesting.go b/requesting.go
index aaaa4f9c..c1ffa697 100644
--- a/requesting.go
+++ b/requesting.go
@@ -160,7 +160,7 @@ func (p *Peer) getDesiredRequestState() (desired desiredRequestState) {
 	}
 	request_strategy.GetRequestablePieces(
 		input,
-		p.t.cl.pieceRequestOrder[p.t.storage.Capacity],
+		p.t.getPieceRequestOrder(),
 		func(ih InfoHash, pieceIndex int) {
 			if ih != p.t.infoHash {
 				return
diff --git a/torrent-piece-request-order.go b/torrent-piece-request-order.go
index e1ab9670..7b1ef978 100644
--- a/torrent-piece-request-order.go
+++ b/torrent-piece-request-order.go
@@ -1,7 +1,51 @@
 package torrent
 
+import (
+	request_strategy "github.com/anacrolix/torrent/request-strategy"
+)
+
 func (t *Torrent) updatePieceRequestOrder(pieceIndex int) {
-	t.cl.pieceRequestOrder[t.storage.Capacity].Update(
+	t.cl.pieceRequestOrder[t.clientPieceRequestOrderKey()].Update(
 		t.pieceRequestOrderKey(pieceIndex),
 		t.requestStrategyPieceOrderState(pieceIndex))
 }
+
+func (t *Torrent) clientPieceRequestOrderKey() interface{} {
+	if t.storage.Capacity == nil {
+		return t
+	}
+	return t.storage.Capacity
+}
+
+func (t *Torrent) deletePieceRequestOrder() {
+	cpro := t.cl.pieceRequestOrder
+	key := t.clientPieceRequestOrderKey()
+	pro := cpro[key]
+	for i := 0; i < t.numPieces(); i++ {
+		pro.Delete(t.pieceRequestOrderKey(i))
+	}
+	if pro.Len() == 0 {
+		delete(cpro, key)
+	}
+}
+
+func (t *Torrent) initPieceRequestOrder() {
+	if t.cl.pieceRequestOrder == nil {
+		t.cl.pieceRequestOrder = make(map[interface{}]*request_strategy.PieceRequestOrder)
+	}
+	key := t.clientPieceRequestOrderKey()
+	cpro := t.cl.pieceRequestOrder
+	if cpro[key] == nil {
+		cpro[key] = request_strategy.NewPieceOrder()
+	}
+}
+
+func (t *Torrent) addRequestOrderPiece(i int) {
+	t.cl.pieceRequestOrder[t.clientPieceRequestOrderKey()].Add(
+		t.pieceRequestOrderKey(i),
+		t.requestStrategyPieceOrderState(i))
+}
+
+func (t *Torrent) getPieceRequestOrder() *request_strategy.PieceRequestOrder {
+	return t.cl.pieceRequestOrder[t.clientPieceRequestOrderKey()]
+}
diff --git a/torrent.go b/torrent.go
index db24603b..ade7b933 100644
--- a/torrent.go
+++ b/torrent.go
@@ -437,12 +437,7 @@ func (t *Torrent) pieceRequestOrderKey(i int) request_strategy.PieceRequestOrder
 
 // This seems to be all the follow-up tasks after info is set, that can't fail.
 func (t *Torrent) onSetInfo() {
-	if t.cl.pieceRequestOrder == nil {
-		t.cl.pieceRequestOrder = make(map[storage.TorrentCapacity]*request_strategy.PieceRequestOrder)
-	}
-	if t.cl.pieceRequestOrder[t.storage.Capacity] == nil {
-		t.cl.pieceRequestOrder[t.storage.Capacity] = request_strategy.NewPieceOrder()
-	}
+	t.initPieceRequestOrder()
 	for i := range t.pieces {
 		p := &t.pieces[i]
 		// Need to add availability before updating piece completion, as that may result in conns
@@ -451,9 +446,7 @@ func (t *Torrent) onSetInfo() {
 			panic(p.availability)
 		}
 		p.availability = int64(t.pieceAvailabilityFromPeers(i))
-		t.cl.pieceRequestOrder[t.storage.Capacity].Add(
-			t.pieceRequestOrderKey(i),
-			t.requestStrategyPieceOrderState(i))
+		t.addRequestOrderPiece(i)
 		t.updatePieceCompletion(pieceIndex(i))
 		if !t.initialPieceCheckDisabled && !p.storageCompletionOk {
 			// t.logger.Printf("piece %s completion unknown, queueing check", p)
@@ -818,12 +811,6 @@ func (t *Torrent) numPiecesCompleted() (num pieceIndex) {
 	return pieceIndex(t._completedPieces.GetCardinality())
 }
 
-func (t *Torrent) deletePieceRequestOrder() {
-	for i := 0; i < t.numPieces(); i++ {
-		t.cl.pieceRequestOrder[t.storage.Capacity].Delete(t.pieceRequestOrderKey(i))
-	}
-}
-
 func (t *Torrent) close(wg *sync.WaitGroup) (err error) {
 	t.closed.Set()
 	if t.storage != nil {
-- 
2.44.0