diff --git a/br/pkg/restore/import_retry.go b/br/pkg/restore/import_retry.go index 6f3b9fc1cca53..4a6bdf1b8afcf 100644 --- a/br/pkg/restore/import_retry.go +++ b/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 diff --git a/br/pkg/restore/import_retry_test.go b/br/pkg/restore/import_retry_test.go index 6f3d8f490ef13..4e885657f998f 100644 --- a/br/pkg/restore/import_retry_test.go +++ b/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) { diff --git a/br/pkg/utils/backoff.go b/br/pkg/utils/backoff.go index bff2490b56650..08df56c1c1e53 100644 --- a/br/pkg/utils/backoff.go +++ b/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 {