txnkv: convert LOCK on untouched kv into PUT (#752)

Signed-off-by: zyguan <[email protected]>
tikv · Mar 29, 2023 · 3512be1 · 3512be1
1 parent 603f10d
commit 3512be1
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 6 deletions.
diff --git a/integration_tests/2pc_test.go b/integration_tests/2pc_test.go
@@ -1747,6 +1747,7 @@ func (s *testCommitterSuite) TestFlagsInMemBufferMutations() {
 func (s *testCommitterSuite) TestSetLockedKeyValue() {
 	ctx := context.Background()
 	k1 := []byte("k1")
+	k2 := []byte("t00000001_i000000001")
 	v1 := []byte("v1")
 	v2 := []byte("v2")
 
@@ -1801,6 +1802,26 @@ func (s *testCommitterSuite) TestSetLockedKeyValue() {
 			checkByOpVals(kvrpcpb.Op_Put, v2),
 			checkByOpVals(kvrpcpb.Op_Put, v2),
 		},
+		{
+			"LockAndSetUnnecessaryKeyWithSameValue",
+			[]func(txn transaction.TxnProbe){
+				func(txn transaction.TxnProbe) { txn.SetLockedKeyValue(k2, v2) },
+				func(txn transaction.TxnProbe) { mustLockKey(txn, k2) },
+				func(txn transaction.TxnProbe) { s.Require().NoError(txn.Set(k2, v2)) },
+			},
+			checkByOpVals(kvrpcpb.Op_Put, v2),
+			checkByOpVals(kvrpcpb.Op_Lock, v2),
+		},
+		{
+			"LockAndSetUnnecessaryKeyWithDiffValue",
+			[]func(txn transaction.TxnProbe){
+				func(txn transaction.TxnProbe) { txn.SetLockedKeyValue(k2, v1) },
+				func(txn transaction.TxnProbe) { mustLockKey(txn, k2) },
+				func(txn transaction.TxnProbe) { s.Require().NoError(txn.Set(k2, v2)) },
+			},
+			checkByOpVals(kvrpcpb.Op_Lock, v2),
+			checkByOpVals(kvrpcpb.Op_Lock, v2),
+		},
 		{
 			"LockAndDelete",
 			[]func(txn transaction.TxnProbe){
@@ -1811,12 +1832,26 @@ func (s *testCommitterSuite) TestSetLockedKeyValue() {
 			checkByOpVals(kvrpcpb.Op_Del, []byte{}),
 			checkByOpVals(kvrpcpb.Op_Del, []byte{}),
 		},
+		{
+			"LockAndDeleteYourWrite",
+			[]func(txn transaction.TxnProbe){
+				func(txn transaction.TxnProbe) { txn.SetLockedKeyValue(k1, v1) },
+				func(txn transaction.TxnProbe) { mustLockKey(txn, k1) },
+				func(txn transaction.TxnProbe) {
+					s.Require().NoError(txn.GetMemBuffer().DeleteWithFlags(k1, kv.SetNewlyInserted))
+				},
+			},
+			checkByOpVals(kvrpcpb.Op_Lock, []byte{}),
+			checkByOpVals(kvrpcpb.Op_Lock, []byte{}),
+		},
 	} {
 		var testAll func(name string, state []bool, actions []func(txn transaction.TxnProbe))
 		testAll = func(name string, state []bool, actions []func(txn transaction.TxnProbe)) {
 			if len(actions) == len(tt.actions) {
 				s.Run("Pessimistic"+name, func() {
 					txn := s.begin()
+					defer txn.Rollback()
+					txn.SetKVFilter(kvFilter{})
 					txn.SetPessimistic(true)
 					for _, action := range actions {
 						action(txn)
@@ -1828,13 +1863,14 @@ func (s *testCommitterSuite) TestSetLockedKeyValue() {
 				})
 				s.Run("Optimistic"+name, func() {
 					txn := s.begin()
+					defer txn.Rollback()
+					txn.SetKVFilter(kvFilter{})
 					for _, action := range actions {
 						action(txn)
 					}
 					c, err := txn.NewCommitter(1)
 					s.Require().NoError(err)
 					tt.checkOptimisitc(c.GetMutations())
-					s.Require().NoError(txn.Rollback())
 				})
 				return
 			}

diff --git a/txnkv/transaction/2pc.go b/txnkv/transaction/2pc.go
@@ -562,11 +562,18 @@ func (c *twoPhaseCommitter) initKeysAndMutations(ctx context.Context) error {
 					if !flags.HasLocked() {
 						continue
 					}
-					// If the key was locked before, we should prewrite the lock even if
-					// the KV needn't be committed according to the filter. Otherwise, we
-					// were forgetting removing pessimistic locks added before.
-					op = kvrpcpb.Op_Lock
-					lockCnt++
+					if val, ok := txn.getValueByLockedKey(key); ok && bytes.Equal(val, value) && c.isPessimistic {
+						// Change the LOCK into PUT if the value of this key has a cached value.
+						cachedValue = val
+						op = kvrpcpb.Op_Put
+						putFromLockCnt++
+					} else {
+						// If the key was locked before, we should prewrite the lock even if
+						// the KV needn't be committed according to the filter. Otherwise, we
+						// were forgetting removing pessimistic locks added before.
+						op = kvrpcpb.Op_Lock
+						lockCnt++
+					}
 				} else {
 					op = kvrpcpb.Op_Put
 					if flags.HasPresumeKeyNotExists() {