store/tikv: support resolve specified lock keys (#10292) (#10793)

store/mockstore/mocktikv/mvcc.go

@@ -48,6 +48,7 @@ type mvccLock struct {
 	op          kvrpcpb.Op
 	ttl         uint64
 	forUpdateTS uint64
+	txnSize     uint64
 }
 
 type mvccEntry struct {

store/mockstore/mocktikv/mvcc_leveldb.go

@@ -630,6 +630,7 @@ func (mvcc *MVCCLevelDB) Prewrite(mutations []*kvrpcpb.Mutation, primary []byte,
 	anyError := false
 	batch := &leveldb.Batch{}
 	errs := make([]error, 0, len(mutations))
+	txnSize := len(mutations)
 	for _, m := range mutations {
 		// If the operation is Insert, check if key is exists at first.
 		var err error
@@ -649,7 +650,7 @@ func (mvcc *MVCCLevelDB) Prewrite(mutations []*kvrpcpb.Mutation, primary []byte,
 				continue
 			}
 		}
-		err = prewriteMutation(mvcc.db, batch, m, startTS, primary, ttl)
+		err = prewriteMutation(mvcc.db, batch, m, startTS, primary, ttl, uint64(txnSize))
 		errs = append(errs, err)
 		if err != nil {
 			anyError = true
@@ -684,7 +685,7 @@ func checkConflictValue(iter *Iterator, key []byte, startTS uint64) error {
 	return nil
 }
 
-func prewriteMutation(db *leveldb.DB, batch *leveldb.Batch, mutation *kvrpcpb.Mutation, startTS uint64, primary []byte, ttl uint64) error {
+func prewriteMutation(db *leveldb.DB, batch *leveldb.Batch, mutation *kvrpcpb.Mutation, startTS uint64, primary []byte, ttl uint64, txnSize uint64) error {
 	startKey := mvccEncode(mutation.Key, lockVer)
 	iter := newIterator(db, &util.Range{
 		Start: startKey,
@@ -723,6 +724,7 @@ func prewriteMutation(db *leveldb.DB, batch *leveldb.Batch, mutation *kvrpcpb.Mu
 		value:   mutation.Value,
 		op:      op,
 		ttl:     ttl,
+		txnSize: txnSize,
 	}
 	writeKey := mvccEncode(mutation.Key, lockVer)
 	writeValue, err := lock.MarshalBinary()

store/tikv/2pc.go

@@ -482,6 +482,7 @@ func (c *twoPhaseCommitter) buildPrewriteRequest(batch batchKeys) *tikvrpc.Reque
 			LockTtl:           c.lockTTL,
 			IsPessimisticLock: isPessimisticLock,
 			ForUpdateTs:       c.forUpdateTS,
+			TxnSize:           uint64(len(batch.keys)),
 		},
 		Context: pb.Context{
 			Priority: c.priority,

store/tikv/lock_resolver.go

@@ -33,6 +33,9 @@ import (
 // ResolvedCacheSize is max number of cached txn status.
 const ResolvedCacheSize = 2048
 
+// bigTxnThreshold : transaction involves keys exceed this threshold can be treated as `big transaction`.
+const bigTxnThreshold = 16
+
 var (
 	tikvLockResolverCountWithBatchResolve             = metrics.TiKVLockResolverCounter.WithLabelValues("batch_resolve")
 	tikvLockResolverCountWithExpired                  = metrics.TiKVLockResolverCounter.WithLabelValues("expired")
@@ -43,6 +46,7 @@ var (
 	tikvLockResolverCountWithQueryTxnStatusCommitted  = metrics.TiKVLockResolverCounter.WithLabelValues("query_txn_status_committed")
 	tikvLockResolverCountWithQueryTxnStatusRolledBack = metrics.TiKVLockResolverCounter.WithLabelValues("query_txn_status_rolled_back")
 	tikvLockResolverCountWithResolveLocks             = metrics.TiKVLockResolverCounter.WithLabelValues("query_resolve_locks")
+	tikvLockResolverCountWithResolveLockLite          = metrics.TiKVLockResolverCounter.WithLabelValues("query_resolve_lock_lite")
 )
 
 // LockResolver resolves locks and also caches resolved txn status.
@@ -125,6 +129,7 @@ type Lock struct {
 	Primary []byte
 	TxnID   uint64
 	TTL     uint64
+	TxnSize uint64
 }
 
 func (l *Lock) String() string {
@@ -137,11 +142,13 @@ func NewLock(l *kvrpcpb.LockInfo) *Lock {
 	if ttl == 0 {
 		ttl = defaultLockTTL
 	}
+	txnSize := l.GetTxnSize()
 	return &Lock{
 		Key:     l.GetKey(),
 		Primary: l.GetPrimaryLock(),
 		TxnID:   l.GetLockVersion(),
 		TTL:     ttl,
+		TxnSize: txnSize,
 	}
 }
 
@@ -388,6 +395,7 @@ func (lr *LockResolver) getTxnStatus(bo *Backoffer, txnID uint64, primary []byte
 
 func (lr *LockResolver) resolveLock(bo *Backoffer, l *Lock, status TxnStatus, cleanRegions map[RegionVerID]struct{}) error {
 	tikvLockResolverCountWithResolveLocks.Inc()
+	cleanWholeRegion := l.TxnSize >= bigTxnThreshold
 	for {
 		loc, err := lr.store.GetRegionCache().LocateKey(bo, l.Key)
 		if err != nil {
@@ -405,6 +413,12 @@ func (lr *LockResolver) resolveLock(bo *Backoffer, l *Lock, status TxnStatus, cl
 		if status.IsCommitted() {
 			req.ResolveLock.CommitVersion = status.CommitTS()
 		}
+		if l.TxnSize < bigTxnThreshold {
+			// Only resolve specified keys when it is a small transaction,
+			// prevent from scanning the whole region in this case.
+			tikvLockResolverCountWithResolveLockLite.Inc()
+			req.ResolveLock.Keys = [][]byte{l.Key}
+		}
 		resp, err := lr.store.SendReq(bo, req, loc.Region, readTimeoutShort)
 		if err != nil {
 			return errors.Trace(err)
@@ -429,7 +443,9 @@ func (lr *LockResolver) resolveLock(bo *Backoffer, l *Lock, status TxnStatus, cl
 			logutil.Logger(context.Background()).Error("resolveLock error", zap.Error(err))
 			return err
 		}
-		cleanRegions[loc.Region] = struct{}{}
+		if cleanWholeRegion {
+			cleanRegions[loc.Region] = struct{}{}
+		}
 		return nil
 	}
 }