mirror of https://github.com/tikv/client-go.git
tikvrpc: avoid data race on `XxxRequest.Context` (#1242)
* tikvrpc: avoid data race on `XxxRequest.Context` Signed-off-by: zyguan <zhongyangguan@gmail.com> * fix grammar of codegen comment Signed-off-by: zyguan <zhongyangguan@gmail.com> * address comments Signed-off-by: zyguan <zhongyangguan@gmail.com> * check diff of go generate Signed-off-by: zyguan <zhongyangguan@gmail.com> * fix a typo Signed-off-by: zyguan <zhongyangguan@gmail.com> --------- Signed-off-by: zyguan <zhongyangguan@gmail.com>
This commit is contained in:
zyguan
GitHub
parent
08aa706635
commit
05aaba6cc6
|
|
@ -44,6 +44,11 @@ jobs:
|
|||
with:
|
||||
go-version: 1.21.6
|
||||
|
||||
- name: Go generate and check diff
|
||||
run: |
|
||||
go generate ./...
|
||||
git diff --exit-code
|
||||
|
||||
- name: Lint
|
||||
uses: golangci/golangci-lint-action@v3
|
||||
with:
|
||||
|
|
|
|||
|
|
@ -2518,7 +2518,7 @@ func TestReplicaReadAccessPathByGenError(t *testing.T) {
|
|||
|
||||
maxAccessErrCnt := 4
|
||||
if israce.RaceEnabled {
|
||||
// When run this test with race, it will take a long time, so we reduce the maxAccessErrCnt to 3 to speed up test to avoid timeout.
|
||||
// When run this test with race, it will take a long time, so we reduce the maxAccessErrCnt to 2 to speed up test to avoid timeout.
|
||||
maxAccessErrCnt = 2
|
||||
}
|
||||
totalValidCaseCount := 0
|
||||
|
|
|
|||
|
|
@ -0,0 +1,384 @@
|
|||
// Code generated by gen.sh. DO NOT EDIT.
|
||||
|
||||
package tikvrpc
|
||||
|
||||
import (
|
||||
"github.com/pingcap/kvproto/pkg/kvrpcpb"
|
||||
)
|
||||
|
||||
func patchCmdCtx(req *Request, cmd CmdType, ctx *kvrpcpb.Context) bool {
|
||||
switch cmd {
|
||||
case CmdGet:
|
||||
if req.rev == 0 {
|
||||
req.Get().Context = ctx
|
||||
} else {
|
||||
cmd := *req.Get()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdScan:
|
||||
if req.rev == 0 {
|
||||
req.Scan().Context = ctx
|
||||
} else {
|
||||
cmd := *req.Scan()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdPrewrite:
|
||||
if req.rev == 0 {
|
||||
req.Prewrite().Context = ctx
|
||||
} else {
|
||||
cmd := *req.Prewrite()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdPessimisticLock:
|
||||
if req.rev == 0 {
|
||||
req.PessimisticLock().Context = ctx
|
||||
} else {
|
||||
cmd := *req.PessimisticLock()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdPessimisticRollback:
|
||||
if req.rev == 0 {
|
||||
req.PessimisticRollback().Context = ctx
|
||||
} else {
|
||||
cmd := *req.PessimisticRollback()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdCommit:
|
||||
if req.rev == 0 {
|
||||
req.Commit().Context = ctx
|
||||
} else {
|
||||
cmd := *req.Commit()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdCleanup:
|
||||
if req.rev == 0 {
|
||||
req.Cleanup().Context = ctx
|
||||
} else {
|
||||
cmd := *req.Cleanup()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdBatchGet:
|
||||
if req.rev == 0 {
|
||||
req.BatchGet().Context = ctx
|
||||
} else {
|
||||
cmd := *req.BatchGet()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdBatchRollback:
|
||||
if req.rev == 0 {
|
||||
req.BatchRollback().Context = ctx
|
||||
} else {
|
||||
cmd := *req.BatchRollback()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdScanLock:
|
||||
if req.rev == 0 {
|
||||
req.ScanLock().Context = ctx
|
||||
} else {
|
||||
cmd := *req.ScanLock()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdResolveLock:
|
||||
if req.rev == 0 {
|
||||
req.ResolveLock().Context = ctx
|
||||
} else {
|
||||
cmd := *req.ResolveLock()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdGC:
|
||||
if req.rev == 0 {
|
||||
req.GC().Context = ctx
|
||||
} else {
|
||||
cmd := *req.GC()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdDeleteRange:
|
||||
if req.rev == 0 {
|
||||
req.DeleteRange().Context = ctx
|
||||
} else {
|
||||
cmd := *req.DeleteRange()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawGet:
|
||||
if req.rev == 0 {
|
||||
req.RawGet().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawGet()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawBatchGet:
|
||||
if req.rev == 0 {
|
||||
req.RawBatchGet().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawBatchGet()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawPut:
|
||||
if req.rev == 0 {
|
||||
req.RawPut().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawPut()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawBatchPut:
|
||||
if req.rev == 0 {
|
||||
req.RawBatchPut().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawBatchPut()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawDelete:
|
||||
if req.rev == 0 {
|
||||
req.RawDelete().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawDelete()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawBatchDelete:
|
||||
if req.rev == 0 {
|
||||
req.RawBatchDelete().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawBatchDelete()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawDeleteRange:
|
||||
if req.rev == 0 {
|
||||
req.RawDeleteRange().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawDeleteRange()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawScan:
|
||||
if req.rev == 0 {
|
||||
req.RawScan().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawScan()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawGetKeyTTL:
|
||||
if req.rev == 0 {
|
||||
req.RawGetKeyTTL().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawGetKeyTTL()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawCompareAndSwap:
|
||||
if req.rev == 0 {
|
||||
req.RawCompareAndSwap().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawCompareAndSwap()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRawChecksum:
|
||||
if req.rev == 0 {
|
||||
req.RawChecksum().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RawChecksum()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdUnsafeDestroyRange:
|
||||
if req.rev == 0 {
|
||||
req.UnsafeDestroyRange().Context = ctx
|
||||
} else {
|
||||
cmd := *req.UnsafeDestroyRange()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRegisterLockObserver:
|
||||
if req.rev == 0 {
|
||||
req.RegisterLockObserver().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RegisterLockObserver()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdCheckLockObserver:
|
||||
if req.rev == 0 {
|
||||
req.CheckLockObserver().Context = ctx
|
||||
} else {
|
||||
cmd := *req.CheckLockObserver()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdRemoveLockObserver:
|
||||
if req.rev == 0 {
|
||||
req.RemoveLockObserver().Context = ctx
|
||||
} else {
|
||||
cmd := *req.RemoveLockObserver()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdPhysicalScanLock:
|
||||
if req.rev == 0 {
|
||||
req.PhysicalScanLock().Context = ctx
|
||||
} else {
|
||||
cmd := *req.PhysicalScanLock()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdCop:
|
||||
if req.rev == 0 {
|
||||
req.Cop().Context = ctx
|
||||
} else {
|
||||
cmd := *req.Cop()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdBatchCop:
|
||||
if req.rev == 0 {
|
||||
req.BatchCop().Context = ctx
|
||||
} else {
|
||||
cmd := *req.BatchCop()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdMvccGetByKey:
|
||||
if req.rev == 0 {
|
||||
req.MvccGetByKey().Context = ctx
|
||||
} else {
|
||||
cmd := *req.MvccGetByKey()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdMvccGetByStartTs:
|
||||
if req.rev == 0 {
|
||||
req.MvccGetByStartTs().Context = ctx
|
||||
} else {
|
||||
cmd := *req.MvccGetByStartTs()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdSplitRegion:
|
||||
if req.rev == 0 {
|
||||
req.SplitRegion().Context = ctx
|
||||
} else {
|
||||
cmd := *req.SplitRegion()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdTxnHeartBeat:
|
||||
if req.rev == 0 {
|
||||
req.TxnHeartBeat().Context = ctx
|
||||
} else {
|
||||
cmd := *req.TxnHeartBeat()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdCheckTxnStatus:
|
||||
if req.rev == 0 {
|
||||
req.CheckTxnStatus().Context = ctx
|
||||
} else {
|
||||
cmd := *req.CheckTxnStatus()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdCheckSecondaryLocks:
|
||||
if req.rev == 0 {
|
||||
req.CheckSecondaryLocks().Context = ctx
|
||||
} else {
|
||||
cmd := *req.CheckSecondaryLocks()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdFlashbackToVersion:
|
||||
if req.rev == 0 {
|
||||
req.FlashbackToVersion().Context = ctx
|
||||
} else {
|
||||
cmd := *req.FlashbackToVersion()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdPrepareFlashbackToVersion:
|
||||
if req.rev == 0 {
|
||||
req.PrepareFlashbackToVersion().Context = ctx
|
||||
} else {
|
||||
cmd := *req.PrepareFlashbackToVersion()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdFlush:
|
||||
if req.rev == 0 {
|
||||
req.Flush().Context = ctx
|
||||
} else {
|
||||
cmd := *req.Flush()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
case CmdBufferBatchGet:
|
||||
if req.rev == 0 {
|
||||
req.BufferBatchGet().Context = ctx
|
||||
} else {
|
||||
cmd := *req.BufferBatchGet()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
|
|
@ -0,0 +1,99 @@
|
|||
#!/bin/bash
|
||||
|
||||
# Copyright 2024 TiKV Authors
|
||||
#
|
||||
# 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,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
output="cmds_generated.go"
|
||||
|
||||
cat <<EOF > $output
|
||||
// Code generated by gen.sh. DO NOT EDIT.
|
||||
|
||||
package tikvrpc
|
||||
|
||||
import (
|
||||
"github.com/pingcap/kvproto/pkg/kvrpcpb"
|
||||
)
|
||||
EOF
|
||||
|
||||
cmds=(
|
||||
Get
|
||||
Scan
|
||||
Prewrite
|
||||
PessimisticLock
|
||||
PessimisticRollback
|
||||
Commit
|
||||
Cleanup
|
||||
BatchGet
|
||||
BatchRollback
|
||||
ScanLock
|
||||
ResolveLock
|
||||
GC
|
||||
DeleteRange
|
||||
RawGet
|
||||
RawBatchGet
|
||||
RawPut
|
||||
RawBatchPut
|
||||
RawDelete
|
||||
RawBatchDelete
|
||||
RawDeleteRange
|
||||
RawScan
|
||||
RawGetKeyTTL
|
||||
RawCompareAndSwap
|
||||
RawChecksum
|
||||
UnsafeDestroyRange
|
||||
RegisterLockObserver
|
||||
CheckLockObserver
|
||||
RemoveLockObserver
|
||||
PhysicalScanLock
|
||||
Cop
|
||||
BatchCop
|
||||
MvccGetByKey
|
||||
MvccGetByStartTs
|
||||
SplitRegion
|
||||
TxnHeartBeat
|
||||
CheckTxnStatus
|
||||
CheckSecondaryLocks
|
||||
FlashbackToVersion
|
||||
PrepareFlashbackToVersion
|
||||
Flush
|
||||
BufferBatchGet
|
||||
)
|
||||
|
||||
cat <<EOF >> $output
|
||||
|
||||
func patchCmdCtx(req *Request, cmd CmdType, ctx *kvrpcpb.Context) bool {
|
||||
switch cmd {
|
||||
EOF
|
||||
|
||||
for cmd in "${cmds[@]}"; do
|
||||
cat <<EOF >> $output
|
||||
case Cmd${cmd}:
|
||||
if req.rev == 0 {
|
||||
req.${cmd}().Context = ctx
|
||||
} else {
|
||||
cmd := *req.${cmd}()
|
||||
cmd.Context = ctx
|
||||
req.Req = &cmd
|
||||
}
|
||||
req.rev++
|
||||
EOF
|
||||
done
|
||||
|
||||
cat <<EOF >> $output
|
||||
default:
|
||||
return false
|
||||
}
|
||||
return true
|
||||
}
|
||||
EOF
|
||||
|
|
@ -85,7 +85,7 @@ const (
|
|||
CmdRawBatchDelete
|
||||
CmdRawDeleteRange
|
||||
CmdRawScan
|
||||
CmdGetKeyTTL
|
||||
CmdRawGetKeyTTL
|
||||
CmdRawCompareAndSwap
|
||||
CmdRawChecksum
|
||||
|
||||
|
|
@ -118,6 +118,11 @@ const (
|
|||
CmdEmpty CmdType = 3072 + iota
|
||||
)
|
||||
|
||||
// CmdType aliases.
|
||||
const (
|
||||
CmdGetKeyTTL = CmdRawGetKeyTTL
|
||||
)
|
||||
|
||||
func (t CmdType) String() string {
|
||||
switch t {
|
||||
case CmdGet:
|
||||
|
|
@ -164,6 +169,10 @@ func (t CmdType) String() string {
|
|||
return "RawScan"
|
||||
case CmdRawChecksum:
|
||||
return "RawChecksum"
|
||||
case CmdRawGetKeyTTL:
|
||||
return "RawGetKeyTTL"
|
||||
case CmdRawCompareAndSwap:
|
||||
return "RawCompareAndSwap"
|
||||
case CmdUnsafeDestroyRange:
|
||||
return "UnsafeDestroyRange"
|
||||
case CmdRegisterLockObserver:
|
||||
|
|
@ -225,6 +234,10 @@ func (t CmdType) String() string {
|
|||
// Request wraps all kv/coprocessor requests.
|
||||
type Request struct {
|
||||
Type CmdType
|
||||
// Req is one of the request type defined in kvrpcpb.
|
||||
//
|
||||
// WARN: It may be read concurrently in batch-send-loop, so you should ONLY modify it via `AttachContext`,
|
||||
// otherwise there could be a risk of data race.
|
||||
Req interface{}
|
||||
kvrpcpb.Context
|
||||
ReadReplicaScope string
|
||||
|
|
@ -244,6 +257,9 @@ type Request struct {
|
|||
ReadType string
|
||||
// InputRequestSource is the input source of the request, if it's not empty, the final RequestSource sent to store will be attached with the retry info.
|
||||
InputRequestSource string
|
||||
|
||||
// rev represents the revision of the request, it's increased when `Req.Context` gets patched.
|
||||
rev uint32
|
||||
}
|
||||
|
||||
// NewRequest returns new kv rpc request.
|
||||
|
|
@ -731,104 +747,31 @@ type MPPStreamResponse struct {
|
|||
Lease
|
||||
}
|
||||
|
||||
//go:generate bash gen.sh
|
||||
|
||||
// AttachContext sets the request context to the request,
|
||||
// return false if encounter unknown request type.
|
||||
// Parameter `rpcCtx` use `kvrpcpb.Context` instead of `*kvrpcpb.Context` to avoid concurrent modification by shallow copy.
|
||||
func AttachContext(req *Request, rpcCtx kvrpcpb.Context) bool {
|
||||
ctx := &rpcCtx
|
||||
cmd := req.Type
|
||||
// CmdCopStream and CmdCop share the same request type.
|
||||
if cmd == CmdCopStream {
|
||||
cmd = CmdCop
|
||||
}
|
||||
if patchCmdCtx(req, cmd, ctx) {
|
||||
return true
|
||||
}
|
||||
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 CmdGetKeyTTL:
|
||||
req.RawGetKeyTTL().Context = ctx
|
||||
case CmdRawCompareAndSwap:
|
||||
req.RawCompareAndSwap().Context = ctx
|
||||
case CmdRawChecksum:
|
||||
req.RawChecksum().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
|
||||
// Dispatching MPP tasks don't need a region context, because it's a request for store but not region.
|
||||
case CmdMPPTask:
|
||||
case CmdMPPConn:
|
||||
case CmdMPPCancel:
|
||||
case CmdMPPAlive:
|
||||
|
||||
case CmdMvccGetByKey:
|
||||
req.MvccGetByKey().Context = ctx
|
||||
case CmdMvccGetByStartTs:
|
||||
req.MvccGetByStartTs().Context = ctx
|
||||
case CmdSplitRegion:
|
||||
req.SplitRegion().Context = ctx
|
||||
// Empty command doesn't need a region context.
|
||||
case CmdEmpty:
|
||||
req.SplitRegion().Context = ctx
|
||||
case CmdTxnHeartBeat:
|
||||
req.TxnHeartBeat().Context = ctx
|
||||
case CmdCheckTxnStatus:
|
||||
req.CheckTxnStatus().Context = ctx
|
||||
case CmdCheckSecondaryLocks:
|
||||
req.CheckSecondaryLocks().Context = ctx
|
||||
case CmdFlashbackToVersion:
|
||||
req.FlashbackToVersion().Context = ctx
|
||||
case CmdPrepareFlashbackToVersion:
|
||||
req.PrepareFlashbackToVersion().Context = ctx
|
||||
case CmdFlush:
|
||||
req.Flush().Context = ctx
|
||||
case CmdBufferBatchGet:
|
||||
req.BufferBatchGet().Context = ctx
|
||||
|
||||
default:
|
||||
return false
|
||||
}
|
||||
|
|
|
|||
|
|
@ -35,8 +35,14 @@
|
|||
package tikvrpc
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"math/rand"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/pingcap/kvproto/pkg/coprocessor"
|
||||
"github.com/pingcap/kvproto/pkg/kvrpcpb"
|
||||
"github.com/pingcap/kvproto/pkg/tikvpb"
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
|
@ -47,3 +53,87 @@ func TestBatchResponse(t *testing.T) {
|
|||
assert.Nil(t, batchResp)
|
||||
assert.NotNil(t, err)
|
||||
}
|
||||
|
||||
// https://github.com/pingcap/tidb/issues/51921
|
||||
func TestTiDB51921(t *testing.T) {
|
||||
for _, r := range []*Request{
|
||||
NewRequest(CmdGet, &kvrpcpb.GetRequest{}),
|
||||
NewRequest(CmdScan, &kvrpcpb.ScanRequest{}),
|
||||
NewRequest(CmdPrewrite, &kvrpcpb.PrewriteRequest{}),
|
||||
NewRequest(CmdPessimisticLock, &kvrpcpb.PessimisticLockRequest{}),
|
||||
NewRequest(CmdPessimisticRollback, &kvrpcpb.PessimisticRollbackRequest{}),
|
||||
NewRequest(CmdCommit, &kvrpcpb.CommitRequest{}),
|
||||
NewRequest(CmdCleanup, &kvrpcpb.CleanupRequest{}),
|
||||
NewRequest(CmdBatchGet, &kvrpcpb.BatchGetRequest{}),
|
||||
NewRequest(CmdBatchRollback, &kvrpcpb.BatchRollbackRequest{}),
|
||||
NewRequest(CmdScanLock, &kvrpcpb.ScanLockRequest{}),
|
||||
NewRequest(CmdResolveLock, &kvrpcpb.ResolveLockRequest{}),
|
||||
NewRequest(CmdGC, &kvrpcpb.GCRequest{}),
|
||||
NewRequest(CmdDeleteRange, &kvrpcpb.DeleteRangeRequest{}),
|
||||
NewRequest(CmdRawGet, &kvrpcpb.RawGetRequest{}),
|
||||
NewRequest(CmdRawBatchGet, &kvrpcpb.RawBatchGetRequest{}),
|
||||
NewRequest(CmdRawPut, &kvrpcpb.RawPutRequest{}),
|
||||
NewRequest(CmdRawBatchPut, &kvrpcpb.RawBatchPutRequest{}),
|
||||
NewRequest(CmdRawDelete, &kvrpcpb.RawDeleteRequest{}),
|
||||
NewRequest(CmdRawBatchDelete, &kvrpcpb.RawBatchDeleteRequest{}),
|
||||
NewRequest(CmdRawDeleteRange, &kvrpcpb.RawDeleteRangeRequest{}),
|
||||
NewRequest(CmdRawScan, &kvrpcpb.RawScanRequest{}),
|
||||
NewRequest(CmdRawGetKeyTTL, &kvrpcpb.RawGetKeyTTLRequest{}),
|
||||
NewRequest(CmdRawCompareAndSwap, &kvrpcpb.RawCASRequest{}),
|
||||
NewRequest(CmdRawChecksum, &kvrpcpb.RawChecksumRequest{}),
|
||||
NewRequest(CmdUnsafeDestroyRange, &kvrpcpb.UnsafeDestroyRangeRequest{}),
|
||||
NewRequest(CmdRegisterLockObserver, &kvrpcpb.RegisterLockObserverRequest{}),
|
||||
NewRequest(CmdCheckLockObserver, &kvrpcpb.CheckLockObserverRequest{}),
|
||||
NewRequest(CmdRemoveLockObserver, &kvrpcpb.RemoveLockObserverRequest{}),
|
||||
NewRequest(CmdPhysicalScanLock, &kvrpcpb.PhysicalScanLockRequest{}),
|
||||
NewRequest(CmdCop, &coprocessor.Request{}),
|
||||
NewRequest(CmdCopStream, &coprocessor.Request{}),
|
||||
NewRequest(CmdBatchCop, &coprocessor.BatchRequest{}),
|
||||
NewRequest(CmdMvccGetByKey, &kvrpcpb.MvccGetByKeyRequest{}),
|
||||
NewRequest(CmdMvccGetByStartTs, &kvrpcpb.MvccGetByStartTsRequest{}),
|
||||
NewRequest(CmdSplitRegion, &kvrpcpb.SplitRegionRequest{}),
|
||||
NewRequest(CmdTxnHeartBeat, &kvrpcpb.TxnHeartBeatRequest{}),
|
||||
NewRequest(CmdCheckTxnStatus, &kvrpcpb.CheckTxnStatusRequest{}),
|
||||
NewRequest(CmdCheckSecondaryLocks, &kvrpcpb.CheckSecondaryLocksRequest{}),
|
||||
NewRequest(CmdFlashbackToVersion, &kvrpcpb.FlashbackToVersionRequest{}),
|
||||
NewRequest(CmdPrepareFlashbackToVersion, &kvrpcpb.PrepareFlashbackToVersionRequest{}),
|
||||
NewRequest(CmdFlush, &kvrpcpb.FlushRequest{}),
|
||||
NewRequest(CmdBufferBatchGet, &kvrpcpb.BufferBatchGetRequest{}),
|
||||
} {
|
||||
req := r
|
||||
t.Run(fmt.Sprintf("%s#%d", req.Type.String(), req.Type), func(t *testing.T) {
|
||||
if req.ToBatchCommandsRequest() == nil {
|
||||
t.Skipf("%s doesn't support batch commands", req.Type.String())
|
||||
}
|
||||
done := make(chan struct{})
|
||||
cmds := make(chan *tikvpb.BatchCommandsRequest_Request, 8)
|
||||
wg := sync.WaitGroup{}
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for {
|
||||
select {
|
||||
case <-done:
|
||||
close(cmds)
|
||||
return
|
||||
default:
|
||||
// mock relocate and retry
|
||||
AttachContext(req, kvrpcpb.Context{RegionId: rand.Uint64()})
|
||||
cmds <- req.ToBatchCommandsRequest()
|
||||
}
|
||||
}
|
||||
}()
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
for cmd := range cmds {
|
||||
// mock send and marshal in batch-send-loop
|
||||
cmd.Marshal()
|
||||
}
|
||||
}()
|
||||
|
||||
time.Sleep(time.Second / 4)
|
||||
close(done)
|
||||
wg.Wait()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue