Add --read-data-subset flag to check command

Signed-off-by: Igor Fedorenko <igor@ifedorenko.com>

cmd/restic/cmd_check.go

@@ -3,6 +3,8 @@ package main
 import (
 	"fmt"
 	"os"
+	"strconv"
+	"strings"
 	"time"
 
 	"github.com/spf13/cobra"
@@ -26,13 +28,17 @@ repository and not use a local cache.
 	RunE: func(cmd *cobra.Command, args []string) error {
 		return runCheck(checkOptions, globalOptions, args)
 	},
+	PreRunE: func(cmd *cobra.Command, args []string) error {
+		return checkFlags(checkOptions)
+	},
 }
 
 // CheckOptions bundles all options for the 'check' command.
 type CheckOptions struct {
-	ReadData    bool
-	CheckUnused bool
-	WithCache   bool
+	ReadData       bool
+	ReadDataSubset string
+	CheckUnused    bool
+	WithCache      bool
 }
 
 var checkOptions CheckOptions
@@ -42,10 +48,45 @@ func init() {
 
 	f := cmdCheck.Flags()
 	f.BoolVar(&checkOptions.ReadData, "read-data", false, "read all data blobs")
+	f.StringVar(&checkOptions.ReadDataSubset, "read-data-subset", "", "read subset of data packs")
 	f.BoolVar(&checkOptions.CheckUnused, "check-unused", false, "find unused blobs")
 	f.BoolVar(&checkOptions.WithCache, "with-cache", false, "use the cache")
 }
 
+func checkFlags(opts CheckOptions) error {
+	if opts.ReadData && opts.ReadDataSubset != "" {
+		return errors.Fatalf("check flags --read-data and --read-data-subset cannot be used together")
+	}
+	if opts.ReadDataSubset != "" {
+		dataSubset, err := stringToIntSlice(opts.ReadDataSubset)
+		if err != nil || len(dataSubset) != 2 {
+			return errors.Fatalf("check flag --read-data-subset must have two positive integer values, e.g. --read-data-subset=1/2")
+		}
+		if dataSubset[0] == 0 || dataSubset[1] == 0 || dataSubset[0] > dataSubset[1] {
+			return errors.Fatalf("check flag --read-data-subset=n/t values must be positive integers, and n <= t, e.g. --read-data-subset=1/2")
+		}
+	}
+
+	return nil
+}
+
+// stringToIntSlice converts string to []uint, using '/' as element separator
+func stringToIntSlice(param string) (split []uint, err error) {
+	if param == "" {
+		return nil, nil
+	}
+	parts := strings.Split(param, "/")
+	result := make([]uint, len(parts))
+	for idx, part := range parts {
+		uintval, err := strconv.ParseUint(part, 10, 0)
+		if err != nil {
+			return nil, err
+		}
+		result[idx] = uint(uintval)
+	}
+	return result, nil
+}
+
 func newReadProgress(gopts GlobalOptions, todo restic.Stat) *restic.Progress {
 	if gopts.Quiet {
 		return nil
@@ -158,13 +199,25 @@ func runCheck(opts CheckOptions, gopts GlobalOptions, args []string) error {
 		}
 	}
 
-	if opts.ReadData {
-		Verbosef("read all data\n")
+	doReadData := func(bucket, totalBuckets uint) {
+		packs := restic.IDSet{}
+		for pack := range chkr.GetPacks() {
+			if (uint(pack[0]) % totalBuckets) == (bucket - 1) {
+				packs.Insert(pack)
+			}
+		}
+		packCount := uint64(len(packs))
 
-		p := newReadProgress(gopts, restic.Stat{Blobs: chkr.CountPacks()})
+		if packCount < chkr.CountPacks() {
+			Verbosef(fmt.Sprintf("read group #%d of %d data packs (out of total %d packs in %d groups)\n", bucket, packCount, chkr.CountPacks(), totalBuckets))
+		} else {
+			Verbosef("read all data\n")
+		}
+
+		p := newReadProgress(gopts, restic.Stat{Blobs: packCount})
 		errChan := make(chan error)
 
-		go chkr.ReadData(gopts.ctx, p, errChan)
+		go chkr.ReadPacks(gopts.ctx, packs, p, errChan)
 
 		for err := range errChan {
 			errorsFound = true
@@ -172,6 +225,14 @@ func runCheck(opts CheckOptions, gopts GlobalOptions, args []string) error {
 		}
 	}
 
+	switch {
+	case opts.ReadData:
+		doReadData(1, 1)
+	case opts.ReadDataSubset != "":
+		dataSubset, _ := stringToIntSlice(opts.ReadDataSubset)
+		doReadData(dataSubset[0], dataSubset[1])
+	}
+
 	if errorsFound {
 		return errors.Fatal("repository contains errors")
 	}