From 5e7a36dc7670eee47b46ed4f2d9e0e4d094ce479 Mon Sep 17 00:00:00 2001 From: Ti Chi Robot Date: Tue, 23 Apr 2024 11:36:10 +0800 Subject: [PATCH] executor: add memory tracker for memtableRetriever (#51455) (#52823) close pingcap/tidb#51456 --- executor/builder.go | 9 ++++++--- executor/infoschema_reader.go | 4 ++++ 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/executor/builder.go b/executor/builder.go index c19bf82946e97..c5ffa80884c32 100644 --- a/executor/builder.go +++ b/executor/builder.go @@ -2003,13 +2003,16 @@ func (b *executorBuilder) buildMemTable(v *plannercore.PhysicalMemTable) Executo strings.ToLower(infoschema.ClusterTableMemoryUsage), strings.ToLower(infoschema.ClusterTableMemoryUsageOpsHistory), strings.ToLower(infoschema.TableResourceGroups): + memTracker := memory.NewTracker(v.ID(), -1) + memTracker.AttachTo(b.ctx.GetSessionVars().StmtCtx.MemTracker) return &MemTableReaderExec{ baseExecutor: newBaseExecutor(b.ctx, v.Schema(), v.ID()), table: v.Table, retriever: &memtableRetriever{ - table: v.Table, - columns: v.Columns, - extractor: v.Extractor, + table: v.Table, + columns: v.Columns, + extractor: v.Extractor, + memTracker: memTracker, }, } case strings.ToLower(infoschema.TableTiDBTrx), diff --git a/executor/infoschema_reader.go b/executor/infoschema_reader.go index 29938e2462fab..e414829d4e969 100644 --- a/executor/infoschema_reader.go +++ b/executor/infoschema_reader.go @@ -87,6 +87,7 @@ type memtableRetriever struct { retrieved bool initialized bool extractor plannercore.MemTablePredicateExtractor + memTracker *memory.Tracker } // retrieve implements the infoschemaRetriever interface @@ -190,6 +191,9 @@ func (e *memtableRetriever) retrieve(ctx context.Context, sctx sessionctx.Contex return nil, err } e.initialized = true + if e.memTracker != nil { + e.memTracker.Consume(calculateDatumsSize(e.rows)) + } } // Adjust the amount of each return