session: fix the memory tracker leak in prepare binary protocol (#44618…

…) (#44639) close #44612
pingcap · Jun 19, 2023 · af1110c · af1110c
1 parent f917cc4
commit af1110c
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 4 deletions.
diff --git a/server/conn_stmt_test.go b/server/conn_stmt_test.go
@@ -407,21 +407,41 @@ func TestCursorDetachMemTracker(t *testing.T) {
 	// testkit also uses `PREPARE` related calls to run statement with arguments.
 	// format the SQL to avoid the interference from testkit.
 	tk.MustExec(fmt.Sprintf("set tidb_mem_quota_query=%d", maxConsumed/2))
-	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 1)
+	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 0)
 	// This query should exceed the memory limitation during `openExecutor`
 	require.Error(t, c.Dispatch(ctx, append(
 		appendUint32([]byte{mysql.ComStmtExecute}, uint32(stmt.ID())),
 		mysql.CursorTypeReadOnly, 0x1, 0x0, 0x0, 0x0,
 	)))
-	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 1)
+	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 0)
 
 	// The next query should succeed
 	tk.MustExec(fmt.Sprintf("set tidb_mem_quota_query=%d", maxConsumed+1))
-	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 1)
+	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 0)
 	// This query should succeed
 	require.NoError(t, c.Dispatch(ctx, append(
 		appendUint32([]byte{mysql.ComStmtExecute}, uint32(stmt.ID())),
 		mysql.CursorTypeReadOnly, 0x1, 0x0, 0x0, 0x0,
 	)))
-	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 1)
+	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 0)
+}
+
+func TestMemoryTrackForPrepareBinaryProtocol(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	srv := CreateMockServer(t, store)
+	srv.SetDomain(dom)
+	defer srv.Close()
+
+	c := CreateMockConn(t, srv).(*mockConn)
+
+	tk := testkit.NewTestKitWithSession(t, store, c.Context().Session)
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t(id_2 int)")
+	for i := 0; i <= 10; i++ {
+		stmt, _, _, err := c.Context().Prepare("select count(id_2) from t")
+		require.NoError(t, err)
+		require.NoError(t, stmt.Close())
+	}
+	require.Len(t, tk.Session().GetSessionVars().MemTracker.GetChildrenForTest(), 0)
 }
diff --git a/session/session.go b/session/session.go
@@ -2447,6 +2447,11 @@ func (s *session) rollbackOnError(ctx context.Context) {
 
 // PrepareStmt is used for executing prepare statement in binary protocol
 func (s *session) PrepareStmt(sql string) (stmtID uint32, paramCount int, fields []*ast.ResultField, err error) {
+	defer func() {
+		if s.sessionVars.StmtCtx != nil {
+			s.sessionVars.StmtCtx.DetachMemDiskTracker()
+		}
+	}()
 	if s.sessionVars.TxnCtx.InfoSchema == nil {
 		// We don't need to create a transaction for prepare statement, just get information schema will do.
 		s.sessionVars.TxnCtx.InfoSchema = domain.GetDomain(s).InfoSchema()