From 8bec188b6abeb2e30dce574528d246bcd0b8ac75 Mon Sep 17 00:00:00 2001
From: Ewan Chou <coocood@gmail.com>
Date: Thu, 21 Jun 2018 17:26:14 +0800
Subject: [PATCH] session: add a global/session variable
 'tidb_disable_txn_auto_retry' (#6872)

---
 session/session.go               | 12 ++++++++++-
 session/session_test.go          | 37 ++++++++++++++++++++++++++++++++
 sessionctx/variable/session.go   |  6 +++++-
 sessionctx/variable/sysvar.go    |  1 +
 sessionctx/variable/tidb_vars.go |  4 ++++
 5 files changed, 58 insertions(+), 2 deletions(-)

diff --git a/session/session.go b/session/session.go
index 570991122c4d0..0f6658db7330a 100644
--- a/session/session.go
+++ b/session/session.go
@@ -346,6 +346,15 @@ func (s *session) doCommitWithRetry(ctx context.Context) error {
 	err := s.doCommit(ctx)
 	if err != nil {
 		commitRetryLimit := s.sessionVars.RetryLimit
+		if s.sessionVars.DisableTxnAutoRetry && !s.sessionVars.InRestrictedSQL {
+			// Do not retry non-autocommit transactions.
+			// For autocommit single statement transactions, the history count is always 1.
+			// For explicit transactions, the statement count is more than 1.
+			history := GetHistory(s)
+			if history.Count() > 1 {
+				commitRetryLimit = 0
+			}
+		}
 		// Don't retry in BatchInsert mode. As a counter-example, insert into t1 select * from t2,
 		// BatchInsert already commit the first batch 1000 rows, then it commit 1000-2000 and retry the statement,
 		// Finally t1 will have more data than t2, with no errors return to user!
@@ -1295,7 +1304,8 @@ const loadCommonGlobalVarsSQL = "select HIGH_PRIORITY * from mysql.global_variab
 	variable.TiDBOptInSubqUnFolding + quoteCommaQuote +
 	variable.TiDBDistSQLScanConcurrency + quoteCommaQuote +
 	variable.TiDBMaxChunkSize + quoteCommaQuote +
-	variable.TiDBRetryLimit + "')"
+	variable.TiDBRetryLimit + quoteCommaQuote +
+	variable.TiDBDisableTxnAutoRetry + "')"
 
 // loadCommonGlobalVariablesIfNeeded loads and applies commonly used global variables for the session.
 func (s *session) loadCommonGlobalVariablesIfNeeded() error {
diff --git a/session/session_test.go b/session/session_test.go
index ed59d46a0ce77..7a4c71d2e1825 100644
--- a/session/session_test.go
+++ b/session/session_test.go
@@ -2056,3 +2056,40 @@ func (s *testSessionSuite) TestCommitRetryCount(c *C) {
 	_, err := tk1.Se.Execute(context.Background(), "commit")
 	c.Assert(err, NotNil)
 }
+
+func (s *testSessionSuite) TestDisableTxnAutoRetry(c *C) {
+	tk1 := testkit.NewTestKitWithInit(c, s.store)
+	tk2 := testkit.NewTestKitWithInit(c, s.store)
+	tk1.MustExec("create table no_retry (id int)")
+	tk1.MustExec("insert into no_retry values (1)")
+	tk1.MustExec("set @@tidb_disable_txn_auto_retry = 1")
+
+	tk1.MustExec("begin")
+	tk1.MustExec("update no_retry set id = 2")
+
+	tk2.MustExec("begin")
+	tk2.MustExec("update no_retry set id = 3")
+	tk2.MustExec("commit")
+
+	// No auto retry because tidb_disable_txn_auto_retry is set to 1.
+	_, err := tk1.Se.Execute(context.Background(), "commit")
+	c.Assert(err, NotNil)
+
+	// session 1 starts a transaction early.
+	// execute a select statement to clear retry history.
+	tk1.MustExec("select 1")
+	tk1.Se.NewTxn()
+	// session 2 update the value.
+	tk2.MustExec("update no_retry set id = 4")
+	// Autocommit update will retry, so it would not fail.
+	tk1.MustExec("update no_retry set id = 5")
+
+	// RestrictedSQL should retry.
+	tk1.Se.GetSessionVars().InRestrictedSQL = true
+	tk1.MustExec("begin")
+
+	tk2.MustExec("update no_retry set id = 6")
+
+	tk1.MustExec("update no_retry set id = 7")
+	tk1.MustExec("commit")
+}
diff --git a/sessionctx/variable/session.go b/sessionctx/variable/session.go
index d1f408886e45b..eb999254de99f 100644
--- a/sessionctx/variable/session.go
+++ b/sessionctx/variable/session.go
@@ -152,7 +152,8 @@ type SessionVars struct {
 	Concurrency
 	MemQuota
 	BatchSize
-	RetryLimit int64
+	RetryLimit          int64
+	DisableTxnAutoRetry bool
 	// UsersLock is a lock for user defined variables.
 	UsersLock sync.RWMutex
 	// Users are user defined variables.
@@ -303,6 +304,7 @@ func NewSessionVars() *SessionVars {
 		AllowAggPushDown:          false,
 		OptimizerSelectivityLevel: DefTiDBOptimizerSelectivityLevel,
 		RetryLimit:                DefTiDBRetryLimit,
+		DisableTxnAutoRetry:       DefTiDBDisableTxnAutoRetry,
 	}
 	vars.Concurrency = Concurrency{
 		IndexLookupConcurrency:     DefIndexLookupConcurrency,
@@ -535,6 +537,8 @@ func (s *SessionVars) SetSystemVar(name string, val string) error {
 		atomic.StoreUint32(&ProcessGeneralLog, uint32(tidbOptPositiveInt32(val, DefTiDBGeneralLog)))
 	case TiDBRetryLimit:
 		s.RetryLimit = tidbOptInt64(val, DefTiDBRetryLimit)
+	case TiDBDisableTxnAutoRetry:
+		s.DisableTxnAutoRetry = TiDBOptOn(val)
 	case TiDBEnableStreaming:
 		s.EnableStreaming = TiDBOptOn(val)
 	case TiDBOptimizerSelectivityLevel:
diff --git a/sessionctx/variable/sysvar.go b/sessionctx/variable/sysvar.go
index 9fce56be7485e..3393243eb2285 100644
--- a/sessionctx/variable/sysvar.go
+++ b/sessionctx/variable/sysvar.go
@@ -642,6 +642,7 @@ var defaultSysVars = []*SysVar{
 	{ScopeGlobal | ScopeSession, TiDBHashAggFinalConcurrency, strconv.Itoa(DefTiDBHashAggFinalConcurrency)},
 	{ScopeGlobal | ScopeSession, TiDBBackoffLockFast, strconv.Itoa(kv.DefBackoffLockFast)},
 	{ScopeGlobal | ScopeSession, TiDBRetryLimit, strconv.Itoa(DefTiDBRetryLimit)},
+	{ScopeGlobal | ScopeSession, TiDBDisableTxnAutoRetry, boolToIntStr(DefTiDBDisableTxnAutoRetry)},
 	{ScopeSession, TiDBOptimizerSelectivityLevel, strconv.Itoa(DefTiDBOptimizerSelectivityLevel)},
 	/* The following variable is defined as session scope but is actually server scope. */
 	{ScopeSession, TiDBGeneralLog, strconv.Itoa(DefTiDBGeneralLog)},
diff --git a/sessionctx/variable/tidb_vars.go b/sessionctx/variable/tidb_vars.go
index d5fffbe2661b3..8627efc52511d 100644
--- a/sessionctx/variable/tidb_vars.go
+++ b/sessionctx/variable/tidb_vars.go
@@ -92,6 +92,9 @@ const (
 	// tidb_retry_limit is the maximun number of retries when committing a transaction.
 	TiDBRetryLimit = "tidb_retry_limit"
 
+	// tidb_disable_txn_auto_retry disables transaction auto retry.
+	TiDBDisableTxnAutoRetry = "tidb_disable_txn_auto_retry"
+
 	// tidb_enable_streaming enables TiDB to use streaming API for coprocessor requests.
 	TiDBEnableStreaming = "tidb_enable_streaming"
 
@@ -198,6 +201,7 @@ const (
 	DefTiDBMemQuotaNestedLoopApply   = 32 << 30 // 32GB.
 	DefTiDBGeneralLog                = 0
 	DefTiDBRetryLimit                = 10
+	DefTiDBDisableTxnAutoRetry       = false
 	DefTiDBHashJoinConcurrency       = 5
 	DefTiDBProjectionConcurrency     = 4
 	DefTiDBOptimizerSelectivityLevel = 0