315 lines
7.6 KiB
Go
315 lines
7.6 KiB
Go
package request_strategy
|
|
|
|
import (
|
|
"encoding/gob"
|
|
"math"
|
|
"testing"
|
|
|
|
"github.com/RoaringBitmap/roaring"
|
|
qt "github.com/frankban/quicktest"
|
|
"github.com/google/go-cmp/cmp"
|
|
)
|
|
|
|
func init() {
|
|
gob.Register(chunkIterRange(0))
|
|
gob.Register(sliceChunksIter{})
|
|
}
|
|
|
|
type chunkIterRange ChunkIndex
|
|
|
|
func (me chunkIterRange) Iter(f func(ChunkIndex)) {
|
|
for offset := ChunkIndex(0); offset < ChunkIndex(me); offset += 1 {
|
|
f(offset)
|
|
}
|
|
}
|
|
|
|
type sliceChunksIter []ChunkIndex
|
|
|
|
func chunkIter(offsets ...ChunkIndex) ChunksIter {
|
|
return sliceChunksIter(offsets)
|
|
}
|
|
|
|
func (offsets sliceChunksIter) Iter(f func(ChunkIndex)) {
|
|
for _, offset := range offsets {
|
|
f(offset)
|
|
}
|
|
}
|
|
|
|
func requestSetFromSlice(rs ...RequestIndex) (ret roaring.Bitmap) {
|
|
ret.AddMany(rs)
|
|
return
|
|
}
|
|
|
|
func init() {
|
|
gob.Register(intPeerId(0))
|
|
}
|
|
|
|
type intPeerId int
|
|
|
|
func (i intPeerId) Uintptr() uintptr {
|
|
return uintptr(i)
|
|
}
|
|
|
|
var hasAllRequests = func() (all roaring.Bitmap) {
|
|
all.AddRange(0, roaring.MaxRange)
|
|
return
|
|
}()
|
|
|
|
func TestStealingFromSlowerPeer(t *testing.T) {
|
|
c := qt.New(t)
|
|
basePeer := Peer{
|
|
MaxRequests: math.MaxInt16,
|
|
DownloadRate: 2,
|
|
}
|
|
basePeer.Pieces.Add(0)
|
|
// Slower than the stealers, but has all requests already.
|
|
stealee := basePeer
|
|
stealee.DownloadRate = 1
|
|
stealee.ExistingRequests = hasAllRequests
|
|
stealee.Id = intPeerId(1)
|
|
firstStealer := basePeer
|
|
firstStealer.Id = intPeerId(2)
|
|
secondStealer := basePeer
|
|
secondStealer.Id = intPeerId(3)
|
|
results := Run(Input{Torrents: []Torrent{{
|
|
ChunksPerPiece: 9,
|
|
Pieces: []Piece{{
|
|
Request: true,
|
|
NumPendingChunks: 5,
|
|
IterPendingChunks: chunkIterRange(5),
|
|
}},
|
|
Peers: []Peer{
|
|
stealee,
|
|
firstStealer,
|
|
secondStealer,
|
|
},
|
|
}}})
|
|
|
|
c.Assert(results, qt.HasLen, 3)
|
|
check := func(p PeerId, l uint64) {
|
|
addressableBm := results[p].Requests
|
|
c.Check(addressableBm.GetCardinality(), qt.ContentEquals, l)
|
|
c.Check(results[p].Interested, qt.Equals, l > 0)
|
|
}
|
|
check(stealee.Id, 1)
|
|
check(firstStealer.Id, 2)
|
|
check(secondStealer.Id, 2)
|
|
}
|
|
|
|
func checkNumRequestsAndInterest(c *qt.C, next PeerNextRequestState, num uint64, interest bool) {
|
|
addressableBm := next.Requests
|
|
c.Check(addressableBm.GetCardinality(), qt.ContentEquals, num)
|
|
c.Check(next.Interested, qt.Equals, interest)
|
|
}
|
|
|
|
func TestStealingFromSlowerPeersBasic(t *testing.T) {
|
|
c := qt.New(t)
|
|
basePeer := Peer{
|
|
MaxRequests: math.MaxInt16,
|
|
DownloadRate: 2,
|
|
}
|
|
basePeer.Pieces.Add(0)
|
|
stealee := basePeer
|
|
stealee.DownloadRate = 1
|
|
stealee.ExistingRequests = hasAllRequests
|
|
stealee.Id = intPeerId(1)
|
|
firstStealer := basePeer
|
|
firstStealer.Id = intPeerId(2)
|
|
secondStealer := basePeer
|
|
secondStealer.Id = intPeerId(3)
|
|
results := Run(Input{Torrents: []Torrent{{
|
|
ChunksPerPiece: 9,
|
|
Pieces: []Piece{{
|
|
Request: true,
|
|
NumPendingChunks: 2,
|
|
IterPendingChunks: chunkIter(0, 1),
|
|
}},
|
|
Peers: []Peer{
|
|
stealee,
|
|
firstStealer,
|
|
secondStealer,
|
|
},
|
|
}}})
|
|
|
|
checkNumRequestsAndInterest(c, results[firstStealer.Id], 1, true)
|
|
checkNumRequestsAndInterest(c, results[secondStealer.Id], 1, true)
|
|
checkNumRequestsAndInterest(c, results[stealee.Id], 0, false)
|
|
}
|
|
|
|
func checkResultsRequestsLen(t *testing.T, reqs roaring.Bitmap, l uint64) {
|
|
qt.Check(t, reqs.GetCardinality(), qt.Equals, l)
|
|
}
|
|
|
|
func TestPeerKeepsExistingIfReasonable(t *testing.T) {
|
|
c := qt.New(t)
|
|
basePeer := Peer{
|
|
MaxRequests: math.MaxInt16,
|
|
DownloadRate: 2,
|
|
}
|
|
basePeer.Pieces.Add(0)
|
|
// Slower than the stealers, but has all requests already.
|
|
stealee := basePeer
|
|
stealee.DownloadRate = 1
|
|
keepReq := RequestIndex(0)
|
|
stealee.ExistingRequests = requestSetFromSlice(keepReq)
|
|
stealee.Id = intPeerId(1)
|
|
firstStealer := basePeer
|
|
firstStealer.Id = intPeerId(2)
|
|
secondStealer := basePeer
|
|
secondStealer.Id = intPeerId(3)
|
|
results := Run(Input{Torrents: []Torrent{{
|
|
ChunksPerPiece: 9,
|
|
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 uint64) {
|
|
checkResultsRequestsLen(t, results[p].Requests, l)
|
|
c.Check(results[p].Interested, qt.Equals, l > 0)
|
|
}
|
|
check(firstStealer.Id, 2)
|
|
check(secondStealer.Id, 1)
|
|
c.Check(
|
|
results[stealee.Id],
|
|
peerNextRequestStateChecker,
|
|
PeerNextRequestState{
|
|
Interested: true,
|
|
Requests: requestSetFromSlice(keepReq),
|
|
},
|
|
)
|
|
}
|
|
|
|
var peerNextRequestStateChecker = qt.CmpEquals(
|
|
cmp.Transformer(
|
|
"bitmap",
|
|
func(bm roaring.Bitmap) []uint32 {
|
|
return bm.ToArray()
|
|
}))
|
|
|
|
func TestDontStealUnnecessarily(t *testing.T) {
|
|
c := qt.New(t)
|
|
basePeer := Peer{
|
|
MaxRequests: math.MaxInt16,
|
|
DownloadRate: 2,
|
|
}
|
|
basePeer.Pieces.AddRange(0, 5)
|
|
// Slower than the stealers, but has all requests already.
|
|
stealee := basePeer
|
|
stealee.DownloadRate = 1
|
|
r := func(i, c RequestIndex) RequestIndex {
|
|
return i*9 + c
|
|
}
|
|
keepReqs := requestSetFromSlice(
|
|
r(3, 2), r(3, 4), r(3, 6), r(3, 8),
|
|
r(4, 0), r(4, 1), r(4, 7), r(4, 8))
|
|
stealee.ExistingRequests = keepReqs
|
|
stealee.Id = intPeerId(1)
|
|
firstStealer := basePeer
|
|
firstStealer.Id = intPeerId(2)
|
|
secondStealer := basePeer
|
|
secondStealer.Id = intPeerId(3)
|
|
secondStealer.Pieces = roaring.Bitmap{}
|
|
secondStealer.Pieces.Add(1)
|
|
secondStealer.Pieces.Add(3)
|
|
results := Run(Input{Torrents: []Torrent{{
|
|
ChunksPerPiece: 9,
|
|
Pieces: []Piece{
|
|
{
|
|
Request: true,
|
|
NumPendingChunks: 0,
|
|
IterPendingChunks: chunkIterRange(9),
|
|
},
|
|
{
|
|
Request: true,
|
|
NumPendingChunks: 7,
|
|
IterPendingChunks: chunkIterRange(7),
|
|
},
|
|
{
|
|
Request: true,
|
|
NumPendingChunks: 0,
|
|
IterPendingChunks: chunkIterRange(0),
|
|
},
|
|
{
|
|
Request: true,
|
|
NumPendingChunks: 9,
|
|
IterPendingChunks: chunkIterRange(9),
|
|
},
|
|
{
|
|
Request: true,
|
|
NumPendingChunks: 9,
|
|
IterPendingChunks: chunkIterRange(9),
|
|
}},
|
|
Peers: []Peer{
|
|
firstStealer,
|
|
stealee,
|
|
secondStealer,
|
|
},
|
|
}}})
|
|
|
|
c.Assert(results, qt.HasLen, 3)
|
|
check := func(p PeerId, l uint64) {
|
|
checkResultsRequestsLen(t, results[p].Requests, l)
|
|
c.Check(results[p].Interested, qt.Equals, l > 0)
|
|
}
|
|
check(firstStealer.Id, 5)
|
|
check(secondStealer.Id, 7+9)
|
|
c.Check(
|
|
results[stealee.Id],
|
|
peerNextRequestStateChecker,
|
|
PeerNextRequestState{
|
|
Interested: true,
|
|
Requests: requestSetFromSlice(r(4, 0), r(4, 1), r(4, 7), r(4, 8)),
|
|
},
|
|
)
|
|
}
|
|
|
|
// This tests a situation where multiple peers had the same existing request, due to "actual" and
|
|
// "next" request states being out of sync. This reasonable occurs when a peer hasn't fully updated
|
|
// its actual request state since the last request strategy run.
|
|
func TestDuplicatePreallocations(t *testing.T) {
|
|
peer := func(id int, downloadRate float64) Peer {
|
|
p := Peer{
|
|
ExistingRequests: hasAllRequests,
|
|
MaxRequests: 2,
|
|
Id: intPeerId(id),
|
|
DownloadRate: downloadRate,
|
|
}
|
|
p.Pieces.AddRange(0, roaring.MaxRange)
|
|
return p
|
|
}
|
|
results := Run(Input{
|
|
Torrents: []Torrent{{
|
|
ChunksPerPiece: 1,
|
|
Pieces: []Piece{{
|
|
Request: true,
|
|
NumPendingChunks: 1,
|
|
IterPendingChunks: chunkIterRange(1),
|
|
}, {
|
|
Request: true,
|
|
NumPendingChunks: 1,
|
|
IterPendingChunks: chunkIterRange(1),
|
|
}},
|
|
Peers: []Peer{
|
|
// The second peer was be marked as the preallocation, clobbering the first. The
|
|
// first peer is preferred, and the piece isn't striped, so it gets preallocated a
|
|
// request, and then gets reallocated from the peer the same request.
|
|
peer(1, 2),
|
|
peer(2, 1),
|
|
},
|
|
}},
|
|
})
|
|
c := qt.New(t)
|
|
req1 := results[intPeerId(1)].Requests
|
|
req2 := results[intPeerId(2)].Requests
|
|
c.Assert(uint64(2), qt.Equals, req1.GetCardinality()+req2.GetCardinality())
|
|
}
|