Merge pull request #10230 from Lyndon-Li/block-uploader-restore-progress

Use source size in progress for block uploader restore

pkg/uploader/block/snapshot.go

@@ -235,12 +235,12 @@ func Restore(ctx context.Context, blkUp Uploader, rep udmrepo.BackupRepo, snapsh
 		return 0, errors.Wrapf(err, "error reset pos of block device %s", dest)
 	}
 
-	size, err := blkUp.Restore(snapshot, destInfo{dev: destDev, path: destPath, size: destSize}, bitmap.Iterator(), uploaderCfg)
+	_, totalSize, err := blkUp.Restore(snapshot, destInfo{dev: destDev, path: destPath, size: destSize}, bitmap.Iterator(), uploaderCfg)
 	if err != nil {
 		return 0, errors.Wrapf(err, "error restoring to block dev %s", destPath)
 	}
 
-	return size, nil
+	return totalSize, nil
 }
 
 func findPreviousSnapshot(ctx context.Context, rep udmrepo.BackupRepo, path string, snapshotTags map[string]string, noLaterThan *time.Time, log logrus.FieldLogger) (udmrepo.Snapshot, error) {

pkg/uploader/block/snapshot_test.go

@@ -46,9 +46,9 @@ func (m *mockUploader) Backup(src sourceInfo, parent udmrepo.ID, iter cbttypes.I
 	return args.Get(0).(udmrepo.Snapshot), args.Get(1).(int64), args.Error(2)
 }
 
-func (m *mockUploader) Restore(snap udmrepo.Snapshot, dest destInfo, iter cbttypes.Iterator, cfg map[string]string) (int64, error) {
+func (m *mockUploader) Restore(snap udmrepo.Snapshot, dest destInfo, iter cbttypes.Iterator, cfg map[string]string) (int64, int64, error) {
 	args := m.Called(snap, dest, iter, cfg)
-	return args.Get(0).(int64), args.Error(1)
+	return args.Get(0).(int64), args.Get(1).(int64), args.Error(2)
 }
 
 func testLog() logrus.FieldLogger {
@@ -574,7 +574,7 @@ func TestRestore(t *testing.T) {
 				repo.On("GetSnapshot", mock.Anything, udmrepo.ID("snap-001")).
 					Return(storedSnap, nil)
 				blkup.On("Restore", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
-					Return(int64(0), errors.New("restore I/O error"))
+					Return(int64(0), int64(0), errors.New("restore I/O error"))
 			},
 			setupOpenDev: func(t *testing.T) *os.File {
 				t.Helper()
@@ -588,7 +588,7 @@ func TestRestore(t *testing.T) {
 				repo.On("GetSnapshot", mock.Anything, udmrepo.ID("snap-001")).
 					Return(storedSnap, nil)
 				blkup.On("Restore", mock.Anything, mock.Anything, mock.Anything, mock.Anything).
-					Return(int64(4096), nil)
+					Return(int64(4096), int64(4096), nil)
 			},
 			setupOpenDev: func(t *testing.T) *os.File {
 				t.Helper()

pkg/uploader/block/uploader.go

@@ -59,7 +59,7 @@ type destInfo struct {
 
 type Uploader interface {
 	Backup(sourceInfo, udmrepo.ID, cbt.Iterator, map[string]string) (udmrepo.Snapshot, int64, error)
-	Restore(udmrepo.Snapshot, destInfo, cbt.Iterator, map[string]string) (int64, error)
+	Restore(udmrepo.Snapshot, destInfo, cbt.Iterator, map[string]string) (int64, int64, error)
 }
 
 type blockUploader struct {
@@ -148,18 +148,18 @@ func (blkup *blockUploader) Backup(source sourceInfo, parentObject udmrepo.ID, b
 	}, backupSize, nil
 }
 
-func (blkup *blockUploader) Restore(snapshot udmrepo.Snapshot, dest destInfo, bitmap cbt.Iterator, configs map[string]string) (int64, error) {
+func (blkup *blockUploader) Restore(snapshot udmrepo.Snapshot, dest destInfo, bitmap cbt.Iterator, configs map[string]string) (int64, int64, error) {
 	if bitmap == nil {
-		return 0, errors.New("bitmap is not available")
+		return 0, 0, errors.New("bitmap is not available")
 	}
 
 	meta, err := blkup.repoWriter.ReadMetadata(blkup.ctx, snapshot.RootObject.ID)
 	if err != nil {
-		return 0, errors.Wrapf(err, "error reading snapshot metadata for %s", snapshot.Description)
+		return 0, 0, errors.Wrapf(err, "error reading snapshot metadata for %s", snapshot.Description)
 	}
 
 	if len(meta.SubObjects) != 1 {
-		return 0, errors.Errorf("unexpected number of bdev object (%d) for snapshot %s", len(meta.SubObjects), snapshot.Description)
+		return 0, 0, errors.Errorf("unexpected number of bdev object (%d) for snapshot %s", len(meta.SubObjects), snapshot.Description)
 	}
 
 	sourceSize, err := getSourceSize(snapshot)
@@ -169,11 +169,11 @@ func (blkup *blockUploader) Restore(snapshot udmrepo.Snapshot, dest destInfo, bi
 	}
 
 	if sourceSize > meta.SubObjects[0].Size {
-		return 0, errors.Errorf("unexpected size (%v vs. %v) for bdev object %s", meta.SubObjects[0].Size, sourceSize, meta.SubObjects[0].Name)
+		return 0, 0, errors.Errorf("unexpected size (%v vs. %v) for bdev object %s", meta.SubObjects[0].Size, sourceSize, meta.SubObjects[0].Name)
 	}
 
 	if sourceSize > dest.size {
-		return 0, errors.Errorf("dest dev(%s) size is too small (%v vs. %v)", dest.path, dest.size, sourceSize)
+		return 0, 0, errors.Errorf("dest dev(%s) size is too small (%v vs. %v)", dest.path, dest.size, sourceSize)
 	}
 
 	reader, err := blkup.repoWriter.OpenObject(blkup.ctx, meta.SubObjects[0].ID, udmrepo.ObjectReadOptions{
@@ -181,16 +181,16 @@ func (blkup *blockUploader) Restore(snapshot udmrepo.Snapshot, dest destInfo, bi
 		PrefetchBudgetMB: 256,
 	})
 	if err != nil {
-		return 0, errors.Wrapf(err, "error opening bdev object %v", meta.SubObjects[0].Name)
+		return 0, 0, errors.Wrapf(err, "error opening bdev object %v", meta.SubObjects[0].Name)
 	}
 	defer reader.Close()
 
 	size, err := blkup.restoreData(reader, dest.dev, bitmap, sourceSize, dest.path)
 	if err != nil {
-		return 0, errors.Wrapf(err, "error restoring bdev object %s to volume %s", meta.SubObjects[0].Name, dest.path)
+		return 0, 0, errors.Wrapf(err, "error restoring bdev object %s to volume %s", meta.SubObjects[0].Name, dest.path)
 	}
 
-	return size, nil
+	return size, sourceSize, nil
 }
 
 func (blkup *blockUploader) backupObject(dev *os.File, dest udmrepo.ObjectWriter, bitmap cbt.Iterator, totalLength int64) (udmrepo.ID, int64, int64, error) {
@@ -444,6 +444,8 @@ func (blkup *blockUploader) restoreData(reader io.ReadSeeker, dest *os.File, bit
 		return written, errors.Wrap(writeErr, "error writing data")
 	}
 
+	blkup.progress.UpdateProgress(&uploader.Progress{BytesDone: totalLength, TotalBytes: totalLength})
+
 	return written, nil
 }
 
@@ -579,7 +581,7 @@ func restoreWriteProc(ctx context.Context, dest *os.File, resultChan chan readRe
 
 		result.resetBuffer(list)
 
-		progress.UpdateProgress(&uploader.Progress{BytesDone: written, TotalBytes: totalLength})
+		progress.UpdateProgress(&uploader.Progress{BytesDone: result.offset + length, TotalBytes: totalLength})
 	}
 
 	result.resetBuffer(list)

pkg/uploader/block/uploader_test.go

@@ -623,7 +623,7 @@ func TestBlockUploaderRestore(t *testing.T) {
 		repoWriter.On("ReadMetadata", mock.Anything, udmrepo.ID("root-id")).Return(nil, errors.New("meta not found"))
 
 		iterMock := cbtmocks.NewIterator(t)
-		_, err := blkup.Restore(udmrepo.Snapshot{RootObject: udmrepo.ObjectMetadata{ID: "root-id"}}, destInfo{}, iterMock, nil)
+		_, _, err := blkup.Restore(udmrepo.Snapshot{RootObject: udmrepo.ObjectMetadata{ID: "root-id"}}, destInfo{}, iterMock, nil)
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "meta not found")
 	})
@@ -685,7 +685,7 @@ func TestBlockUploaderRestore(t *testing.T) {
 		iterMock.On("Next").Return(uint64(0), false)
 		iterMock.On("BlockSize").Return(uint(1048576))
 
-		written, err := blkup.Restore(snap, dest, iterMock, nil)
+		written, _, err := blkup.Restore(snap, dest, iterMock, nil)
 		require.NoError(t, err)
 		assert.Equal(t, int64(1048576), written)
 	})
@@ -709,7 +709,7 @@ func TestBlockUploaderRestore(t *testing.T) {
 		dest := destInfo{size: 4194304, path: "/dev/target"}
 		iterMock := cbtmocks.NewIterator(t)
 
-		_, err := blkup.Restore(snap, dest, iterMock, nil)
+		_, _, err := blkup.Restore(snap, dest, iterMock, nil)
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "unexpected size (1048576 vs. 2097152) for bdev object bdev")
 	})
@@ -733,7 +733,7 @@ func TestBlockUploaderRestore(t *testing.T) {
 		dest := destInfo{size: 512, path: "/dev/small"}
 		iterMock := cbtmocks.NewIterator(t)
 
-		_, err := blkup.Restore(snap, dest, iterMock, nil)
+		_, _, err := blkup.Restore(snap, dest, iterMock, nil)
 		require.Error(t, err)
 		assert.Contains(t, err.Error(), "dest dev(/dev/small) size is too small")
 	})