Support setting assertions and receiving assertion errors for Prewrite requests

error/error.go

@@ -231,6 +231,15 @@ func (e *ErrTokenLimit) Error() string {
 	return fmt.Sprintf("Store token is up to the limit, store id = %d.", e.StoreID)
 }
 
+// ErrAssertionFailed is the error that assertion on data failed.
+type ErrAssertionFailed struct {
+	*kvrpcpb.AssertionFailed
+}
+
+func (e *ErrAssertionFailed) Error() string {
+	return fmt.Sprintf("assertion failed { %s }", e.AssertionFailed.String())
+}
+
 // ExtractKeyErr extracts a KeyError.
 func ExtractKeyErr(keyErr *kvrpcpb.KeyError) error {
 	if val, err := util.EvalFailpoint("mockRetryableErrorResp"); err == nil {
@@ -248,6 +257,10 @@ func ExtractKeyErr(keyErr *kvrpcpb.KeyError) error {
 		return errors.WithStack(&ErrRetryable{Retryable: keyErr.Retryable})
 	}
 
+	if keyErr.AssertionFailed != nil {
+		return &ErrAssertionFailed{AssertionFailed: keyErr.AssertionFailed}
+	}
+
 	if keyErr.Abort != "" {
 		err := errors.Errorf("tikv aborts txn: %s", keyErr.GetAbort())
 		logutil.BgLogger().Warn("2PC failed", zap.Error(err))

go.mod

@@ -17,7 +17,7 @@ require (
 	github.com/opentracing/opentracing-go v1.1.0
 	github.com/pingcap/failpoint v0.0.0-20210316064728-7acb0f0a3dfd
 	github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989
-	github.com/pingcap/kvproto v0.0.0-20211122024046-03abd340988f
+	github.com/pingcap/kvproto v0.0.0-20211202065422-a412f7a319c3
 	github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.5.1

go.sum

@@ -78,7 +78,6 @@ github.com/fatih/color v1.7.0/go.mod h1:Zm6kSWBoL9eyXnKyktHP6abPY2pDugNf5Kwzbycv
 github.com/fatih/color v1.9.0/go.mod h1:eQcE1qtQxscV5RaZvpXrrb8Drkc3/DdQ+uUYCNjL+zU=
 github.com/fatih/structtag v1.2.0/go.mod h1:mBJUNpUnHmRKrKlQQlmCrh5PuhftFbNv8Ys4/aAZl94=
 github.com/fogleman/gg v1.3.0/go.mod h1:R/bRT+9gY/C5z7JzPU0zXsXHKM4/ayA+zqcVNZzPa1k=
-github.com/fsnotify/fsnotify v1.4.7 h1:IXs+QLmnXW2CcXuY+8Mzv/fWEsPGWxqefPtCP5CnV9I=
 github.com/fsnotify/fsnotify v1.4.7/go.mod h1:jwhsz4b93w/PPRr/qN1Yymfu8t87LnFCMoQvtojpjFo=
 github.com/ghodss/yaml v1.0.0/go.mod h1:4dBDuWmgqj2HViK6kFavaiC9ZROes6MMH2rRYeMEF04=
 github.com/gin-contrib/gzip v0.0.1/go.mod h1:fGBJBCdt6qCZuCAOwWuFhBB4OOq9EFqlo5dEaFhhu5w=
@@ -183,7 +182,6 @@ github.com/jinzhu/now v1.1.1/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/
 github.com/jinzhu/now v1.1.2/go.mod h1:d3SSVoowX0Lcu0IBviAWJpolVfI5UJVZZ7cO71lE/z8=
 github.com/jmespath/go-jmespath v0.4.0 h1:BEgLn5cpjn8UN1mAw4NjwDrS35OdebyEtFe+9YPoQUg=
 github.com/jmespath/go-jmespath v0.4.0/go.mod h1:T8mJZnbsbmF+m6zOOFylbeCJqk5+pHWvzYPziyZiYoo=
-github.com/jmespath/go-jmespath/internal/testify v1.5.1 h1:shLQSRRSCCPj3f2gpwzGwWFoC7ycTf1rcQZHOlsJ6N8=
 github.com/jmespath/go-jmespath/internal/testify v1.5.1/go.mod h1:L3OGu8Wl2/fWfCI6z80xFu9LTZmf1ZRjMHUOPmWr69U=
 github.com/joho/godotenv v1.3.0/go.mod h1:7hK45KPybAkOC6peb+G5yklZfMxEjkZhHbwpqxOKXbg=
 github.com/jonboulle/clockwork v0.1.0 h1:VKV+ZcuP6l3yW9doeqz6ziZGgcynBVQO+obU0+0hcPo=
@@ -283,8 +281,8 @@ github.com/pingcap/goleveldb v0.0.0-20191226122134-f82aafb29989/go.mod h1:O17Xtb
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
 github.com/pingcap/kvproto v0.0.0-20200411081810-b85805c9476c/go.mod h1:IOdRDPLyda8GX2hE/jO7gqaCV/PNFh8BZQCQZXfIOqI=
 github.com/pingcap/kvproto v0.0.0-20210819164333-bd5706b9d9f2/go.mod h1:IOdRDPLyda8GX2hE/jO7gqaCV/PNFh8BZQCQZXfIOqI=
-github.com/pingcap/kvproto v0.0.0-20211122024046-03abd340988f h1:hjInxK1Ie6CYx7Jy2pYnBdEnWI8jIfr423l9Yh6LRy8=
-github.com/pingcap/kvproto v0.0.0-20211122024046-03abd340988f/go.mod h1:IOdRDPLyda8GX2hE/jO7gqaCV/PNFh8BZQCQZXfIOqI=
+github.com/pingcap/kvproto v0.0.0-20211202065422-a412f7a319c3 h1:PZMewlB2b5JC0pMteDmWNBtAgmJv2ih1+z+WifdVcEk=
+github.com/pingcap/kvproto v0.0.0-20211202065422-a412f7a319c3/go.mod h1:IOdRDPLyda8GX2hE/jO7gqaCV/PNFh8BZQCQZXfIOqI=
 github.com/pingcap/log v0.0.0-20191012051959-b742a5d432e9/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v0.0.0-20200511115504-543df19646ad/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=
 github.com/pingcap/log v0.0.0-20210317133921-96f4fcab92a4/go.mod h1:4rbK1p9ILyIfb6hU7OG2CiWSqMXnp3JMbiaVJ6mvoY8=

integration_tests/2pc_test.go

@@ -746,6 +746,39 @@ func (s *testCommitterSuite) TestPessimisticLockReturnValues() {
 	s.Equal(lockCtx.Values[string(key2)].Value, key2)
 }
 
+func (s *testCommitterSuite) TestPessimisticLockCheckExistence() {
+	key := []byte("key")
+	key2 := []byte("key2")
+	txn := s.begin()
+	s.Nil(txn.Set(key, key))
+	s.Nil(txn.Commit(context.Background()))
+
+	txn = s.begin()
+	txn.SetPessimistic(true)
+	lockCtx := &kv.LockCtx{ForUpdateTS: txn.StartTS(), WaitStartTime: time.Now()}
+	lockCtx.InitCheckExistence(2)
+	s.Nil(txn.LockKeys(context.Background(), lockCtx, key, key2))
+	s.Len(lockCtx.Values, 2)
+	s.Empty(lockCtx.Values[string(key)].Value)
+	s.True(lockCtx.Values[string(key)].Exists)
+	s.Empty(lockCtx.Values[string(key2)].Value)
+	s.False(lockCtx.Values[string(key2)].Exists)
+	s.Nil(txn.Rollback())
+
+	txn = s.begin()
+	txn.SetPessimistic(true)
+	lockCtx = &kv.LockCtx{ForUpdateTS: txn.StartTS(), WaitStartTime: time.Now()}
+	lockCtx.InitCheckExistence(2)
+	lockCtx.InitReturnValues(2)
+	s.Nil(txn.LockKeys(context.Background(), lockCtx, key, key2))
+	s.Len(lockCtx.Values, 2)
+	s.Equal(lockCtx.Values[string(key)].Value, key)
+	s.True(lockCtx.Values[string(key)].Exists)
+	s.Empty(lockCtx.Values[string(key2)].Value)
+	s.False(lockCtx.Values[string(key2)].Exists)
+	s.Nil(txn.Rollback())
+}
+
 // TestElapsedTTL tests that elapsed time is correct even if ts physical time is greater than local time.
 func (s *testCommitterSuite) TestElapsedTTL() {
 	key := []byte("key")
@@ -1215,7 +1248,7 @@ func (s *testCommitterSuite) TestResolveMixed() {
 	// stop txn ttl manager and remove primary key, make the other keys left behind
 	committer.CloseTTLManager()
 	muts := transaction.NewPlainMutations(1)
-	muts.Push(kvrpcpb.Op_Lock, pk, nil, true)
+	muts.Push(kvrpcpb.Op_Lock, pk, nil, true, false, false)
 	err = committer.PessimisticRollbackMutations(context.Background(), &muts)
 	s.Nil(err)
 
@@ -1598,3 +1631,56 @@ func (s *testCommitterSuite) TestNewlyInsertedMemDBFlag() {
 	err = txn.Commit(ctx)
 	s.Nil(err)
 }
+
+func (s *testCommitterSuite) TestFlagsInMemBufferMutations() {
+	// Get a MemDB object from a transaction object.
+	db := s.begin().GetMemBuffer()
+
+	// A helper for iterating all cases.
+	forEachCase := func(f func(op kvrpcpb.Op, key []byte, value []byte, index int, isPessimisticLock, assertExist, assertNotExist bool)) {
+		keyIndex := 0
+		for _, op := range []kvrpcpb.Op{kvrpcpb.Op_Put, kvrpcpb.Op_Del, kvrpcpb.Op_CheckNotExists} {
+			for flags := 0; flags < (1 << 3); flags++ {
+				key := []byte(fmt.Sprintf("k%05d", keyIndex))
+				value := []byte(fmt.Sprintf("v%05d", keyIndex))
+
+				// `flag` Iterates all combinations of flags in binary.
+				isPessimisticLock := (flags & 0x4) != 0
+				assertExist := (flags & 0x2) != 0
+				assertNotExist := (flags & 0x1) != 0
+
+				f(op, key, value, keyIndex, isPessimisticLock, assertExist, assertNotExist)
+				keyIndex++
+			}
+		}
+	}
+
+	// Put some keys to the MemDB
+	forEachCase(func(op kvrpcpb.Op, key []byte, value []byte, i int, isPessimisticLock, assertExist, assertNotExist bool) {
+		if op == kvrpcpb.Op_Put {
+			err := db.Set(key, value)
+			s.Nil(err)
+		} else if op == kvrpcpb.Op_Del {
+			err := db.Delete(key)
+			s.Nil(err)
+		} else {
+			db.UpdateFlags(key, kv.SetPresumeKeyNotExists)
+		}
+	})
+
+	// Create memBufferMutations object and add keys with flags to it.
+	mutations := transaction.NewMemBufferMutationsProbe(db.Len(), db)
+
+	forEachCase(func(op kvrpcpb.Op, key []byte, value []byte, i int, isPessimisticLock, assertExist, assertNotExist bool) {
+		handle := db.IterWithFlags(key, nil).Handle()
+		mutations.Push(op, isPessimisticLock, assertExist, assertNotExist, handle)
+	})
+
+	forEachCase(func(op kvrpcpb.Op, key []byte, value []byte, i int, isPessimisticLock, assertExist, assertNotExist bool) {
+		s.Equal(key, mutations.GetKey(i))
+		s.Equal(op, mutations.GetOp(i))
+		s.Equal(isPessimisticLock, mutations.IsPessimisticLock(i))
+		s.Equal(assertExist, mutations.IsAssertExists(i))
+		s.Equal(assertNotExist, mutations.IsAssertNotExist(i))
+	})
+}

integration_tests/assertion_test.go

@@ -0,0 +1,209 @@
+// Copyright 2021 TiKV Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package tikv_test
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/pingcap/errors"
+	"github.com/pingcap/failpoint"
+	"github.com/pingcap/kvproto/pkg/kvrpcpb"
+	"github.com/stretchr/testify/suite"
+	tikverr "github.com/tikv/client-go/v2/error"
+	"github.com/tikv/client-go/v2/kv"
+	"github.com/tikv/client-go/v2/testutils"
+	"github.com/tikv/client-go/v2/tikv"
+	"github.com/tikv/client-go/v2/txnkv/transaction"
+)
+
+func TestAssertion(t *testing.T) {
+	suite.Run(t, new(testAssertionSuite))
+}
+
+type testAssertionSuite struct {
+	suite.Suite
+	cluster testutils.Cluster
+	store   tikv.StoreProbe
+}
+
+func (s *testAssertionSuite) SetupTest() {
+	s.store = tikv.StoreProbe{KVStore: NewTestStore(s.T())}
+}
+
+func (s *testAssertionSuite) TearDownTest() {
+	s.store.Close()
+}
+
+type mockAmender struct{}
+
+func (*mockAmender) AmendTxn(ctx context.Context, startInfoSchema transaction.SchemaVer, change *transaction.RelatedSchemaChange, mutations transaction.CommitterMutations) (transaction.CommitterMutations, error) {
+	return nil, nil
+}
+
+func (s *testAssertionSuite) testAssertionImpl(keyPrefix byte, pessimistic bool, lockKeys bool, assertionLevel kvrpcpb.AssertionLevel, enableAmend bool) {
+	if assertionLevel != kvrpcpb.AssertionLevel_Strict {
+		s.Nil(failpoint.Enable("tikvclient/assertionSkipCheckFromPrewrite", "return"))
+		defer func() {
+			s.Nil(failpoint.Disable("tikvclient/assertionSkipCheckFromPrewrite"))
+		}()
+	}
+	if assertionLevel != kvrpcpb.AssertionLevel_Fast {
+		s.Nil(failpoint.Enable("tikvclient/assertionSkipCheckFromLock", "return"))
+		defer func() {
+			s.Nil(failpoint.Disable("tikvclient/assertionSkipCheckFromLock"))
+		}()
+	}
+
+	// Compose the key
+	k := func(i byte) []byte {
+		return []byte{keyPrefix, 'k', i}
+	}
+
+	// Prepare some data. Make k1, k3, k7 exist.
+	prepareTxn, err := s.store.Begin()
+	s.Nil(err)
+	err = prepareTxn.Set(k(1), []byte("v1"))
+	s.Nil(err)
+	err = prepareTxn.Set(k(3), []byte("v3"))
+	s.Nil(err)
+	err = prepareTxn.Set(k(7), []byte("v7"))
+	s.Nil(err)
+	err = prepareTxn.Commit(context.Background())
+	s.Nil(err)
+	prepareStartTS := prepareTxn.GetCommitter().GetStartTS()
+	prepareCommitTS := prepareTxn.GetCommitTS()
+
+	// A helper to perform a complete transaction. When multiple keys are passed in, assertion will be set on only
+	// the last key.
+	doTxn := func(lastAssertion kv.FlagsOp, keys ...[]byte) (uint64, error) {
+		txn, err := s.store.Begin()
+		s.Nil(err)
+		txn.SetAssertionLevel(assertionLevel)
+		txn.SetPessimistic(pessimistic)
+		if enableAmend {
+			txn.SetSchemaAmender(&mockAmender{})
+		}
+		if lockKeys {
+			lockCtx := kv.NewLockCtx(txn.StartTS(), 1000, time.Now())
+			lockCtx.InitCheckExistence(1)
+			err = txn.LockKeys(context.Background(), lockCtx, keys...)
+			s.Nil(err)
+		} else if pessimistic {
+			// Since we don't want to lock the keys to be tested, set another key as the primary.
+			err = txn.LockKeysWithWaitTime(context.Background(), 10000, []byte("primary"))
+			s.Nil(err)
+		}
+		for _, key := range keys {
+			err = txn.Set(key, append([]byte{'v'}, key...))
+			s.Nil(err)
+		}
+		txn.GetMemBuffer().UpdateFlags(keys[len(keys)-1], lastAssertion)
+		err = txn.Commit(context.Background())
+		startTS := txn.GetCommitter().GetStartTS()
+		return startTS, err
+	}
+
+	checkAssertionFailError := func(err error, startTS uint64, key []byte, assertion kvrpcpb.Assertion, existingStartTS uint64, existingCommitTS uint64) {
+		assertionFailed, ok := errors.Cause(err).(*tikverr.ErrAssertionFailed)
+		s.True(ok)
+		s.Equal(startTS, assertionFailed.StartTs)
+		s.Equal(key, assertionFailed.Key)
+		s.Equal(assertion, assertionFailed.Assertion)
+		s.Equal(existingStartTS, assertionFailed.ExistingStartTs)
+		s.Equal(existingCommitTS, assertionFailed.ExistingCommitTs)
+	}
+
+	if assertionLevel == kvrpcpb.AssertionLevel_Strict && !enableAmend {
+		// Single key.
+		_, err = doTxn(kv.SetAssertExist, k(1))
+		s.Nil(err)
+		_, err = doTxn(kv.SetAssertNotExist, k(2))
+		s.Nil(err)
+		startTS, err := doTxn(kv.SetAssertNotExist, k(3))
+		s.NotNil(err)
+		checkAssertionFailError(err, startTS, k(3), kvrpcpb.Assertion_NotExist, prepareStartTS, prepareCommitTS)
+		startTS, err = doTxn(kv.SetAssertExist, k(4))
+		s.NotNil(err)
+		checkAssertionFailError(err, startTS, k(4), kvrpcpb.Assertion_Exist, 0, 0)
+
+		// Multiple keys
+		startTS, err = doTxn(kv.SetAssertNotExist, k(5), k(6), k(7))
+		s.NotNil(err)
+		checkAssertionFailError(err, startTS, k(7), kvrpcpb.Assertion_NotExist, prepareStartTS, prepareCommitTS)
+		startTS, err = doTxn(kv.SetAssertExist, k(8), k(9), k(10))
+		s.NotNil(err)
+		checkAssertionFailError(err, startTS, k(10), kvrpcpb.Assertion_Exist, 0, 0)
+	} else if assertionLevel == kvrpcpb.AssertionLevel_Fast && pessimistic && lockKeys && !enableAmend {
+		// Different from STRICT level, the already-existing version's startTS and commitTS cannot be fetched.
+
+		// Single key.
+		_, err = doTxn(kv.SetAssertExist, k(1))
+		s.Nil(err)
+		_, err = doTxn(kv.SetAssertNotExist, k(2))
+		s.Nil(err)
+		startTS, err := doTxn(kv.SetAssertNotExist, k(3))
+		s.NotNil(err)
+		checkAssertionFailError(err, startTS, k(3), kvrpcpb.Assertion_NotExist, 0, 0)
+		startTS, err = doTxn(kv.SetAssertExist, k(4))
+		s.NotNil(err)
+		checkAssertionFailError(err, startTS, k(4), kvrpcpb.Assertion_Exist, 0, 0)
+
+		// Multiple keys
+		startTS, err = doTxn(kv.SetAssertNotExist, k(5), k(6), k(7))
+		s.NotNil(err)
+		checkAssertionFailError(err, startTS, k(7), kvrpcpb.Assertion_NotExist, 0, 0)
+		startTS, err = doTxn(kv.SetAssertExist, k(8), k(9), k(10))
+		s.NotNil(err)
+		checkAssertionFailError(err, startTS, k(10), kvrpcpb.Assertion_Exist, 0, 0)
+	} else {
+		// Nothing will be detected.
+
+		// Single key.
+		_, err = doTxn(kv.SetAssertExist, k(1))
+		s.Nil(err)
+		_, err = doTxn(kv.SetAssertNotExist, k(2))
+		s.Nil(err)
+		_, err := doTxn(kv.SetAssertNotExist, k(3))
+		s.Nil(err)
+		_, err = doTxn(kv.SetAssertExist, k(4))
+		s.Nil(err)
+
+		// Multiple keys
+		_, err = doTxn(kv.SetAssertNotExist, k(5), k(6), k(7))
+		s.Nil(err)
+		_, err = doTxn(kv.SetAssertExist, k(8), k(9), k(10))
+		s.Nil(err)
+	}
+}
+
+func (s *testAssertionSuite) TestPrewriteAssertion() {
+	s.testAssertionImpl('a', false, false, kvrpcpb.AssertionLevel_Strict, false)
+	s.testAssertionImpl('b', true, false, kvrpcpb.AssertionLevel_Strict, false)
+	s.testAssertionImpl('c', true, true, kvrpcpb.AssertionLevel_Strict, false)
+	s.testAssertionImpl('a', false, false, kvrpcpb.AssertionLevel_Strict, true)
+	s.testAssertionImpl('b', true, false, kvrpcpb.AssertionLevel_Strict, true)
+	s.testAssertionImpl('c', true, true, kvrpcpb.AssertionLevel_Strict, true)
+}
+
+func (s *testAssertionSuite) TestFastAssertion() {
+	s.testAssertionImpl('a', false, false, kvrpcpb.AssertionLevel_Fast, false)
+	s.testAssertionImpl('b', true, false, kvrpcpb.AssertionLevel_Fast, false)
+	s.testAssertionImpl('c', true, true, kvrpcpb.AssertionLevel_Fast, false)
+	s.testAssertionImpl('a', false, false, kvrpcpb.AssertionLevel_Fast, true)
+	s.testAssertionImpl('b', true, false, kvrpcpb.AssertionLevel_Fast, true)
+	s.testAssertionImpl('c', true, true, kvrpcpb.AssertionLevel_Fast, true)
+}