Get request rebalancing working optimally!

[btrtrc.git] / request-strategy / order_test.go

diff --git a/request-strategy/order_test.go b/request-strategy/order_test.go
index 2096140386078b533ffaf2fa2f3c14d163f238e1..213b6f52674e27457ec0e49d2ed48c127e3b4e58 100644 (file)
--- a/request-strategy/order_test.go
+++ b/request-strategy/order_test.go
@@ -5,6 +5,7 @@ import (
        "testing"
 
        pp "github.com/anacrolix/torrent/peer_protocol"
+       "github.com/bradfitz/iter"
        qt "github.com/frankban/quicktest"
 )
 
@@ -12,6 +13,14 @@ func r(i pieceIndex, begin int) Request {
        return Request{pp.Integer(i), ChunkSpec{pp.Integer(begin), 1}}
 }
 
+func chunkIterRange(end int) func(func(ChunkSpec)) {
+       return func(f func(ChunkSpec)) {
+               for offset := range iter.N(end) {
+                       f(ChunkSpec{pp.Integer(offset), 1})
+               }
+       }
+}
+
 func chunkIter(offsets ...int) func(func(ChunkSpec)) {
        return func(f func(ChunkSpec)) {
                for _, offset := range offsets {
@@ -59,7 +68,7 @@ func TestStealingFromSlowerPeer(t *testing.T) {
                Pieces: []Piece{{
                        Request:           true,
                        NumPendingChunks:  5,
-                       IterPendingChunks: chunkIter(0, 1, 2, 3, 4),
+                       IterPendingChunks: chunkIterRange(5),
                }},
                Peers: []Peer{
                        stealee,
@@ -170,3 +179,51 @@ func TestPeerKeepsExistingIfReasonable(t *testing.T) {
                Requests:   requestSetFromSlice(keepReq),
        })
 }
+
+func TestDontStealUnnecessarily(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
+       keepReqs := requestSetFromSlice(r(0, 0), r(0, 1), r(0, 2))
+       stealee.HasExistingRequest = func(r Request) bool {
+               _, ok := keepReqs[r]
+               return ok
+       }
+       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:  9,
+                       IterPendingChunks: chunkIterRange(9),
+               }},
+               Peers: []Peer{
+                       firstStealer,
+                       stealee,
+                       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, 3)
+       check(secondStealer.Id, 3)
+       c.Check(results[stealee.Id], qt.ContentEquals, PeerNextRequestState{
+               Interested: true,
+               Requests:   keepReqs,
+       })
+}