From 07ba6e921028c5eb4186439f983997b8f8cfba79 Mon Sep 17 00:00:00 2001
From: Matt Joiner <anacrolix@gmail.com>
Date: Thu, 13 May 2021 13:50:41 +1000
Subject: [PATCH] New tests and fixes for them

Not complete. There's still a request-stealing balancing issue, but it's functional for now.
---
 request-strategy/order.go      | 12 +++--
 request-strategy/order_test.go | 92 +++++++++++++++++++++++++++++++++-
 2 files changed, 99 insertions(+), 5 deletions(-)

diff --git a/request-strategy/order.go b/request-strategy/order.go
index f0a5eaef..307e47ef 100644
--- a/request-strategy/order.go
+++ b/request-strategy/order.go
@@ -1,12 +1,12 @@
 package request_strategy
 
 import (
+	"math"
 	"sort"
 
 	"github.com/anacrolix/multiless"
 	pp "github.com/anacrolix/torrent/peer_protocol"
 	"github.com/anacrolix/torrent/types"
-	"github.com/davecgh/go-spew/spew"
 )
 
 type (
@@ -164,24 +164,28 @@ func (requestOrder *ClientPieceOrder) DoRequests(torrents []*Torrent) map[PeerId
 		if f := torrentPiece.IterPendingChunks; f != nil {
 			f(func(chunk types.ChunkSpec) {
 				req := Request{pp.Integer(p.index), chunk}
-				pendingChunksRemaining--
+				defer func() { pendingChunksRemaining-- }()
 				sortPeersForPiece()
-				spew.Dump(peersForPiece)
 				skipped := 0
 				// Try up to the number of peers that could legitimately receive the request equal to
 				// the number of chunks left. This should ensure that only the best peers serve the last
 				// few chunks in a piece.
+				lowestNumRequestsInPiece := math.MaxInt16
 				for _, peer := range peersForPiece {
 					if !peer.canFitRequest() || !peer.HasPiece(p.index) || (!peer.pieceAllowedFastOrDefault(p.index) && peer.Choking) {
 						continue
 					}
-					if skipped >= pendingChunksRemaining {
+					if skipped+1 >= pendingChunksRemaining {
 						break
 					}
 					if f := peer.HasExistingRequest; f == nil || !f(req) {
 						skipped++
+						lowestNumRequestsInPiece = peer.requestsInPiece
 						continue
 					}
+					if peer.requestsInPiece > lowestNumRequestsInPiece {
+						break
+					}
 					if !peer.pieceAllowedFastOrDefault(p.index) {
 						// We must stay interested for this.
 						peer.nextState.Interested = true
diff --git a/request-strategy/order_test.go b/request-strategy/order_test.go
index c972b301..20961403 100644
--- a/request-strategy/order_test.go
+++ b/request-strategy/order_test.go
@@ -34,7 +34,50 @@ func (i intPeerId) Uintptr() uintptr {
 	return uintptr(i)
 }
 
-func TestStealingFromSlowerPeers(t *testing.T) {
+func TestStealingFromSlowerPeer(t *testing.T) {
+	c := qt.New(t)
+	order := ClientPieceOrder{}
+	basePeer := Peer{
+		HasPiece: func(i pieceIndex) bool {
+			return true
+		},
+		MaxRequests:  math.MaxInt16,
+		DownloadRate: 2,
+	}
+	// Slower than the stealers, but has all requests already.
+	stealee := basePeer
+	stealee.DownloadRate = 1
+	stealee.HasExistingRequest = func(r Request) bool {
+		return true
+	}
+	stealee.Id = intPeerId(1)
+	firstStealer := basePeer
+	firstStealer.Id = intPeerId(2)
+	secondStealer := basePeer
+	secondStealer.Id = intPeerId(3)
+	results := order.DoRequests([]*Torrent{{
+		Pieces: []Piece{{
+			Request:           true,
+			NumPendingChunks:  5,
+			IterPendingChunks: chunkIter(0, 1, 2, 3, 4),
+		}},
+		Peers: []Peer{
+			stealee,
+			firstStealer,
+			secondStealer,
+		},
+	}})
+	c.Assert(results, qt.HasLen, 3)
+	check := func(p PeerId, l int) {
+		c.Check(results[p].Requests, qt.HasLen, l)
+		c.Check(results[p].Interested, qt.Equals, l > 0)
+	}
+	check(stealee.Id, 1)
+	check(firstStealer.Id, 2)
+	check(secondStealer.Id, 2)
+}
+
+func TestStealingFromSlowerPeersBasic(t *testing.T) {
 	c := qt.New(t)
 	order := ClientPieceOrder{}
 	basePeer := Peer{
@@ -80,3 +123,50 @@ func TestStealingFromSlowerPeers(t *testing.T) {
 		},
 	})
 }
+
+func TestPeerKeepsExistingIfReasonable(t *testing.T) {
+	c := qt.New(t)
+	order := ClientPieceOrder{}
+	basePeer := Peer{
+		HasPiece: func(i pieceIndex) bool {
+			return true
+		},
+		MaxRequests:  math.MaxInt16,
+		DownloadRate: 2,
+	}
+	// Slower than the stealers, but has all requests already.
+	stealee := basePeer
+	stealee.DownloadRate = 1
+	keepReq := r(0, 0)
+	stealee.HasExistingRequest = func(r Request) bool {
+		return r == keepReq
+	}
+	stealee.Id = intPeerId(1)
+	firstStealer := basePeer
+	firstStealer.Id = intPeerId(2)
+	secondStealer := basePeer
+	secondStealer.Id = intPeerId(3)
+	results := order.DoRequests([]*Torrent{{
+		Pieces: []Piece{{
+			Request:           true,
+			NumPendingChunks:  4,
+			IterPendingChunks: chunkIter(0, 1, 3, 4),
+		}},
+		Peers: []Peer{
+			stealee,
+			firstStealer,
+			secondStealer,
+		},
+	}})
+	c.Assert(results, qt.HasLen, 3)
+	check := func(p PeerId, l int) {
+		c.Check(results[p].Requests, qt.HasLen, l)
+		c.Check(results[p].Interested, qt.Equals, l > 0)
+	}
+	check(firstStealer.Id, 2)
+	check(secondStealer.Id, 1)
+	c.Check(results[stealee.Id], qt.ContentEquals, PeerNextRequestState{
+		Interested: true,
+		Requests:   requestSetFromSlice(keepReq),
+	})
+}
-- 
2.48.1