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 }