spilling data didn't clean for cursor fetch

[guo-shaoge]

Bug Report

Please answer these questions before submitting your issue. Thanks!

1. Minimal reproduce step (Required)

drop table if exists t1;
CREATE TABLE `t1` (`c1` varchar(100));
insert into t1 (c1)values("201801");
insert into t1 (c1)values("201801");
insert into t1 (c1)values("xxxxx");
insert into t1 (select * from t1);

-- run the last insert multiple times until you have 786432 rows in t1
set global tidb_mem_quota_query =10000000;

change tidb config to set tmp-storage-quota = 16500000

import java.sql.*;

public class MySQLCursorFetch {
    public static void main(String[] args) {
        String url = "jdbc:mysql://127.0.0.1:41501/test?useCursorFetch=true&defaultFetchSize=100";
        String user = "root";
        String password = "";

        try {
            Class.forName("com.mysql.cj.jdbc.Driver");
            Connection conn = DriverManager.getConnection(url, user, password);

            String query = "SELECT * FROM test.t1";
            PreparedStatement pstmt = conn.prepareStatement(query);
            pstmt.setFetchSize(100);

            ResultSet rs = pstmt.executeQuery();

            while (rs.next()) {
                String columnData = rs.getString("c1"); 
                System.out.println("Processing row: " + columnData);
            }

            rs.close();
            pstmt.close();
            conn.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}

java -cp .:/home/guojiangtao/tar_gz/jdbc/mysql-connector-j-9.2.0/mysql-connector-j-9.2.0.jar ./jdbc.java (note: you need set up you java env and download mysql-connector-java first)

2. What did you expect to see? (Required)

return error like the following and the spilling files have been deleted.

java -cp .:/home/guojiangtao/tar_gz/jdbc/mysql-connector-j-9.2.0/mysql-connector-j-9.2.0.jar ./jdbc.java
java.sql.SQLException: Out Of Quota For Local Temporary Space!
	at com.mysql.cj.jdbc.exceptions.SQLError.createSQLException(SQLError.java:121)
	at com.mysql.cj.jdbc.exceptions.SQLExceptionsMapping.translateException(SQLExceptionsMapping.java:114)
	at com.mysql.cj.jdbc.ServerPreparedStatement.serverExecute(ServerPreparedStatement.java:612)
	at com.mysql.cj.jdbc.ServerPreparedStatement.executeInternal(ServerPreparedStatement.java:320)
	at com.mysql.cj.jdbc.ClientPreparedStatement.executeQuery(ClientPreparedStatement.java:1058)
	at MySQLCursorFetch.main(jdbc.java:22)
	at java.base/jdk.internal.reflect.DirectMethodHandleAccessor.invoke(DirectMethodHandleAccessor.java:103)
	at java.base/java.lang.reflect.Method.invoke(Method.java:580)
	at jdk.compiler/com.sun.tools.javac.launcher.Main.execute(Main.java:484)
	at jdk.compiler/com.sun.tools.javac.launcher.Main.run(Main.java:208)
	at jdk.compiler/com.sun.tools.javac.launcher.Main.main(Main.java:135)

3. What did you see instead (Required)

spilling file didn't clean

4. What is your TiDB version? (Required)

v6.5.7

[guo-shaoge]

Looks like rowContainer.Add() panic in L411, and it will nerver be closed.

tidb/pkg/server/conn_stmt.go

Lines 400 to 415 in 8dc50b3

	for {
	chk := rs.NewChunk(nil)
	if err = rs.Next(ctx, chk); err != nil {
	return false, err
	}
	rowCount := chk.NumRows()
	if rowCount == 0 {
	break
	}
	err = rowContainer.Add(chk)
	if err != nil {
	return false, err
	}
	}

And in normal process, rowContainer will be closed when all data has been fetched(in L492)

tidb/server/conn_stmt.go

Lines 486 to 497 in cd073ea

	_, err = cc.writeResultset(ctx, rs, true, cc.ctx.Status(), int(fetchSize))
	// if the iterator reached the end before writing result, we could say the `FETCH` command will send EOF
	if rs.GetRowContainerReader().Current() == rs.GetRowContainerReader().End() {
	// also reset the statement when the cursor reaches the end
	// don't overwrite the `err` in outer scope, to avoid redundant `Reset()` in `defer` statement (though, it's not
	// a big problem, as the `Reset()` function call is idempotent.)
	err := stmt.Reset()
	if err != nil {
	logutil.Logger(ctx).Error("Fail to reset statement when FETCH command reaches the end. May cause resource leak",
	zap.NamedError("error", err))
	}
	}

[YangKeao]

func TestCursorExceedQuota(t *testing.T) {
	cfg := util2.NewTestConfig()
	require.NoError(t, os.RemoveAll(cfg.TempStoragePath))
	require.NoError(t, os.MkdirAll(cfg.TempStoragePath, 0o755))

	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 := 10000
	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, stmt.Close())
	require.NoError(t, conn.Close())
	// MySQL Driver doesn't wait for the connection to be closed.
	time.Sleep(1 * time.Second)

	tempStoragePath := cfg.TempStoragePath
	files, err := os.ReadDir(tempStoragePath)
	require.NoError(t, err)
	require.Empty(t, files)
}

This test in pkg/server/tests/cursor/ can reproduce this issue. Let me fix it.