github.com/KinWaiYuen/client-go/v2@v2.5.4/internal/mockstore/mocktikv/rpc.go (about)

     1  // Copyright 2021 TiKV Authors
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  // NOTE: The code in this file is based on code from the
    16  // TiDB project, licensed under the Apache License v 2.0
    17  //
    18  // https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/mockstore/mocktikv/rpc.go
    19  //
    20  
    21  // Copyright 2016 PingCAP, Inc.
    22  //
    23  // Licensed under the Apache License, Version 2.0 (the "License");
    24  // you may not use this file except in compliance with the License.
    25  // You may obtain a copy of the License at
    26  //
    27  //     http://www.apache.org/licenses/LICENSE-2.0
    28  //
    29  // Unless required by applicable law or agreed to in writing, software
    30  // distributed under the License is distributed on an "AS IS" BASIS,
    31  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    32  // See the License for the specific language governing permissions and
    33  // limitations under the License.
    34  
    35  package mocktikv
    36  
    37  import (
    38  	"bytes"
    39  	"context"
    40  	"math"
    41  	"strconv"
    42  	"time"
    43  
    44  	"github.com/KinWaiYuen/client-go/v2/tikvrpc"
    45  	"github.com/KinWaiYuen/client-go/v2/util"
    46  	"github.com/opentracing/opentracing-go"
    47  	"github.com/pingcap/errors"
    48  	"github.com/pingcap/kvproto/pkg/coprocessor"
    49  	"github.com/pingcap/kvproto/pkg/debugpb"
    50  	"github.com/pingcap/kvproto/pkg/errorpb"
    51  	"github.com/pingcap/kvproto/pkg/kvrpcpb"
    52  	"github.com/pingcap/kvproto/pkg/metapb"
    53  	"github.com/pingcap/parser/terror"
    54  )
    55  
    56  // For gofail injection.
    57  var undeterminedErr = terror.ErrResultUndetermined
    58  
    59  const requestMaxSize = 8 * 1024 * 1024
    60  
    61  func checkGoContext(ctx context.Context) error {
    62  	select {
    63  	case <-ctx.Done():
    64  		return ctx.Err()
    65  	default:
    66  		return nil
    67  	}
    68  }
    69  
    70  func convertToKeyError(err error) *kvrpcpb.KeyError {
    71  	if locked, ok := errors.Cause(err).(*ErrLocked); ok {
    72  		return &kvrpcpb.KeyError{
    73  			Locked: &kvrpcpb.LockInfo{
    74  				Key:             locked.Key.Raw(),
    75  				PrimaryLock:     locked.Primary,
    76  				LockVersion:     locked.StartTS,
    77  				LockTtl:         locked.TTL,
    78  				TxnSize:         locked.TxnSize,
    79  				LockType:        locked.LockType,
    80  				LockForUpdateTs: locked.ForUpdateTS,
    81  			},
    82  		}
    83  	}
    84  	if alreadyExist, ok := errors.Cause(err).(*ErrKeyAlreadyExist); ok {
    85  		return &kvrpcpb.KeyError{
    86  			AlreadyExist: &kvrpcpb.AlreadyExist{
    87  				Key: alreadyExist.Key,
    88  			},
    89  		}
    90  	}
    91  	if writeConflict, ok := errors.Cause(err).(*ErrConflict); ok {
    92  		return &kvrpcpb.KeyError{
    93  			Conflict: &kvrpcpb.WriteConflict{
    94  				Key:              writeConflict.Key,
    95  				ConflictTs:       writeConflict.ConflictTS,
    96  				ConflictCommitTs: writeConflict.ConflictCommitTS,
    97  				StartTs:          writeConflict.StartTS,
    98  			},
    99  		}
   100  	}
   101  	if dead, ok := errors.Cause(err).(*ErrDeadlock); ok {
   102  		return &kvrpcpb.KeyError{
   103  			Deadlock: &kvrpcpb.Deadlock{
   104  				LockTs:          dead.LockTS,
   105  				LockKey:         dead.LockKey,
   106  				DeadlockKeyHash: dead.DealockKeyHash,
   107  			},
   108  		}
   109  	}
   110  	if retryable, ok := errors.Cause(err).(ErrRetryable); ok {
   111  		return &kvrpcpb.KeyError{
   112  			Retryable: retryable.Error(),
   113  		}
   114  	}
   115  	if expired, ok := errors.Cause(err).(*ErrCommitTSExpired); ok {
   116  		return &kvrpcpb.KeyError{
   117  			CommitTsExpired: &expired.CommitTsExpired,
   118  		}
   119  	}
   120  	if tmp, ok := errors.Cause(err).(*ErrTxnNotFound); ok {
   121  		return &kvrpcpb.KeyError{
   122  			TxnNotFound: &tmp.TxnNotFound,
   123  		}
   124  	}
   125  	return &kvrpcpb.KeyError{
   126  		Abort: err.Error(),
   127  	}
   128  }
   129  
   130  func convertToKeyErrors(errs []error) []*kvrpcpb.KeyError {
   131  	var keyErrors = make([]*kvrpcpb.KeyError, 0)
   132  	for _, err := range errs {
   133  		if err != nil {
   134  			keyErrors = append(keyErrors, convertToKeyError(err))
   135  		}
   136  	}
   137  	return keyErrors
   138  }
   139  
   140  func convertToPbPairs(pairs []Pair) []*kvrpcpb.KvPair {
   141  	kvPairs := make([]*kvrpcpb.KvPair, 0, len(pairs))
   142  	for _, p := range pairs {
   143  		var kvPair *kvrpcpb.KvPair
   144  		if p.Err == nil {
   145  			kvPair = &kvrpcpb.KvPair{
   146  				Key:   p.Key,
   147  				Value: p.Value,
   148  			}
   149  		} else {
   150  			kvPair = &kvrpcpb.KvPair{
   151  				Error: convertToKeyError(p.Err),
   152  			}
   153  		}
   154  		kvPairs = append(kvPairs, kvPair)
   155  	}
   156  	return kvPairs
   157  }
   158  
   159  // kvHandler mocks tikv's side handler behavior. In general, you may assume
   160  // TiKV just translate the logic from Go to Rust.
   161  type kvHandler struct {
   162  	*Session
   163  }
   164  
   165  func (h kvHandler) handleKvGet(req *kvrpcpb.GetRequest) *kvrpcpb.GetResponse {
   166  	if !h.checkKeyInRegion(req.Key) {
   167  		panic("KvGet: key not in region")
   168  	}
   169  
   170  	val, err := h.mvccStore.Get(req.Key, req.GetVersion(), h.isolationLevel, req.Context.GetResolvedLocks())
   171  	if err != nil {
   172  		return &kvrpcpb.GetResponse{
   173  			Error: convertToKeyError(err),
   174  		}
   175  	}
   176  	return &kvrpcpb.GetResponse{
   177  		Value: val,
   178  	}
   179  }
   180  
   181  func (h kvHandler) handleKvScan(req *kvrpcpb.ScanRequest) *kvrpcpb.ScanResponse {
   182  	endKey := MvccKey(h.endKey).Raw()
   183  	var pairs []Pair
   184  	if !req.Reverse {
   185  		if !h.checkKeyInRegion(req.GetStartKey()) {
   186  			panic("KvScan: startKey not in region")
   187  		}
   188  		if len(req.EndKey) > 0 && (len(endKey) == 0 || bytes.Compare(NewMvccKey(req.EndKey), h.endKey) < 0) {
   189  			endKey = req.EndKey
   190  		}
   191  		pairs = h.mvccStore.Scan(req.GetStartKey(), endKey, int(req.GetLimit()), req.GetVersion(), h.isolationLevel, req.Context.ResolvedLocks)
   192  	} else {
   193  		// TiKV use range [end_key, start_key) for reverse scan.
   194  		// Should use the req.EndKey to check in region.
   195  		if !h.checkKeyInRegion(req.GetEndKey()) {
   196  			panic("KvScan: startKey not in region")
   197  		}
   198  
   199  		// TiKV use range [end_key, start_key) for reverse scan.
   200  		// So the req.StartKey actually is the end_key.
   201  		if len(req.StartKey) > 0 && (len(endKey) == 0 || bytes.Compare(NewMvccKey(req.StartKey), h.endKey) < 0) {
   202  			endKey = req.StartKey
   203  		}
   204  
   205  		pairs = h.mvccStore.ReverseScan(req.EndKey, endKey, int(req.GetLimit()), req.GetVersion(), h.isolationLevel, req.Context.ResolvedLocks)
   206  	}
   207  
   208  	return &kvrpcpb.ScanResponse{
   209  		Pairs: convertToPbPairs(pairs),
   210  	}
   211  }
   212  
   213  func (h kvHandler) handleKvPrewrite(req *kvrpcpb.PrewriteRequest) *kvrpcpb.PrewriteResponse {
   214  	regionID := req.Context.RegionId
   215  	h.cluster.handleDelay(req.StartVersion, regionID)
   216  
   217  	for _, m := range req.Mutations {
   218  		if !h.checkKeyInRegion(m.Key) {
   219  			panic("KvPrewrite: key not in region")
   220  		}
   221  	}
   222  	errs := h.mvccStore.Prewrite(req)
   223  	return &kvrpcpb.PrewriteResponse{
   224  		Errors: convertToKeyErrors(errs),
   225  	}
   226  }
   227  
   228  func (h kvHandler) handleKvPessimisticLock(req *kvrpcpb.PessimisticLockRequest) *kvrpcpb.PessimisticLockResponse {
   229  	for _, m := range req.Mutations {
   230  		if !h.checkKeyInRegion(m.Key) {
   231  			panic("KvPessimisticLock: key not in region")
   232  		}
   233  	}
   234  	startTS := req.StartVersion
   235  	regionID := req.Context.RegionId
   236  	h.cluster.handleDelay(startTS, regionID)
   237  	return h.mvccStore.PessimisticLock(req)
   238  }
   239  
   240  func simulateServerSideWaitLock(errs []error) {
   241  	for _, err := range errs {
   242  		if _, ok := err.(*ErrLocked); ok {
   243  			time.Sleep(time.Millisecond * 5)
   244  			break
   245  		}
   246  	}
   247  }
   248  
   249  func (h kvHandler) handleKvPessimisticRollback(req *kvrpcpb.PessimisticRollbackRequest) *kvrpcpb.PessimisticRollbackResponse {
   250  	for _, key := range req.Keys {
   251  		if !h.checkKeyInRegion(key) {
   252  			panic("KvPessimisticRollback: key not in region")
   253  		}
   254  	}
   255  	errs := h.mvccStore.PessimisticRollback(req.Keys, req.StartVersion, req.ForUpdateTs)
   256  	return &kvrpcpb.PessimisticRollbackResponse{
   257  		Errors: convertToKeyErrors(errs),
   258  	}
   259  }
   260  
   261  func (h kvHandler) handleKvCommit(req *kvrpcpb.CommitRequest) *kvrpcpb.CommitResponse {
   262  	for _, k := range req.Keys {
   263  		if !h.checkKeyInRegion(k) {
   264  			panic("KvCommit: key not in region")
   265  		}
   266  	}
   267  	var resp kvrpcpb.CommitResponse
   268  	err := h.mvccStore.Commit(req.Keys, req.GetStartVersion(), req.GetCommitVersion())
   269  	if err != nil {
   270  		resp.Error = convertToKeyError(err)
   271  	}
   272  	return &resp
   273  }
   274  
   275  func (h kvHandler) handleKvCleanup(req *kvrpcpb.CleanupRequest) *kvrpcpb.CleanupResponse {
   276  	if !h.checkKeyInRegion(req.Key) {
   277  		panic("KvCleanup: key not in region")
   278  	}
   279  	var resp kvrpcpb.CleanupResponse
   280  	err := h.mvccStore.Cleanup(req.Key, req.GetStartVersion(), req.GetCurrentTs())
   281  	if err != nil {
   282  		if commitTS, ok := errors.Cause(err).(ErrAlreadyCommitted); ok {
   283  			resp.CommitVersion = uint64(commitTS)
   284  		} else {
   285  			resp.Error = convertToKeyError(err)
   286  		}
   287  	}
   288  	return &resp
   289  }
   290  
   291  func (h kvHandler) handleKvCheckTxnStatus(req *kvrpcpb.CheckTxnStatusRequest) *kvrpcpb.CheckTxnStatusResponse {
   292  	if !h.checkKeyInRegion(req.PrimaryKey) {
   293  		panic("KvCheckTxnStatus: key not in region")
   294  	}
   295  	var resp kvrpcpb.CheckTxnStatusResponse
   296  	ttl, commitTS, action, err := h.mvccStore.CheckTxnStatus(req.GetPrimaryKey(), req.GetLockTs(), req.GetCallerStartTs(), req.GetCurrentTs(), req.GetRollbackIfNotExist(), req.ResolvingPessimisticLock)
   297  	if err != nil {
   298  		resp.Error = convertToKeyError(err)
   299  	} else {
   300  		resp.LockTtl, resp.CommitVersion, resp.Action = ttl, commitTS, action
   301  	}
   302  	return &resp
   303  }
   304  
   305  func (h kvHandler) handleTxnHeartBeat(req *kvrpcpb.TxnHeartBeatRequest) *kvrpcpb.TxnHeartBeatResponse {
   306  	if !h.checkKeyInRegion(req.PrimaryLock) {
   307  		panic("KvTxnHeartBeat: key not in region")
   308  	}
   309  	var resp kvrpcpb.TxnHeartBeatResponse
   310  	ttl, err := h.mvccStore.TxnHeartBeat(req.PrimaryLock, req.StartVersion, req.AdviseLockTtl)
   311  	if err != nil {
   312  		resp.Error = convertToKeyError(err)
   313  	}
   314  	resp.LockTtl = ttl
   315  	return &resp
   316  }
   317  
   318  func (h kvHandler) handleKvBatchGet(req *kvrpcpb.BatchGetRequest) *kvrpcpb.BatchGetResponse {
   319  	for _, k := range req.Keys {
   320  		if !h.checkKeyInRegion(k) {
   321  			panic("KvBatchGet: key not in region")
   322  		}
   323  	}
   324  	pairs := h.mvccStore.BatchGet(req.Keys, req.GetVersion(), h.isolationLevel, req.Context.GetResolvedLocks())
   325  	return &kvrpcpb.BatchGetResponse{
   326  		Pairs: convertToPbPairs(pairs),
   327  	}
   328  }
   329  
   330  func (h kvHandler) handleMvccGetByKey(req *kvrpcpb.MvccGetByKeyRequest) *kvrpcpb.MvccGetByKeyResponse {
   331  	debugger, ok := h.mvccStore.(MVCCDebugger)
   332  	if !ok {
   333  		return &kvrpcpb.MvccGetByKeyResponse{
   334  			Error: "not implement",
   335  		}
   336  	}
   337  
   338  	if !h.checkKeyInRegion(req.Key) {
   339  		panic("MvccGetByKey: key not in region")
   340  	}
   341  	var resp kvrpcpb.MvccGetByKeyResponse
   342  	resp.Info = debugger.MvccGetByKey(req.Key)
   343  	return &resp
   344  }
   345  
   346  func (h kvHandler) handleMvccGetByStartTS(req *kvrpcpb.MvccGetByStartTsRequest) *kvrpcpb.MvccGetByStartTsResponse {
   347  	debugger, ok := h.mvccStore.(MVCCDebugger)
   348  	if !ok {
   349  		return &kvrpcpb.MvccGetByStartTsResponse{
   350  			Error: "not implement",
   351  		}
   352  	}
   353  	var resp kvrpcpb.MvccGetByStartTsResponse
   354  	resp.Info, resp.Key = debugger.MvccGetByStartTS(req.StartTs)
   355  	return &resp
   356  }
   357  
   358  func (h kvHandler) handleKvBatchRollback(req *kvrpcpb.BatchRollbackRequest) *kvrpcpb.BatchRollbackResponse {
   359  	err := h.mvccStore.Rollback(req.Keys, req.StartVersion)
   360  	if err != nil {
   361  		return &kvrpcpb.BatchRollbackResponse{
   362  			Error: convertToKeyError(err),
   363  		}
   364  	}
   365  	return &kvrpcpb.BatchRollbackResponse{}
   366  }
   367  
   368  func (h kvHandler) handleKvScanLock(req *kvrpcpb.ScanLockRequest) *kvrpcpb.ScanLockResponse {
   369  	startKey := MvccKey(h.startKey).Raw()
   370  	endKey := MvccKey(h.endKey).Raw()
   371  	locks, err := h.mvccStore.ScanLock(startKey, endKey, req.GetMaxVersion())
   372  	if err != nil {
   373  		return &kvrpcpb.ScanLockResponse{
   374  			Error: convertToKeyError(err),
   375  		}
   376  	}
   377  	return &kvrpcpb.ScanLockResponse{
   378  		Locks: locks,
   379  	}
   380  }
   381  
   382  func (h kvHandler) handleKvResolveLock(req *kvrpcpb.ResolveLockRequest) *kvrpcpb.ResolveLockResponse {
   383  	startKey := MvccKey(h.startKey).Raw()
   384  	endKey := MvccKey(h.endKey).Raw()
   385  	err := h.mvccStore.ResolveLock(startKey, endKey, req.GetStartVersion(), req.GetCommitVersion())
   386  	if err != nil {
   387  		return &kvrpcpb.ResolveLockResponse{
   388  			Error: convertToKeyError(err),
   389  		}
   390  	}
   391  	return &kvrpcpb.ResolveLockResponse{}
   392  }
   393  
   394  func (h kvHandler) handleKvGC(req *kvrpcpb.GCRequest) *kvrpcpb.GCResponse {
   395  	startKey := MvccKey(h.startKey).Raw()
   396  	endKey := MvccKey(h.endKey).Raw()
   397  	err := h.mvccStore.GC(startKey, endKey, req.GetSafePoint())
   398  	if err != nil {
   399  		return &kvrpcpb.GCResponse{
   400  			Error: convertToKeyError(err),
   401  		}
   402  	}
   403  	return &kvrpcpb.GCResponse{}
   404  }
   405  
   406  func (h kvHandler) handleKvDeleteRange(req *kvrpcpb.DeleteRangeRequest) *kvrpcpb.DeleteRangeResponse {
   407  	if !h.checkKeyInRegion(req.StartKey) {
   408  		panic("KvDeleteRange: key not in region")
   409  	}
   410  	var resp kvrpcpb.DeleteRangeResponse
   411  	err := h.mvccStore.DeleteRange(req.StartKey, req.EndKey)
   412  	if err != nil {
   413  		resp.Error = err.Error()
   414  	}
   415  	return &resp
   416  }
   417  
   418  func (h kvHandler) handleKvRawGet(req *kvrpcpb.RawGetRequest) *kvrpcpb.RawGetResponse {
   419  	rawKV, ok := h.mvccStore.(RawKV)
   420  	if !ok {
   421  		return &kvrpcpb.RawGetResponse{
   422  			Error: "not implemented",
   423  		}
   424  	}
   425  	return &kvrpcpb.RawGetResponse{
   426  		Value: rawKV.RawGet(req.GetKey()),
   427  	}
   428  }
   429  
   430  func (h kvHandler) handleKvRawBatchGet(req *kvrpcpb.RawBatchGetRequest) *kvrpcpb.RawBatchGetResponse {
   431  	rawKV, ok := h.mvccStore.(RawKV)
   432  	if !ok {
   433  		// TODO should we add error ?
   434  		return &kvrpcpb.RawBatchGetResponse{
   435  			RegionError: &errorpb.Error{
   436  				Message: "not implemented",
   437  			},
   438  		}
   439  	}
   440  	values := rawKV.RawBatchGet(req.Keys)
   441  	kvPairs := make([]*kvrpcpb.KvPair, len(values))
   442  	for i, key := range req.Keys {
   443  		kvPairs[i] = &kvrpcpb.KvPair{
   444  			Key:   key,
   445  			Value: values[i],
   446  		}
   447  	}
   448  	return &kvrpcpb.RawBatchGetResponse{
   449  		Pairs: kvPairs,
   450  	}
   451  }
   452  
   453  func (h kvHandler) handleKvRawPut(req *kvrpcpb.RawPutRequest) *kvrpcpb.RawPutResponse {
   454  	rawKV, ok := h.mvccStore.(RawKV)
   455  	if !ok {
   456  		return &kvrpcpb.RawPutResponse{
   457  			Error: "not implemented",
   458  		}
   459  	}
   460  	rawKV.RawPut(req.GetKey(), req.GetValue())
   461  	return &kvrpcpb.RawPutResponse{}
   462  }
   463  
   464  func (h kvHandler) handleKvRawBatchPut(req *kvrpcpb.RawBatchPutRequest) *kvrpcpb.RawBatchPutResponse {
   465  	rawKV, ok := h.mvccStore.(RawKV)
   466  	if !ok {
   467  		return &kvrpcpb.RawBatchPutResponse{
   468  			Error: "not implemented",
   469  		}
   470  	}
   471  	keys := make([][]byte, 0, len(req.Pairs))
   472  	values := make([][]byte, 0, len(req.Pairs))
   473  	for _, pair := range req.Pairs {
   474  		keys = append(keys, pair.Key)
   475  		values = append(values, pair.Value)
   476  	}
   477  	rawKV.RawBatchPut(keys, values)
   478  	return &kvrpcpb.RawBatchPutResponse{}
   479  }
   480  
   481  func (h kvHandler) handleKvRawDelete(req *kvrpcpb.RawDeleteRequest) *kvrpcpb.RawDeleteResponse {
   482  	rawKV, ok := h.mvccStore.(RawKV)
   483  	if !ok {
   484  		return &kvrpcpb.RawDeleteResponse{
   485  			Error: "not implemented",
   486  		}
   487  	}
   488  	rawKV.RawDelete(req.GetKey())
   489  	return &kvrpcpb.RawDeleteResponse{}
   490  }
   491  
   492  func (h kvHandler) handleKvRawBatchDelete(req *kvrpcpb.RawBatchDeleteRequest) *kvrpcpb.RawBatchDeleteResponse {
   493  	rawKV, ok := h.mvccStore.(RawKV)
   494  	if !ok {
   495  		return &kvrpcpb.RawBatchDeleteResponse{
   496  			Error: "not implemented",
   497  		}
   498  	}
   499  	rawKV.RawBatchDelete(req.Keys)
   500  	return &kvrpcpb.RawBatchDeleteResponse{}
   501  }
   502  
   503  func (h kvHandler) handleKvRawDeleteRange(req *kvrpcpb.RawDeleteRangeRequest) *kvrpcpb.RawDeleteRangeResponse {
   504  	rawKV, ok := h.mvccStore.(RawKV)
   505  	if !ok {
   506  		return &kvrpcpb.RawDeleteRangeResponse{
   507  			Error: "not implemented",
   508  		}
   509  	}
   510  	rawKV.RawDeleteRange(req.GetStartKey(), req.GetEndKey())
   511  	return &kvrpcpb.RawDeleteRangeResponse{}
   512  }
   513  
   514  func (h kvHandler) handleKvRawScan(req *kvrpcpb.RawScanRequest) *kvrpcpb.RawScanResponse {
   515  	rawKV, ok := h.mvccStore.(RawKV)
   516  	if !ok {
   517  		errStr := "not implemented"
   518  		return &kvrpcpb.RawScanResponse{
   519  			RegionError: &errorpb.Error{
   520  				Message: errStr,
   521  			},
   522  		}
   523  	}
   524  
   525  	var pairs []Pair
   526  	if req.Reverse {
   527  		lowerBound := h.startKey
   528  		if bytes.Compare(req.EndKey, lowerBound) > 0 {
   529  			lowerBound = req.EndKey
   530  		}
   531  		pairs = rawKV.RawReverseScan(
   532  			req.StartKey,
   533  			lowerBound,
   534  			int(req.GetLimit()),
   535  		)
   536  	} else {
   537  		upperBound := h.endKey
   538  		if len(req.EndKey) > 0 && (len(upperBound) == 0 || bytes.Compare(req.EndKey, upperBound) < 0) {
   539  			upperBound = req.EndKey
   540  		}
   541  		pairs = rawKV.RawScan(
   542  			req.StartKey,
   543  			upperBound,
   544  			int(req.GetLimit()),
   545  		)
   546  	}
   547  
   548  	return &kvrpcpb.RawScanResponse{
   549  		Kvs: convertToPbPairs(pairs),
   550  	}
   551  }
   552  
   553  func (h kvHandler) handleSplitRegion(req *kvrpcpb.SplitRegionRequest) *kvrpcpb.SplitRegionResponse {
   554  	keys := req.GetSplitKeys()
   555  	resp := &kvrpcpb.SplitRegionResponse{Regions: make([]*metapb.Region, 0, len(keys)+1)}
   556  	for i, key := range keys {
   557  		k := NewMvccKey(key)
   558  		region, _ := h.cluster.GetRegionByKey(k)
   559  		if bytes.Equal(region.GetStartKey(), key) {
   560  			continue
   561  		}
   562  		if i == 0 {
   563  			// Set the leftmost region.
   564  			resp.Regions = append(resp.Regions, region)
   565  		}
   566  		newRegionID, newPeerIDs := h.cluster.AllocID(), h.cluster.AllocIDs(len(region.Peers))
   567  		newRegion := h.cluster.SplitRaw(region.GetId(), newRegionID, k, newPeerIDs, newPeerIDs[0])
   568  		resp.Regions = append(resp.Regions, newRegion)
   569  	}
   570  	return resp
   571  }
   572  
   573  // Client is a client that sends RPC.
   574  // This is same with tikv.Client, define again for avoid circle import.
   575  type Client interface {
   576  	// Close should release all data.
   577  	Close() error
   578  	// SendRequest sends Request.
   579  	SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error)
   580  }
   581  
   582  // CoprRPCHandler is the interface to handle coprocessor RPC commands.
   583  type CoprRPCHandler interface {
   584  	HandleCmdCop(reqCtx *kvrpcpb.Context, session *Session, r *coprocessor.Request) *coprocessor.Response
   585  	HandleBatchCop(ctx context.Context, reqCtx *kvrpcpb.Context, session *Session, r *coprocessor.BatchRequest, timeout time.Duration) (*tikvrpc.BatchCopStreamResponse, error)
   586  	HandleCopStream(ctx context.Context, reqCtx *kvrpcpb.Context, session *Session, r *coprocessor.Request, timeout time.Duration) (*tikvrpc.CopStreamResponse, error)
   587  	Close()
   588  }
   589  
   590  // RPCClient sends kv RPC calls to mock cluster. RPCClient mocks the behavior of
   591  // a rpc client at tikv's side.
   592  type RPCClient struct {
   593  	Cluster     *Cluster
   594  	MvccStore   MVCCStore
   595  	coprHandler CoprRPCHandler
   596  }
   597  
   598  // NewRPCClient creates an RPCClient.
   599  // Note that close the RPCClient may close the underlying MvccStore.
   600  func NewRPCClient(cluster *Cluster, mvccStore MVCCStore, coprHandler CoprRPCHandler) *RPCClient {
   601  	return &RPCClient{
   602  		Cluster:     cluster,
   603  		MvccStore:   mvccStore,
   604  		coprHandler: coprHandler,
   605  	}
   606  }
   607  
   608  func (c *RPCClient) getAndCheckStoreByAddr(addr string) (*metapb.Store, error) {
   609  	stores, err := c.Cluster.GetAndCheckStoreByAddr(addr)
   610  	if err != nil {
   611  		return nil, err
   612  	}
   613  	if len(stores) == 0 {
   614  		return nil, errors.New("connect fail")
   615  	}
   616  	for _, store := range stores {
   617  		if store.GetState() != metapb.StoreState_Offline &&
   618  			store.GetState() != metapb.StoreState_Tombstone {
   619  			return store, nil
   620  		}
   621  	}
   622  	return nil, errors.New("connection refused")
   623  }
   624  
   625  func (c *RPCClient) checkArgs(ctx context.Context, addr string) (*Session, error) {
   626  	if err := checkGoContext(ctx); err != nil {
   627  		return nil, err
   628  	}
   629  
   630  	store, err := c.getAndCheckStoreByAddr(addr)
   631  	if err != nil {
   632  		return nil, err
   633  	}
   634  	session := &Session{
   635  		cluster:   c.Cluster,
   636  		mvccStore: c.MvccStore,
   637  		// set store id for current request
   638  		storeID: store.GetId(),
   639  	}
   640  	return session, nil
   641  }
   642  
   643  // SendRequest sends a request to mock cluster.
   644  func (c *RPCClient) SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) {
   645  	if span := opentracing.SpanFromContext(ctx); span != nil && span.Tracer() != nil {
   646  		span1 := span.Tracer().StartSpan("RPCClient.SendRequest", opentracing.ChildOf(span.Context()))
   647  		defer span1.Finish()
   648  		ctx = opentracing.ContextWithSpan(ctx, span1)
   649  	}
   650  
   651  	// increase coverage for mock tikv
   652  	_ = req.Type.String()
   653  	_ = req.ToBatchCommandsRequest()
   654  
   655  	reqCtx := &req.Context
   656  	resp := &tikvrpc.Response{}
   657  
   658  	session, err := c.checkArgs(ctx, addr)
   659  	if err != nil {
   660  		return nil, err
   661  	}
   662  	switch req.Type {
   663  	case tikvrpc.CmdGet:
   664  		r := req.Get()
   665  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   666  			resp.Resp = &kvrpcpb.GetResponse{RegionError: err}
   667  			return resp, nil
   668  		}
   669  		resp.Resp = kvHandler{session}.handleKvGet(r)
   670  	case tikvrpc.CmdScan:
   671  		r := req.Scan()
   672  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   673  			resp.Resp = &kvrpcpb.ScanResponse{RegionError: err}
   674  			return resp, nil
   675  		}
   676  		resp.Resp = kvHandler{session}.handleKvScan(r)
   677  
   678  	case tikvrpc.CmdPrewrite:
   679  		if val, err := util.EvalFailpoint("rpcAllowedOnAlmostFull"); err == nil {
   680  			switch val.(string) {
   681  			case "true":
   682  				if req.Context.DiskFullOpt != kvrpcpb.DiskFullOpt_AllowedOnAlmostFull {
   683  					return &tikvrpc.Response{
   684  						Resp: &kvrpcpb.PrewriteResponse{
   685  							RegionError: &errorpb.Error{
   686  								DiskFull: &errorpb.DiskFull{StoreId: []uint64{1, 10}, Reason: "disk almost full"},
   687  							},
   688  						},
   689  					}, nil
   690  				}
   691  			}
   692  		}
   693  
   694  		if val, err := util.EvalFailpoint("rpcPrewriteResult"); err == nil {
   695  			switch val.(string) {
   696  			case "notLeader":
   697  				return &tikvrpc.Response{
   698  					Resp: &kvrpcpb.PrewriteResponse{RegionError: &errorpb.Error{NotLeader: &errorpb.NotLeader{}}},
   699  				}, nil
   700  			}
   701  		}
   702  
   703  		r := req.Prewrite()
   704  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   705  			resp.Resp = &kvrpcpb.PrewriteResponse{RegionError: err}
   706  			return resp, nil
   707  		}
   708  		resp.Resp = kvHandler{session}.handleKvPrewrite(r)
   709  	case tikvrpc.CmdPessimisticLock:
   710  		r := req.PessimisticLock()
   711  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   712  			resp.Resp = &kvrpcpb.PessimisticLockResponse{RegionError: err}
   713  			return resp, nil
   714  		}
   715  		resp.Resp = kvHandler{session}.handleKvPessimisticLock(r)
   716  	case tikvrpc.CmdPessimisticRollback:
   717  		r := req.PessimisticRollback()
   718  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   719  			resp.Resp = &kvrpcpb.PessimisticRollbackResponse{RegionError: err}
   720  			return resp, nil
   721  		}
   722  		resp.Resp = kvHandler{session}.handleKvPessimisticRollback(r)
   723  	case tikvrpc.CmdCommit:
   724  		if val, err := util.EvalFailpoint("rpcAllowedOnAlmostFull"); err == nil {
   725  			switch val.(string) {
   726  			case "true":
   727  				if req.Context.DiskFullOpt != kvrpcpb.DiskFullOpt_AllowedOnAlmostFull {
   728  					return &tikvrpc.Response{
   729  						Resp: &kvrpcpb.CommitResponse{
   730  							RegionError: &errorpb.Error{
   731  								DiskFull: &errorpb.DiskFull{StoreId: []uint64{1, 10}, Reason: "disk almost full"},
   732  							},
   733  						},
   734  					}, nil
   735  				}
   736  			}
   737  		}
   738  
   739  		if val, err := util.EvalFailpoint("rpcCommitResult"); err == nil {
   740  			switch val.(string) {
   741  			case "timeout":
   742  				return nil, errors.New("timeout")
   743  			case "notLeader":
   744  				return &tikvrpc.Response{
   745  					Resp: &kvrpcpb.CommitResponse{RegionError: &errorpb.Error{NotLeader: &errorpb.NotLeader{}}},
   746  				}, nil
   747  			case "keyError":
   748  				return &tikvrpc.Response{
   749  					Resp: &kvrpcpb.CommitResponse{Error: &kvrpcpb.KeyError{}},
   750  				}, nil
   751  			}
   752  		}
   753  
   754  		r := req.Commit()
   755  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   756  			resp.Resp = &kvrpcpb.CommitResponse{RegionError: err}
   757  			return resp, nil
   758  		}
   759  		resp.Resp = kvHandler{session}.handleKvCommit(r)
   760  		if val, err := util.EvalFailpoint("rpcCommitTimeout"); err == nil {
   761  			if val.(bool) {
   762  				return nil, undeterminedErr
   763  			}
   764  		}
   765  	case tikvrpc.CmdCleanup:
   766  		r := req.Cleanup()
   767  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   768  			resp.Resp = &kvrpcpb.CleanupResponse{RegionError: err}
   769  			return resp, nil
   770  		}
   771  		resp.Resp = kvHandler{session}.handleKvCleanup(r)
   772  	case tikvrpc.CmdCheckTxnStatus:
   773  		r := req.CheckTxnStatus()
   774  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   775  			resp.Resp = &kvrpcpb.CheckTxnStatusResponse{RegionError: err}
   776  			return resp, nil
   777  		}
   778  		resp.Resp = kvHandler{session}.handleKvCheckTxnStatus(r)
   779  	case tikvrpc.CmdTxnHeartBeat:
   780  		r := req.TxnHeartBeat()
   781  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   782  			resp.Resp = &kvrpcpb.TxnHeartBeatResponse{RegionError: err}
   783  			return resp, nil
   784  		}
   785  		resp.Resp = kvHandler{session}.handleTxnHeartBeat(r)
   786  	case tikvrpc.CmdBatchGet:
   787  		r := req.BatchGet()
   788  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   789  			resp.Resp = &kvrpcpb.BatchGetResponse{RegionError: err}
   790  			return resp, nil
   791  		}
   792  		resp.Resp = kvHandler{session}.handleKvBatchGet(r)
   793  	case tikvrpc.CmdBatchRollback:
   794  		r := req.BatchRollback()
   795  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   796  			resp.Resp = &kvrpcpb.BatchRollbackResponse{RegionError: err}
   797  			return resp, nil
   798  		}
   799  		resp.Resp = kvHandler{session}.handleKvBatchRollback(r)
   800  	case tikvrpc.CmdScanLock:
   801  		r := req.ScanLock()
   802  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   803  			resp.Resp = &kvrpcpb.ScanLockResponse{RegionError: err}
   804  			return resp, nil
   805  		}
   806  		resp.Resp = kvHandler{session}.handleKvScanLock(r)
   807  	case tikvrpc.CmdResolveLock:
   808  		r := req.ResolveLock()
   809  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   810  			resp.Resp = &kvrpcpb.ResolveLockResponse{RegionError: err}
   811  			return resp, nil
   812  		}
   813  		resp.Resp = kvHandler{session}.handleKvResolveLock(r)
   814  	case tikvrpc.CmdGC:
   815  		r := req.GC()
   816  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   817  			resp.Resp = &kvrpcpb.GCResponse{RegionError: err}
   818  			return resp, nil
   819  		}
   820  		resp.Resp = kvHandler{session}.handleKvGC(r)
   821  	case tikvrpc.CmdDeleteRange:
   822  		r := req.DeleteRange()
   823  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   824  			resp.Resp = &kvrpcpb.DeleteRangeResponse{RegionError: err}
   825  			return resp, nil
   826  		}
   827  		resp.Resp = kvHandler{session}.handleKvDeleteRange(r)
   828  	case tikvrpc.CmdRawGet:
   829  		r := req.RawGet()
   830  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   831  			resp.Resp = &kvrpcpb.RawGetResponse{RegionError: err}
   832  			return resp, nil
   833  		}
   834  		resp.Resp = kvHandler{session}.handleKvRawGet(r)
   835  	case tikvrpc.CmdRawBatchGet:
   836  		r := req.RawBatchGet()
   837  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   838  			resp.Resp = &kvrpcpb.RawBatchGetResponse{RegionError: err}
   839  			return resp, nil
   840  		}
   841  		resp.Resp = kvHandler{session}.handleKvRawBatchGet(r)
   842  	case tikvrpc.CmdRawPut:
   843  		r := req.RawPut()
   844  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   845  			resp.Resp = &kvrpcpb.RawPutResponse{RegionError: err}
   846  			return resp, nil
   847  		}
   848  		resp.Resp = kvHandler{session}.handleKvRawPut(r)
   849  	case tikvrpc.CmdRawBatchPut:
   850  		r := req.RawBatchPut()
   851  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   852  			resp.Resp = &kvrpcpb.RawBatchPutResponse{RegionError: err}
   853  			return resp, nil
   854  		}
   855  		resp.Resp = kvHandler{session}.handleKvRawBatchPut(r)
   856  	case tikvrpc.CmdRawDelete:
   857  		r := req.RawDelete()
   858  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   859  			resp.Resp = &kvrpcpb.RawDeleteResponse{RegionError: err}
   860  			return resp, nil
   861  		}
   862  		resp.Resp = kvHandler{session}.handleKvRawDelete(r)
   863  	case tikvrpc.CmdRawBatchDelete:
   864  		r := req.RawBatchDelete()
   865  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   866  			resp.Resp = &kvrpcpb.RawBatchDeleteResponse{RegionError: err}
   867  		}
   868  		resp.Resp = kvHandler{session}.handleKvRawBatchDelete(r)
   869  	case tikvrpc.CmdRawDeleteRange:
   870  		r := req.RawDeleteRange()
   871  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   872  			resp.Resp = &kvrpcpb.RawDeleteRangeResponse{RegionError: err}
   873  			return resp, nil
   874  		}
   875  		resp.Resp = kvHandler{session}.handleKvRawDeleteRange(r)
   876  	case tikvrpc.CmdRawScan:
   877  		r := req.RawScan()
   878  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   879  			resp.Resp = &kvrpcpb.RawScanResponse{RegionError: err}
   880  			return resp, nil
   881  		}
   882  		resp.Resp = kvHandler{session}.handleKvRawScan(r)
   883  	case tikvrpc.CmdUnsafeDestroyRange:
   884  		panic("unimplemented")
   885  	case tikvrpc.CmdRegisterLockObserver:
   886  		return nil, errors.New("unimplemented")
   887  	case tikvrpc.CmdCheckLockObserver:
   888  		return nil, errors.New("unimplemented")
   889  	case tikvrpc.CmdRemoveLockObserver:
   890  		return nil, errors.New("unimplemented")
   891  	case tikvrpc.CmdPhysicalScanLock:
   892  		return nil, errors.New("unimplemented")
   893  	case tikvrpc.CmdCop:
   894  		if c.coprHandler == nil {
   895  			return nil, errors.New("unimplemented")
   896  		}
   897  		session.rawStartKey = MvccKey(session.startKey).Raw()
   898  		session.rawEndKey = MvccKey(session.endKey).Raw()
   899  		resp.Resp = c.coprHandler.HandleCmdCop(reqCtx, session, req.Cop())
   900  	case tikvrpc.CmdBatchCop:
   901  		if value, err := util.EvalFailpoint("BatchCopCancelled"); err == nil {
   902  			if value.(bool) {
   903  				return nil, context.Canceled
   904  			}
   905  		}
   906  
   907  		if value, err := util.EvalFailpoint("BatchCopRpcErr"); err != nil {
   908  			if value.(string) == addr {
   909  				return nil, errors.New("rpc error")
   910  			}
   911  		}
   912  		if c.coprHandler == nil {
   913  			return nil, errors.New("unimplemented")
   914  		}
   915  		batchResp, err := c.coprHandler.HandleBatchCop(ctx, reqCtx, session, req.BatchCop(), timeout)
   916  		if err != nil {
   917  			return nil, errors.Trace(err)
   918  		}
   919  		resp.Resp = batchResp
   920  	case tikvrpc.CmdCopStream:
   921  		if c.coprHandler == nil {
   922  			return nil, errors.New("unimplemented")
   923  		}
   924  		session.rawStartKey = MvccKey(session.startKey).Raw()
   925  		session.rawEndKey = MvccKey(session.endKey).Raw()
   926  		streamResp, err := c.coprHandler.HandleCopStream(ctx, reqCtx, session, req.Cop(), timeout)
   927  		if err != nil {
   928  			return nil, errors.Trace(err)
   929  		}
   930  		resp.Resp = streamResp
   931  	case tikvrpc.CmdMvccGetByKey:
   932  		r := req.MvccGetByKey()
   933  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   934  			resp.Resp = &kvrpcpb.MvccGetByKeyResponse{RegionError: err}
   935  			return resp, nil
   936  		}
   937  		resp.Resp = kvHandler{session}.handleMvccGetByKey(r)
   938  	case tikvrpc.CmdMvccGetByStartTs:
   939  		r := req.MvccGetByStartTs()
   940  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   941  			resp.Resp = &kvrpcpb.MvccGetByStartTsResponse{RegionError: err}
   942  			return resp, nil
   943  		}
   944  		resp.Resp = kvHandler{session}.handleMvccGetByStartTS(r)
   945  	case tikvrpc.CmdSplitRegion:
   946  		r := req.SplitRegion()
   947  		if err := session.checkRequest(reqCtx, r.Size()); err != nil {
   948  			resp.Resp = &kvrpcpb.SplitRegionResponse{RegionError: err}
   949  			return resp, nil
   950  		}
   951  		resp.Resp = kvHandler{session}.handleSplitRegion(r)
   952  	// DebugGetRegionProperties is for fast analyze in mock tikv.
   953  	case tikvrpc.CmdDebugGetRegionProperties:
   954  		r := req.DebugGetRegionProperties()
   955  		region, _ := c.Cluster.GetRegion(r.RegionId)
   956  		var reqCtx kvrpcpb.Context
   957  		scanResp := kvHandler{session}.handleKvScan(&kvrpcpb.ScanRequest{
   958  			Context:  &reqCtx,
   959  			StartKey: MvccKey(region.StartKey).Raw(),
   960  			EndKey:   MvccKey(region.EndKey).Raw(),
   961  			Version:  math.MaxUint64,
   962  			Limit:    math.MaxUint32})
   963  		resp.Resp = &debugpb.GetRegionPropertiesResponse{
   964  			Props: []*debugpb.Property{{
   965  				Name:  "mvcc.num_rows",
   966  				Value: strconv.Itoa(len(scanResp.Pairs)),
   967  			}}}
   968  	default:
   969  		return nil, errors.Errorf("unsupported this request type %v", req.Type)
   970  	}
   971  	return resp, nil
   972  }
   973  
   974  // Close closes the client.
   975  func (c *RPCClient) Close() error {
   976  	if c.coprHandler != nil {
   977  		c.coprHandler.Close()
   978  	}
   979  
   980  	var err error
   981  	if c.MvccStore != nil {
   982  		err = c.MvccStore.Close()
   983  		if err != nil {
   984  			return err
   985  		}
   986  	}
   987  	return nil
   988  }