session: fix rollback on killed txn may panic (#64001)

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",

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()
+}

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()
 }