Skip to content

Commit 8b7057a

Browse files
ti-chi-botti-chi-bot
authored
session: fix select for update statement can't get stmt-count-limit error (#48412) (#48468)
close #48411
1 parent 98fe03d commit 8b7057a

File tree

5 files changed

+114
-5
lines changed

5 files changed

+114
-5
lines changed

‎server/server_test.go

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,7 @@ import (
3737
"github.com/pingcap/errors"
3838
"github.com/pingcap/failpoint"
3939
"github.com/pingcap/log"
40+
"github.com/pingcap/tidb/config"
4041
"github.com/pingcap/tidb/errno"
4142
"github.com/pingcap/tidb/kv"
4243
tmysql "github.com/pingcap/tidb/parser/mysql"
@@ -2569,3 +2570,76 @@ func TestIssue46197(t *testing.T) {
25692570
path := testdata.ConvertRowsToStrings(tk.MustQuery("select @@tidb_last_plan_replayer_token").Rows())
25702571
require.NoError(t, os.Remove(filepath.Join(replayer.GetPlanReplayerDirName(), path[0])))
25712572
}
2573+
2574+
func (cli *testServerClient) RunTestStmtCountLimit(t *testing.T) {
2575+
originalStmtCountLimit := config.GetGlobalConfig().Performance.StmtCountLimit
2576+
config.UpdateGlobal(func(conf *config.Config) {
2577+
conf.Performance.StmtCountLimit = 3
2578+
})
2579+
defer func() {
2580+
config.UpdateGlobal(func(conf *config.Config) {
2581+
conf.Performance.StmtCountLimit = originalStmtCountLimit
2582+
})
2583+
}()
2584+
2585+
cli.runTests(t, nil, func(dbt *testkit.DBTestKit) {
2586+
dbt.MustExec("create table t (id int key);")
2587+
dbt.MustExec("set @@tidb_disable_txn_auto_retry=0;")
2588+
dbt.MustExec("set autocommit=0;")
2589+
dbt.MustExec("begin optimistic;")
2590+
dbt.MustExec("insert into t values (1);")
2591+
dbt.MustExec("insert into t values (2);")
2592+
_, err := dbt.GetDB().Query("select * from t for update;")
2593+
require.Error(t, err)
2594+
require.Equal(t, "Error 1105 (HY000): statement count 4 exceeds the transaction limitation, transaction has been rollback, autocommit = false", err.Error())
2595+
dbt.MustExec("insert into t values (3);")
2596+
dbt.MustExec("commit;")
2597+
rows := dbt.MustQuery("select * from t;")
2598+
var id int
2599+
count := 0
2600+
for rows.Next() {
2601+
rows.Scan(&id)
2602+
count++
2603+
}
2604+
require.NoError(t, rows.Close())
2605+
require.Equal(t, 3, id)
2606+
require.Equal(t, 1, count)
2607+
2608+
dbt.MustExec("delete from t;")
2609+
dbt.MustExec("commit;")
2610+
dbt.MustExec("set @@tidb_disable_txn_auto_retry=0;")
2611+
dbt.MustExec("set autocommit=0;")
2612+
dbt.MustExec("begin optimistic;")
2613+
dbt.MustExec("insert into t values (1);")
2614+
dbt.MustExec("insert into t values (2);")
2615+
_, err = dbt.GetDB().Exec("insert into t values (3);")
2616+
require.Error(t, err)
2617+
require.Equal(t, "Error 1105 (HY000): statement count 4 exceeds the transaction limitation, transaction has been rollback, autocommit = false", err.Error())
2618+
dbt.MustExec("commit;")
2619+
rows = dbt.MustQuery("select count(*) from t;")
2620+
for rows.Next() {
2621+
rows.Scan(&count)
2622+
}
2623+
require.NoError(t, rows.Close())
2624+
require.Equal(t, 0, count)
2625+
2626+
dbt.MustExec("delete from t;")
2627+
dbt.MustExec("commit;")
2628+
dbt.MustExec("set @@tidb_batch_commit=1;")
2629+
dbt.MustExec("set @@tidb_disable_txn_auto_retry=0;")
2630+
dbt.MustExec("set autocommit=0;")
2631+
dbt.MustExec("begin optimistic;")
2632+
dbt.MustExec("insert into t values (1);")
2633+
dbt.MustExec("insert into t values (2);")
2634+
dbt.MustExec("insert into t values (3);")
2635+
dbt.MustExec("insert into t values (4);")
2636+
dbt.MustExec("insert into t values (5);")
2637+
dbt.MustExec("commit;")
2638+
rows = dbt.MustQuery("select count(*) from t;")
2639+
for rows.Next() {
2640+
rows.Scan(&count)
2641+
}
2642+
require.NoError(t, rows.Close())
2643+
require.Equal(t, 5, count)
2644+
})
2645+
}

‎server/tidb_test.go

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1122,6 +1122,11 @@ func TestSumAvg(t *testing.T) {
11221122
ts.runTestSumAvg(t)
11231123
}
11241124

1125+
func TestStmtCountLimit(t *testing.T) {
1126+
ts := createTidbTestSuite(t)
1127+
ts.RunTestStmtCountLimit(t)
1128+
}
1129+
11251130
func TestNullFlag(t *testing.T) {
11261131
ts := createTidbTestSuite(t)
11271132

‎session/session.go

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2399,6 +2399,14 @@ func runStmt(ctx context.Context, se *session, s sqlexec.Statement) (rs sqlexec.
23992399
if err != nil {
24002400
return nil, err
24012401
}
2402+
if sessVars.TxnCtx.CouldRetry && !s.IsReadOnly(sessVars) {
2403+
// Only when the txn is could retry and the statement is not read only, need to do stmt-count-limit check,
2404+
// otherwise, the stmt won't be add into stmt history, and also don't need check.
2405+
// About `stmt-count-limit`, see more in https://docs.pingcap.com/tidb/stable/tidb-configuration-file#stmt-count-limit
2406+
if err := checkStmtLimit(ctx, se, false); err != nil {
2407+
return nil, err
2408+
}
2409+
}
24022410

24032411
rs, err = s.Exec(ctx)
24042412
se.updateTelemetryMetric(s.(*executor.ExecStmt))

‎session/sessiontest/session_test.go

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -968,6 +968,16 @@ func TestBatchCommit(t *testing.T) {
968968
tk.MustExec("insert into t values (7)")
969969
tk1.MustQuery("select * from t").Check(testkit.Rows("5", "6", "7"))
970970

971+
tk.MustExec("delete from t")
972+
tk.MustExec("commit")
973+
tk.MustExec("begin")
974+
tk.MustExec("explain analyze insert into t values (5)")
975+
tk1.MustQuery("select * from t").Check(testkit.Rows())
976+
tk.MustExec("explain analyze insert into t values (6)")
977+
tk1.MustQuery("select * from t").Check(testkit.Rows())
978+
tk.MustExec("explain analyze insert into t values (7)")
979+
tk1.MustQuery("select * from t").Check(testkit.Rows("5", "6", "7"))
980+
971981
// The session is still in transaction.
972982
tk.MustExec("insert into t values (8)")
973983
tk1.MustQuery("select * from t").Check(testkit.Rows("5", "6", "7"))

‎session/tidb.go

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -271,7 +271,7 @@ func finishStmt(ctx context.Context, se *session, meetsErr error, sql sqlexec.St
271271
if err != nil {
272272
return err
273273
}
274-
return checkStmtLimit(ctx, se)
274+
return checkStmtLimit(ctx, se, true)
275275
}
276276

277277
func autoCommitAfterStmt(ctx context.Context, se *session, meetsErr error, sql sqlexec.Statement) error {
@@ -305,18 +305,29 @@ func autoCommitAfterStmt(ctx context.Context, se *session, meetsErr error, sql s
305305
return nil
306306
}
307307

308-
func checkStmtLimit(ctx context.Context, se *session) error {
308+
func checkStmtLimit(ctx context.Context, se *session, isFinish bool) error {
309309
// If the user insert, insert, insert ... but never commit, TiDB would OOM.
310310
// So we limit the statement count in a transaction here.
311311
var err error
312312
sessVars := se.GetSessionVars()
313313
history := GetHistory(se)
314-
if history.Count() > int(config.GetGlobalConfig().Performance.StmtCountLimit) {
314+
stmtCount := history.Count()
315+
if !isFinish {
316+
// history stmt count + current stmt, since current stmt is not finish, it has not add to history.
317+
stmtCount++
318+
}
319+
if stmtCount > int(config.GetGlobalConfig().Performance.StmtCountLimit) {
315320
if !sessVars.BatchCommit {
316321
se.RollbackTxn(ctx)
317-
return errors.Errorf("statement count %d exceeds the transaction limitation, autocommit = %t",
318-
history.Count(), sessVars.IsAutocommit())
322+
return errors.Errorf("statement count %d exceeds the transaction limitation, transaction has been rollback, autocommit = %t",
323+
stmtCount, sessVars.IsAutocommit())
324+
}
325+
if !isFinish {
326+
// if the stmt is not finish execute, then just return, since some work need to be done such as StmtCommit.
327+
return nil
319328
}
329+
// If the stmt is finish execute, and exceed the StmtCountLimit, and BatchCommit is true,
330+
// then commit the current transaction and create a new transaction.
320331
err = sessiontxn.NewTxn(ctx, se)
321332
// The transaction does not committed yet, we need to keep it in transaction.
322333
// The last history could not be "commit"/"rollback" statement.
@@ -328,6 +339,7 @@ func checkStmtLimit(ctx context.Context, se *session) error {
328339
}
329340

330341
// GetHistory get all stmtHistory in current txn. Exported only for test.
342+
// If stmtHistory is nil, will create a new one for current txn.
331343
func GetHistory(ctx sessionctx.Context) *StmtHistory {
332344
hist, ok := ctx.GetSessionVars().TxnCtx.History.(*StmtHistory)
333345
if ok {

0 commit comments

Comments
 (0)