diff --git a/pkg/distsql/request_builder.go b/pkg/distsql/request_builder.go index 313dbdd79c44d..c9f63b74f48a3 100644 --- a/pkg/distsql/request_builder.go +++ b/pkg/distsql/request_builder.go @@ -20,6 +20,7 @@ import ( "sort" "sync/atomic" "time" + "unsafe" "github.com/pingcap/errors" "github.com/pingcap/failpoint" @@ -749,6 +750,9 @@ func indexRangesToKVWithoutSplit(dctx *distsqlctx.DistSQLContext, tids []int64, krs[i] = make([]kv.KeyRange, 0, len(ranges)) } + if memTracker != nil { + memTracker.Consume(int64(unsafe.Sizeof(kv.KeyRange{})) * int64(len(ranges))) + } const checkSignalStep = 8 var estimatedMemUsage int64 // encodeIndexKey and EncodeIndexSeekKey is time-consuming, thus we need to @@ -777,6 +781,9 @@ func indexRangesToKVWithoutSplit(dctx *distsqlctx.DistSQLContext, tids []int64, if interruptSignal != nil && interruptSignal.Load().(bool) { return kv.NewPartitionedKeyRanges(nil), nil } + if memTracker != nil { + memTracker.HandleKillSignal() + } } } return kv.NewPartitionedKeyRanges(krs), nil diff --git a/pkg/executor/distsql.go b/pkg/executor/distsql.go index f3aaf6f306fba..ef1f426adaa5b 100644 --- a/pkg/executor/distsql.go +++ b/pkg/executor/distsql.go @@ -540,6 +540,14 @@ func (e *IndexLookUpExecutor) Open(ctx context.Context) error { return err } } + + if e.memTracker != nil { + e.memTracker.Reset() + } else { + e.memTracker = memory.NewTracker(e.ID(), -1) + } + e.memTracker.AttachTo(e.stmtMemTracker) + err = e.buildTableKeyRanges() if err != nil { return err @@ -570,7 +578,7 @@ func (e *IndexLookUpExecutor) buildTableKeyRanges() (err error) { if e.index.ID == -1 { kvRange, err = distsql.CommonHandleRangesToKVRanges(dctx, []int64{physicalID}, ranges) } else { - kvRange, err = distsql.IndexRangesToKVRanges(dctx, physicalID, e.index.ID, ranges) + kvRange, err = distsql.IndexRangesToKVRangesWithInterruptSignal(dctx, physicalID, e.index.ID, ranges, e.memTracker, nil) } if err != nil { return err @@ -583,7 +591,7 @@ func (e *IndexLookUpExecutor) buildTableKeyRanges() (err error) { if e.index.ID == -1 { kvRanges, err = distsql.CommonHandleRangesToKVRanges(dctx, []int64{physicalID}, e.ranges) } else { - kvRanges, err = distsql.IndexRangesToKVRanges(dctx, physicalID, e.index.ID, e.ranges) + kvRanges, err = distsql.IndexRangesToKVRangesWithInterruptSignal(dctx, physicalID, e.index.ID, e.ranges, e.memTracker, nil) } e.kvRanges = kvRanges.FirstPartitionRange() } diff --git a/pkg/executor/executor_pkg_test.go b/pkg/executor/executor_pkg_test.go index b0d2133af38a3..b965d240b1732 100644 --- a/pkg/executor/executor_pkg_test.go +++ b/pkg/executor/executor_pkg_test.go @@ -131,7 +131,7 @@ func TestBuildKvRangesForIndexJoinWithoutCwcAndWithMemoryTracker(t *testing.T) { } require.Equal(t, 2*bytesConsumed1, bytesConsumed2) - require.Equal(t, int64(20760), bytesConsumed1) + require.Equal(t, int64(25560), bytesConsumed1) } func generateIndexRange(vals ...int64) *ranger.Range {