dragonfly/client/daemon/peer/piece_manager_test.go

817 lines
21 KiB
Go

/*
* Copyright 2020 The Dragonfly Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package peer
import (
"bytes"
"context"
"crypto/md5"
"encoding/hex"
"fmt"
"io"
"math"
"net/http"
"net/http/httptest"
"os"
"testing"
"time"
"github.com/go-http-utils/headers"
testifyassert "github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"go.uber.org/atomic"
"go.uber.org/mock/gomock"
"golang.org/x/time/rate"
commonv1 "d7y.io/api/v2/pkg/apis/common/v1"
schedulerv1 "d7y.io/api/v2/pkg/apis/scheduler/v1"
"d7y.io/dragonfly/v2/client/config"
"d7y.io/dragonfly/v2/client/daemon/storage"
"d7y.io/dragonfly/v2/client/daemon/test"
clientutil "d7y.io/dragonfly/v2/client/util"
logger "d7y.io/dragonfly/v2/internal/dflog"
"d7y.io/dragonfly/v2/internal/util"
"d7y.io/dragonfly/v2/pkg/digest"
nethttp "d7y.io/dragonfly/v2/pkg/net/http"
_ "d7y.io/dragonfly/v2/pkg/rpc/dfdaemon/server"
"d7y.io/dragonfly/v2/pkg/source"
"d7y.io/dragonfly/v2/pkg/source/clients/httpprotocol"
)
func TestPieceManager_DownloadSource(t *testing.T) {
assert := testifyassert.New(t)
ctrl := gomock.NewController(t)
source.UnRegister("http")
require.Nil(t, source.Register("http", httpprotocol.NewHTTPSourceClient(), httpprotocol.Adapter))
defer source.UnRegister("http")
testBytes, err := os.ReadFile(test.File)
assert.Nil(err, "load test file")
var (
peerID = "peer0"
taskID = "task0"
output = "../test/testdata/test.output"
)
pieceDownloadTimeout := 30 * time.Second
storageManager, _ := storage.NewStorageManager(
config.SimpleLocalTaskStoreStrategy,
&config.StorageOption{
DataPath: t.TempDir(),
TaskExpireTime: clientutil.Duration{
Duration: -1 * time.Second,
},
}, func(request storage.CommonTaskRequest) {}, os.FileMode(0700))
hash := md5.New()
hash.Write(testBytes)
digest := digest.New(digest.AlgorithmMD5, hex.EncodeToString(hash.Sum(nil)[:16]))
testCases := []struct {
name string
pieceSize uint32
withContentLength bool
checkDigest bool
recordDownloadTime bool
bandwidth clientutil.Size
concurrentOption *config.ConcurrentOption
}{
{
name: "multiple pieces with content length, check digest",
pieceSize: 1024,
checkDigest: true,
withContentLength: true,
},
{
name: "multiple pieces with content length",
pieceSize: 1024,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
},
{
name: "multiple pieces with content length, concurrent download with 2 goroutines",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 2,
ThresholdSize: clientutil.Size{
Limit: 1024,
},
},
},
{
name: "multiple pieces with content length, concurrent download with 4 goroutines",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 4,
ThresholdSize: clientutil.Size{
Limit: 1024,
},
},
},
{
name: "multiple pieces with content length, concurrent download with 8 goroutines",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 8,
ThresholdSize: clientutil.Size{
Limit: 1024,
},
},
},
{
name: "multiple pieces with content length, concurrent download with 16 goroutines",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 16,
ThresholdSize: clientutil.Size{
Limit: 1024,
},
},
},
{
name: "multiple pieces with content length, single-thread download with download bandwidth 2KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 2048},
},
{
name: "multiple pieces with content length, concurrent download with 2 goroutines, download bandwidth 2KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 2048},
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 2,
ThresholdSize: clientutil.Size{
Limit: 1024 * 1024,
},
ThresholdSpeed: 8192,
},
},
{
name: "multiple pieces with content length, concurrent download with 4 goroutines, download bandwidth 2KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 2048},
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 4,
ThresholdSize: clientutil.Size{
Limit: 1024 * 1024,
},
ThresholdSpeed: 8192,
},
},
{
name: "multiple pieces with content length, concurrent download with 8 goroutines, download bandwidth 2KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 2048},
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 8,
ThresholdSize: clientutil.Size{
Limit: 1024 * 1024,
},
ThresholdSpeed: 8192,
},
},
{
name: "multiple pieces with content length, concurrent download with 16 goroutines, download bandwidth 2KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 2048},
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 16,
ThresholdSize: clientutil.Size{
Limit: 1024 * 1024,
},
ThresholdSpeed: 8192,
},
},
{
name: "multiple pieces with content length, single-thread download with download bandwidth 4KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 4096},
},
{
name: "multiple pieces with content length, concurrent download with download bandwidth 4KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 4096},
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 4,
ThresholdSize: clientutil.Size{
Limit: 1024 * 1024,
},
ThresholdSpeed: 8192,
},
},
{
name: "multiple pieces with content length, single-thread download with download bandwidth 8KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 8192},
},
{
name: "multiple pieces with content length, concurrent download with download bandwidth 8KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 8192},
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 4,
ThresholdSize: clientutil.Size{
Limit: 1024 * 1024,
},
ThresholdSpeed: 8192,
},
},
{
name: "multiple pieces with content length, single-thread download with download bandwidth 16KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 16384},
},
{
name: "multiple pieces with content length, concurrent download with download bandwidth 16KB/s",
pieceSize: 2048,
checkDigest: false,
withContentLength: true,
recordDownloadTime: true,
bandwidth: clientutil.Size{Limit: 16384},
concurrentOption: &config.ConcurrentOption{
GoroutineCount: 4,
ThresholdSize: clientutil.Size{
Limit: 1024 * 1024,
},
ThresholdSpeed: 8192,
},
},
{
name: "multiple pieces without content length, check digest",
pieceSize: 1024,
checkDigest: true,
withContentLength: false,
},
{
name: "multiple pieces without content length",
pieceSize: 1024,
checkDigest: false,
withContentLength: false,
},
{
name: "one pieces with content length case 1",
pieceSize: uint32(len(testBytes)),
withContentLength: true,
},
{
name: "one pieces without content length case 1",
pieceSize: uint32(len(testBytes)),
withContentLength: false,
},
{
name: "one pieces with content length case 2",
pieceSize: uint32(len(testBytes)) + 1,
withContentLength: true,
},
{
name: "one pieces without content length case 2",
pieceSize: uint32(len(testBytes)) + 1,
withContentLength: false,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
/********** prepare test start **********/
mockPeerTask := NewMockTask(ctrl)
var (
totalPieces = &atomic.Int32{}
taskStorage storage.TaskStorageDriver
)
mockPeerTask.EXPECT().SetContentLength(gomock.Any()).AnyTimes().DoAndReturn(
func(arg0 int64) error {
return nil
})
mockPeerTask.EXPECT().SetTotalPieces(gomock.Any()).AnyTimes().DoAndReturn(
func(arg0 int32) {
totalPieces.Store(arg0)
})
mockPeerTask.EXPECT().GetTotalPieces().AnyTimes().DoAndReturn(
func() int32 {
return totalPieces.Load()
})
mockPeerTask.EXPECT().GetPeerID().AnyTimes().DoAndReturn(
func() string {
return peerID
})
mockPeerTask.EXPECT().GetTaskID().AnyTimes().DoAndReturn(
func() string {
return taskID
})
mockPeerTask.EXPECT().GetStorage().AnyTimes().DoAndReturn(
func() storage.TaskStorageDriver {
return taskStorage
})
mockPeerTask.EXPECT().AddTraffic(gomock.Any()).AnyTimes().DoAndReturn(func(int642 uint64) {})
mockPeerTask.EXPECT().ReportPieceResult(gomock.Any(), gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn(
func(*DownloadPieceRequest, *DownloadPieceResult, error) {
})
mockPeerTask.EXPECT().PublishPieceInfo(gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn(
func(pieceNum int32, size uint32) {
})
mockPeerTask.EXPECT().Context().AnyTimes().DoAndReturn(func() context.Context {
return context.Background()
})
mockPeerTask.EXPECT().Log().AnyTimes().DoAndReturn(func() *logger.SugaredLoggerOnWith {
return logger.With("test case", tc.name)
})
taskStorage, err = storageManager.RegisterTask(context.Background(),
&storage.RegisterTaskRequest{
PeerTaskMetadata: storage.PeerTaskMetadata{
PeerID: mockPeerTask.GetPeerID(),
TaskID: mockPeerTask.GetTaskID(),
},
DesiredLocation: output,
ContentLength: int64(len(testBytes)),
})
assert.Nil(err)
defer storageManager.CleanUp()
defer os.Remove(output)
/********** prepare test end **********/
ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
if tc.withContentLength {
w.Header().Set(headers.ContentLength,
fmt.Sprintf("%d", len(testBytes)))
}
var (
buf *bytes.Buffer
l int
)
rg := r.Header.Get(headers.Range)
if rg == "" {
buf = bytes.NewBuffer(testBytes)
l = len(testBytes)
} else {
parsedRange, err := nethttp.ParseRange(rg, int64(len(testBytes)))
assert.Nil(err)
w.Header().Set(headers.ContentRange,
fmt.Sprintf("bytes %d-%d/%d",
parsedRange[0].Start,
parsedRange[0].Start+parsedRange[0].Length-1,
len(testBytes)))
w.WriteHeader(http.StatusPartialContent)
buf = bytes.NewBuffer(testBytes[parsedRange[0].Start : parsedRange[0].Start+parsedRange[0].Length])
l = int(parsedRange[0].Length)
}
if tc.bandwidth.Limit > 0 {
ctx := context.Background()
limiter := rate.NewLimiter(tc.bandwidth.Limit, int(tc.bandwidth.Limit))
limiter.AllowN(time.Now(), int(tc.bandwidth.Limit))
maxPieceNum := util.ComputePieceCount(int64(l), uint32(tc.bandwidth.Limit))
for pieceNum := int32(0); pieceNum < maxPieceNum; pieceNum++ {
size := int(tc.bandwidth.Limit)
offset := uint64(pieceNum) * uint64(size)
// calculate piece size for last piece
if int64(offset)+int64(size) > int64(l) {
size = int(int64(l) - int64(offset))
}
err := limiter.WaitN(ctx, size)
assert.Nil(err)
n, err := io.CopyN(w, buf, int64(size))
if err != nil {
// broken pipe due to auto switch to concurrent download, which will close conn
return
}
assert.Equal(int64(size), n)
}
} else {
n, err := io.Copy(w, buf)
assert.Nil(err)
assert.Equal(int64(l), n)
}
}))
defer ts.Close()
pm, err := NewPieceManager(pieceDownloadTimeout, WithConcurrentOption(tc.concurrentOption))
assert.Nil(err)
pm.(*pieceManager).computePieceSize = func(length int64) uint32 {
return tc.pieceSize
}
request := &schedulerv1.PeerTaskRequest{
Url: ts.URL,
UrlMeta: &commonv1.UrlMeta{
Digest: "",
Range: "",
Header: nil,
},
}
if tc.checkDigest {
request.UrlMeta.Digest = digest.String()
}
var start time.Time
if tc.recordDownloadTime {
start = time.Now()
}
err = pm.DownloadSource(context.Background(), mockPeerTask, request, nil)
assert.Nil(err)
if tc.recordDownloadTime {
elapsed := time.Since(start)
log := mockPeerTask.Log()
log.Infof("download took %s", elapsed)
}
err = storageManager.Store(context.Background(),
&storage.StoreRequest{
CommonTaskRequest: storage.CommonTaskRequest{
PeerID: peerID,
TaskID: taskID,
Destination: output,
},
})
assert.Nil(err)
outputBytes, err := os.ReadFile(output)
assert.Nil(err, "load output file")
if string(testBytes) != string(outputBytes) {
assert.Equal(string(testBytes), string(outputBytes), "output and desired output must match")
}
})
}
}
func TestDetectBackSourceError(t *testing.T) {
assert := testifyassert.New(t)
testCases := []struct {
name string
genError func() error
detectError bool
isBackSourceError bool
}{
{
name: "is back source error - connect error",
genError: func() error {
_, err := http.Get("http://127.0.0.1:12345")
return err
},
detectError: true,
isBackSourceError: true,
},
{
name: "is back source error - timeout",
genError: func() error {
client := http.Client{
Timeout: 10 * time.Millisecond,
}
request, _ := http.NewRequest(http.MethodGet, "http://127.0.0.2:12345", nil)
_, err := client.Do(request)
return err
},
detectError: true,
isBackSourceError: true,
},
{
name: "not back source error",
genError: func() error {
return fmt.Errorf("test")
},
detectError: true,
isBackSourceError: false,
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
err := tc.genError()
err = detectBackSourceError(err)
if tc.detectError {
assert.NotNil(err)
} else {
assert.Nil(err)
}
assert.Equal(tc.isBackSourceError, isBackSourceError(err))
})
}
}
func TestPieceGroup(t *testing.T) {
assert := testifyassert.New(t)
testCases := []struct {
name string
parsedRange *nethttp.Range
startPieceNum int32
pieceSize uint32
con int32
pieceGroups []pieceGroup
}{
{
name: "100-200-2",
pieceSize: 100,
parsedRange: &nethttp.Range{Start: 0, Length: 200},
con: 2,
pieceGroups: []pieceGroup{
{
start: 0,
end: 0,
startByte: 0,
endByte: 99,
},
{
start: 1,
end: 1,
startByte: 100,
endByte: 199,
},
},
},
{
name: "100-300-2",
pieceSize: 100,
parsedRange: &nethttp.Range{Start: 0, Length: 300},
con: 2,
pieceGroups: []pieceGroup{
{
start: 0,
end: 1,
startByte: 0,
endByte: 199,
},
{
start: 2,
end: 2,
startByte: 200,
endByte: 299,
},
},
},
{
name: "100-500-4",
pieceSize: 100,
parsedRange: &nethttp.Range{Start: 0, Length: 500},
con: 4,
pieceGroups: []pieceGroup{
{
start: 0,
end: 1,
startByte: 0,
endByte: 199,
},
{
start: 2,
end: 2,
startByte: 200,
endByte: 299,
},
{
start: 3,
end: 3,
startByte: 300,
endByte: 399,
},
{
start: 4,
end: 4,
startByte: 400,
endByte: 499,
},
},
},
{
name: "100-600-4",
pieceSize: 100,
parsedRange: &nethttp.Range{Start: 0, Length: 600},
con: 4,
pieceGroups: []pieceGroup{
{
start: 0,
end: 1,
startByte: 0,
endByte: 199,
},
{
start: 2,
end: 3,
startByte: 200,
endByte: 399,
},
{
start: 4,
end: 4,
startByte: 400,
endByte: 499,
},
{
start: 5,
end: 5,
startByte: 500,
endByte: 599,
},
},
},
{
name: "100-700-4",
pieceSize: 100,
parsedRange: &nethttp.Range{Start: 0, Length: 700},
con: 4,
pieceGroups: []pieceGroup{
{
start: 0,
end: 1,
startByte: 0,
endByte: 199,
},
{
start: 2,
end: 3,
startByte: 200,
endByte: 399,
},
{
start: 4,
end: 5,
startByte: 400,
endByte: 599,
},
{
start: 6,
end: 6,
startByte: 600,
endByte: 699,
},
},
},
{
name: "1-100-700-4",
pieceSize: 100,
startPieceNum: 1,
parsedRange: &nethttp.Range{Start: 90, Length: 707}, // last byte: 90 + 706 = 796
con: 4,
pieceGroups: []pieceGroup{
{
start: 1,
end: 2,
startByte: 190,
endByte: 389,
},
{
start: 3,
end: 4,
startByte: 390,
endByte: 589,
},
{
start: 5,
end: 6,
startByte: 590,
endByte: 789,
},
{
start: 7,
end: 7,
startByte: 790,
endByte: 796,
},
},
},
{
name: "100-1100-4",
pieceSize: 100,
parsedRange: &nethttp.Range{Start: 0, Length: 1100},
con: 4,
pieceGroups: []pieceGroup{
{
start: 0,
end: 2,
startByte: 0,
endByte: 299,
},
{
start: 3,
end: 5,
startByte: 300,
endByte: 599,
},
{
start: 6,
end: 8,
startByte: 600,
endByte: 899,
},
{
start: 9,
end: 10,
startByte: 900,
endByte: 1099,
},
},
},
{
name: "from-real-e2e-test",
pieceSize: 4194304,
startPieceNum: 1,
parsedRange: &nethttp.Range{Start: 984674, Length: 20638941},
con: 4,
pieceGroups: []pieceGroup{
{
start: 1,
end: 1,
startByte: 5178978,
endByte: 9373281,
},
{
start: 2,
end: 2,
startByte: 9373282,
endByte: 13567585,
},
{
start: 3,
end: 3,
startByte: 13567586,
endByte: 17761889,
},
{
start: 4,
end: 4,
startByte: 17761890,
endByte: 21623614,
},
},
},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
pieceCount := int32(math.Ceil(float64(tc.parsedRange.Length) / float64(tc.pieceSize)))
pieceCountToDownload := pieceCount - tc.startPieceNum
minPieceCountPerGroup := pieceCountToDownload / tc.con
reminderPieces := pieceCountToDownload % tc.con
for i := int32(0); i < tc.con; i++ {
tc.pieceGroups[i].pieceSize = tc.pieceSize
tc.pieceGroups[i].parsedRange = tc.parsedRange
}
var pieceGroups []pieceGroup
for i := int32(0); i < tc.con; i++ {
pg := newPieceGroup(i, reminderPieces, tc.startPieceNum, minPieceCountPerGroup, tc.pieceSize, tc.parsedRange)
pieceGroups = append(pieceGroups, *pg)
}
assert.Equal(tc.pieceGroups, pieceGroups, "piece group should equal")
})
}
}