From 4cb200ddcbbbceb7cfaf51a82c6d7432cbfb9af8 Mon Sep 17 00:00:00 2001
From: Ti Chi Robot <ti-community-prow-bot@tidb.io>
Date: Fri, 27 Oct 2023 11:23:04 +0800
Subject: [PATCH] statistics: check Killed in the GenJSONTableFromStats
 (#47778) (#47826)

close pingcap/tidb#47779
---
 pkg/statistics/handle/storage/json.go         | 23 +++++++++++++++----
 .../handle/storage/stats_read_writer.go       |  9 ++++----
 pkg/statistics/handle/util/util.go            | 14 +++++++++++
 3 files changed, 38 insertions(+), 8 deletions(-)

diff --git a/pkg/statistics/handle/storage/json.go b/pkg/statistics/handle/storage/json.go
index af90c56d3d657..8d4fff42246ee 100644
--- a/pkg/statistics/handle/storage/json.go
+++ b/pkg/statistics/handle/storage/json.go
@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"encoding/json"
 	"io"
+	"sync/atomic"
 	"time"
 
 	"github.com/klauspost/compress/gzip"
@@ -31,6 +32,7 @@ import (
 	"github.com/pingcap/tidb/pkg/types"
 	compressutil "github.com/pingcap/tidb/pkg/util/compress"
 	"github.com/pingcap/tidb/pkg/util/logutil"
+	"github.com/pingcap/tidb/pkg/util/memory"
 	"go.uber.org/zap"
 )
 
@@ -88,7 +90,10 @@ func dumpJSONCol(hist *statistics.Histogram, cmsketch *statistics.CMSketch, topn
 }
 
 // GenJSONTableFromStats generate jsonTable from tableInfo and stats
-func GenJSONTableFromStats(dbName string, tableInfo *model.TableInfo, tbl *statistics.Table) (*util.JSONTable, error) {
+func GenJSONTableFromStats(sctx sessionctx.Context, dbName string, tableInfo *model.TableInfo, tbl *statistics.Table) (*util.JSONTable, error) {
+	tracker := memory.NewTracker(memory.LabelForAnalyzeMemory, -1)
+	tracker.AttachTo(sctx.GetSessionVars().MemTracker)
+	defer tracker.Detach()
 	jsonTbl := &util.JSONTable{
 		DatabaseName: dbName,
 		TableName:    tableInfo.Name.L,
@@ -104,11 +109,21 @@ func GenJSONTableFromStats(dbName string, tableInfo *model.TableInfo, tbl *stati
 		if err != nil {
 			return nil, errors.Trace(err)
 		}
-		jsonTbl.Columns[col.Info.Name.L] = dumpJSONCol(hist, col.CMSketch, col.TopN, col.FMSketch, &col.StatsVer)
+		proto := dumpJSONCol(hist, col.CMSketch, col.TopN, col.FMSketch, &col.StatsVer)
+		tracker.Consume(proto.TotalMemoryUsage())
+		if atomic.LoadUint32(&sctx.GetSessionVars().Killed) == 1 {
+			return nil, errors.Trace(statistics.ErrQueryInterrupted)
+		}
+		jsonTbl.Columns[col.Info.Name.L] = proto
+		col.FMSketch.DestroyAndPutToPool()
 	}
-
 	for _, idx := range tbl.Indices {
-		jsonTbl.Indices[idx.Info.Name.L] = dumpJSONCol(&idx.Histogram, idx.CMSketch, idx.TopN, nil, &idx.StatsVer)
+		proto := dumpJSONCol(&idx.Histogram, idx.CMSketch, idx.TopN, nil, &idx.StatsVer)
+		tracker.Consume(proto.TotalMemoryUsage())
+		if atomic.LoadUint32(&sctx.GetSessionVars().Killed) == 1 {
+			return nil, errors.Trace(statistics.ErrQueryInterrupted)
+		}
+		jsonTbl.Indices[idx.Info.Name.L] = proto
 	}
 	jsonTbl.ExtStats = dumpJSONExtendedStats(tbl.ExtendedStats)
 	return jsonTbl, nil
diff --git a/pkg/statistics/handle/storage/stats_read_writer.go b/pkg/statistics/handle/storage/stats_read_writer.go
index 7feb85f7794e1..bf9f84b7bd2ec 100644
--- a/pkg/statistics/handle/storage/stats_read_writer.go
+++ b/pkg/statistics/handle/storage/stats_read_writer.go
@@ -493,17 +493,18 @@ func (s *statsReadWriter) TableStatsToJSON(dbName string, tableInfo *model.Table
 	if err != nil || tbl == nil {
 		return nil, err
 	}
+	var jsonTbl *util.JSONTable
 	err = util.CallWithSCtx(s.statsHandler.SPool(), func(sctx sessionctx.Context) error {
 		tbl.Version, tbl.ModifyCount, tbl.RealtimeCount, err = StatsMetaByTableIDFromStorage(sctx, physicalID, snapshot)
+		if err != nil {
+			return err
+		}
+		jsonTbl, err = GenJSONTableFromStats(sctx, dbName, tableInfo, tbl)
 		return err
 	})
 	if err != nil {
 		return nil, err
 	}
-	jsonTbl, err := GenJSONTableFromStats(dbName, tableInfo, tbl)
-	if err != nil {
-		return nil, err
-	}
 	return jsonTbl, nil
 }
 
diff --git a/pkg/statistics/handle/util/util.go b/pkg/statistics/handle/util/util.go
index 202791b6a2305..50add779777d0 100644
--- a/pkg/statistics/handle/util/util.go
+++ b/pkg/statistics/handle/util/util.go
@@ -280,3 +280,17 @@ type JSONColumn struct {
 	LastUpdateVersion uint64  `json:"last_update_version"`
 	Correlation       float64 `json:"correlation"`
 }
+
+// TotalMemoryUsage returns the total memory usage of this column.
+func (col *JSONColumn) TotalMemoryUsage() (size int64) {
+	if col.Histogram != nil {
+		size += int64(col.Histogram.Size())
+	}
+	if col.CMSketch != nil {
+		size += int64(col.CMSketch.Size())
+	}
+	if col.FMSketch != nil {
+		size += int64(col.FMSketch.Size())
+	}
+	return size
+}