address comment

store/tikv/2pc.go

@@ -725,7 +725,10 @@ func (action actionPessimisticLock) handleSingleBatch(c *twoPhaseCommitter, bo *
 		// When a txn runs into LockKeys() and backoff here, it has no chance to call
 		// executor.Next() and check the killed flag.
 		if action.killed != nil {
-			if atomic.CompareAndSwapUint32(action.killed, 1, 0) {
+			// Do not reset the killed flag here!
+			// actionPessimisticLock runs on each region parallelly, we have to consider that
+			// the error may be dropped.
+			if atomic.LoadUint32(action.killed) == 1 {
 				return ErrQueryInterrupted
 			}
 		}

store/tikv/error.go

@@ -61,6 +61,7 @@ func init() {
 		mysql.ErrTiKVServerBusy:      mysql.ErrTiKVServerBusy,
 		mysql.ErrGCTooEarly:          mysql.ErrGCTooEarly,
 		mysql.ErrTruncatedWrongValue: mysql.ErrTruncatedWrongValue,
+		mysql.ErrQueryInterrupted:    mysql.ErrQueryInterrupted,
 	}
 	terror.ErrClassToMySQLCodes[terror.ClassTiKV] = tikvMySQLErrCodes
 }

store/tikv/txn.go

@@ -17,6 +17,7 @@ import (
 	"context"
 	"fmt"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/dgryski/go-farm"
@@ -425,6 +426,12 @@ func (txn *tikvTxn) LockKeys(ctx context.Context, killed *uint32, forUpdateTS ui
 		// concurrently execute on multiple regions may lead to deadlock.
 		txn.committer.isFirstLock = len(txn.lockKeys) == 0 && len(keys) == 1
 		err := txn.committer.pessimisticLockKeys(bo, killed, keys)
+		if killed != nil {
+			// If the kill signal is received during waiting for pessimisticLock,
+			// pessimisticLockKeys would handle the error but it doesn't reset the flag.
+			// We need to reset the killed flag here.
+			atomic.CompareAndSwapUint32(killed, 1, 0)
+		}
 		if err != nil {
 			for _, key := range keys {
 				txn.us.DeleteKeyExistErrInfo(key)