store/tikv: unset pessimistic txn primary key when statement failed (pingcap#10867) (pingcap#10894)

Author: coocood

session/pessimistic_test.go

@@ -232,3 +232,15 @@ func (s *testPessimisticSuite) TestSingleStatementRollback(c *C) {
 	tk.MustExec("commit")
 	syncCh <- true
 }
+
+func (s *testPessimisticSuite) TestFirstStatementFail(c *C) {
+	tk := testkit.NewTestKitWithInit(c, s.store)
+	tk.MustExec("drop table if exists first")
+	tk.MustExec("create table first (k int unique)")
+	tk.MustExec("insert first values (1)")
+	tk.MustExec("begin pessimistic")
+	_, err := tk.Exec("insert first values (1)")
+	c.Assert(err, NotNil)
+	tk.MustExec("insert first values (2)")
+	tk.MustExec("commit")
+}

store/mockstore/mocktikv/mock_tikv_test.go

@@ -94,7 +94,12 @@ func (s *testMockTiKVSuite) mustGetRC(c *C, key string, ts uint64, expect string
 }
 
 func (s *testMockTiKVSuite) mustPutOK(c *C, key, value string, startTS, commitTS uint64) {
-	errs := s.store.Prewrite(putMutations(key, value), []byte(key), startTS, 0)
+	req := &kvrpcpb.PrewriteRequest{
+		Mutations:    putMutations(key, value),
+		PrimaryLock:  []byte(key),
+		StartVersion: startTS,
+	}
+	errs := s.store.Prewrite(req)
 	for _, err := range errs {
 		c.Assert(err, IsNil)
 	}
@@ -109,7 +114,12 @@ func (s *testMockTiKVSuite) mustDeleteOK(c *C, key string, startTS, commitTS uin
 			Key: []byte(key),
 		},
 	}
-	errs := s.store.Prewrite(mutations, []byte(key), startTS, 0)
+	req := &kvrpcpb.PrewriteRequest{
+		Mutations:    mutations,
+		PrimaryLock:  []byte(key),
+		StartVersion: startTS,
+	}
+	errs := s.store.Prewrite(req)
 	for _, err := range errs {
 		c.Assert(err, IsNil)
 	}
@@ -146,7 +156,12 @@ func (s *testMockTiKVSuite) mustRangeReverseScanOK(c *C, start, end string, limi
 }
 
 func (s *testMockTiKVSuite) mustPrewriteOK(c *C, mutations []*kvrpcpb.Mutation, primary string, startTS uint64) {
-	errs := s.store.Prewrite(mutations, []byte(primary), startTS, 0)
+	req := &kvrpcpb.PrewriteRequest{
+		Mutations:    mutations,
+		PrimaryLock:  []byte(primary),
+		StartVersion: startTS,
+	}
+	errs := s.store.Prewrite(req)
 	for _, err := range errs {
 		c.Assert(err, IsNil)
 	}
@@ -412,7 +427,12 @@ func (s *testMockTiKVSuite) TestCommitConflict(c *C) {
 	// A prewrite.
 	s.mustPrewriteOK(c, putMutations("x", "A"), "x", 5)
 	// B prewrite and find A's lock.
-	errs := s.store.Prewrite(putMutations("x", "B"), []byte("x"), 10, 0)
+	req := &kvrpcpb.PrewriteRequest{
+		Mutations:    putMutations("x", "B"),
+		PrimaryLock:  []byte("x"),
+		StartVersion: 10,
+	}
+	errs := s.store.Prewrite(req)
 	c.Assert(errs[0], NotNil)
 	// B find rollback A because A exist too long.
 	s.mustRollbackOK(c, [][]byte{[]byte("x")}, 5)
@@ -470,17 +490,27 @@ func (s *testMockTiKVSuite) TestBatchResolveLock(c *C) {
 
 func (s *testMockTiKVSuite) TestRollbackAndWriteConflict(c *C) {
 	s.mustPutOK(c, "test", "test", 1, 3)
-
-	errs := s.store.Prewrite(putMutations("lock", "lock", "test", "test1"), []byte("test"), 2, 2)
+	req := &kvrpcpb.PrewriteRequest{
+		Mutations:    putMutations("lock", "lock", "test", "test1"),
+		PrimaryLock:  []byte("test"),
+		StartVersion: 2,
+		LockTtl:      2,
+	}
+	errs := s.store.Prewrite(req)
 	s.mustWriteWriteConflict(c, errs, 1)
 
 	s.mustPutOK(c, "test", "test2", 5, 8)
 
 	// simulate `getTxnStatus` for txn 2.
 	err := s.store.Cleanup([]byte("test"), 2)
 	c.Assert(err, IsNil)
-
-	errs = s.store.Prewrite(putMutations("test", "test3"), []byte("test"), 6, 1)
+	req = &kvrpcpb.PrewriteRequest{
+		Mutations:    putMutations("test", "test3"),
+		PrimaryLock:  []byte("test"),
+		StartVersion: 6,
+		LockTtl:      1,
+	}
+	errs = s.store.Prewrite(req)
 	s.mustWriteWriteConflict(c, errs, 0)
 }
 

store/mockstore/mocktikv/mvcc.go

@@ -434,7 +434,7 @@ type MVCCStore interface {
 	BatchGet(ks [][]byte, startTS uint64, isoLevel kvrpcpb.IsolationLevel) []Pair
 	PessimisticLock(mutations []*kvrpcpb.Mutation, primary []byte, startTS, forUpdateTS uint64, ttl uint64) []error
 	PessimisticRollback(keys [][]byte, startTS, forUpdateTS uint64) []error
-	Prewrite(mutations []*kvrpcpb.Mutation, primary []byte, startTS, ttl uint64) []error
+	Prewrite(req *kvrpcpb.PrewriteRequest) []error
 	Commit(keys [][]byte, startTS, commitTS uint64) error
 	Rollback(keys [][]byte, startTS uint64) error
 	Cleanup(key []byte, startTS uint64) error

store/mockstore/mocktikv/mvcc_leveldb.go

@@ -554,7 +554,7 @@ func (mvcc *MVCCLevelDB) pessimisticLockMutation(batch *leveldb.Batch, mutation
 		}
 		return nil
 	}
-	if err = checkConflictValue(iter, mutation.Key, forUpdateTS); err != nil {
+	if err = checkConflictValue(iter, mutation, forUpdateTS); err != nil {
 		return err
 	}
 
@@ -623,15 +623,19 @@ func pessimisticRollbackKey(db *leveldb.DB, batch *leveldb.Batch, key []byte, st
 }
 
 // Prewrite implements the MVCCStore interface.
-func (mvcc *MVCCLevelDB) Prewrite(mutations []*kvrpcpb.Mutation, primary []byte, startTS uint64, ttl uint64) []error {
+func (mvcc *MVCCLevelDB) Prewrite(req *kvrpcpb.PrewriteRequest) []error {
+	mutations := req.Mutations
+	primary := req.PrimaryLock
+	startTS := req.StartVersion
+	ttl := req.LockTtl
 	mvcc.mu.Lock()
 	defer mvcc.mu.Unlock()
 
 	anyError := false
 	batch := &leveldb.Batch{}
 	errs := make([]error, 0, len(mutations))
 	txnSize := len(mutations)
-	for _, m := range mutations {
+	for i, m := range mutations {
 		// If the operation is Insert, check if key is exists at first.
 		var err error
 		if m.GetOp() == kvrpcpb.Op_Insert {
@@ -650,7 +654,8 @@ func (mvcc *MVCCLevelDB) Prewrite(mutations []*kvrpcpb.Mutation, primary []byte,
 				continue
 			}
 		}
-		err = prewriteMutation(mvcc.db, batch, m, startTS, primary, ttl, uint64(txnSize))
+		isPessimisticLock := len(req.IsPessimisticLock) > 0 && req.IsPessimisticLock[i]
+		err = prewriteMutation(mvcc.db, batch, m, startTS, primary, ttl, uint64(txnSize), isPessimisticLock)
 		errs = append(errs, err)
 		if err != nil {
 			anyError = true
@@ -666,26 +671,34 @@ func (mvcc *MVCCLevelDB) Prewrite(mutations []*kvrpcpb.Mutation, primary []byte,
 	return errs
 }
 
-func checkConflictValue(iter *Iterator, key []byte, startTS uint64) error {
+func checkConflictValue(iter *Iterator, m *kvrpcpb.Mutation, startTS uint64) error {
 	dec := valueDecoder{
-		expectKey: key,
+		expectKey: m.Key,
 	}
 	ok, err := dec.Decode(iter)
 	if err != nil {
 		return errors.Trace(err)
 	}
+	if !ok {
+		return nil
+	}
 	// Note that it's a write conflict here, even if the value is a rollback one.
-	if ok && dec.value.commitTS >= startTS {
+	if dec.value.commitTS >= startTS {
 		return &ErrConflict{
 			StartTS:    startTS,
 			ConflictTS: dec.value.commitTS,
-			Key:        key,
+			Key:        m.Key,
+		}
+	}
+	if m.Op == kvrpcpb.Op_PessimisticLock && m.Assertion == kvrpcpb.Assertion_NotExist {
+		return &ErrKeyAlreadyExist{
+			Key: m.Key,
 		}
 	}
 	return nil
 }
 
-func prewriteMutation(db *leveldb.DB, batch *leveldb.Batch, mutation *kvrpcpb.Mutation, startTS uint64, primary []byte, ttl uint64, txnSize uint64) error {
+func prewriteMutation(db *leveldb.DB, batch *leveldb.Batch, mutation *kvrpcpb.Mutation, startTS uint64, primary []byte, ttl uint64, txnSize uint64, isPessimisticLock bool) error {
 	startKey := mvccEncode(mutation.Key, lockVer)
 	iter := newIterator(db, &util.Range{
 		Start: startKey,
@@ -708,7 +721,10 @@ func prewriteMutation(db *leveldb.DB, batch *leveldb.Batch, mutation *kvrpcpb.Mu
 		}
 		// Overwrite the pessimistic lock.
 	} else {
-		err = checkConflictValue(iter, mutation.Key, startTS)
+		if isPessimisticLock {
+			return ErrAbort("pessimistic lock not found")
+		}
+		err = checkConflictValue(iter, mutation, startTS)
 		if err != nil {
 			return err
 		}

store/mockstore/mocktikv/rpc.go

@@ -292,7 +292,7 @@ func (h *rpcHandler) handleKvPrewrite(req *kvrpcpb.PrewriteRequest) *kvrpcpb.Pre
 			panic("KvPrewrite: key not in region")
 		}
 	}
-	errs := h.mvccStore.Prewrite(req.Mutations, req.PrimaryLock, req.GetStartVersion(), req.GetLockTtl())
+	errs := h.mvccStore.Prewrite(req)
 	return &kvrpcpb.PrewriteResponse{
 		Errors: convertToKeyErrors(errs),
 	}

store/tikv/2pc_test.go

@@ -461,6 +461,28 @@ func (s *testCommitterSuite) TestPessimisticPrewriteRequest(c *C) {
 	c.Assert(req.Prewrite.ForUpdateTs, Equals, uint64(100))
 }
 
+func (s *testCommitterSuite) TestUnsetPrimaryKey(c *C) {
+	// This test checks that the isPessimisticLock field is set in the request even when no keys are pessimistic lock.
+	key := kv.Key("key")
+	txn := s.begin(c)
+	c.Assert(txn.Set(key, key), IsNil)
+	c.Assert(txn.Commit(context.Background()), IsNil)
+
+	txn = s.begin(c)
+	txn.SetOption(kv.Pessimistic, true)
+	txn.SetOption(kv.PresumeKeyNotExists, nil)
+	_, _ = txn.us.Get(key)
+	c.Assert(txn.Set(key, key), IsNil)
+	txn.DelOption(kv.PresumeKeyNotExists)
+	err := txn.LockKeys(context.Background(), txn.startTS, key)
+	c.Assert(err, NotNil)
+	c.Assert(txn.Delete(key), IsNil)
+	key2 := kv.Key("key2")
+	c.Assert(txn.Set(key2, key2), IsNil)
+	err = txn.Commit(context.Background())
+	c.Assert(err, IsNil)
+}
+
 func (s *testCommitterSuite) TestPessimisticLockedKeysDedup(c *C) {
 	txn := s.begin(c)
 	txn.SetOption(kv.Pessimistic, true)

store/tikv/txn.go

@@ -373,7 +373,11 @@ func (txn *tikvTxn) LockKeys(ctx context.Context, forUpdateTS uint64, keysInput
 			if err != nil {
 				return err
 			}
+		}
+		var assignedPrimaryKey bool
+		if txn.committer.primaryKey == nil {
 			txn.committer.primaryKey = keys[0]
+			assignedPrimaryKey = true
 		}
 
 		bo := NewBackoffer(ctx, pessimisticLockMaxBackoff).WithVars(txn.vars)
@@ -391,6 +395,10 @@ func (txn *tikvTxn) LockKeys(ctx context.Context, forUpdateTS uint64, keysInput
 				// Sleep a little, wait for the other transaction that blocked by this transaction to acquire the lock.
 				time.Sleep(time.Millisecond * 5)
 			}
+			if assignedPrimaryKey {
+				// unset the primary key if we assigned primary key when failed to lock it.
+				txn.committer.primaryKey = nil
+			}
 			return err
 		}
 	}