fix: small size task failed due to digest error (#886)

Signed-off-by: Jim Ma <[email protected]>
dragonflyoss · Jun 28, 2023 · 5b016d9 · 5b016d9
1 parent 6866ecc
commit 5b016d9
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/client/daemon/peer/peertask_base.go b/client/daemon/peer/peertask_base.go
@@ -39,6 +39,7 @@ import (
 	dfclient "d7y.io/dragonfly/v2/pkg/rpc/dfdaemon/client"
 	"d7y.io/dragonfly/v2/pkg/rpc/scheduler"
 	schedulerclient "d7y.io/dragonfly/v2/pkg/rpc/scheduler/client"
+	"d7y.io/dragonfly/v2/pkg/util/digestutils"
 )
 
 const (
@@ -347,7 +348,7 @@ func (pt *peerTask) pullSinglePiece(cleanUnfinishedFunc func()) {
 
 	pt.contentLength.Store(int64(pt.singlePiece.PieceInfo.RangeSize))
 	pt.SetTotalPieces(1)
-	pt.SetPieceMd5Sign(pt.singlePiece.PieceInfo.PieceMd5)
+	pt.SetPieceMd5Sign(digestutils.Sha256(pt.singlePiece.PieceInfo.PieceMd5))
 	if err := pt.callback.Init(pt); err != nil {
 		pt.failedReason = err.Error()
 		pt.failedCode = base.Code_ClientError

diff --git a/client/daemon/peer/piece_manager.go b/client/daemon/peer/piece_manager.go
@@ -405,6 +405,7 @@ func (pm *pieceManager) DownloadSource(ctx context.Context, pt Task, request *sc
 			// last piece, piece size maybe 0
 			if n < int64(size) {
 				contentLength = int64(pieceNum)*int64(pieceSize) + n
+				pt.SetTotalPieces(int32(math.Ceil(float64(contentLength) / float64(pieceSize))))
 				if err := pm.storageManager.UpdateTask(ctx,
 					&storage.UpdateTaskRequest{
 						PeerTaskMetadata: storage.PeerTaskMetadata{
@@ -416,7 +417,6 @@ func (pm *pieceManager) DownloadSource(ctx context.Context, pt Task, request *sc
 					}); err != nil {
 					log.Errorf("update task failed %s", err)
 				}
-				pt.SetTotalPieces(pieceNum + 1)
 				return pt.SetContentLength(contentLength)
 			}
 		}