From 3854bb8711c43fd6881a0ca995b7cd40dfe80c84 Mon Sep 17 00:00:00 2001
From: Shenghui Wu <793703860@qq.com>
Date: Wed, 23 Feb 2022 18:55:42 +0800
Subject: [PATCH] cherry pick #32554 to release-5.4

Signed-off-by: ti-srebot <ti-srebot@pingcap.com>
---
 session/session.go       |  6 ++---
 util/chunk/alloc.go      |  3 +++
 util/chunk/alloc_test.go | 51 ++++++++++++++++++++++++++++++++++++++++
 util/chunk/codec.go      |  3 +++
 util/chunk/column.go     |  2 ++
 5 files changed, 61 insertions(+), 4 deletions(-)

diff --git a/session/session.go b/session/session.go
index 7089008367636..41a26f5704a54 100644
--- a/session/session.go
+++ b/session/session.go
@@ -1142,8 +1142,7 @@ func createSessionFunc(store kv.Storage) pools.Factory {
 		}
 		se.sessionVars.CommonGlobalLoaded = true
 		se.sessionVars.InRestrictedSQL = true
-		// TODO: Remove this line after fixing https://github.com/pingcap/tidb/issues/30880
-		// Chunk RPC protocol may have memory leak issue not solved.
+		// Internal session uses default format to prevent memory leak problem.
 		se.sessionVars.EnableChunkRPC = false
 		return se, nil
 	}
@@ -1165,8 +1164,7 @@ func createSessionWithDomainFunc(store kv.Storage) func(*domain.Domain) (pools.R
 		}
 		se.sessionVars.CommonGlobalLoaded = true
 		se.sessionVars.InRestrictedSQL = true
-		// TODO: Remove this line after fixing https://github.com/pingcap/tidb/issues/30880
-		// Chunk RPC protocol may have memory leak issue not solved.
+		// Internal session uses default format to prevent memory leak problem.
 		se.sessionVars.EnableChunkRPC = false
 		return se, nil
 	}
diff --git a/util/chunk/alloc.go b/util/chunk/alloc.go
index c2bfae8a4f600..c426e4a2753e3 100644
--- a/util/chunk/alloc.go
+++ b/util/chunk/alloc.go
@@ -116,6 +116,9 @@ func (alloc *poolColumnAllocator) init() {
 }
 
 func (alloc *poolColumnAllocator) put(col *Column) {
+	if col.avoidReusing {
+		return
+	}
 	typeSize := col.typeSize()
 	if typeSize <= 0 {
 		return
diff --git a/util/chunk/alloc_test.go b/util/chunk/alloc_test.go
index 59c097b09d1bd..9687c7bae725f 100644
--- a/util/chunk/alloc_test.go
+++ b/util/chunk/alloc_test.go
@@ -151,3 +151,54 @@ func TestNoDuplicateColumnReuse(t *testing.T) {
 		}
 	}
 }
+
+func TestAvoidColumnReuse(t *testing.T) {
+	// For issue: https://github.com/pingcap/tidb/issues/31981
+	// Some chunk columns are references to rpc message.
+	// So when reusing Chunk, we should ignore them.
+
+	fieldTypes := []*types.FieldType{
+		{Tp: mysql.TypeVarchar},
+		{Tp: mysql.TypeJSON},
+		{Tp: mysql.TypeFloat},
+		{Tp: mysql.TypeNewDecimal},
+		{Tp: mysql.TypeDouble},
+		{Tp: mysql.TypeLonglong},
+		{Tp: mysql.TypeTimestamp},
+		{Tp: mysql.TypeDatetime},
+	}
+	alloc := NewAllocator()
+	for i := 0; i < maxFreeChunks+10; i++ {
+		chk := alloc.Alloc(fieldTypes, 5, 10)
+		for _, col := range chk.columns {
+			col.avoidReusing = true
+		}
+	}
+	alloc.Reset()
+
+	a := alloc.columnAlloc
+	// Make sure no duplicated column in the pool.
+	for _, p := range a.pool {
+		require.True(t, p.empty())
+	}
+
+	// test decoder will set avoid reusing flag.
+	chk := alloc.Alloc(fieldTypes, 5, 1024)
+	for i := 0; i <= 10; i++ {
+		for _, col := range chk.columns {
+			col.AppendNull()
+		}
+	}
+	codec := &Codec{fieldTypes}
+	buf := codec.Encode(chk)
+
+	decoder := NewDecoder(
+		NewChunkWithCapacity(fieldTypes, 0),
+		fieldTypes,
+	)
+	decoder.Reset(buf)
+	decoder.ReuseIntermChk(chk)
+	for _, col := range chk.columns {
+		require.True(t, col.avoidReusing)
+	}
+}
diff --git a/util/chunk/codec.go b/util/chunk/codec.go
index ca872e7797bb9..9f2d4bb948d82 100644
--- a/util/chunk/codec.go
+++ b/util/chunk/codec.go
@@ -140,6 +140,9 @@ func (c *Codec) decodeColumn(buffer []byte, col *Column, ordinal int) (remained
 
 	// decode data.
 	col.data = buffer[:numDataBytes:numDataBytes]
+	// The column reference the data of the grpc response, the memory of the grpc message cannot be GCed if we reuse
+	// this column. Thus, we set `avoidReusing` to true.
+	col.avoidReusing = true
 	return buffer[numDataBytes:]
 }
 
diff --git a/util/chunk/column.go b/util/chunk/column.go
index 2f8795e1ed06c..8d0a0c469ce2c 100644
--- a/util/chunk/column.go
+++ b/util/chunk/column.go
@@ -66,6 +66,8 @@ type Column struct {
 	offsets    []int64 // used for varLen column. Row i starts from data[offsets[i]]
 	data       []byte
 	elemBuf    []byte
+
+	avoidReusing bool // avoid reusing the Column by allocator
 }
 
 // ColumnAllocator defines an allocator for Column.