Merge pull request #7512 from qiuming-best/support-parallel-restore

Make parallel restore configurable

pkg/uploader/kopia/snapshot.go

@@ -401,6 +401,8 @@ func Restore(ctx context.Context, rep repo.RepositoryWriter, progress *Progress,
 		IgnorePermissionErrors: true,
 	}
 
+	restoreConcurrency := runtime.NumCPU()
+
 	if len(uploaderCfg) > 0 {
 		writeSparseFiles, err := uploaderutil.GetWriteSparseFiles(uploaderCfg)
 		if err != nil {
@@ -409,9 +411,17 @@ func Restore(ctx context.Context, rep repo.RepositoryWriter, progress *Progress,
 		if writeSparseFiles {
 			fsOutput.WriteSparseFiles = true
 		}
+
+		concurrency, err := uploaderutil.GetRestoreConcurrency(uploaderCfg)
+		if err != nil {
+			return 0, 0, errors.Wrap(err, "failed to get parallel restore uploader config")
+		}
+		if concurrency > 0 {
+			restoreConcurrency = concurrency
+		}
 	}
 
-	log.Debugf("Restore filesystem output %v", fsOutput)
+	log.Debugf("Restore filesystem output %v, concurrency %d", fsOutput, restoreConcurrency)
 
 	err = fsOutput.Init(ctx)
 	if err != nil {
@@ -426,7 +436,7 @@ func Restore(ctx context.Context, rep repo.RepositoryWriter, progress *Progress,
 	}
 
 	stat, err := restoreEntryFunc(kopiaCtx, rep, output, rootEntry, restore.Options{
-		Parallel:               runtime.NumCPU(),
+		Parallel:               restoreConcurrency,
 		RestoreDirEntryAtDepth: math.MaxInt32,
 		Cancel:                 cancleCh,
 		ProgressCallback: func(ctx context.Context, stats restore.Stats) {

pkg/uploader/provider/restic.go

@@ -246,5 +246,9 @@ func (rp *resticProvider) parseRestoreExtraFlags(uploaderCfg map[string]string)
 		extraFlags = append(extraFlags, "--sparse")
 	}
 
+	if restoreConcurrency, err := uploaderutil.GetRestoreConcurrency(uploaderCfg); err == nil && restoreConcurrency > 0 {
+		return extraFlags, errors.New("restic does not support parallel restore")
+	}
+
 	return extraFlags, nil
 }

pkg/uploader/provider/restic_test.go

@@ -434,6 +434,13 @@ func TestParseUploaderConfig(t *testing.T) {
 			},
 			expectedFlags: []string{},
 		},
+		{
+			name: "RestoreConcorrency",
+			uploaderConfig: map[string]string{
+				"Parallel": "5",
+			},
+			expectedFlags: []string{},
+		},
 	}
 
 	for _, testCase := range testCases {

pkg/uploader/util/uploader_config.go

@@ -27,6 +27,7 @@ import (
 const (
 	ParallelFilesUpload = "ParallelFilesUpload"
 	WriteSparseFiles    = "WriteSparseFiles"
+	RestoreConcurrency  = "ParallelFilesDownload"
 )
 
 func StoreBackupConfig(config *velerov1api.UploaderConfigForBackup) map[string]string {
@@ -42,6 +43,10 @@ func StoreRestoreConfig(config *velerov1api.UploaderConfigForRestore) map[string
 	} else {
 		data[WriteSparseFiles] = strconv.FormatBool(false)
 	}
+
+	if config.ParallelFilesDownload > 0 {
+		data[RestoreConcurrency] = strconv.Itoa(config.ParallelFilesDownload)
+	}
 	return data
 }
 
@@ -68,3 +73,15 @@ func GetWriteSparseFiles(uploaderCfg map[string]string) (bool, error) {
 	}
 	return false, nil
 }
+
+func GetRestoreConcurrency(uploaderCfg map[string]string) (int, error) {
+	restoreConcurrency, ok := uploaderCfg[RestoreConcurrency]
+	if ok {
+		restoreConcurrencyInt, err := strconv.Atoi(restoreConcurrency)
+		if err != nil {
+			return 0, errors.Wrap(err, "failed to parse RestoreConcurrency config")
+		}
+		return restoreConcurrencyInt, nil
+	}
+	return 0, nil
+}

pkg/uploader/util/uploader_config_test.go

@@ -78,6 +78,16 @@ func TestStoreRestoreConfig(t *testing.T) {
 				WriteSparseFiles: "false", // Assuming default value is false for nil case
 			},
 		},
+		{
+			name: "Parallel is set",
+			config: &velerov1api.UploaderConfigForRestore{
+				ParallelFilesDownload: 5,
+			},
+			expectedData: map[string]string{
+				RestoreConcurrency: "5",
+				WriteSparseFiles:   "false",
+			},
+		},
 	}
 
 	for _, tc := range testCases {
@@ -180,3 +190,53 @@ func TestGetWriteSparseFiles(t *testing.T) {
 		})
 	}
 }
+
+func TestGetRestoreConcurrency(t *testing.T) {
+	testCases := []struct {
+		Name             string
+		UploaderCfg      map[string]string
+		ExpectedResult   int
+		ExpectedError    bool
+		ExpectedErrorMsg string
+	}{
+		{
+			Name:           "Valid Configuration",
+			UploaderCfg:    map[string]string{RestoreConcurrency: "10"},
+			ExpectedResult: 10,
+			ExpectedError:  false,
+		},
+		{
+			Name:           "Missing Configuration",
+			UploaderCfg:    map[string]string{},
+			ExpectedResult: 0,
+			ExpectedError:  false,
+		},
+		{
+			Name:             "Invalid Configuration",
+			UploaderCfg:      map[string]string{RestoreConcurrency: "not_an_integer"},
+			ExpectedResult:   0,
+			ExpectedError:    true,
+			ExpectedErrorMsg: "failed to parse RestoreConcurrency config: strconv.Atoi: parsing \"not_an_integer\": invalid syntax",
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.Name, func(t *testing.T) {
+			result, err := GetRestoreConcurrency(tc.UploaderCfg)
+
+			if tc.ExpectedError {
+				if err.Error() != tc.ExpectedErrorMsg {
+					t.Errorf("Expected error message %s, but got %s", tc.ExpectedErrorMsg, err.Error())
+				}
+			} else {
+				if err != nil {
+					t.Errorf("Expected no error, but got %v", err)
+				}
+			}
+
+			if result != tc.ExpectedResult {
+				t.Errorf("Expected result %d, but got %d", tc.ExpectedResult, result)
+			}
+		})
+	}
+}