You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

952 lines
30 KiB

// Copyright 2017 PingCAP, Inc.
//
// 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,
// See the License for the specific language governing permissions and
// limitations under the License.
package tikvrpc
import (
"context"
"fmt"
"sync/atomic"
"time"
"github.com/pingcap/errors"
"github.com/pingcap/kvproto/pkg/coprocessor"
"github.com/pingcap/kvproto/pkg/debugpb"
"github.com/pingcap/kvproto/pkg/errorpb"
"github.com/pingcap/kvproto/pkg/kvrpcpb"
"github.com/pingcap/kvproto/pkg/metapb"
"github.com/pingcap/kvproto/pkg/tikvpb"
"github.com/pingcap/tidb/kv"
)
// CmdType represents the concrete request type in Request or response type in Response.
type CmdType uint16
// CmdType values.
const (
CmdGet CmdType = 1 + iota
CmdScan
CmdPrewrite
CmdCommit
CmdCleanup
CmdBatchGet
CmdBatchRollback
CmdScanLock
CmdResolveLock
CmdGC
CmdDeleteRange
CmdPessimisticLock
CmdPessimisticRollback
CmdTxnHeartBeat
CmdCheckTxnStatus
CmdRawGet CmdType = 256 + iota
CmdRawBatchGet
CmdRawPut
CmdRawBatchPut
CmdRawDelete
CmdRawBatchDelete
CmdRawDeleteRange
CmdRawScan
CmdUnsafeDestroyRange
CmdRegisterLockObserver
CmdCheckLockObserver
CmdRemoveLockObserver
CmdPhysicalScanLock
CmdCop CmdType = 512 + iota
CmdCopStream
CmdBatchCop
CmdMvccGetByKey CmdType = 1024 + iota
CmdMvccGetByStartTs
CmdSplitRegion
CmdDebugGetRegionProperties CmdType = 2048 + iota
CmdEmpty CmdType = 3072 + iota
)
func (t CmdType) String() string {
switch t {
case CmdGet:
return "Get"
case CmdScan:
return "Scan"
case CmdPrewrite:
return "Prewrite"
case CmdPessimisticLock:
return "PessimisticLock"
case CmdPessimisticRollback:
return "PessimisticRollback"
case CmdCommit:
return "Commit"
case CmdCleanup:
return "Cleanup"
case CmdBatchGet:
return "BatchGet"
case CmdBatchRollback:
return "BatchRollback"
case CmdScanLock:
return "ScanLock"
case CmdResolveLock:
return "ResolveLock"
case CmdGC:
return "GC"
case CmdDeleteRange:
return "DeleteRange"
case CmdRawGet:
return "RawGet"
case CmdRawBatchGet:
return "RawBatchGet"
case CmdRawPut:
return "RawPut"
case CmdRawBatchPut:
return "RawBatchPut"
case CmdRawDelete:
return "RawDelete"
case CmdRawBatchDelete:
return "RawBatchDelete"
case CmdRawDeleteRange:
return "RawDeleteRange"
case CmdRawScan:
return "RawScan"
case CmdUnsafeDestroyRange:
return "UnsafeDestroyRange"
case CmdRegisterLockObserver:
return "RegisterLockObserver"
case CmdCheckLockObserver:
return "CheckLockObserver"
case CmdRemoveLockObserver:
return "RemoveLockObserver"
case CmdPhysicalScanLock:
return "PhysicalScanLock"
case CmdCop:
return "Cop"
case CmdCopStream:
return "CopStream"
case CmdBatchCop:
return "BatchCop"
case CmdMvccGetByKey:
return "MvccGetByKey"
case CmdMvccGetByStartTs:
return "MvccGetByStartTS"
case CmdSplitRegion:
return "SplitRegion"
case CmdCheckTxnStatus:
return "CheckTxnStatus"
case CmdDebugGetRegionProperties:
return "DebugGetRegionProperties"
case CmdTxnHeartBeat:
return "TxnHeartBeat"
}
return "Unknown"
}
// Request wraps all kv/coprocessor requests.
type Request struct {
Type CmdType
req interface{}
kvrpcpb.Context
ReplicaReadType kv.ReplicaReadType // dirrerent from `kvrpcpb.Context.ReplicaRead`
ReplicaReadSeed *uint32 // pointer to follower read seed in snapshot/coprocessor
StoreTp kv.StoreType
}
// NewRequest returns new kv rpc request.
func NewRequest(typ CmdType, pointer interface{}, ctxs ...kvrpcpb.Context) *Request {
if len(ctxs) > 0 {
return &Request{
Type: typ,
req: pointer,
Context: ctxs[0],
}
}
return &Request{
Type: typ,
req: pointer,
}
}
// NewReplicaReadRequest returns new kv rpc request with replica read.
func NewReplicaReadRequest(typ CmdType, pointer interface{}, replicaReadType kv.ReplicaReadType, replicaReadSeed *uint32, ctxs ...kvrpcpb.Context) *Request {
req := NewRequest(typ, pointer, ctxs...)
req.ReplicaRead = replicaReadType.IsFollowerRead()
req.ReplicaReadType = replicaReadType
req.ReplicaReadSeed = replicaReadSeed
return req
}
// Get returns GetRequest in request.
func (req *Request) Get() *kvrpcpb.GetRequest {
return req.req.(*kvrpcpb.GetRequest)
}
// Scan returns ScanRequest in request.
func (req *Request) Scan() *kvrpcpb.ScanRequest {
return req.req.(*kvrpcpb.ScanRequest)
}
// Prewrite returns PrewriteRequest in request.
func (req *Request) Prewrite() *kvrpcpb.PrewriteRequest {
return req.req.(*kvrpcpb.PrewriteRequest)
}
// Commit returns CommitRequest in request.
func (req *Request) Commit() *kvrpcpb.CommitRequest {
return req.req.(*kvrpcpb.CommitRequest)
}
// Cleanup returns CleanupRequest in request.
func (req *Request) Cleanup() *kvrpcpb.CleanupRequest {
return req.req.(*kvrpcpb.CleanupRequest)
}
// BatchGet returns BatchGetRequest in request.
func (req *Request) BatchGet() *kvrpcpb.BatchGetRequest {
return req.req.(*kvrpcpb.BatchGetRequest)
}
// BatchRollback returns BatchRollbackRequest in request.
func (req *Request) BatchRollback() *kvrpcpb.BatchRollbackRequest {
return req.req.(*kvrpcpb.BatchRollbackRequest)
}
// ScanLock returns ScanLockRequest in request.
func (req *Request) ScanLock() *kvrpcpb.ScanLockRequest {
return req.req.(*kvrpcpb.ScanLockRequest)
}
// ResolveLock returns ResolveLockRequest in request.
func (req *Request) ResolveLock() *kvrpcpb.ResolveLockRequest {
return req.req.(*kvrpcpb.ResolveLockRequest)
}
// GC returns GCRequest in request.
func (req *Request) GC() *kvrpcpb.GCRequest {
return req.req.(*kvrpcpb.GCRequest)
}
// DeleteRange returns DeleteRangeRequest in request.
func (req *Request) DeleteRange() *kvrpcpb.DeleteRangeRequest {
return req.req.(*kvrpcpb.DeleteRangeRequest)
}
// RawGet returns RawGetRequest in request.
func (req *Request) RawGet() *kvrpcpb.RawGetRequest {
return req.req.(*kvrpcpb.RawGetRequest)
}
// RawBatchGet returns RawBatchGetRequest in request.
func (req *Request) RawBatchGet() *kvrpcpb.RawBatchGetRequest {
return req.req.(*kvrpcpb.RawBatchGetRequest)
}
// RawPut returns RawPutRequest in request.
func (req *Request) RawPut() *kvrpcpb.RawPutRequest {
return req.req.(*kvrpcpb.RawPutRequest)
}
// RawBatchPut returns RawBatchPutRequest in request.
func (req *Request) RawBatchPut() *kvrpcpb.RawBatchPutRequest {
return req.req.(*kvrpcpb.RawBatchPutRequest)
}
// RawDelete returns PrewriteRequest in request.
func (req *Request) RawDelete() *kvrpcpb.RawDeleteRequest {
return req.req.(*kvrpcpb.RawDeleteRequest)
}
// RawBatchDelete returns RawBatchDeleteRequest in request.
func (req *Request) RawBatchDelete() *kvrpcpb.RawBatchDeleteRequest {
return req.req.(*kvrpcpb.RawBatchDeleteRequest)
}
// RawDeleteRange returns RawDeleteRangeRequest in request.
func (req *Request) RawDeleteRange() *kvrpcpb.RawDeleteRangeRequest {
return req.req.(*kvrpcpb.RawDeleteRangeRequest)
}
// RawScan returns RawScanRequest in request.
func (req *Request) RawScan() *kvrpcpb.RawScanRequest {
return req.req.(*kvrpcpb.RawScanRequest)
}
// UnsafeDestroyRange returns UnsafeDestroyRangeRequest in request.
func (req *Request) UnsafeDestroyRange() *kvrpcpb.UnsafeDestroyRangeRequest {
return req.req.(*kvrpcpb.UnsafeDestroyRangeRequest)
}
// RegisterLockObserver returns RegisterLockObserverRequest in request.
func (req *Request) RegisterLockObserver() *kvrpcpb.RegisterLockObserverRequest {
return req.req.(*kvrpcpb.RegisterLockObserverRequest)
}
// CheckLockObserver returns CheckLockObserverRequest in request.
func (req *Request) CheckLockObserver() *kvrpcpb.CheckLockObserverRequest {
return req.req.(*kvrpcpb.CheckLockObserverRequest)
}
// RemoveLockObserver returns RemoveLockObserverRequest in request.
func (req *Request) RemoveLockObserver() *kvrpcpb.RemoveLockObserverRequest {
return req.req.(*kvrpcpb.RemoveLockObserverRequest)
}
// PhysicalScanLock returns PhysicalScanLockRequest in request.
func (req *Request) PhysicalScanLock() *kvrpcpb.PhysicalScanLockRequest {
return req.req.(*kvrpcpb.PhysicalScanLockRequest)
}
// Cop returns coprocessor request in request.
func (req *Request) Cop() *coprocessor.Request {
return req.req.(*coprocessor.Request)
}
// BatchCop returns coprocessor request in request.
func (req *Request) BatchCop() *coprocessor.BatchRequest {
return req.req.(*coprocessor.BatchRequest)
}
// MvccGetByKey returns MvccGetByKeyRequest in request.
func (req *Request) MvccGetByKey() *kvrpcpb.MvccGetByKeyRequest {
return req.req.(*kvrpcpb.MvccGetByKeyRequest)
}
// MvccGetByStartTs returns MvccGetByStartTsRequest in request.
func (req *Request) MvccGetByStartTs() *kvrpcpb.MvccGetByStartTsRequest {
return req.req.(*kvrpcpb.MvccGetByStartTsRequest)
}
// SplitRegion returns SplitRegionRequest in request.
func (req *Request) SplitRegion() *kvrpcpb.SplitRegionRequest {
return req.req.(*kvrpcpb.SplitRegionRequest)
}
// PessimisticLock returns PessimisticLockRequest in request.
func (req *Request) PessimisticLock() *kvrpcpb.PessimisticLockRequest {
return req.req.(*kvrpcpb.PessimisticLockRequest)
}
// PessimisticRollback returns PessimisticRollbackRequest in request.
func (req *Request) PessimisticRollback() *kvrpcpb.PessimisticRollbackRequest {
return req.req.(*kvrpcpb.PessimisticRollbackRequest)
}
// DebugGetRegionProperties returns GetRegionPropertiesRequest in request.
func (req *Request) DebugGetRegionProperties() *debugpb.GetRegionPropertiesRequest {
return req.req.(*debugpb.GetRegionPropertiesRequest)
}
// Empty returns BatchCommandsEmptyRequest in request.
func (req *Request) Empty() *tikvpb.BatchCommandsEmptyRequest {
return req.req.(*tikvpb.BatchCommandsEmptyRequest)
}
// CheckTxnStatus returns CheckTxnStatusRequest in request.
func (req *Request) CheckTxnStatus() *kvrpcpb.CheckTxnStatusRequest {
return req.req.(*kvrpcpb.CheckTxnStatusRequest)
}
// TxnHeartBeat returns TxnHeartBeatRequest in request.
func (req *Request) TxnHeartBeat() *kvrpcpb.TxnHeartBeatRequest {
return req.req.(*kvrpcpb.TxnHeartBeatRequest)
}
// ToBatchCommandsRequest converts the request to an entry in BatchCommands request.
func (req *Request) ToBatchCommandsRequest() *tikvpb.BatchCommandsRequest_Request {
switch req.Type {
case CmdGet:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Get{Get: req.Get()}}
case CmdScan:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Scan{Scan: req.Scan()}}
case CmdPrewrite:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Prewrite{Prewrite: req.Prewrite()}}
case CmdCommit:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Commit{Commit: req.Commit()}}
case CmdCleanup:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Cleanup{Cleanup: req.Cleanup()}}
case CmdBatchGet:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_BatchGet{BatchGet: req.BatchGet()}}
case CmdBatchRollback:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_BatchRollback{BatchRollback: req.BatchRollback()}}
case CmdScanLock:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_ScanLock{ScanLock: req.ScanLock()}}
case CmdResolveLock:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_ResolveLock{ResolveLock: req.ResolveLock()}}
case CmdGC:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_GC{GC: req.GC()}}
case CmdDeleteRange:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_DeleteRange{DeleteRange: req.DeleteRange()}}
case CmdRawGet:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_RawGet{RawGet: req.RawGet()}}
case CmdRawBatchGet:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_RawBatchGet{RawBatchGet: req.RawBatchGet()}}
case CmdRawPut:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_RawPut{RawPut: req.RawPut()}}
case CmdRawBatchPut:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_RawBatchPut{RawBatchPut: req.RawBatchPut()}}
case CmdRawDelete:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_RawDelete{RawDelete: req.RawDelete()}}
case CmdRawBatchDelete:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_RawBatchDelete{RawBatchDelete: req.RawBatchDelete()}}
case CmdRawDeleteRange:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_RawDeleteRange{RawDeleteRange: req.RawDeleteRange()}}
case CmdRawScan:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_RawScan{RawScan: req.RawScan()}}
case CmdCop:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Coprocessor{Coprocessor: req.Cop()}}
case CmdPessimisticLock:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_PessimisticLock{PessimisticLock: req.PessimisticLock()}}
case CmdPessimisticRollback:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_PessimisticRollback{PessimisticRollback: req.PessimisticRollback()}}
case CmdEmpty:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_Empty{Empty: req.Empty()}}
case CmdCheckTxnStatus:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_CheckTxnStatus{CheckTxnStatus: req.CheckTxnStatus()}}
case CmdTxnHeartBeat:
return &tikvpb.BatchCommandsRequest_Request{Cmd: &tikvpb.BatchCommandsRequest_Request_TxnHeartBeat{TxnHeartBeat: req.TxnHeartBeat()}}
}
return nil
}
// IsDebugReq check whether the req is debug req.
func (req *Request) IsDebugReq() bool {
switch req.Type {
case CmdDebugGetRegionProperties:
return true
}
return false
}
// Response wraps all kv/coprocessor responses.
type Response struct {
Resp interface{}
}
// FromBatchCommandsResponse converts a BatchCommands response to Response.
func FromBatchCommandsResponse(res *tikvpb.BatchCommandsResponse_Response) (*Response, error) {
if res.GetCmd() == nil {
return nil, errors.New("Unknown command response")
}
switch res := res.GetCmd().(type) {
case *tikvpb.BatchCommandsResponse_Response_Get:
return &Response{Resp: res.Get}, nil
case *tikvpb.BatchCommandsResponse_Response_Scan:
return &Response{Resp: res.Scan}, nil
case *tikvpb.BatchCommandsResponse_Response_Prewrite:
return &Response{Resp: res.Prewrite}, nil
case *tikvpb.BatchCommandsResponse_Response_Commit:
return &Response{Resp: res.Commit}, nil
case *tikvpb.BatchCommandsResponse_Response_Cleanup:
return &Response{Resp: res.Cleanup}, nil
case *tikvpb.BatchCommandsResponse_Response_BatchGet:
return &Response{Resp: res.BatchGet}, nil
case *tikvpb.BatchCommandsResponse_Response_BatchRollback:
return &Response{Resp: res.BatchRollback}, nil
case *tikvpb.BatchCommandsResponse_Response_ScanLock:
return &Response{Resp: res.ScanLock}, nil
case *tikvpb.BatchCommandsResponse_Response_ResolveLock:
return &Response{Resp: res.ResolveLock}, nil
case *tikvpb.BatchCommandsResponse_Response_GC:
return &Response{Resp: res.GC}, nil
case *tikvpb.BatchCommandsResponse_Response_DeleteRange:
return &Response{Resp: res.DeleteRange}, nil
case *tikvpb.BatchCommandsResponse_Response_RawGet:
return &Response{Resp: res.RawGet}, nil
case *tikvpb.BatchCommandsResponse_Response_RawBatchGet:
return &Response{Resp: res.RawBatchGet}, nil
case *tikvpb.BatchCommandsResponse_Response_RawPut:
return &Response{Resp: res.RawPut}, nil
case *tikvpb.BatchCommandsResponse_Response_RawBatchPut:
return &Response{Resp: res.RawBatchPut}, nil
case *tikvpb.BatchCommandsResponse_Response_RawDelete:
return &Response{Resp: res.RawDelete}, nil
case *tikvpb.BatchCommandsResponse_Response_RawBatchDelete:
return &Response{Resp: res.RawBatchDelete}, nil
case *tikvpb.BatchCommandsResponse_Response_RawDeleteRange:
return &Response{Resp: res.RawDeleteRange}, nil
case *tikvpb.BatchCommandsResponse_Response_RawScan:
return &Response{Resp: res.RawScan}, nil
case *tikvpb.BatchCommandsResponse_Response_Coprocessor:
return &Response{Resp: res.Coprocessor}, nil
case *tikvpb.BatchCommandsResponse_Response_PessimisticLock:
return &Response{Resp: res.PessimisticLock}, nil
case *tikvpb.BatchCommandsResponse_Response_PessimisticRollback:
return &Response{Resp: res.PessimisticRollback}, nil
case *tikvpb.BatchCommandsResponse_Response_Empty:
return &Response{Resp: res.Empty}, nil
case *tikvpb.BatchCommandsResponse_Response_TxnHeartBeat:
return &Response{Resp: res.TxnHeartBeat}, nil
case *tikvpb.BatchCommandsResponse_Response_CheckTxnStatus:
return &Response{Resp: res.CheckTxnStatus}, nil
}
panic("unreachable")
}
// CopStreamResponse combinates tikvpb.Tikv_CoprocessorStreamClient and the first Recv() result together.
// In streaming API, get grpc stream client may not involve any network packet, then region error have
// to be handled in Recv() function. This struct facilitates the error handling.
type CopStreamResponse struct {
tikvpb.Tikv_CoprocessorStreamClient
*coprocessor.Response // The first result of Recv()
Timeout time.Duration
Lease // Shared by this object and a background goroutine.
}
// BatchCopStreamResponse comprises the BatchCoprocessorClient , the first result and timeout detector.
type BatchCopStreamResponse struct {
tikvpb.Tikv_BatchCoprocessorClient
*coprocessor.BatchResponse
Timeout time.Duration
Lease // Shared by this object and a background goroutine.
}
// SetContext set the Context field for the given req to the specified ctx.
func SetContext(req *Request, region *metapb.Region, peer *metapb.Peer) error {
ctx := &req.Context
if region != nil {
ctx.RegionId = region.Id
ctx.RegionEpoch = region.RegionEpoch
}
ctx.Peer = peer
switch req.Type {
case CmdGet:
req.Get().Context = ctx
case CmdScan:
req.Scan().Context = ctx
case CmdPrewrite:
req.Prewrite().Context = ctx
case CmdPessimisticLock:
req.PessimisticLock().Context = ctx
case CmdPessimisticRollback:
req.PessimisticRollback().Context = ctx
case CmdCommit:
req.Commit().Context = ctx
case CmdCleanup:
req.Cleanup().Context = ctx
case CmdBatchGet:
req.BatchGet().Context = ctx
case CmdBatchRollback:
req.BatchRollback().Context = ctx
case CmdScanLock:
req.ScanLock().Context = ctx
case CmdResolveLock:
req.ResolveLock().Context = ctx
case CmdGC:
req.GC().Context = ctx
case CmdDeleteRange:
req.DeleteRange().Context = ctx
case CmdRawGet:
req.RawGet().Context = ctx
case CmdRawBatchGet:
req.RawBatchGet().Context = ctx
case CmdRawPut:
req.RawPut().Context = ctx
case CmdRawBatchPut:
req.RawBatchPut().Context = ctx
case CmdRawDelete:
req.RawDelete().Context = ctx
case CmdRawBatchDelete:
req.RawBatchDelete().Context = ctx
case CmdRawDeleteRange:
req.RawDeleteRange().Context = ctx
case CmdRawScan:
req.RawScan().Context = ctx
case CmdUnsafeDestroyRange:
req.UnsafeDestroyRange().Context = ctx
case CmdRegisterLockObserver:
req.RegisterLockObserver().Context = ctx
case CmdCheckLockObserver:
req.CheckLockObserver().Context = ctx
case CmdRemoveLockObserver:
req.RemoveLockObserver().Context = ctx
case CmdPhysicalScanLock:
req.PhysicalScanLock().Context = ctx
case CmdCop:
req.Cop().Context = ctx
case CmdCopStream:
req.Cop().Context = ctx
case CmdBatchCop:
req.BatchCop().Context = ctx
case CmdMvccGetByKey:
req.MvccGetByKey().Context = ctx
case CmdMvccGetByStartTs:
req.MvccGetByStartTs().Context = ctx
case CmdSplitRegion:
req.SplitRegion().Context = ctx
case CmdEmpty:
req.SplitRegion().Context = ctx
case CmdTxnHeartBeat:
req.TxnHeartBeat().Context = ctx
case CmdCheckTxnStatus:
req.CheckTxnStatus().Context = ctx
default:
return fmt.Errorf("invalid request type %v", req.Type)
}
return nil
}
// GenRegionErrorResp returns corresponding Response with specified RegionError
// according to the given req.
func GenRegionErrorResp(req *Request, e *errorpb.Error) (*Response, error) {
var p interface{}
resp := &Response{}
switch req.Type {
case CmdGet:
p = &kvrpcpb.GetResponse{
RegionError: e,
}
case CmdScan:
p = &kvrpcpb.ScanResponse{
RegionError: e,
}
case CmdPrewrite:
p = &kvrpcpb.PrewriteResponse{
RegionError: e,
}
case CmdPessimisticLock:
p = &kvrpcpb.PessimisticLockResponse{
RegionError: e,
}
case CmdPessimisticRollback:
p = &kvrpcpb.PessimisticRollbackResponse{
RegionError: e,
}
case CmdCommit:
p = &kvrpcpb.CommitResponse{
RegionError: e,
}
case CmdCleanup:
p = &kvrpcpb.CleanupResponse{
RegionError: e,
}
case CmdBatchGet:
p = &kvrpcpb.BatchGetResponse{
RegionError: e,
}
case CmdBatchRollback:
p = &kvrpcpb.BatchRollbackResponse{
RegionError: e,
}
case CmdScanLock:
p = &kvrpcpb.ScanLockResponse{
RegionError: e,
}
case CmdResolveLock:
p = &kvrpcpb.ResolveLockResponse{
RegionError: e,
}
case CmdGC:
p = &kvrpcpb.GCResponse{
RegionError: e,
}
case CmdDeleteRange:
p = &kvrpcpb.DeleteRangeResponse{
RegionError: e,
}
case CmdRawGet:
p = &kvrpcpb.RawGetResponse{
RegionError: e,
}
case CmdRawBatchGet:
p = &kvrpcpb.RawBatchGetResponse{
RegionError: e,
}
case CmdRawPut:
p = &kvrpcpb.RawPutResponse{
RegionError: e,
}
case CmdRawBatchPut:
p = &kvrpcpb.RawBatchPutResponse{
RegionError: e,
}
case CmdRawDelete:
p = &kvrpcpb.RawDeleteResponse{
RegionError: e,
}
case CmdRawBatchDelete:
p = &kvrpcpb.RawBatchDeleteResponse{
RegionError: e,
}
case CmdRawDeleteRange:
p = &kvrpcpb.RawDeleteRangeResponse{
RegionError: e,
}
case CmdRawScan:
p = &kvrpcpb.RawScanResponse{
RegionError: e,
}
case CmdUnsafeDestroyRange:
p = &kvrpcpb.UnsafeDestroyRangeResponse{
RegionError: e,
}
case CmdCop:
p = &coprocessor.Response{
RegionError: e,
}
case CmdCopStream:
p = &CopStreamResponse{
Response: &coprocessor.Response{
RegionError: e,
},
}
case CmdMvccGetByKey:
p = &kvrpcpb.MvccGetByKeyResponse{
RegionError: e,
}
case CmdMvccGetByStartTs:
p = &kvrpcpb.MvccGetByStartTsResponse{
RegionError: e,
}
case CmdSplitRegion:
p = &kvrpcpb.SplitRegionResponse{
RegionError: e,
}
case CmdEmpty:
case CmdTxnHeartBeat:
p = &kvrpcpb.TxnHeartBeatResponse{
RegionError: e,
}
case CmdCheckTxnStatus:
p = &kvrpcpb.CheckTxnStatusResponse{
RegionError: e,
}
default:
return nil, fmt.Errorf("invalid request type %v", req.Type)
}
resp.Resp = p
return resp, nil
}
type getRegionError interface {
GetRegionError() *errorpb.Error
}
// GetRegionError returns the RegionError of the underlying concrete response.
func (resp *Response) GetRegionError() (*errorpb.Error, error) {
if resp.Resp == nil {
return nil, nil
}
err, ok := resp.Resp.(getRegionError)
if !ok {
if _, isEmpty := resp.Resp.(*tikvpb.BatchCommandsEmptyResponse); isEmpty {
return nil, nil
}
return nil, fmt.Errorf("invalid response type %v", resp)
}
return err.GetRegionError(), nil
}
// CallRPC launches a rpc call.
// ch is needed to implement timeout for coprocessor streaing, the stream object's
// cancel function will be sent to the channel, together with a lease checked by a background goroutine.
func CallRPC(ctx context.Context, client tikvpb.TikvClient, req *Request) (*Response, error) {
resp := &Response{}
var err error
switch req.Type {
case CmdGet:
resp.Resp, err = client.KvGet(ctx, req.Get())
case CmdScan:
resp.Resp, err = client.KvScan(ctx, req.Scan())
case CmdPrewrite:
resp.Resp, err = client.KvPrewrite(ctx, req.Prewrite())
case CmdPessimisticLock:
resp.Resp, err = client.KvPessimisticLock(ctx, req.PessimisticLock())
case CmdPessimisticRollback:
resp.Resp, err = client.KVPessimisticRollback(ctx, req.PessimisticRollback())
case CmdCommit:
resp.Resp, err = client.KvCommit(ctx, req.Commit())
case CmdCleanup:
resp.Resp, err = client.KvCleanup(ctx, req.Cleanup())
case CmdBatchGet:
resp.Resp, err = client.KvBatchGet(ctx, req.BatchGet())
case CmdBatchRollback:
resp.Resp, err = client.KvBatchRollback(ctx, req.BatchRollback())
case CmdScanLock:
resp.Resp, err = client.KvScanLock(ctx, req.ScanLock())
case CmdResolveLock:
resp.Resp, err = client.KvResolveLock(ctx, req.ResolveLock())
case CmdGC:
resp.Resp, err = client.KvGC(ctx, req.GC())
case CmdDeleteRange:
resp.Resp, err = client.KvDeleteRange(ctx, req.DeleteRange())
case CmdRawGet:
resp.Resp, err = client.RawGet(ctx, req.RawGet())
case CmdRawBatchGet:
resp.Resp, err = client.RawBatchGet(ctx, req.RawBatchGet())
case CmdRawPut:
resp.Resp, err = client.RawPut(ctx, req.RawPut())
case CmdRawBatchPut:
resp.Resp, err = client.RawBatchPut(ctx, req.RawBatchPut())
case CmdRawDelete:
resp.Resp, err = client.RawDelete(ctx, req.RawDelete())
case CmdRawBatchDelete:
resp.Resp, err = client.RawBatchDelete(ctx, req.RawBatchDelete())
case CmdRawDeleteRange:
resp.Resp, err = client.RawDeleteRange(ctx, req.RawDeleteRange())
case CmdRawScan:
resp.Resp, err = client.RawScan(ctx, req.RawScan())
case CmdUnsafeDestroyRange:
resp.Resp, err = client.UnsafeDestroyRange(ctx, req.UnsafeDestroyRange())
case CmdRegisterLockObserver:
resp.Resp, err = client.RegisterLockObserver(ctx, req.RegisterLockObserver())
case CmdCheckLockObserver:
resp.Resp, err = client.CheckLockObserver(ctx, req.CheckLockObserver())
case CmdRemoveLockObserver:
resp.Resp, err = client.RemoveLockObserver(ctx, req.RemoveLockObserver())
case CmdPhysicalScanLock:
resp.Resp, err = client.PhysicalScanLock(ctx, req.PhysicalScanLock())
case CmdCop:
resp.Resp, err = client.Coprocessor(ctx, req.Cop())
case CmdCopStream:
var streamClient tikvpb.Tikv_CoprocessorStreamClient
streamClient, err = client.CoprocessorStream(ctx, req.Cop())
resp.Resp = &CopStreamResponse{
Tikv_CoprocessorStreamClient: streamClient,
}
case CmdBatchCop:
var streamClient tikvpb.Tikv_BatchCoprocessorClient
streamClient, err = client.BatchCoprocessor(ctx, req.BatchCop())
resp.Resp = &BatchCopStreamResponse{
Tikv_BatchCoprocessorClient: streamClient,
}
case CmdMvccGetByKey:
resp.Resp, err = client.MvccGetByKey(ctx, req.MvccGetByKey())
case CmdMvccGetByStartTs:
resp.Resp, err = client.MvccGetByStartTs(ctx, req.MvccGetByStartTs())
case CmdSplitRegion:
resp.Resp, err = client.SplitRegion(ctx, req.SplitRegion())
case CmdEmpty:
resp.Resp, err = &tikvpb.BatchCommandsEmptyResponse{}, nil
case CmdCheckTxnStatus:
resp.Resp, err = client.KvCheckTxnStatus(ctx, req.CheckTxnStatus())
case CmdTxnHeartBeat:
resp.Resp, err = client.KvTxnHeartBeat(ctx, req.TxnHeartBeat())
default:
return nil, errors.Errorf("invalid request type: %v", req.Type)
}
if err != nil {
return nil, errors.Trace(err)
}
return resp, nil
}
// CallDebugRPC launches a debug rpc call.
func CallDebugRPC(ctx context.Context, client debugpb.DebugClient, req *Request) (*Response, error) {
resp := &Response{}
var err error
switch req.Type {
case CmdDebugGetRegionProperties:
resp.Resp, err = client.GetRegionProperties(ctx, req.DebugGetRegionProperties())
default:
return nil, errors.Errorf("invalid request type: %v", req.Type)
}
return resp, err
}
// Lease is used to implement grpc stream timeout.
type Lease struct {
Cancel context.CancelFunc
deadline int64 // A time.UnixNano value, if time.Now().UnixNano() > deadline, cancel() would be called.
}
// Recv overrides the stream client Recv() function.
func (resp *CopStreamResponse) Recv() (*coprocessor.Response, error) {
deadline := time.Now().Add(resp.Timeout).UnixNano()
atomic.StoreInt64(&resp.Lease.deadline, deadline)
ret, err := resp.Tikv_CoprocessorStreamClient.Recv()
atomic.StoreInt64(&resp.Lease.deadline, 0) // Stop the lease check.
return ret, errors.Trace(err)
}
// Close closes the CopStreamResponse object.
func (resp *CopStreamResponse) Close() {
atomic.StoreInt64(&resp.Lease.deadline, 1)
// We also call cancel here because CheckStreamTimeoutLoop
// is not guaranteed to cancel all items when it exits.
if resp.Lease.Cancel != nil {
resp.Lease.Cancel()
}
}
// Recv overrides the stream client Recv() function.
func (resp *BatchCopStreamResponse) Recv() (*coprocessor.BatchResponse, error) {
deadline := time.Now().Add(resp.Timeout).UnixNano()
atomic.StoreInt64(&resp.Lease.deadline, deadline)
ret, err := resp.Tikv_BatchCoprocessorClient.Recv()
atomic.StoreInt64(&resp.Lease.deadline, 0) // Stop the lease check.
return ret, errors.Trace(err)
}
// Close closes the CopStreamResponse object.
func (resp *BatchCopStreamResponse) Close() {
atomic.StoreInt64(&resp.Lease.deadline, 1)
// We also call cancel here because CheckStreamTimeoutLoop
// is not guaranteed to cancel all items when it exits.
if resp.Lease.Cancel != nil {
resp.Lease.Cancel()
}
}
// CheckStreamTimeoutLoop runs periodically to check is there any stream request timeouted.
// Lease is an object to track stream requests, call this function with "go CheckStreamTimeoutLoop()"
// It is not guaranteed to call every Lease.Cancel() putting into channel when exits.
// If grpc-go supports SetDeadline(https://github.com/grpc/grpc-go/issues/2917), we can stop using this method.
func CheckStreamTimeoutLoop(ch <-chan *Lease, done <-chan struct{}) {
ticker := time.NewTicker(200 * time.Millisecond)
defer ticker.Stop()
array := make([]*Lease, 0, 1024)
for {
select {
case <-done:
drainLoop:
// Try my best cleaning the channel to make SendRequest which is blocking by it continues.
for {
select {
case <-ch:
default:
break drainLoop
}
}
return
case item := <-ch:
array = append(array, item)
case now := <-ticker.C:
array = keepOnlyActive(array, now.UnixNano())
}
}
}
// keepOnlyActive removes completed items, call cancel function for timeout items.
func keepOnlyActive(array []*Lease, now int64) []*Lease {
idx := 0
for i := 0; i < len(array); i++ {
item := array[i]
deadline := atomic.LoadInt64(&item.deadline)
if deadline == 0 || deadline > now {
array[idx] = array[i]
idx++
} else {
item.Cancel()
}
}
return array[:idx]
}