From 09d3fa9d882c4c57fe692572ba615de48779c9b9 Mon Sep 17 00:00:00 2001
From: MyonKeminta <9948422+MyonKeminta@users.noreply.github.com>
Date: Thu, 9 Feb 2023 10:19:59 +0800
Subject: [PATCH] txn: Fix the issue that analyze panics when aggressive
 locking is enabled globally (#41196)

close pingcap/tidb#41194
---
 session/bootstrap.go                         |  8 ++++++++
 sessiontxn/isolation/base.go                 |  5 ++++-
 sessiontxn/isolation/repeatable_read_test.go | 15 +++++++++++++++
 statistics/handle/handle.go                  |  4 ++++
 4 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/session/bootstrap.go b/session/bootstrap.go
index a18947be5230c..9a9e96b4acd63 100644
--- a/session/bootstrap.go
+++ b/session/bootstrap.go
@@ -29,6 +29,7 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/bindinfo"
 	"github.com/pingcap/tidb/config"
 	"github.com/pingcap/tidb/domain"
@@ -2490,6 +2491,13 @@ func doDMLWorks(s Session) {
 		case variable.TiDBEnableMutationChecker:
 			vVal = variable.On
 		}
+
+		failpoint.Inject("enableAggressiveLockingOnBootstrap", func() {
+			if v.Name == variable.TiDBPessimisticTransactionAggressiveLocking {
+				vVal = variable.On
+			}
+		})
+
 		// sanitize k and vVal
 		value := fmt.Sprintf(`("%s", "%s")`, sqlexec.EscapeString(k), sqlexec.EscapeString(vVal))
 		values = append(values, value)
diff --git a/sessiontxn/isolation/base.go b/sessiontxn/isolation/base.go
index dd1e2881b389e..85fcfde5f3828 100644
--- a/sessiontxn/isolation/base.go
+++ b/sessiontxn/isolation/base.go
@@ -510,7 +510,10 @@ func (p *basePessimisticTxnContextProvider) OnHandlePessimisticStmtStart(ctx con
 	if err := p.baseTxnContextProvider.OnHandlePessimisticStmtStart(ctx); err != nil {
 		return err
 	}
-	if p.sctx.GetSessionVars().PessimisticTransactionAggressiveLocking && p.txn != nil {
+	if p.sctx.GetSessionVars().PessimisticTransactionAggressiveLocking &&
+		p.txn != nil &&
+		p.sctx.GetSessionVars().ConnectionID != 0 &&
+		!p.sctx.GetSessionVars().InRestrictedSQL {
 		if err := p.txn.StartAggressiveLocking(); err != nil {
 			return err
 		}
diff --git a/sessiontxn/isolation/repeatable_read_test.go b/sessiontxn/isolation/repeatable_read_test.go
index 085b64c34cc38..a2c0641ff456a 100644
--- a/sessiontxn/isolation/repeatable_read_test.go
+++ b/sessiontxn/isolation/repeatable_read_test.go
@@ -700,3 +700,18 @@ func TestRRWaitTSTimeInSlowLog(t *testing.T) {
 	require.NotEqual(t, waitTS1, waitTS3)
 	require.NotEqual(t, waitTS2, waitTS3)
 }
+
+func TestIssue41194(t *testing.T) {
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/session/enableAggressiveLockingOnBootstrap", "return"))
+	defer func() {
+		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/session/enableAggressiveLockingOnBootstrap"))
+	}()
+
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+
+	tk.MustExec("use test")
+	tk.MustExec("create table t (id int primary key, v int)")
+	tk.MustExec("insert into t values (1, 1), (2, 2), (3, 3)")
+	tk.MustExec("analyze table t")
+}
diff --git a/statistics/handle/handle.go b/statistics/handle/handle.go
index 7743c2338dcfa..b2c472c2d2c20 100644
--- a/statistics/handle/handle.go
+++ b/statistics/handle/handle.go
@@ -1566,6 +1566,10 @@ func SaveTableStatsToStorage(sctx sessionctx.Context, results *statistics.Analyz
 	if err != nil {
 		return err
 	}
+	err = rs.Close()
+	if err != nil {
+		return err
+	}
 	var curCnt, curModifyCnt int64
 	if len(rows) > 0 {
 		snapshot := rows[0].GetUint64(0)