clean rowContainer even if the execution paniced

Signed-off-by: Yang Keao <[email protected]>

Author: YangKeao

pkg/server/conn_stmt.go

@@ -385,8 +385,14 @@ func (cc *clientConn) executeWithCursor(ctx context.Context, stmt PreparedStatem
 		action := memory.NewActionWithPriority(rowContainer.ActionSpill(), memory.DefCursorFetchSpillPriority)
 		vars.MemTracker.FallbackOldAndSetNewAction(action)
 	}
+	// store the rowContainer in the statement right after it's created, so that even if the logic in defer is not triggered,
+	// the rowContainer will be released when the statement is closed.
+	stmt.StoreRowContainer(rowContainer)
 	defer func() {
 		if err != nil {
+			// if the execution panic, it'll not reach this branch. The `rowContainer` will be released in the `stmt.Close`.
+			stmt.StoreRowContainer(nil)
+
 			rowContainer.GetMemTracker().Detach()
 			rowContainer.GetDiskTracker().Detach()
 			errCloseRowContainer := rowContainer.Close()
@@ -424,7 +430,6 @@ func (cc *clientConn) executeWithCursor(ctx context.Context, stmt PreparedStatem
 	if cl, ok := crs.(resultset.FetchNotifier); ok {
 		cl.OnFetchReturned()
 	}
-	stmt.StoreRowContainer(rowContainer)
 
 	err = cc.writeExecuteResultWithCursor(ctx, stmt, crs)
 	return false, err
@@ -448,16 +453,13 @@ func (cc *clientConn) executeWithLazyCursor(ctx context.Context, stmt PreparedSt
 
 // writeExecuteResultWithCursor will store the `ResultSet` in `stmt` and send the column info to the client. The logic is shared between
 // lazy cursor fetch and normal(eager) cursor fetch.
-func (cc *clientConn) writeExecuteResultWithCursor(ctx context.Context, stmt PreparedStatement, rs resultset.CursorResultSet) error {
-	var err error
-
+func (cc *clientConn) writeExecuteResultWithCursor(ctx context.Context, stmt PreparedStatement, rs resultset.CursorResultSet) (err error) {
 	stmt.StoreResultSet(rs)
 	stmt.SetCursorActive(true)
 	defer func() {
 		if err != nil {
 			// the resultSet and rowContainer have been closed in former "defer" statement.
 			stmt.StoreResultSet(nil)
-			stmt.StoreRowContainer(nil)
 			stmt.SetCursorActive(false)
 		}
 	}()

pkg/server/tests/cursor/BUILD.bazel

@@ -8,12 +8,14 @@ go_test(
         "main_test.go",
     ],
     flaky = True,
-    shard_count = 8,
+    shard_count = 9,
     deps = [
         "//pkg/config",
+        "//pkg/executor",
         "//pkg/metrics",
         "//pkg/parser/mysql",
         "//pkg/server",
+        "//pkg/server/internal/util",
         "//pkg/server/tests/servertestkit",
         "//pkg/store/mockstore/unistore",
         "//pkg/testkit",

pkg/server/tests/cursor/cursor_test.go

@@ -21,13 +21,17 @@ import (
 	"fmt"
 	"io"
 	"math/rand"
+	"os"
 	"testing"
+	"time"
 
 	mysqlcursor "github.com/YangKeao/go-mysql-driver"
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/config"
+	"github.com/pingcap/tidb/pkg/executor"
 	tmysql "github.com/pingcap/tidb/pkg/parser/mysql"
 	server2 "github.com/pingcap/tidb/pkg/server"
+	util2 "github.com/pingcap/tidb/pkg/server/internal/util"
 	"github.com/pingcap/tidb/pkg/server/tests/servertestkit"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/stretchr/testify/require"
@@ -506,3 +510,51 @@ outerLoop:
 		}
 	}
 }
+
+func TestCursorExceedQuota(t *testing.T) {
+	cfg := util2.NewTestConfig()
+	cfg.TempStoragePath = t.TempDir()
+
+	cfg.Port = 0
+	cfg.Status.StatusPort = 0
+	cfg.TempStorageQuota = 1000
+	executor.GlobalDiskUsageTracker.SetBytesLimit(cfg.TempStorageQuota)
+	ts := servertestkit.CreateTidbTestSuiteWithCfg(t, cfg)
+
+	mysqldriver := &mysqlcursor.MySQLDriver{}
+	rawConn, err := mysqldriver.Open(ts.GetDSNWithCursor(10))
+	require.NoError(t, err)
+	conn := rawConn.(mysqlcursor.Connection)
+
+	_, err = conn.ExecContext(context.Background(), "drop table if exists t1", nil)
+	require.NoError(t, err)
+	_, err = conn.ExecContext(context.Background(), "CREATE TABLE `t1` (`c1` varchar(100));", nil)
+	require.NoError(t, err)
+	rowCount := 1000
+	for i := 0; i < rowCount; i++ {
+		_, err = conn.ExecContext(context.Background(), "insert into t1 (c1) values ('201801');", nil)
+		require.NoError(t, err)
+	}
+
+	_, err = conn.ExecContext(context.Background(), "set tidb_mem_quota_query = 1;", nil)
+	require.NoError(t, err)
+	_, err = conn.ExecContext(context.Background(), "set global tidb_enable_tmp_storage_on_oom = 'ON'", nil)
+	require.NoError(t, err)
+
+	rawStmt, err := conn.Prepare("SELECT * FROM test.t1")
+	require.NoError(t, err)
+	stmt := rawStmt.(mysqlcursor.Statement)
+
+	_, err = stmt.QueryContext(context.Background(), nil)
+	require.Error(t, err)
+	require.Contains(t, err.Error(), "Out Of Quota For Local Temporary Space!")
+
+	require.NoError(t, conn.Close())
+
+	time.Sleep(time.Second)
+
+	tempStoragePath := cfg.TempStoragePath
+	files, err := os.ReadDir(tempStoragePath)
+	require.NoError(t, err)
+	require.Empty(t, files)
+}