[close #757] Support verifying primary for check_txn_status

dev/proto.sh

@@ -21,9 +21,9 @@ if [ -d $proto_dir ]; then
 fi
 
 repos=("https://github.com/pingcap/kvproto" "https://github.com/pingcap/raft-rs" "https://github.com/pingcap/tipb")
-commits=(3056ca36e6f2a71a9fc7ba7135e6b119fd977553 b9891b673573fad77ebcf9bbe0969cf945841926 c4d518eb1d60c21f05b028b36729e64610346dac)
+commits=(971c6f2217153da745d51886f94d7b701dec9799 b9891b673573fad77ebcf9bbe0969cf945841926 c4d518eb1d60c21f05b028b36729e64610346dac)
 
-for i in "${!repos[@]}"; do 
+for i in "${!repos[@]}"; do
 	repo_name=$(basename ${repos[$i]})
 	git_command="git -C $repo_name"
 

src/main/java/org/tikv/txn/LockResolverClientV4.java

@@ -52,6 +52,7 @@
 import org.tikv.kvproto.TikvGrpc;
 import org.tikv.kvproto.TikvGrpc.TikvBlockingStub;
 import org.tikv.kvproto.TikvGrpc.TikvFutureStub;
+import org.tikv.txn.exception.PrimaryMismatchException;
 import org.tikv.txn.exception.TxnNotFoundException;
 import org.tikv.txn.exception.WriteConflictException;
 
@@ -114,7 +115,22 @@
     Set<Long> pushed = new HashSet<>(locks.size());
 
     for (Lock l : locks) {
-      TxnStatus status = getTxnStatusFromLock(bo, l, callerStartTS);
+
+      TxnStatus status = null;
+
+      try {
+        status = getTxnStatusFromLock(bo, l, callerStartTS);
+      } catch (PrimaryMismatchException e) {
+        if (l.getLockType() != Kvrpcpb.Op.PessimisticLock) {
+          logger.info(
+              String.format(
+                  "unexpected primaryMismatch error occurred on a non-pessimistic lock, lock: %s",
+                  l));
+          throw e;
+        }
+        // Pessimistic rollback the pessimistic lock as it points to an invalid primary.
+        status = new TxnStatus();
+      }
 
       if (status.getTtl() == 0) {
         Set<RegionVerID> cleanRegion =
@@ -231,7 +247,13 @@
     while (true) {
       try {
         return getTxnStatus(
-            bo, lock.getTxnID(), lock.getPrimary(), callerStartTS, currentTS, rollbackIfNotExist);
+            bo,
+            lock.getTxnID(),
+            lock.getPrimary(),
+            callerStartTS,
+            currentTS,
+            rollbackIfNotExist,
+            lock);
       } catch (TxnNotFoundException e) {
         // If the error is something other than txnNotFoundErr, throw the error (network
         // unavailable, tikv down, backoff timeout etc) to the caller.
@@ -271,12 +293,16 @@
       ByteString primary,
       Long callerStartTS,
       Long currentTS,
-      boolean rollbackIfNotExist) {
+      boolean rollbackIfNotExist,
+      Lock lock) {
     TxnStatus status = getResolved(txnID);
     if (status != null) {
       return status;
     }
 
+    boolean resolvingPessimisticLock =
+        (lock != null && lock.getLockType() == Kvrpcpb.Op.PessimisticLock);
+
     // CheckTxnStatus may meet the following cases:
     // 1. LOCK
     // 1.1 Lock expired -- orphan lock, fail to update TTL, crash recovery etc.
@@ -295,6 +321,7 @@
               .setCallerStartTs(callerStartTS)
               .setCurrentTs(currentTS)
               .setRollbackIfNotExist(rollbackIfNotExist)
+              .setVerifyIsPrimary(true)
               .build();
         };
 
@@ -338,6 +365,12 @@
           throw new TxnNotFoundException();
         }
 
+        if (resolvingPessimisticLock && keyError.hasPrimaryMismatch()) {
+          logger.info(
+              String.format("getTxnStatus was called on secondary lock. err: %s", keyError));
+          throw new PrimaryMismatchException();
+        }
+
         logger.error(String.format("unexpected cleanup err: %s, tid: %d", keyError, txnID));
         throw new KeyException(keyError);
       }

src/main/java/org/tikv/txn/exception/PrimaryMismatchException.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2023 TiKV Project 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 org.tikv.txn.exception;
+
+public class PrimaryMismatchException extends RuntimeException {}