restore: enhance error handling

pkg/restore/backoff.go

@@ -11,19 +11,19 @@ import (
 )
 
 var (
-	errNotLeader           = errors.New("not leader")
-	errKeyNotInRegion      = errors.New("key not in region")
-	errResp                = errors.New("response error")
-	errRewriteRuleNotFound = errors.New("rewrite rule not found")
-	errRangeIsEmpty        = errors.New("range is empty")
-	errGrpc                = errors.New("gRPC error")
+	errNotLeader           = errors.NewNoStackError("not leader")
+	errKeyNotInRegion      = errors.NewNoStackError("key not in region")
+	errResp                = errors.NewNoStackError("response error")
+	errRewriteRuleNotFound = errors.NewNoStackError("rewrite rule not found")
+	errRangeIsEmpty        = errors.NewNoStackError("range is empty")
+	errGrpc                = errors.NewNoStackError("gRPC error")
 
 	// TODO: add `error` field to `DownloadResponse` for distinguish the errors of gRPC
 	// and the errors of request
-	errBadFormat      = errors.New("bad format")
-	errWrongKeyPrefix = errors.New("wrong key prefix")
-	errFileCorrupted  = errors.New("file corrupted")
-	errCannotRead     = errors.New("cannot read externel storage")
+	errBadFormat      = errors.NewNoStackError("bad format")
+	errWrongKeyPrefix = errors.NewNoStackError("wrong key prefix")
+	errFileCorrupted  = errors.NewNoStackError("file corrupted")
+	errCannotRead     = errors.NewNoStackError("cannot read externel storage")
 )
 
 const (

pkg/restore/backoff_test.go

@@ -1,6 +1,7 @@
 package restore
 
 import (
+	"context"
 	"time"
 
 	. "github.com/pingcap/check"
@@ -27,7 +28,7 @@ func (s *testBackofferSuite) TearDownSuite(c *C) {
 
 func (s *testBackofferSuite) TestImporterBackoffer(c *C) {
 	var counter int
-	err := utils.WithRetry(func() error {
+	err := utils.WithRetry(context.Background(), func() error {
 		defer func() { counter++ }()
 		switch counter {
 		case 0:
@@ -48,7 +49,7 @@ func (s *testBackofferSuite) TestImporterBackoffer(c *C) {
 		delayTime:    time.Nanosecond,
 		maxDelayTime: time.Nanosecond,
 	}
-	err = utils.WithRetry(func() error {
+	err = utils.WithRetry(context.Background(), func() error {
 		defer func() { counter++ }()
 		return errResp
 	}, &backoffer)

pkg/restore/client.go

@@ -132,7 +132,7 @@ func (rc *Client) ResetTS(pdAddrs []string) error {
 	restoreTS := rc.backupMeta.GetEndVersion()
 	log.Info("reset pd timestamp", zap.Uint64("ts", restoreTS))
 	i := 0
-	return utils.WithRetry(func() error {
+	return utils.WithRetry(rc.ctx, func() error {
 		idx := i % len(pdAddrs)
 		return utils.ResetTS(pdAddrs[idx], restoreTS)
 	}, newResetTSBackoffer())

pkg/restore/import.go

@@ -159,7 +159,7 @@ func (importer *FileImporter) Import(file *backup.File, rewriteRules *RewriteRul
 		zap.Stringer("file", file),
 		zap.Binary("startKey", startKey),
 		zap.Binary("endKey", endKey))
-	err = utils.WithRetry(func() error {
+	err = utils.WithRetry(importer.ctx, func() error {
 		ctx, cancel := context.WithTimeout(importer.ctx, importScanResgionTime)
 		defer cancel()
 		// Scan regions covered by the file range
@@ -173,7 +173,7 @@ func (importer *FileImporter) Import(file *backup.File, rewriteRules *RewriteRul
 			info := regionInfo
 			// Try to download file.
 			var downloadMeta *import_sstpb.SSTMeta
-			err1 = utils.WithRetry(func() error {
+			err1 = utils.WithRetry(importer.ctx, func() error {
 				var e error
 				downloadMeta, e = importer.downloadSST(info, file, rewriteRules)
 				return e

pkg/utils/retry.go

@@ -1,6 +1,9 @@
 package utils
 
-import "time"
+import (
+	"context"
+	"time"
+)
 
 // RetryableFunc presents a retryable opreation
 type RetryableFunc func() error
@@ -15,6 +18,7 @@ type Backoffer interface {
 
 // WithRetry retrys a given operation with a backoff policy
 func WithRetry(
+	ctx context.Context,
 	retryableFunc RetryableFunc,
 	backoffer Backoffer,
 ) error {
@@ -23,7 +27,12 @@ func WithRetry(
 		err := retryableFunc()
 		if err != nil {
 			lastErr = err
-			time.Sleep(backoffer.NextBackoff(err))
+			select {
+			case <-ctx.Done():
+				return lastErr
+			default:
+				time.Sleep(backoffer.NextBackoff(err))
+			}
 		} else {
 			return nil
 		}