planner, executor: fix haven't track the memory usage of PointGet/BatchPointGet

executor/batch_point_get.go

@@ -39,6 +39,7 @@ import (
  "github.com/pingcap/tidb/util/hack"
  "github.com/pingcap/tidb/util/logutil/consistency"
  "github.com/pingcap/tidb/util/math"
+ "github.com/pingcap/tidb/util/memory"
  "github.com/pingcap/tidb/util/rowcodec"
  "github.com/tikv/client-go/v2/txnkv/txnsnapshot"
 )
@@ -82,6 +83,8 @@ type BatchPointGetExec struct {
  snapshot kv.Snapshot
  stats *runtimeStatsWithSnapshot
  cacheTable kv.MemBuffer
+
+ memTracker *memory.Tracker
 }
 
 // buildVirtualColumnInfo saves virtual column indices and sort them in definition order
@@ -165,6 +168,9 @@ func (e *BatchPointGetExec) Open(context.Context) error {
  }
  e.snapshot = snapshot
  e.batchGetter = batchGetter
+
+ e.memTracker = memory.NewTracker(e.id, -1)
+ e.memTracker.AttachTo(e.ctx.GetSessionVars().StmtCtx.MemTracker)
  return nil
 }
 
@@ -212,6 +218,12 @@ func MockNewCacheTableSnapShot(snapshot kv.Snapshot, memBuffer kv.MemBuffer) *ca
 
 // Close implements the Executor interface.
 func (e *BatchPointGetExec) Close() error {
+
+ if e.memTracker != nil {
+ e.memTracker.Consume(-e.memTracker.BytesConsumed())
+ e.memTracker = nil
+ }
+
  if e.runtimeStats != nil && e.snapshot != nil {
  e.snapshot.SetOption(kv.CollectRuntimeStats, nil)
  }
@@ -271,6 +283,7 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error {
  dedup := make(map[hack.MutableString]struct{})
  toFetchIndexKeys := make([]kv.Key, 0, len(e.idxVals))
  for _, idxVals := range e.idxVals {
+ e.memTracker.Consume(types.EstimatedMemUsage(idxVals, 1))
  // For all x, 'x IN (null)' evaluate to null, so the query get no result.
  if datumsContainNull(idxVals) {
  continue
@@ -498,7 +511,9 @@ func (e *BatchPointGetExec) initialize(ctx context.Context) error {
  continue
  }
  e.values = append(e.values, val)
+ e.memTracker.Consume(int64(cap(val)))
  handles = append(handles, e.handles[i])
+ e.memTracker.Consume(e.handles[i].MemoryUsage())
  if e.lock && rc {
  existKeys = append(existKeys, key)
  // when e.handles is set in builder directly, index should be primary key and the plan is CommonHandleRead

executor/point_get.go

@@ -39,6 +39,7 @@ import (
  "github.com/pingcap/tidb/util/codec"
  "github.com/pingcap/tidb/util/execdetails"
  "github.com/pingcap/tidb/util/logutil/consistency"
+ "github.com/pingcap/tidb/util/memory"
  "github.com/pingcap/tidb/util/rowcodec"
  "github.com/tikv/client-go/v2/txnkv/txnsnapshot"
 )
@@ -103,6 +104,8 @@ type PointGetExecutor struct {
 
  stats *runtimeStatsWithSnapshot
  cacheTable kv.MemBuffer
+
+ memTracker *memory.Tracker
 }
 
 // Init set fields needed for PointGetExecutor reuse, this does NOT change baseExecutor field
@@ -174,6 +177,8 @@ func (e *PointGetExecutor) Open(context.Context) error {
  if readReplicaType.IsFollowerRead() {
  e.snapshot.SetOption(kv.ReplicaRead, readReplicaType)
  }
+ e.memTracker = memory.NewTracker(e.id, -1)
+ e.memTracker.AttachTo(e.ctx.GetSessionVars().StmtCtx.MemTracker)
  e.snapshot.SetOption(kv.TaskID, e.ctx.GetSessionVars().StmtCtx.TaskID)
  e.snapshot.SetOption(kv.ReadReplicaScope, e.readReplicaScope)
  e.snapshot.SetOption(kv.IsStalenessReadOnly, e.isStaleness)
@@ -198,6 +203,11 @@ func (e *PointGetExecutor) Open(context.Context) error {
 
 // Close implements the Executor interface.
 func (e *PointGetExecutor) Close() error {
+ if e.memTracker != nil {
+ e.memTracker.Consume(-e.memTracker.BytesConsumed())
+ e.memTracker = nil
+ }
+
  if e.runtimeStats != nil && e.snapshot != nil {
  e.snapshot.SetOption(kv.CollectRuntimeStats, nil)
  }
@@ -248,14 +258,15 @@ func (e *PointGetExecutor) Next(ctx context.Context, req *chunk.Chunk) error {
  if err != nil && !kv.ErrNotExist.Equal(err) {
  return err
  }
+ e.memTracker.Consume(int64(cap(e.idxKey)))
 
  e.handleVal, err = e.get(ctx, e.idxKey)
  if err != nil {
  if !kv.ErrNotExist.Equal(err) {
  return err
  }
  }
-
+ e.memTracker.Consume(int64(cap(e.handleVal)))
  // try lock the index key if isolation level is not read consistency
  // also lock key if read consistency read a value
  if !e.ctx.GetSessionVars().IsPessimisticReadConsistency() || len(e.handleVal) > 0 {
@@ -301,7 +312,7 @@ func (e *PointGetExecutor) Next(ctx context.Context, req *chunk.Chunk) error {
  })
  }
  }
-
+ e.memTracker.Consume(e.handle.MemoryUsage())
  key := tablecodec.EncodeRowKeyWithHandle(tblID, e.handle)
  val, err := e.getAndLock(ctx, key)
  if err != nil {

kv/key.go

@@ -157,6 +157,8 @@
  Data() ([]types.Datum, error)
  // String implements the fmt.Stringer interface.
  String() string
+ // MemoryUsage returns the total memory usage of a Handle in bytes.
+ MemoryUsage() int64
 }
 
 // IntHandle implement the Handle interface for int64 type handle.
@@ -228,6 +230,11 @@
  return strconv.FormatInt(int64(ih), 10)
 }
 
+// MemoryUsage implements the Handle interface.
+func (ih IntHandle) MemoryUsage() int64 {
+ return 8
+}
+
 // CommonHandle implements the Handle interface for non-int64 type handle.
 type CommonHandle struct {
  encoded []byte
@@ -347,6 +354,11 @@
  return fmt.Sprintf("{%s}", strings.Join(strs, ", "))
 }
 
+// MemoryUsage implements the Handle interface.
+func (ch *CommonHandle) MemoryUsage() int64 {
+ return int64(cap(ch.encoded))
+}
+
 // HandleMap is the map for Handle.
 type HandleMap struct {
  ints map[int64]interface{}

session/session_test.go

@@ -3458,6 +3458,32 @@ func (s *testSessionSuite2) TestStmtHints(c *C) {
  c.Assert(tk.Se.GetSessionVars().GetReplicaRead(), Equals, kv.ReplicaReadFollower)
 }
 
+// Test memory track in POINT_GET/BATCH_POINT_GET for issue 21653
+func (s *testSessionSuite2) TestPointGetMemoryTracking(c *C) {
+ tk := testkit.NewTestKit(c, s.store)
+ tk.MustExec("use test")
+ tk.MustExec("drop table if exists t1;")
+ tk.MustExec("create table t1(a int primary key);")
+ config.UpdateGlobal(func(conf *config.Config) {
+ conf.OOMAction = config.OOMActionCancel
+ })
+ tk.MustExec("insert into t1 values (1);")
+ tk.MustQuery("select * from t1 where a = 1;").Check(testkit.Rows("1"))
+ tk.MustExec("set tidb_mem_quota_query=1;")
+ err := tk.QueryToErr("select * from t1 where a = 1;")
+ c.Assert(err, NotNil)
+ c.Assert(err.Error(), Matches, "Out Of Memory Quota!.*")
+ err = tk.QueryToErr("select * from t1 where a in (1,2,3);")
+ c.Assert(err, NotNil)
+ c.Assert(err.Error(), Matches, "Out Of Memory Quota!.*")
+ tk.MustExec("set tidb_mem_quota_query=1024;")
+ tk.MustQuery("select * from t1 where a = 1;").Check(testkit.Rows("1"))
+ err = tk.QueryToErr("select * from t1 where a = 1;")
+ c.Check(err, IsNil)
+ err = tk.QueryToErr("select * from t1 where a in (1,2,3);")
+ c.Check(err, IsNil)
+}
+
 func (s *testSessionSuite3) TestPessimisticLockOnPartition(c *C) {
  // This test checks that 'select ... for update' locks the partition instead of the table.
  // Cover a bug that table ID is used to encode the lock key mistakenly.