diff --git a/pkg/session/test/txn/BUILD.bazel b/pkg/session/test/txn/BUILD.bazel index 456c700f3123a..fc7a7053e3a6a 100644 --- a/pkg/session/test/txn/BUILD.bazel +++ b/pkg/session/test/txn/BUILD.bazel @@ -9,7 +9,7 @@ go_test( ], flaky = True, race = "on", - shard_count = 10, + shard_count = 11, deps = [ "//pkg/config", "//pkg/kv", @@ -20,6 +20,8 @@ go_test( "//pkg/testkit/testmain", "//pkg/testkit/testsetup", "//pkg/util/dbterror/plannererrors", + "//pkg/util/memory", + "//pkg/util/sqlkiller", "@com_github_pingcap_failpoint//:failpoint", "@com_github_stretchr_testify//require", "@com_github_tikv_client_go_v2//tikv", diff --git a/pkg/session/test/txn/txn_test.go b/pkg/session/test/txn/txn_test.go index 9a40530c76427..bce0dea1a4e44 100644 --- a/pkg/session/test/txn/txn_test.go +++ b/pkg/session/test/txn/txn_test.go @@ -30,6 +30,8 @@ import ( "github.com/pingcap/tidb/pkg/parser/terror" "github.com/pingcap/tidb/pkg/testkit" "github.com/pingcap/tidb/pkg/util/dbterror/plannererrors" + "github.com/pingcap/tidb/pkg/util/memory" + "github.com/pingcap/tidb/pkg/util/sqlkiller" "github.com/stretchr/testify/require" ) @@ -579,3 +581,25 @@ func TestMemBufferCleanupMemoryLeak(t *testing.T) { } tk.MustExec("commit") } + +func TestPanicOnRollbackKilledTxn(t *testing.T) { + store := testkit.CreateMockStore(t) + tk := testkit.NewTestKit(t, store) + tk.MustExec("use test") + tk.MustExec("drop table if exists t") + tk.MustExec("create table t(id int)") + tk.MustExec("begin pessimistic") + tk.MustExec("insert into t values(1);") + tk.MustExec("insert into t select * from t") + tk.MustExec("insert into t select * from t") + tk.MustExec("insert into t select * from t") + tk.MustExec("insert into t select * from t") + tk.MustExec("insert into t select * from t") + tk.MustExec("insert into t select * from t") + mockTracker := memory.NewTracker(-1, -1) + mockTracker.IsRootTrackerOfSess = true + mockTracker.Killer = &sqlkiller.SQLKiller{} + tk.Session().GetSessionVars().MemTracker.AttachTo(mockTracker) + mockTracker.Killer.SendKillSignal(sqlkiller.QueryInterrupted) + tk.Session().Close() +} diff --git a/pkg/session/txn.go b/pkg/session/txn.go index 7fcac80bf218d..2635048d498ac 100644 --- a/pkg/session/txn.go +++ b/pkg/session/txn.go @@ -442,6 +442,8 @@ func (txn *LazyTxn) Rollback() error { txn.mu.Unlock() // mockSlowRollback is used to mock a rollback which takes a long time failpoint.Inject("mockSlowRollback", func(_ failpoint.Value) {}) + // When rolling back a txn, swap with a dummy hook to avoid operations on an invalid memory tracker. + txn.SetMemoryFootprintChangeHook(func(uint64) {}) return txn.Transaction.Rollback() }