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
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]
|
|
}
|
|
|