This is an automated cherry-pick of pingcap#56847

Signed-off-by: ti-chi-bot <ti-community-prow-bot@tidb.io>

pkg/executor/test/txn/BUILD.bazel

@@ -8,13 +8,18 @@ go_test(
         "txn_test.go",
     ],
     flaky = True,
+<<<<<<< HEAD
     shard_count = 9,
+=======
+    shard_count = 11,
+>>>>>>> 38187e5b1d6 (txn: fix issue innodb_lock_wait_timeout doesn't work in some case (#56847))
     deps = [
         "//pkg/config",
         "//pkg/errno",
         "//pkg/meta/autoid",
         "//pkg/testkit",
         "//pkg/testkit/testfailpoint",
+        "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_stretchr_testify//require",
         "@com_github_tikv_client_go_v2//tikv",
         "@io_opencensus_go//stats/view",

pkg/executor/test/txn/txn_test.go

@@ -15,13 +15,15 @@
 package txn_test
 
 import (
+	"context"
 	"fmt"
 	"strconv"
 	"strings"
 	"sync"
 	"testing"
 	"time"
 
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/errno"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/pkg/testkit/testfailpoint"
@@ -706,3 +708,94 @@ func TestColumnNotMatchError(t *testing.T) {
 	tk.MustExec("delete from t where id=1")
 	tk.MustGetErrCode("commit", errno.ErrInfoSchemaChanged)
 }
+<<<<<<< HEAD
+=======
+
+func TestSavepointWithForeignKey(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table parent (id int primary key)")
+	tk.MustExec("create table child (id int primary key, parent_id int, foreign key (parent_id) references parent(id))")
+
+	tk.MustExec("begin")
+	// rollback insertion to child table
+	tk.MustExec("savepoint sp1")
+	tk.MustExec("insert into parent values (1)")
+	tk.MustExec("savepoint sp2")
+	tk.MustExec("insert into child values (1, 1)")
+	tk.MustExec("rollback to sp2")
+	tk.MustQuery("select * from child").Check(testkit.Rows())
+	// rollback insertion to parent table
+	tk.MustExec("rollback to sp1")
+	require.Contains(t, tk.ExecToErr("insert into child values (1, 1)").Error(), "a foreign key constraint fails")
+	tk.MustQuery("select * from parent").Check(testkit.Rows())
+	tk.MustExec("rollback")
+
+	// A known limitation: the lock of the parent table is not released after rollback to savepoint.
+	tk.MustExec("insert into parent values (1)")
+
+	tk.MustExec("begin")
+	tk.MustExec("savepoint sp1")
+	tk.MustExec("insert into child values (1, 1)")
+	// Now, the row in parent table is locked
+	tk2 := testkit.NewTestKit(t, store)
+	tk2.MustExec("use test")
+	var wg sync.WaitGroup
+	wg.Add(1)
+	var updateTime, beforeCommitTime time.Time
+	go func() {
+		defer wg.Done()
+
+		// This SQL will be blocked until the transaction is committed. It's a known limitation that the pessimistic lock
+		// is not released after rollback to savepoint. If the lock can be released after rollback in the future, this test
+		// can be modified.
+		tk2.MustExec("update parent set id = 2 where id = 1")
+		updateTime = time.Now()
+	}()
+	tk.MustExec("rollback to sp1")
+	time.Sleep(500 * time.Millisecond)
+	beforeCommitTime = time.Now()
+	tk.MustExec("commit")
+
+	wg.Wait()
+	require.Greater(t, updateTime, beforeCommitTime)
+}
+
+func TestInnodbLockWaitTimeout(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 auto_increment, k int,c varchar(255), unique index idx(id))")
+	tk.MustExec("insert into t (k,c) values (1,'abcdefg');")
+	for i := 0; i < 8; i++ {
+		tk.MustExec("insert into t (k,c) select k,c from t;")
+	}
+	tk.MustExec("update t set k= id, c = id")
+	tk.MustExec("split table t by (0), (50), (100);")
+	tk.MustExec("split table t index idx by (0), (50), (100);")
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/pkg/store/mockstore/unistore/tikv/pessimisticLockReturnWriteConflict", `return(true)`))
+	defer func() {
+		require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/pkg/store/mockstore/unistore/tikv/pessimisticLockReturnWriteConflict"))
+	}()
+	tk.MustExec("set @@innodb_lock_wait_timeout=1")
+	isolations := []string{"REPEATABLE READ", "READ COMMITTED"}
+	for _, isolation := range isolations {
+		tk.MustExec("SET SESSION TRANSACTION ISOLATION LEVEL " + isolation)
+		tk.MustExec("begin")
+		start := time.Now()
+		ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
+		res, err := tk.ExecWithContext(ctx, "update t use index (idx) set k=k+1 where id >0;")
+		cancel()
+		if res != nil {
+			require.NoError(t, res.Close())
+		}
+		require.Error(t, err)
+		msg := fmt.Sprintf("cost: %v", time.Since(start))
+		require.Equal(t, "lock wait timeout", err.Error(), msg)
+		require.Less(t, time.Since(start), time.Second*2)
+		tk.MustExec("commit")
+	}
+}
+>>>>>>> 38187e5b1d6 (txn: fix issue innodb_lock_wait_timeout doesn't work in some case (#56847))

pkg/sessiontxn/isolation/readcommitted.go

@@ -231,6 +231,11 @@ func (p *PessimisticRCTxnContextProvider) handleAfterPessimisticLockError(ctx co
 			return sessiontxn.ErrorAction(err)
 		}
 	} else if terror.ErrorEqual(kv.ErrWriteConflict, lockErr) {
+		sessVars := p.sctx.GetSessionVars()
+		waitTime := time.Since(sessVars.StmtCtx.GetLockWaitStartTime())
+		if waitTime.Milliseconds() >= sessVars.LockWaitTimeout {
+			return sessiontxn.ErrorAction(tikverr.ErrLockWaitTimeout)
+		}
 		logutil.Logger(p.ctx).Debug("pessimistic write conflict, retry statement",
 			zap.Uint64("txn", txnCtx.StartTS),
 			zap.Uint64("forUpdateTS", txnCtx.GetForUpdateTS()),

pkg/sessiontxn/isolation/repeatable_read.go

@@ -258,6 +258,11 @@ func (p *PessimisticRRTxnContextProvider) handleAfterPessimisticLockError(ctx co
 			return sessiontxn.ErrorAction(err)
 		}
 	} else if terror.ErrorEqual(kv.ErrWriteConflict, lockErr) {
+		waitTime := time.Since(sessVars.StmtCtx.GetLockWaitStartTime())
+		if waitTime.Milliseconds() >= sessVars.LockWaitTimeout {
+			return sessiontxn.ErrorAction(tikverr.ErrLockWaitTimeout)
+		}
+
 		// Always update forUpdateTS by getting a new timestamp from PD.
 		// If we use the conflict commitTS as the new forUpdateTS and async commit
 		// is used, the commitTS of this transaction may exceed the max timestamp

pkg/store/mockstore/unistore/tikv/BUILD.bazel

@@ -39,6 +39,7 @@ go_library(
         "@com_github_pingcap_badger//:badger",
         "@com_github_pingcap_badger//y",
         "@com_github_pingcap_errors//:errors",
+        "@com_github_pingcap_failpoint//:failpoint",
         "@com_github_pingcap_kvproto//pkg/coprocessor",
         "@com_github_pingcap_kvproto//pkg/deadlock",
         "@com_github_pingcap_kvproto//pkg/errorpb",

pkg/store/mockstore/unistore/tikv/server.go

@@ -21,6 +21,7 @@ import (
 	"time"
 
 	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
 	"github.com/pingcap/kvproto/pkg/coprocessor"
 	deadlockPb "github.com/pingcap/kvproto/pkg/deadlock"
 	"github.com/pingcap/kvproto/pkg/errorpb"
@@ -193,6 +194,18 @@ func (svr *Server) KvScan(ctx context.Context, req *kvrpcpb.ScanRequest) (*kvrpc
 
 // KvPessimisticLock implements the tikvpb.TikvServer interface.
 func (svr *Server) KvPessimisticLock(ctx context.Context, req *kvrpcpb.PessimisticLockRequest) (*kvrpcpb.PessimisticLockResponse, error) {
+	failpoint.Inject("pessimisticLockReturnWriteConflict", func(val failpoint.Value) {
+		if val.(bool) {
+			time.Sleep(time.Millisecond * 100)
+			err := &kverrors.ErrConflict{
+				StartTS:          req.GetForUpdateTs(),
+				ConflictTS:       req.GetForUpdateTs() + 1,
+				ConflictCommitTS: req.GetForUpdateTs() + 2,
+			}
+			failpoint.Return(&kvrpcpb.PessimisticLockResponse{Errors: []*kvrpcpb.KeyError{convertToKeyError(err)}}, nil)
+		}
+	})
+
 	reqCtx, err := newRequestCtx(svr, req.Context, "PessimisticLock")
 	if err != nil {
 		return &kvrpcpb.PessimisticLockResponse{Errors: []*kvrpcpb.KeyError{convertToKeyError(err)}}, nil