func (h *rpcHandler) handleRequest(req *kvrpcpb.Request) *kvrpcpb.Response {
resp := &kvrpcpb.Response{
Type: req.Type,
}
if err := h.checkContext(req.GetContext()); err != nil {
resp.RegionError = err
return resp
}
switch req.GetType() {
case kvrpcpb.MessageType_CmdGet:
resp.CmdGetResp = h.onGet(req.CmdGetReq)
case kvrpcpb.MessageType_CmdScan:
resp.CmdScanResp = h.onScan(req.CmdScanReq)
case kvrpcpb.MessageType_CmdPrewrite:
resp.CmdPrewriteResp = h.onPrewrite(req.CmdPrewriteReq)
case kvrpcpb.MessageType_CmdCommit:
resp.CmdCommitResp = h.onCommit(req.CmdCommitReq)
case kvrpcpb.MessageType_CmdCleanup:
resp.CmdCleanupResp = h.onCleanup(req.CmdCleanupReq)
case kvrpcpb.MessageType_CmdCommitThenGet:
resp.CmdCommitGetResp = h.onCommitThenGet(req.CmdCommitGetReq)
case kvrpcpb.MessageType_CmdRollbackThenGet:
resp.CmdRbGetResp = h.onRollbackThenGet(req.CmdRbGetReq)
case kvrpcpb.MessageType_CmdBatchGet:
resp.CmdBatchGetResp = h.onBatchGet(req.CmdBatchGetReq)
}
return resp
}
func (h *rpcHandler) handleRequest(req *kvrpcpb.Request) *kvrpcpb.Response {
var resp kvrpcpb.Response
if err := h.checkContext(req.GetContext()); err != nil {
resp.RegionError = err
return &resp
}
switch req.GetType() {
case kvrpcpb.MessageType_CmdGet:
resp.CmdGetResp = h.onGet(req.CmdGetReq)
case kvrpcpb.MessageType_CmdScan:
resp.CmdScanResp = h.onScan(req.CmdScanReq)
case kvrpcpb.MessageType_CmdPrewrite:
resp.CmdPrewriteResp = h.onPrewrite(req.CmdPrewriteReq)
case kvrpcpb.MessageType_CmdCommit:
resp.CmdCommitResp = h.onCommit(req.CmdCommitReq)
case kvrpcpb.MessageType_CmdCleanup:
resp.CmdCleanupResp = h.onCleanup(req.CmdCleanupReq)
case kvrpcpb.MessageType_CmdBatchGet:
resp.CmdBatchGetResp = h.onBatchGet(req.CmdBatchGetReq)
case kvrpcpb.MessageType_CmdScanLock:
resp.CmdResolveLockResp = h.onResolveLock(req.CmdResolveLockReq)
case kvrpcpb.MessageType_CmdResolveLock:
resp.CmdResolveLockResp = h.onResolveLock(req.CmdResolveLockReq)
case kvrpcpb.MessageType_CmdRawGet:
resp.CmdRawGetResp = h.onRawGet(req.CmdRawGetReq)
case kvrpcpb.MessageType_CmdRawPut:
resp.CmdRawPutResp = h.onRawPut(req.CmdRawPutReq)
case kvrpcpb.MessageType_CmdRawDelete:
resp.CmdRawDeleteResp = h.onRawDelete(req.CmdRawDeleteReq)
}
resp.Type = req.Type
return &resp
}
// Server read message then close, so `Send` will return retry error.
func (s *testClientSuite) TestRetryReadThenClose(c *C) {
l := startServer(":61236", c, readThenCloseRequest)
defer l.Close()
cli := newRPCClient()
req := new(pb.Request)
req.Type = pb.MessageType_CmdGet
resp, err := cli.SendKVReq(":61236", req, readTimeoutShort)
c.Assert(err, NotNil)
c.Assert(resp, IsNil)
}
// One normally `Send`.
func (s *testClientSuite) TestSendBySelf(c *C) {
l := startServer(":61234", c, handleRequest)
defer l.Close()
cli := newRPCClient()
req := new(pb.Request)
req.Type = pb.MessageType_CmdGet
getReq := new(pb.CmdGetRequest)
getReq.Key = []byte("a")
ver := uint64(0)
getReq.Version = ver
req.CmdGetReq = getReq
resp, err := cli.SendKVReq(":61234", req, readTimeoutShort)
c.Assert(err, IsNil)
c.Assert(req.GetType(), Equals, resp.GetType())
}
// sendKVReq sends req to tikv server. It will retry internally to find the right
// region leader if i) fails to establish a connection to server or ii) server
// returns `NotLeader`.
func (s *tikvStore) SendKVReq(bo *Backoffer, req *pb.Request, regionID RegionVerID) (*pb.Response, error) {
for {
region := s.regionCache.GetRegionByVerID(regionID)
if region == nil {
// If the region is not found in cache, it must be out
// of date and already be cleaned up. We can skip the
// RPC by returning RegionError directly.
return &pb.Response{
Type: req.GetType(),
RegionError: &errorpb.Error{StaleEpoch: &errorpb.StaleEpoch{}},
}, nil
}
req.Context = region.GetContext()
resp, err := s.client.SendKVReq(region.GetAddress(), req)
if err != nil {
s.regionCache.NextPeer(region.VerID())
err = bo.Backoff(boTiKVRPC, errors.Errorf("send tikv request error: %v, ctx: %s, try next peer later", err, req.Context))
if err != nil {
return nil, errors.Trace(err)
}
continue
}
if regionErr := resp.GetRegionError(); regionErr != nil {
// Retry if error is `NotLeader`.
if notLeader := regionErr.GetNotLeader(); notLeader != nil {
log.Warnf("tikv reports `NotLeader`: %s, ctx: %s, retry later", notLeader, req.Context)
s.regionCache.UpdateLeader(region.VerID(), notLeader.GetLeader().GetId())
if notLeader.GetLeader() == nil {
err = bo.Backoff(boRegionMiss, errors.Errorf("not leader: %v, ctx: %s", notLeader, req.Context))
if err != nil {
return nil, errors.Trace(err)
}
}
continue
}
// For other errors, we only drop cache here.
// Because caller may need to re-split the request.
log.Warnf("tikv reports region error: %s, ctx: %s", resp.GetRegionError(), req.Context)
s.regionCache.DropRegion(region.VerID())
return resp, nil
}
if resp.GetType() != req.GetType() {
return nil, errors.Trace(errMismatch(resp, req))
}
return resp, nil
}
}
// SendKVReq sends a KV request to tikv server.
func (s *RegionRequestSender) SendKVReq(req *kvrpcpb.Request, regionID RegionVerID, timeout time.Duration) (*kvrpcpb.Response, error) {
for {
select {
case <-s.bo.ctx.Done():
return nil, errors.Trace(s.bo.ctx.Err())
default:
}
ctx, err := s.regionCache.GetRPCContext(s.bo, regionID)
if err != nil {
return nil, errors.Trace(err)
}
if ctx == nil {
// If the region is not found in cache, it must be out
// of date and already be cleaned up. We can skip the
// RPC by returning RegionError directly.
return &kvrpcpb.Response{
Type: req.GetType(),
RegionError: &errorpb.Error{StaleEpoch: &errorpb.StaleEpoch{}},
}, nil
}
resp, retry, err := s.sendKVReqToRegion(ctx, req, timeout)
if err != nil {
return nil, errors.Trace(err)
}
if retry {
continue
}
if regionErr := resp.GetRegionError(); regionErr != nil {
retry, err := s.onRegionError(ctx, regionErr)
if err != nil {
return nil, errors.Trace(err)
}
if retry {
continue
}
return resp, nil
}
if resp.GetType() != req.GetType() {
return nil, errors.Trace(errMismatch(resp, req))
}
return resp, nil
}
}
func (s *RegionRequestSender) sendKVReqToRegion(ctx *RPCContext, req *kvrpcpb.Request, timeout time.Duration) (resp *kvrpcpb.Response, retry bool, err error) {
req.Context = ctx.KVCtx
resp, err = s.client.SendKVReq(ctx.Addr, req, timeout)
if err != nil {
if e := s.onSendFail(ctx, err); e != nil {
return nil, false, errors.Trace(e)
}
return nil, true, nil
}
return
}
// errMismatch if response mismatches request return error.
func errMismatch(resp *pb.Response, req *pb.Request) error {
return errors.Errorf("message type mismatches, response[%s] request[%s]",
resp.GetType(), req.GetType())
}