executor: make the memory tracker of Jsonobjectagg more accurate

executor/aggfuncs/func_json_objectagg.go

@@ -27,6 +27,8 @@ import (
 const (
 	// DefPartialResult4JsonObjectAgg is the size of partialResult4JsonObject
 	DefPartialResult4JsonObjectAgg = int64(unsafe.Sizeof(partialResult4JsonObjectAgg{}))
+	// DefMapStringInterfaceBucketSize = bucketSize*(1+unsafe.Sizeof(string) + unsafe.Sizeof(interface{}))+2*ptrSize
+	DefMapStringInterfaceBucketSize = 8*(1+16+16) + 16
 )
 
 type jsonObjectAgg struct {
@@ -35,17 +37,20 @@ type jsonObjectAgg struct {
 
 type partialResult4JsonObjectAgg struct {
 	entries map[string]interface{}
+	bInMap  int // indicate there are 2^bInMap buckets in entries.
 }
 
 func (e *jsonObjectAgg) AllocPartialResult() (pr PartialResult, memDelta int64) {
 	p := partialResult4JsonObjectAgg{}
 	p.entries = make(map[string]interface{})
-	return PartialResult(&p), DefPartialResult4JsonObjectAgg
+	p.bInMap = 0
+	return PartialResult(&p), DefPartialResult4JsonObjectAgg + (1<<p.bInMap)*DefMapStringInterfaceBucketSize
 }
 
 func (e *jsonObjectAgg) ResetPartialResult(pr PartialResult) {
 	p := (*partialResult4JsonObjectAgg)(pr)
 	p.entries = make(map[string]interface{})
+	p.bInMap = 0
 }
 
 func (e *jsonObjectAgg) AppendFinalResult2Chunk(sctx sessionctx.Context, pr PartialResult, chk *chunk.Chunk) error {
@@ -105,8 +110,10 @@ func (e *jsonObjectAgg) UpdatePartialResult(sctx sessionctx.Context, rowsInGroup
 		switch x := realVal.(type) {
 		case nil, bool, int64, uint64, float64, string, json.BinaryJSON, *types.MyDecimal, []uint8, types.Time, types.Duration:
 			if _, ok := p.entries[keyString]; !ok {
-				memDelta += int64(len(keyString))
-				memDelta += getValMemDelta(realVal)
+				memDelta += int64(len(keyString)) + getValMemDelta(realVal)
+				if len(p.entries)+1 > (1<<p.bInMap)*loadFactorNum/loadFactorDen {
+					memDelta += (1 << p.bInMap) * DefMapStringInterfaceBucketSize
+				}
 			}
 			p.entries[keyString] = realVal
 		default:
@@ -150,6 +157,10 @@ func (e *jsonObjectAgg) MergePartialResult(sctx sessionctx.Context, src, dst Par
 	// and only the last value encountered is used with that key in the returned object
 	for k, v := range p1.entries {
 		p2.entries[k] = v
+		memDelta += int64(len(k)) + getValMemDelta(v)
+		if len(p2.entries)+1 > (1<<p2.bInMap)*loadFactorNum/loadFactorDen {
+			memDelta += (1 << p2.bInMap) * DefMapStringInterfaceBucketSize
+		}
 	}
 	return 0, nil
 }

executor/aggfuncs/func_json_objectagg_test.go

@@ -141,8 +141,8 @@ func (s *testSuite) TestMemJsonObjectagg(c *C) {
 		}
 
 		tests := []multiArgsAggMemTest{
-			buildMultiArgsAggMemTester(ast.AggFuncJsonObjectAgg, argTypes, mysql.TypeJSON, numRows, aggfuncs.DefPartialResult4JsonObjectAgg, defaultMultiArgsMemDeltaGens, true),
-			buildMultiArgsAggMemTester(ast.AggFuncJsonObjectAgg, argTypes, mysql.TypeJSON, numRows, aggfuncs.DefPartialResult4JsonObjectAgg, defaultMultiArgsMemDeltaGens, false),
+			buildMultiArgsAggMemTester(ast.AggFuncJsonObjectAgg, argTypes, mysql.TypeJSON, numRows, aggfuncs.DefPartialResult4JsonObjectAgg+aggfuncs.DefMapStringInterfaceBucketSize, defaultMultiArgsMemDeltaGens, true),
+			buildMultiArgsAggMemTester(ast.AggFuncJsonObjectAgg, argTypes, mysql.TypeJSON, numRows, aggfuncs.DefPartialResult4JsonObjectAgg+aggfuncs.DefMapStringInterfaceBucketSize, defaultMultiArgsMemDeltaGens, false),
 		}
 		for _, test := range tests {
 			s.testMultiArgsAggMemFunc(c, test)

executor/aggfuncs/set_with_memory_usage.go

@@ -0,0 +1,21 @@
+// Copyright 2021 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package aggfuncs
+
+const (
+	// Maximum average load of a bucket that triggers growth is 6.5.
+	// Represent as loadFactorNum/loadFactDen, to allow integer math.
+	loadFactorNum = 13
+	loadFactorDen = 2
+)