From d83fcd04bf1f324174328588b04df5381542e270 Mon Sep 17 00:00:00 2001
From: Yilin Chen <sticnarf@gmail.com>
Date: Wed, 2 Mar 2022 20:03:17 +0800
Subject: [PATCH] Fix incorrect LockNoWait value (#435)

Signed-off-by: Yilin Chen <sticnarf@gmail.com>
---
 integration_tests/lock_test.go | 37 ++++++++++++++++++++++++++++++++++
 kv/kv.go                       |  5 +++--
 2 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/integration_tests/lock_test.go b/integration_tests/lock_test.go
index b3c1ea8962..6b45059e7f 100644
--- a/integration_tests/lock_test.go
+++ b/integration_tests/lock_test.go
@@ -856,3 +856,40 @@ func (s *testLockSuite) TestStartHeartBeatAfterLockingPrimary() {
 
 	s.Nil(txn.Rollback())
 }
+
+func (s *testLockSuite) TestLockWaitTimeLimit() {
+	k1 := []byte("k1")
+	k2 := []byte("k2")
+
+	txn1, err := s.store.Begin()
+	s.Nil(err)
+	txn1.SetPessimistic(true)
+
+	// lock the primary key
+	lockCtx := &kv.LockCtx{ForUpdateTS: txn1.StartTS(), WaitStartTime: time.Now()}
+	err = txn1.LockKeys(context.Background(), lockCtx, k1, k2)
+	s.Nil(err)
+
+	txn2, err := s.store.Begin()
+	s.Nil(err)
+	txn2.SetPessimistic(true)
+
+	// test no wait
+	lockCtx = kv.NewLockCtx(txn2.StartTS(), kv.LockNoWait, time.Now())
+	err = txn2.LockKeys(context.Background(), lockCtx, k1)
+	// cannot acquire lock immediately thus error
+	s.Equal(tikverr.ErrLockAcquireFailAndNoWaitSet.Error(), err.Error())
+	// Default wait time is 1s, we use 500ms as an upper bound
+	s.Less(time.Since(lockCtx.WaitStartTime), 500*time.Millisecond)
+
+	// test for wait limited time (200ms)
+	lockCtx = kv.NewLockCtx(txn2.StartTS(), 200, time.Now())
+	err = txn2.LockKeys(context.Background(), lockCtx, k2)
+	// cannot acquire lock in time thus error
+	s.Equal(tikverr.ErrLockWaitTimeout.Error(), err.Error())
+	s.GreaterOrEqual(time.Since(lockCtx.WaitStartTime), 200*time.Millisecond)
+	s.Less(time.Since(lockCtx.WaitStartTime), 800*time.Millisecond)
+
+	s.Nil(txn1.Rollback())
+	s.Nil(txn2.Rollback())
+}
diff --git a/kv/kv.go b/kv/kv.go
index 12bf56c8e3..565c1ae99c 100644
--- a/kv/kv.go
+++ b/kv/kv.go
@@ -37,10 +37,11 @@ type ReturnedValue struct {
 
 // Used for pessimistic lock wait time
 // these two constants are special for lock protocol with tikv
-// math.MaxInt64 means always wait, 0 means nowait, others meaning lock wait in milliseconds
+// math.MaxInt64 means always wait, -1 means nowait, 0 means the default wait duration in TiKV,
+// others meaning lock wait in milliseconds
 const (
 	LockAlwaysWait = int64(math.MaxInt64)
-	LockNoWait     = int64(0)
+	LockNoWait     = int64(-1)
 )
 
 type lockWaitTimeInMs struct {