chore: parameterize tests in peer task (#994)
Signed-off-by: Jim Ma <majinjing3@gmail.com>
This commit is contained in:
parent
88278be464
commit
c87be98799
|
|
@ -58,6 +58,7 @@ type dummyPeerPacketStream struct {
|
|||
}
|
||||
|
||||
func (d *dummyPeerPacketStream) Recv() (pp *scheduler.PeerPacket, err error) {
|
||||
// TODO set base.Code_SchedNeedBackSource in *scheduler.PeerPacket instead of error
|
||||
return nil, dferrors.New(base.Code_SchedNeedBackSource, "")
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -1,278 +0,0 @@
|
|||
/*
|
||||
* 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"
|
||||
"fmt"
|
||||
"io"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/golang/mock/gomock"
|
||||
testifyassert "github.com/stretchr/testify/assert"
|
||||
testifyrequire "github.com/stretchr/testify/require"
|
||||
|
||||
"d7y.io/dragonfly/v2/client/clientutil"
|
||||
"d7y.io/dragonfly/v2/client/config"
|
||||
"d7y.io/dragonfly/v2/client/daemon/test"
|
||||
"d7y.io/dragonfly/v2/pkg/rpc/base"
|
||||
"d7y.io/dragonfly/v2/pkg/rpc/scheduler"
|
||||
"d7y.io/dragonfly/v2/pkg/source"
|
||||
"d7y.io/dragonfly/v2/pkg/source/httpprotocol"
|
||||
sourceMock "d7y.io/dragonfly/v2/pkg/source/mock"
|
||||
)
|
||||
|
||||
func TestFilePeerTask_BackSource_WithContentLength(t *testing.T) {
|
||||
assert := testifyassert.New(t)
|
||||
require := testifyrequire.New(t)
|
||||
ctrl := gomock.NewController(t)
|
||||
|
||||
testBytes, err := os.ReadFile(test.File)
|
||||
assert.Nil(err, "load test file")
|
||||
|
||||
var (
|
||||
pieceParallelCount = int32(4)
|
||||
pieceSize = 1024
|
||||
|
||||
mockContentLength = len(testBytes)
|
||||
//mockPieceCount = int(math.Ceil(float64(mockContentLength) / float64(pieceSize)))
|
||||
|
||||
peerID = "peer-0"
|
||||
taskID = "task-0"
|
||||
|
||||
output = "../test/testdata/test.output"
|
||||
url = "http://localhost/test/data"
|
||||
)
|
||||
defer os.Remove(output)
|
||||
|
||||
schedulerClient, storageManager := setupPeerTaskManagerComponents(
|
||||
ctrl,
|
||||
componentsOption{
|
||||
taskID: taskID,
|
||||
contentLength: int64(mockContentLength),
|
||||
pieceSize: uint32(pieceSize),
|
||||
pieceParallelCount: pieceParallelCount,
|
||||
backSource: true,
|
||||
content: testBytes,
|
||||
})
|
||||
defer storageManager.CleanUp()
|
||||
|
||||
downloader := NewMockPieceDownloader(ctrl)
|
||||
downloader.EXPECT().DownloadPiece(gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn(func(task *DownloadPieceRequest) (io.Reader, io.Closer, error) {
|
||||
rc := io.NopCloser(
|
||||
bytes.NewBuffer(
|
||||
testBytes[task.piece.RangeStart : task.piece.RangeStart+uint64(task.piece.RangeSize)],
|
||||
))
|
||||
return rc, rc, nil
|
||||
})
|
||||
|
||||
sourceClient := sourceMock.NewMockResourceClient(ctrl)
|
||||
source.UnRegister("http")
|
||||
require.Nil(source.Register("http", sourceClient, httpprotocol.Adapter))
|
||||
defer source.UnRegister("http")
|
||||
sourceClient.EXPECT().GetContentLength(gomock.Any()).DoAndReturn(
|
||||
func(request *source.Request) (int64, error) {
|
||||
if request.URL.String() == url {
|
||||
return int64(len(testBytes)), nil
|
||||
}
|
||||
return -1, fmt.Errorf("unexpect url: %s", request.URL.String())
|
||||
})
|
||||
sourceClient.EXPECT().Download(gomock.Any()).DoAndReturn(
|
||||
func(request *source.Request) (*source.Response, error) {
|
||||
if request.URL.String() == url {
|
||||
return source.NewResponse(io.NopCloser(bytes.NewBuffer(testBytes))), nil
|
||||
}
|
||||
return nil, fmt.Errorf("unexpect url: %s", request.URL.String())
|
||||
})
|
||||
|
||||
ptm := &peerTaskManager{
|
||||
calculateDigest: true,
|
||||
host: &scheduler.PeerHost{
|
||||
Ip: "127.0.0.1",
|
||||
},
|
||||
conductorLock: &sync.Mutex{},
|
||||
runningPeerTasks: sync.Map{},
|
||||
pieceManager: &pieceManager{
|
||||
calculateDigest: true,
|
||||
storageManager: storageManager,
|
||||
pieceDownloader: downloader,
|
||||
computePieceSize: func(contentLength int64) uint32 {
|
||||
return uint32(pieceSize)
|
||||
},
|
||||
},
|
||||
storageManager: storageManager,
|
||||
schedulerClient: schedulerClient,
|
||||
schedulerOption: config.SchedulerOption{
|
||||
ScheduleTimeout: clientutil.Duration{Duration: 10 * time.Minute},
|
||||
},
|
||||
}
|
||||
req := &FileTaskRequest{
|
||||
PeerTaskRequest: scheduler.PeerTaskRequest{
|
||||
Url: "http://localhost/test/data",
|
||||
UrlMeta: &base.UrlMeta{
|
||||
Tag: "d7y-test",
|
||||
},
|
||||
PeerId: peerID,
|
||||
PeerHost: &scheduler.PeerHost{},
|
||||
},
|
||||
Output: output,
|
||||
}
|
||||
ctx := context.Background()
|
||||
_, pt, err := ptm.newFileTask(ctx, req, 0)
|
||||
assert.Nil(err, "new file peer task")
|
||||
|
||||
progress, err := pt.Start(ctx)
|
||||
assert.Nil(err, "start file peer task")
|
||||
|
||||
var p *FileTaskProgress
|
||||
for p = range progress {
|
||||
assert.True(p.State.Success)
|
||||
if p.PeerTaskDone {
|
||||
p.DoneCallback()
|
||||
break
|
||||
}
|
||||
}
|
||||
assert.NotNil(p)
|
||||
assert.True(p.PeerTaskDone)
|
||||
|
||||
outputBytes, err := os.ReadFile(output)
|
||||
assert.Nil(err, "load output file")
|
||||
assert.Equal(testBytes, outputBytes, "output and desired output must match")
|
||||
}
|
||||
|
||||
func TestFilePeerTask_BackSource_WithoutContentLength(t *testing.T) {
|
||||
assert := testifyassert.New(t)
|
||||
require := testifyrequire.New(t)
|
||||
ctrl := gomock.NewController(t)
|
||||
|
||||
testBytes, err := os.ReadFile(test.File)
|
||||
assert.Nil(err, "load test file")
|
||||
|
||||
var (
|
||||
pieceParallelCount = int32(4)
|
||||
pieceSize = 1024
|
||||
|
||||
mockContentLength = len(testBytes)
|
||||
//mockPieceCount = int(math.Ceil(float64(mockContentLength) / float64(pieceSize)))
|
||||
|
||||
peerID = "peer-0"
|
||||
taskID = "task-0"
|
||||
|
||||
output = "../test/testdata/test.output"
|
||||
url = "http://localhost/test/data"
|
||||
)
|
||||
defer os.Remove(output)
|
||||
|
||||
schedulerClient, storageManager := setupPeerTaskManagerComponents(
|
||||
ctrl,
|
||||
componentsOption{
|
||||
taskID: taskID,
|
||||
contentLength: int64(mockContentLength),
|
||||
pieceSize: uint32(pieceSize),
|
||||
pieceParallelCount: pieceParallelCount,
|
||||
backSource: true,
|
||||
content: testBytes,
|
||||
})
|
||||
defer storageManager.CleanUp()
|
||||
|
||||
downloader := NewMockPieceDownloader(ctrl)
|
||||
downloader.EXPECT().DownloadPiece(gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn(
|
||||
func(ctx context.Context, task *DownloadPieceRequest) (io.Reader, io.Closer, error) {
|
||||
rc := io.NopCloser(
|
||||
bytes.NewBuffer(
|
||||
testBytes[task.piece.RangeStart : task.piece.RangeStart+uint64(task.piece.RangeSize)],
|
||||
))
|
||||
return rc, rc, nil
|
||||
})
|
||||
|
||||
sourceClient := sourceMock.NewMockResourceClient(ctrl)
|
||||
source.UnRegister("http")
|
||||
require.Nil(source.Register("http", sourceClient, httpprotocol.Adapter))
|
||||
defer source.UnRegister("http")
|
||||
sourceClient.EXPECT().GetContentLength(gomock.Any()).DoAndReturn(
|
||||
func(request *source.Request) (int64, error) {
|
||||
if request.URL.String() == url {
|
||||
return -1, nil
|
||||
}
|
||||
return -1, fmt.Errorf("unexpect url: %s", request.URL.String())
|
||||
})
|
||||
sourceClient.EXPECT().Download(gomock.Any()).DoAndReturn(
|
||||
func(request *source.Request) (*source.Response, error) {
|
||||
if request.URL.String() == url {
|
||||
return source.NewResponse(io.NopCloser(bytes.NewBuffer(testBytes))), nil
|
||||
}
|
||||
return nil, fmt.Errorf("unexpect url: %s", request.URL.String())
|
||||
})
|
||||
|
||||
ptm := &peerTaskManager{
|
||||
calculateDigest: true,
|
||||
host: &scheduler.PeerHost{
|
||||
Ip: "127.0.0.1",
|
||||
},
|
||||
conductorLock: &sync.Mutex{},
|
||||
runningPeerTasks: sync.Map{},
|
||||
pieceManager: &pieceManager{
|
||||
calculateDigest: true,
|
||||
storageManager: storageManager,
|
||||
pieceDownloader: downloader,
|
||||
computePieceSize: func(contentLength int64) uint32 {
|
||||
return uint32(pieceSize)
|
||||
},
|
||||
},
|
||||
storageManager: storageManager,
|
||||
schedulerClient: schedulerClient,
|
||||
schedulerOption: config.SchedulerOption{
|
||||
ScheduleTimeout: clientutil.Duration{Duration: 10 * time.Minute},
|
||||
},
|
||||
}
|
||||
req := &FileTaskRequest{
|
||||
PeerTaskRequest: scheduler.PeerTaskRequest{
|
||||
Url: "http://localhost/test/data",
|
||||
UrlMeta: &base.UrlMeta{
|
||||
Tag: "d7y-test",
|
||||
},
|
||||
PeerId: peerID,
|
||||
PeerHost: &scheduler.PeerHost{},
|
||||
},
|
||||
Output: output,
|
||||
}
|
||||
ctx := context.Background()
|
||||
_, pt, err := ptm.newFileTask(ctx, req, 0)
|
||||
assert.Nil(err, "new file peer task")
|
||||
|
||||
progress, err := pt.Start(ctx)
|
||||
assert.Nil(err, "start file peer task")
|
||||
|
||||
var p *FileTaskProgress
|
||||
for p = range progress {
|
||||
assert.True(p.State.Success)
|
||||
if p.PeerTaskDone {
|
||||
p.DoneCallback()
|
||||
break
|
||||
}
|
||||
}
|
||||
assert.NotNil(p)
|
||||
assert.True(p.PeerTaskDone)
|
||||
|
||||
outputBytes, err := os.ReadFile(output)
|
||||
assert.Nil(err, "load output file")
|
||||
assert.Equal(testBytes, outputBytes, "output and desired output must match")
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
|
|
@ -1,236 +0,0 @@
|
|||
/*
|
||||
* 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"
|
||||
"io"
|
||||
"os"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/golang/mock/gomock"
|
||||
testifyassert "github.com/stretchr/testify/assert"
|
||||
testifyrequire "github.com/stretchr/testify/require"
|
||||
|
||||
"d7y.io/dragonfly/v2/client/clientutil"
|
||||
"d7y.io/dragonfly/v2/client/config"
|
||||
"d7y.io/dragonfly/v2/client/daemon/test"
|
||||
"d7y.io/dragonfly/v2/pkg/rpc/base"
|
||||
"d7y.io/dragonfly/v2/pkg/rpc/scheduler"
|
||||
"d7y.io/dragonfly/v2/pkg/source"
|
||||
"d7y.io/dragonfly/v2/pkg/source/httpprotocol"
|
||||
sourceMock "d7y.io/dragonfly/v2/pkg/source/mock"
|
||||
)
|
||||
|
||||
func TestStreamPeerTask_BackSource_WithContentLength(t *testing.T) {
|
||||
assert := testifyassert.New(t)
|
||||
require := testifyrequire.New(t)
|
||||
ctrl := gomock.NewController(t)
|
||||
|
||||
testBytes, err := os.ReadFile(test.File)
|
||||
assert.Nil(err, "load test file")
|
||||
|
||||
var (
|
||||
pieceParallelCount = int32(4)
|
||||
pieceSize = 1024
|
||||
|
||||
mockContentLength = len(testBytes)
|
||||
//mockPieceCount = int(math.Ceil(float64(mockContentLength) / float64(pieceSize)))
|
||||
|
||||
peerID = "peer-back-source-with-content-length"
|
||||
taskID = "task-back-source-with-content-length"
|
||||
|
||||
url = "http://localhost/test/data"
|
||||
)
|
||||
schedulerClient, storageManager := setupPeerTaskManagerComponents(
|
||||
ctrl,
|
||||
componentsOption{
|
||||
taskID: taskID,
|
||||
backSource: true,
|
||||
contentLength: int64(mockContentLength),
|
||||
pieceSize: uint32(pieceSize),
|
||||
pieceParallelCount: pieceParallelCount,
|
||||
content: testBytes,
|
||||
})
|
||||
defer storageManager.CleanUp()
|
||||
|
||||
downloader := NewMockPieceDownloader(ctrl)
|
||||
downloader.EXPECT().DownloadPiece(gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn(
|
||||
func(ctx context.Context, task *DownloadPieceRequest) (io.Reader, io.Closer, error) {
|
||||
rc := io.NopCloser(
|
||||
bytes.NewBuffer(
|
||||
testBytes[task.piece.RangeStart : task.piece.RangeStart+uint64(task.piece.RangeSize)],
|
||||
))
|
||||
return rc, rc, nil
|
||||
})
|
||||
|
||||
request, err := source.NewRequest(url)
|
||||
assert.Nil(err, "create request")
|
||||
sourceClient := sourceMock.NewMockResourceClient(ctrl)
|
||||
source.UnRegister("http")
|
||||
require.Nil(source.Register("http", sourceClient, httpprotocol.Adapter))
|
||||
defer source.UnRegister("http")
|
||||
sourceClient.EXPECT().GetContentLength(source.RequestEq(request.URL.String())).DoAndReturn(
|
||||
func(request *source.Request) (int64, error) {
|
||||
return int64(len(testBytes)), nil
|
||||
})
|
||||
sourceClient.EXPECT().Download(source.RequestEq(request.URL.String())).DoAndReturn(
|
||||
func(request *source.Request) (*source.Response, error) {
|
||||
return source.NewResponse(io.NopCloser(bytes.NewBuffer(testBytes))), nil
|
||||
})
|
||||
|
||||
ptm := &peerTaskManager{
|
||||
calculateDigest: true,
|
||||
host: &scheduler.PeerHost{
|
||||
Ip: "127.0.0.1",
|
||||
},
|
||||
conductorLock: &sync.Mutex{},
|
||||
runningPeerTasks: sync.Map{},
|
||||
pieceManager: &pieceManager{
|
||||
calculateDigest: true,
|
||||
storageManager: storageManager,
|
||||
pieceDownloader: downloader,
|
||||
computePieceSize: func(contentLength int64) uint32 {
|
||||
return uint32(pieceSize)
|
||||
},
|
||||
},
|
||||
storageManager: storageManager,
|
||||
schedulerClient: schedulerClient,
|
||||
schedulerOption: config.SchedulerOption{
|
||||
ScheduleTimeout: clientutil.Duration{Duration: 10 * time.Minute},
|
||||
},
|
||||
}
|
||||
req := &scheduler.PeerTaskRequest{
|
||||
Url: url,
|
||||
UrlMeta: &base.UrlMeta{
|
||||
Tag: "d7y-test",
|
||||
},
|
||||
PeerId: peerID,
|
||||
PeerHost: &scheduler.PeerHost{},
|
||||
}
|
||||
ctx := context.Background()
|
||||
pt, err := ptm.newStreamTask(ctx, req)
|
||||
assert.Nil(err, "new stream peer task")
|
||||
|
||||
rc, _, err := pt.Start(ctx)
|
||||
assert.Nil(err, "start stream peer task")
|
||||
|
||||
outputBytes, err := io.ReadAll(rc)
|
||||
assert.Nil(err, "load read data")
|
||||
assert.Equal(testBytes, outputBytes, "output and desired output must match")
|
||||
}
|
||||
|
||||
func TestStreamPeerTask_BackSource_WithoutContentLength(t *testing.T) {
|
||||
assert := testifyassert.New(t)
|
||||
require := testifyrequire.New(t)
|
||||
ctrl := gomock.NewController(t)
|
||||
|
||||
testBytes, err := os.ReadFile(test.File)
|
||||
assert.Nil(err, "load test file")
|
||||
|
||||
var (
|
||||
pieceParallelCount = int32(4)
|
||||
pieceSize = 1024
|
||||
|
||||
mockContentLength = len(testBytes)
|
||||
//mockPieceCount = int(math.Ceil(float64(mockContentLength) / float64(pieceSize)))
|
||||
|
||||
peerID = "peer-back-source-without-content-length"
|
||||
taskID = "task-back-source-without-content-length"
|
||||
|
||||
url = "http://localhost/test/data"
|
||||
)
|
||||
schedulerClient, storageManager := setupPeerTaskManagerComponents(
|
||||
ctrl,
|
||||
componentsOption{
|
||||
taskID: taskID,
|
||||
contentLength: int64(mockContentLength),
|
||||
pieceSize: uint32(pieceSize),
|
||||
pieceParallelCount: pieceParallelCount,
|
||||
backSource: true,
|
||||
content: testBytes,
|
||||
})
|
||||
defer storageManager.CleanUp()
|
||||
|
||||
downloader := NewMockPieceDownloader(ctrl)
|
||||
downloader.EXPECT().DownloadPiece(gomock.Any(), gomock.Any()).AnyTimes().DoAndReturn(
|
||||
func(ctx context.Context, task *DownloadPieceRequest) (io.Reader, io.Closer, error) {
|
||||
rc := io.NopCloser(
|
||||
bytes.NewBuffer(
|
||||
testBytes[task.piece.RangeStart : task.piece.RangeStart+uint64(task.piece.RangeSize)],
|
||||
))
|
||||
return rc, rc, nil
|
||||
})
|
||||
|
||||
sourceClient := sourceMock.NewMockResourceClient(ctrl)
|
||||
source.UnRegister("http")
|
||||
require.Nil(source.Register("http", sourceClient, httpprotocol.Adapter))
|
||||
defer source.UnRegister("http")
|
||||
request, err := source.NewRequest(url)
|
||||
assert.Nil(err, "create reqeust")
|
||||
sourceClient.EXPECT().GetContentLength(source.RequestEq(request.URL.String())).DoAndReturn(
|
||||
func(request *source.Request) (int64, error) {
|
||||
return -1, nil
|
||||
})
|
||||
sourceClient.EXPECT().Download(source.RequestEq(request.URL.String())).DoAndReturn(
|
||||
func(request *source.Request) (*source.Response, error) {
|
||||
return source.NewResponse(io.NopCloser(bytes.NewBuffer(testBytes))), nil
|
||||
})
|
||||
|
||||
ptm := &peerTaskManager{
|
||||
calculateDigest: true,
|
||||
host: &scheduler.PeerHost{
|
||||
Ip: "127.0.0.1",
|
||||
},
|
||||
conductorLock: &sync.Mutex{},
|
||||
runningPeerTasks: sync.Map{},
|
||||
pieceManager: &pieceManager{
|
||||
calculateDigest: true,
|
||||
storageManager: storageManager,
|
||||
pieceDownloader: downloader,
|
||||
computePieceSize: func(contentLength int64) uint32 {
|
||||
return uint32(pieceSize)
|
||||
},
|
||||
},
|
||||
storageManager: storageManager,
|
||||
schedulerClient: schedulerClient,
|
||||
schedulerOption: config.SchedulerOption{
|
||||
ScheduleTimeout: clientutil.Duration{Duration: 10 * time.Minute},
|
||||
},
|
||||
}
|
||||
req := &scheduler.PeerTaskRequest{
|
||||
Url: url,
|
||||
UrlMeta: &base.UrlMeta{
|
||||
Tag: "d7y-test",
|
||||
},
|
||||
PeerId: peerID,
|
||||
PeerHost: &scheduler.PeerHost{},
|
||||
}
|
||||
ctx := context.Background()
|
||||
pt, err := ptm.newStreamTask(ctx, req)
|
||||
assert.Nil(err, "new stream peer task")
|
||||
|
||||
rc, _, err := pt.Start(ctx)
|
||||
assert.Nil(err, "start stream peer task")
|
||||
|
||||
outputBytes, err := io.ReadAll(rc)
|
||||
assert.Nil(err, "load read data")
|
||||
assert.Equal(testBytes, outputBytes, "output and desired output must match")
|
||||
}
|
||||
|
|
@ -131,11 +131,11 @@ func WithTransport(rt http.RoundTripper) func(*pieceDownloader) error {
|
|||
}
|
||||
}
|
||||
|
||||
func (p *pieceDownloader) DownloadPiece(ctx context.Context, d *DownloadPieceRequest) (io.Reader, io.Closer, error) {
|
||||
resp, err := p.httpClient.Do(buildDownloadPieceHTTPRequest(ctx, d))
|
||||
func (p *pieceDownloader) DownloadPiece(ctx context.Context, req *DownloadPieceRequest) (io.Reader, io.Closer, error) {
|
||||
resp, err := p.httpClient.Do(buildDownloadPieceHTTPRequest(ctx, req))
|
||||
if err != nil {
|
||||
logger.Errorf("task id: %s, piece num: %d, dst: %s, download piece failed: %s",
|
||||
d.TaskID, d.piece.PieceNum, d.DstAddr, err)
|
||||
req.TaskID, req.piece.PieceNum, req.DstAddr, err)
|
||||
return nil, nil, &pieceDownloadError{err: err, connectionError: true}
|
||||
}
|
||||
if resp.StatusCode > 299 {
|
||||
|
|
@ -143,17 +143,17 @@ func (p *pieceDownloader) DownloadPiece(ctx context.Context, d *DownloadPieceReq
|
|||
_ = resp.Body.Close()
|
||||
return nil, nil, &pieceDownloadError{err: err, connectionError: false, status: resp.Status, statusCode: resp.StatusCode}
|
||||
}
|
||||
r := resp.Body.(io.Reader)
|
||||
c := resp.Body.(io.Closer)
|
||||
if d.CalcDigest {
|
||||
d.log.Debugf("calculate digest for piece %d, digest: %s", d.piece.PieceNum, d.piece.PieceMd5)
|
||||
r = digestutils.NewDigestReader(d.log, io.LimitReader(resp.Body, int64(d.piece.RangeSize)), d.piece.PieceMd5)
|
||||
reader, closer := resp.Body.(io.Reader), resp.Body.(io.Closer)
|
||||
if req.CalcDigest {
|
||||
req.log.Debugf("calculate digest for piece %d, digest: %s", req.piece.PieceNum, req.piece.PieceMd5)
|
||||
reader = digestutils.NewDigestReader(req.log, io.LimitReader(resp.Body, int64(req.piece.RangeSize)), req.piece.PieceMd5)
|
||||
}
|
||||
return r, c, nil
|
||||
return reader, closer, nil
|
||||
}
|
||||
|
||||
func buildDownloadPieceHTTPRequest(ctx context.Context, d *DownloadPieceRequest) *http.Request {
|
||||
b := strings.Builder{}
|
||||
// FIXME switch to https when tls enabled
|
||||
b.WriteString("http://")
|
||||
b.WriteString(d.DstAddr)
|
||||
b.WriteString(upload.PeerDownloadHTTPPathPrefix)
|
||||
|
|
|
|||
|
|
@ -443,6 +443,10 @@ func (pm *pieceManager) downloadUnknownLengthSource(ctx context.Context, pt Task
|
|||
pt.ReportPieceResult(request, result, err)
|
||||
return err
|
||||
}
|
||||
// content length is aligning at piece size
|
||||
if result.Size == 0 {
|
||||
break
|
||||
}
|
||||
pt.PublishPieceInfo(pieceNum, uint32(result.Size))
|
||||
pt.ReportPieceResult(request, result, nil)
|
||||
break
|
||||
|
|
|
|||
Loading…
Reference in New Issue