session: check max_execution_time when get tso

pkg/session/session.go

@@ -2065,7 +2065,11 @@ func (s *session) ExecuteStmt(ctx context.Context, stmtNode ast.StmtNode) (sqlex
 
 	// Uncorrelated subqueries will execute once when building plan, so we reset process info before building plan.
 	s.currentPlan = nil // reset current plan
-	s.SetProcessInfo(stmtNode.Text(), time.Now(), cmdByte, 0)
+	var maxExecutionTime uint64
+	if intest.InTest {
+		maxExecutionTime = sessVars.MaxExecutionTime
+	}
+	s.SetProcessInfo(stmtNode.Text(), time.Now(), cmdByte, maxExecutionTime)
 	s.txn.onStmtStart(digest.String())
 	defer sessiontxn.GetTxnManager(s).OnStmtEnd()
 	defer s.txn.onStmtEnd()

pkg/session/test/txn/txn_test.go

@@ -28,8 +28,10 @@ import (
 	"github.com/pingcap/tidb/pkg/parser/auth"
 	"github.com/pingcap/tidb/pkg/parser/mysql"
 	"github.com/pingcap/tidb/pkg/parser/terror"
+	"github.com/pingcap/tidb/pkg/server"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/util/dbterror/plannererrors"
+	"github.com/pingcap/tidb/pkg/util/logutil"
 	"github.com/stretchr/testify/require"
 )
 
@@ -509,6 +511,34 @@ func TestInTrans(t *testing.T) {
 	require.False(t, txn.Valid())
 }
 
+func TestGetTsoSlow(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	sv := server.CreateMockServer(t, store)
+	sv.SetDomain(dom)
+	defer sv.Close()
+
+	conn1 := server.CreateMockConn(t, sv)
+	tk := testkit.NewTestKitWithSession(t, store, conn1.Context().Session)
+
+	go dom.ExpensiveQueryHandle().SetSessionManager(sv).Run()
+
+	tk.MustExec("use test")
+	tk.MustExec("drop table if exists t;")
+	tk.MustExec("create table t (id int)")
+	tk.MustExec("insert t values (1), (2)")
+
+	tk.MustExec("set @@max_execution_time = 500")
+
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/session/mock_get_tso_slow", "return(3)"))
+	defer func() {
+		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/session/mock_get_tso_slow"))
+	}()
+
+	err := tk.ExecToErr("select * from t")
+	require.NotNil(t, err)
+	require.Contains(t, err.Error(), "Query execution was interrupted, maximum statement execution time exceeded")
+}
+
 func TestMemBufferSnapshotRead(t *testing.T) {
 	store := testkit.CreateMockStore(t)
 

pkg/session/txn.go

@@ -608,20 +608,56 @@ func (txn *LazyTxn) Wait(ctx context.Context, sctx sessionctx.Context) (kv.Trans
 		defer func(begin time.Time) {
 			sctx.GetSessionVars().DurationWaitTS = time.Since(begin)
 		}(time.Now())
+		if err := txn.waitWithSQLKiller(ctx, sctx); err != nil {
+			return txn, err
+		}
+		txn.lazyUniquenessCheckEnabled = !sctx.GetSessionVars().ConstraintCheckInPlacePessimistic
+	}
+	return txn, nil
+}
+
+func (txn *LazyTxn) waitWithSQLKiller(ctx context.Context, sctx sessionctx.Context) error {
+	finishCh := make(chan struct{})
+	defer close(finishCh)
 
+	tsoCh := make(chan error)
+	go func() {
 		// Transaction is lazy initialized.
 		// PrepareTxnCtx is called to get a tso future, makes s.txn a pending txn,
 		// If Txn() is called later, wait for the future to get a valid txn.
-		if err := txn.changePendingToValid(ctx, sctx); err != nil {
+		var err error
+		failpoint.Inject("mock_get_tso_slow", func(v failpoint.Value) {
+			t := v.(int)
+			time.Sleep(time.Duration(t) * time.Second)
+		})
+		if err = txn.changePendingToValid(ctx, sctx); err != nil {
 			logutil.BgLogger().Error("active transaction fail",
 				zap.Error(err))
 			txn.cleanup()
 			sctx.GetSessionVars().TxnCtx.StartTS = 0
-			return txn, err
 		}
-		txn.lazyUniquenessCheckEnabled = !sctx.GetSessionVars().ConstraintCheckInPlacePessimistic
+		select {
+		case tsoCh <- err:
+		case <-finishCh:
+		}
+		close(tsoCh)
+	}()
+
+	ticker := time.NewTicker(100 * time.Millisecond)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			if err := sctx.GetSessionVars().SQLKiller.HandleSignal(); err != nil {
+				return err
+			}
+		case err := <-tsoCh:
+			return err
+		case <-ctx.Done():
+			return ctx.Err()
+		}
 	}
-	return txn, nil
 }
 
 // KeyNeedToLock returns true if the key need to lock.