Fix the false positive assertion error caused by loss of pessimistic locks (#596)

Signed-off-by: ekexium <eke@fastmail.com>

integration_tests/assertion_test.go

@@ -17,6 +17,7 @@ package tikv_test
 import (
 	"context"
 	"fmt"
+	"strings"
 	"testing"
 	"time"
 
@@ -276,7 +277,7 @@ func (s *testAssertionSuite) TestAssertionErrorLessPriorToOtherError() {
 
 		err = tx.Commit(ctx)
 		s.NotNil(err)
-		s.IsType(&tikverr.ErrAssertionFailed{}, errors.Cause(err))
+		s.NotContains(strings.ToLower(err.Error()), "assertion")
 	}
 
 	testOnce([]byte("kr1"), []byte("ki1"), false, false, false)

txnkv/transaction/2pc.go

@@ -179,6 +179,9 @@ type twoPhaseCommitter struct {
 
 	// The total number of kv request after batch split.
 	prewriteTotalReqNum int
+
+	// assertion error happened when initializing mutations, could be false positive if pessimistic lock is lost
+	stashedAssertionError error
 }
 
 type memBufferMutations struct {
@@ -622,8 +625,17 @@ func (c *twoPhaseCommitter) initKeysAndMutations(ctx context.Context) error {
 				// Do not exit immediately here. To rollback the pessimistic locks (if any), we need to finish
 				// collecting all the keys.
 				// Keep only the first assertion error.
+
+				// assertion errors is treated differently from other errors. If there is an assertion error,
+				// it's probably cause by loss of pessimistic locks, so we can't directly return the assertion error.
+				// Instead, we stash the error, forbid async commit and 1PC, then let the prewrite continue.
+				// If the prewrite requests all succeed, the assertion error is returned, otherwise return the error
+				// from the prewrite phase.
 				if err1 != nil && assertionError == nil {
 					assertionError = errors.WithStack(err1)
+					c.stashedAssertionError = assertionError
+					c.txn.enableAsyncCommit = false
+					c.txn.enable1PC = false
 				}
 			}
 
@@ -688,10 +700,6 @@ func (c *twoPhaseCommitter) initKeysAndMutations(ctx context.Context) error {
 	c.resourceGroupTagger = txn.resourceGroupTagger
 	c.setDetail(commitDetail)
 
-	if assertionError != nil {
-		return assertionError
-	}
-
 	return nil
 }
 
@@ -1507,6 +1515,16 @@ func (c *twoPhaseCommitter) execute(ctx context.Context) (err error) {
 		return err
 	}
 
+	// return assertion error found in TiDB after prewrite succeeds to prevent false positive. Note this is only visible
+	// when async commit or 1PC is disabled.
+	if c.stashedAssertionError != nil {
+		if c.isAsyncCommit() || c.isOnePC() {
+			// should be unreachable
+			panic("tidb-side assertion error should forbids async commit or 1PC")
+		}
+		return c.stashedAssertionError
+	}
+
 	// strip check_not_exists keys that no need to commit.
 	c.stripNoNeedCommitKeys()