br: continue to retry when pitr met the error about memory is limited (#41982)

br/pkg/restore/import_retry.go

@@ -4,9 +4,11 @@ package restore
 
 import (
 	"context"
+	"strings"
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/kvproto/pkg/errorpb"
 	"github.com/pingcap/kvproto/pkg/import_sstpb"
 	"github.com/pingcap/kvproto/pkg/metapb"
@@ -85,6 +87,21 @@ func (o *OverRegionsInRangeController) tryFindLeader(ctx context.Context, region
 
 // handleInRegionError handles the error happens internal in the region. Update the region info, and perform a suitable backoff.
 func (o *OverRegionsInRangeController) handleInRegionError(ctx context.Context, result RPCResult, region *split.RegionInfo) (cont bool) {
+	if result.StoreError.GetServerIsBusy() != nil {
+		if strings.Contains(result.StoreError.GetMessage(), "memory is limited") {
+			sleepDuration := 15 * time.Second
+
+			failpoint.Inject("hint-memory-is-limited", func(val failpoint.Value) {
+				if val.(bool) {
+					logutil.CL(ctx).Debug("failpoint hint-memory-is-limited injected.")
+					sleepDuration = 100 * time.Microsecond
+				}
+			})
+			time.Sleep(sleepDuration)
+			return true
+		}
+	}
+
 	if nl := result.StoreError.GetNotLeader(); nl != nil {
 		if nl.Leader != nil {
 			region.Leader = nl.Leader

br/pkg/restore/import_retry_test.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	backuppb "github.com/pingcap/kvproto/pkg/brpb"
 	"github.com/pingcap/kvproto/pkg/errorpb"
 	"github.com/pingcap/kvproto/pkg/import_sstpb"
@@ -163,6 +164,48 @@ func TestServerIsBusy(t *testing.T) {
 	require.NoError(t, err)
 	assertRegions(t, idEqualsTo2Regions, "aay", "bba")
 	assertRegions(t, meetRegions, "", "aay", "bba", "bbh", "cca", "")
+	require.Equal(t, rs.RetryTimes(), 1)
+}
+
+func TestServerIsBusyWithMemoryIsLimited(t *testing.T) {
+	_ = failpoint.Enable("github.com/pingcap/tidb/br/pkg/restore/hint-memory-is-limited", "return(true)")
+	defer func() {
+		_ = failpoint.Disable("github.com/pingcap/tidb/br/pkg/restore/hint-memory-is-limited")
+	}()
+
+	// region: [, aay), [aay, bba), [bba, bbh), [bbh, cca), [cca, )
+	cli := initTestClient(false)
+	rs := utils.InitialRetryState(2, 0, 0)
+	ctl := restore.OverRegionsInRange([]byte(""), []byte(""), cli, &rs)
+	ctx := context.Background()
+
+	serverIsBusy := errorpb.Error{
+		Message: "memory is limited",
+		ServerIsBusy: &errorpb.ServerIsBusy{
+			Reason: "",
+		},
+	}
+	// record the regions we didn't touch.
+	meetRegions := []*split.RegionInfo{}
+	// record all regions we meet with id == 2.
+	idEqualsTo2Regions := []*split.RegionInfo{}
+	theFirstRun := true
+	err := ctl.Run(ctx, func(ctx context.Context, r *split.RegionInfo) restore.RPCResult {
+		if theFirstRun && r.Region.Id == 2 {
+			idEqualsTo2Regions = append(idEqualsTo2Regions, r)
+			theFirstRun = false
+			return restore.RPCResult{
+				StoreError: &serverIsBusy,
+			}
+		}
+		meetRegions = append(meetRegions, r)
+		return restore.RPCResultOK()
+	})
+
+	require.NoError(t, err)
+	assertRegions(t, idEqualsTo2Regions, "aay", "bba")
+	assertRegions(t, meetRegions, "", "aay", "bba", "bbh", "cca", "")
+	require.Equal(t, rs.RetryTimes(), 0)
 }
 
 func printRegion(name string, infos []*split.RegionInfo) {

br/pkg/utils/backoff.go

@@ -82,6 +82,12 @@ func (rs *RetryState) RecordRetry() {
 	rs.retryTimes++
 }
 
+// RetryTimes returns the retry times.
+// usage: unit test.
+func (rs *RetryState) RetryTimes() int {
+	return rs.retryTimes
+}
+
 // Attempt implements the `Backoffer`.
 // TODO: Maybe use this to replace the `exponentialBackoffer` (which is nearly homomorphic to this)?
 func (rs *RetryState) Attempt() int {