From 541ca02f1b606d6ea5e7a7a12ecbe4c54ce9019e Mon Sep 17 00:00:00 2001 From: Jim Ma Date: Thu, 9 Dec 2021 17:27:48 +0800 Subject: [PATCH 1/2] fix: total pieces not set when back source Signed-off-by: Jim Ma --- client/daemon/peer/piece_manager.go | 1 + 1 file changed, 1 insertion(+) diff --git a/client/daemon/peer/piece_manager.go b/client/daemon/peer/piece_manager.go index 4e5f41e305d..59e4c836340 100644 --- a/client/daemon/peer/piece_manager.go +++ b/client/daemon/peer/piece_manager.go @@ -414,6 +414,7 @@ func (pm *pieceManager) DownloadSource(ctx context.Context, pt Task, request *sc }, ContentLength: contentLength, GenPieceDigest: true, + TotalPieces: pt.GetTotalPieces(), }); err != nil { log.Errorf("update task failed %s", err) } From 47f1c849821eab96a52156778cace825b75f9ef0 Mon Sep 17 00:00:00 2001 From: Jim Ma Date: Thu, 9 Dec 2021 17:43:09 +0800 Subject: [PATCH 2/2] fix: test Signed-off-by: Jim Ma --- client/daemon/peer/piece_manager_test.go | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/client/daemon/peer/piece_manager_test.go b/client/daemon/peer/piece_manager_test.go index df7a2e49f11..2c6f059c1b0 100644 --- a/client/daemon/peer/piece_manager_test.go +++ b/client/daemon/peer/piece_manager_test.go @@ -33,6 +33,7 @@ import ( "github.com/golang/mock/gomock" testifyassert "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" + "go.uber.org/atomic" "d7y.io/dragonfly/v2/client/clientutil" "d7y.io/dragonfly/v2/client/config" @@ -83,7 +84,7 @@ func TestPieceManager_DownloadSource(t *testing.T) { checkDigest bool }{ { - name: "multiple pieces with content length", + name: "multiple pieces with content length, check digest", pieceSize: 1024, checkDigest: true, withContentLength: true, @@ -95,7 +96,7 @@ func TestPieceManager_DownloadSource(t *testing.T) { withContentLength: true, }, { - name: "multiple pieces without content length", + name: "multiple pieces without content length, check digest", pieceSize: 1024, checkDigest: true, withContentLength: false, @@ -131,12 +132,18 @@ func TestPieceManager_DownloadSource(t *testing.T) { t.Run(tc.name, func(t *testing.T) { /********** prepare test start **********/ mockPeerTask := NewMockTask(ctrl) + var totalPieces = &atomic.Int32{} 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 {