From f600a5c09f6ecbb5878703829425cd3055ef9a15 Mon Sep 17 00:00:00 2001
From: Shenghui Wu <793703860@qq.com>
Date: Tue, 30 Aug 2022 11:44:23 +0800
Subject: [PATCH] executor: check memory limit before expand string set in
 HashAgg (#37298)

ref pingcap/tidb#35635
---
 executor/aggregate.go             |  1 +
 util/set/set_with_memory_usage.go | 14 ++++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/executor/aggregate.go b/executor/aggregate.go
index f492c0c49e7bd..f505049149e65 100644
--- a/executor/aggregate.go
+++ b/executor/aggregate.go
@@ -418,6 +418,7 @@ func (e *HashAggExec) initForParallelExec(_ sessionctx.Context) {
 		}
 		// There is a bucket in the empty partialResultsMap.
 		e.memTracker.Consume(hack.DefBucketMemoryUsageForMapStrToSlice*(1<<w.BInMap) + setSize)
+		groupSet.SetTracker(e.memTracker)
 		if e.stats != nil {
 			w.stats = &AggWorkerStat{}
 			e.stats.FinalStats = append(e.stats.FinalStats, w.stats)
diff --git a/util/set/set_with_memory_usage.go b/util/set/set_with_memory_usage.go
index 3e48f90c604d8..513d9e3057d76 100644
--- a/util/set/set_with_memory_usage.go
+++ b/util/set/set_with_memory_usage.go
@@ -16,12 +16,17 @@ package set
 
 import (
 	"github.com/pingcap/tidb/util/hack"
+	"github.com/pingcap/tidb/util/memory"
 )
 
 // StringSetWithMemoryUsage is a string set with memory usage.
 type StringSetWithMemoryUsage struct {
 	StringSet
 	bInMap int64
+
+	// For tracking large memory usage in time.
+	// If tracker is non-nil, memDelta will track immediately and reset to 0. Otherwise, memDelta will return and lazy track.
+	tracker *memory.Tracker
 }
 
 // NewStringSetWithMemoryUsage builds a string set.
@@ -44,10 +49,19 @@ func (s *StringSetWithMemoryUsage) Insert(val string) (memDelta int64) {
 	if s.Count() > (1<<s.bInMap)*hack.LoadFactorNum/hack.LoadFactorDen {
 		memDelta = hack.DefBucketMemoryUsageForSetString * (1 << s.bInMap)
 		s.bInMap++
+		if s.tracker != nil {
+			s.tracker.Consume(memDelta)
+			memDelta = 0
+		}
 	}
 	return memDelta
 }
 
+// SetTracker sets memory tracker for StringSetWithMemoryUsage
+func (s *StringSetWithMemoryUsage) SetTracker(t *memory.Tracker) {
+	s.tracker = t
+}
+
 // Float64SetWithMemoryUsage is a float64 set with memory usage.
 type Float64SetWithMemoryUsage struct {
 	Float64Set