From 74c5b700c1a875af2091cd3b88c89e419e74e1cc Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Tue, 9 Jun 2026 22:37:59 +0800 Subject: [PATCH 01/33] transaction: Support file based transaction Signed-off-by: Ping Yu --- config/client.go | 21 + integration_tests/txn_file_test.go | 287 ++++++ kv/variables.go | 15 +- metrics/metrics.go | 44 + metrics/shortcuts.go | 18 + tikv/split_region.go | 55 ++ txnkv/transaction/2pc.go | 42 +- txnkv/transaction/2pc_test.go | 55 ++ txnkv/transaction/txn.go | 24 + txnkv/transaction/txn_file.go | 1312 ++++++++++++++++++++++++++++ txnkv/transaction/txn_file_test.go | 295 +++++++ txnkv/txnlock/lock_resolver.go | 15 +- 12 files changed, 2177 insertions(+), 6 deletions(-) create mode 100644 integration_tests/txn_file_test.go create mode 100644 txnkv/transaction/txn_file.go create mode 100644 txnkv/transaction/txn_file_test.go diff --git a/config/client.go b/config/client.go index c7d434bc34..429ba21628 100644 --- a/config/client.go +++ b/config/client.go @@ -121,6 +121,21 @@ type TiKVClient struct { // RUV2 is the RU v2 TiKV-side weights used to calculate TiKV RU values from ExecDetailsV2.RuV2. RUV2 RUV2TiKVConfig `toml:"ru-v2" json:"ru-v2"` + + // TxnChunkWriterAddr is the address of the txn chunk writer for file-based txn. + TxnChunkWriterAddr string `toml:"txn-chunk-writer-addr" json:"txn-chunk-writer-addr"` + // TxnChunkWriterConcurrency is the concurrency to request the txn chunk writer for file-based txn. + TxnChunkWriterConcurrency uint `toml:"txn-chunk-writer-concurrency" json:"txn-chunk-writer-concurrency"` + // TxnChunkMaxSize is the maximum size of a txn chunk of file-based txn. + TxnChunkMaxSize uint64 `toml:"txn-chunk-max-size" json:"txn-chunk-max-size"` + // TxnFileMinMutationSize is the minimum size of mutations to use file-based txn. + TxnFileMinMutationSize uint64 `toml:"txn-file-min-mutation-size" json:"txn-file-min-mutation-size"` + // TxnFileRUDiscountRatio is the discount ratio of resource unit for file-based txn. + // Will be ignored if it's <= 0 or >= 1. + TxnFileRUDiscountRatio float64 `toml:"txn-file-ru-discount-ratio" json:"txn-file-ru-discount-ratio"` + // TxnFileRequestSourceWhitelist is the whitelist of request source types (RequestSource.RequestSourceType) that can use file-based txn. + // For internal requests only. External requests can always use file-based txn. + TxnFileRequestSourceWhitelist []string `toml:"txn-file-request-source-whitelist" json:"txn-file-request-source-whitelist"` } // RUV2TiKVConfig is the configuration for RU v2 TiKV-side weight calculation. @@ -232,6 +247,12 @@ func DefaultTiKVClient() TiKVClient { MaxConcurrencyRequestLimit: DefMaxConcurrencyRequestLimit, EnableReplicaSelectorV2: true, RUV2: DefaultRUV2TiKVConfig(), + + TxnChunkWriterConcurrency: 4, + TxnChunkMaxSize: 128 * 1024 * 1024, + TxnFileMinMutationSize: 16 * 1024 * 1024, + TxnFileRUDiscountRatio: 0.125, // filed-based txn costs 1/8 RU of normal txn. + TxnFileRequestSourceWhitelist: []string{}, } } diff --git a/integration_tests/txn_file_test.go b/integration_tests/txn_file_test.go new file mode 100644 index 0000000000..e25a4b7a10 --- /dev/null +++ b/integration_tests/txn_file_test.go @@ -0,0 +1,287 @@ +// Copyright 2021 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. + +// NOTE: The code in this file is based on code from the +// TiDB project, licensed under the Apache License v 2.0 +// +// https://github.com/pingcap/tidb/tree/cc5e161ac06827589c4966674597c137cc9e809c/store/tikv/tests/prewrite_test.go +// + +// Copyright 2020 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, +// 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. + +package tikv_test + +import ( + "context" + "encoding/json" + "net/http" + "net/http/httptest" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/pingcap/failpoint" + "github.com/pingcap/kvproto/pkg/kvrpcpb" + "github.com/stretchr/testify/require" + "github.com/tikv/client-go/v2/config" + "github.com/tikv/client-go/v2/kv" + "github.com/tikv/client-go/v2/testutils" + "github.com/tikv/client-go/v2/tikv" + "github.com/tikv/client-go/v2/tikvrpc" + "github.com/tikv/client-go/v2/txnkv/transaction" +) + +func TestTxnFilePrewriteTxnSize(t *testing.T) { + require := require.New(t) + const maxChunkSize = 1024 + + var chunkIDCounter atomic.Uint64 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + id := chunkIDCounter.Add(1) + resp, _ := json.Marshal(map[string]uint64{"chunk_id": id}) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(resp) + })) + defer srv.Close() + + origCfg := config.GetGlobalConfig() + newCfg := *origCfg + newCfg.TiKVClient.TxnChunkWriterAddr = srv.Listener.Addr().String() + newCfg.TiKVClient.TxnChunkMaxSize = maxChunkSize + newCfg.TiKVClient.TxnFileMinMutationSize = 1 + config.StoreGlobalConfig(&newCfg) + defer config.StoreGlobalConfig(origCfg) + + client, cluster, pdClient, err := testutils.NewMockTiKV("", nil) + require.Nil(err) + _, _, regionID := testutils.BootstrapWithSingleStore(cluster) + store, err := tikv.NewTestTiKVStore(client, pdClient, nil, nil, 0) + require.Nil(err) + defer store.Close() + + type capturedPrewrite struct { + txnFileChunks []uint64 + txnSize uint64 + } + var mu sync.Mutex + var captured []capturedPrewrite + + hook := func(req *tikvrpc.Request) { + if req.Type != tikvrpc.CmdPrewrite { + return + } + inner := req.Req.(*kvrpcpb.PrewriteRequest) + if len(inner.TxnFileChunks) == 0 { + return + } + chunks := make([]uint64, len(inner.TxnFileChunks)) + copy(chunks, inner.TxnFileChunks) + mu.Lock() + captured = append(captured, capturedPrewrite{ + txnFileChunks: chunks, + txnSize: inner.TxnSize, + }) + mu.Unlock() + } + + require.Nil(failpoint.Enable("tikvclient/beforeSendReqToRegion", "return")) + defer failpoint.Disable("tikvclient/beforeSendReqToRegion") + ctx := context.WithValue(context.Background(), "sendReqToRegionHook", hook) + + commitTxn := func(keys [][]byte) { + tx, err := store.Begin() + require.Nil(err) + txn := transaction.TxnProbe{KVTxn: tx} + + vars := *kv.DefaultVars + vars.TxnFileMinMutationSize = 1 + txn.SetVars(&vars) + + for _, key := range keys { + val := make([]byte, 64) + require.Nil(txn.Set(key, val)) + } + + // The mock environment is only used to inspect outgoing txn-file prewrite requests. + // Commit may fail later because the mock stack does not fully model txn-file follow-up behavior. + _ = txn.Commit(ctx) + } + + assertCaptured := func(expectedRequests int, expectedTxnSize uint64) { + mu.Lock() + defer mu.Unlock() + require.GreaterOrEqual(len(captured), expectedRequests) + for _, c := range captured { + require.NotEmpty(c.txnFileChunks) + require.Equal(expectedTxnSize, c.txnSize) + } + } + + // Single-region case: exact txn size should match the mutation count. + commitTxn([][]byte{[]byte("a"), []byte("b"), []byte("c"), []byte("d"), []byte("e")}) + assertCaptured(1, 5) + + // Split the single region into two. With the default large chunk size, one chunk spans both + // regions, so each region batch should conservatively reuse the full chunk entry count. + newRegionID := cluster.AllocID() + newPeerID := cluster.AllocID() + cluster.Split(regionID, newRegionID, []byte("m"), []uint64{newPeerID}, newPeerID) + + mu.Lock() + captured = nil + mu.Unlock() + + commitTxn([][]byte{[]byte("a"), []byte("b"), []byte("x"), []byte("y"), []byte("z")}) + assertCaptured(2, 5) +} + +func TestTxnFilePrewriteTxnSizeAfterRegionRegroup(t *testing.T) { + require := require.New(t) + const maxChunkSize = 1024 + + var chunkIDCounter atomic.Uint64 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + id := chunkIDCounter.Add(1) + resp, _ := json.Marshal(map[string]uint64{"chunk_id": id}) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(resp) + })) + defer srv.Close() + + origCfg := config.GetGlobalConfig() + newCfg := *origCfg + newCfg.TiKVClient.TxnChunkWriterAddr = srv.Listener.Addr().String() + newCfg.TiKVClient.TxnChunkMaxSize = maxChunkSize + newCfg.TiKVClient.TxnFileMinMutationSize = 1 + config.StoreGlobalConfig(&newCfg) + defer config.StoreGlobalConfig(origCfg) + + client, cluster, pdClient, err := testutils.NewMockTiKV("", nil) + require.Nil(err) + _, peerID, regionID := testutils.BootstrapWithSingleStore(cluster) + store, err := tikv.NewTestTiKVStore(client, pdClient, nil, nil, 0) + require.Nil(err) + defer store.Close() + + type capturedPrewrite struct { + txnFileChunks []uint64 + txnSize uint64 + isRetry bool + regionErr bool + } + var mu sync.Mutex + var captured []capturedPrewrite + + hook := func(req *tikvrpc.Request, resp *tikvrpc.Response, sendErr error) { + if req.Type != tikvrpc.CmdPrewrite { + return + } + inner, ok := req.Req.(*kvrpcpb.PrewriteRequest) + if !ok || len(inner.TxnFileChunks) == 0 { + return + } + if sendErr != nil { + return + } + chunks := make([]uint64, len(inner.TxnFileChunks)) + copy(chunks, inner.TxnFileChunks) + var regionErr bool + if resp != nil { + if respRegionErr, err := resp.GetRegionError(); err == nil && respRegionErr != nil { + regionErr = true + } + } + mu.Lock() + captured = append(captured, capturedPrewrite{ + txnFileChunks: chunks, + txnSize: inner.TxnSize, + isRetry: req.Context.IsRetryRequest, + regionErr: regionErr, + }) + mu.Unlock() + } + + require.Nil(failpoint.Enable("tikvclient/mockRetrySendReqToRegion", "1*return(true)->return(false)")) + defer failpoint.Disable("tikvclient/mockRetrySendReqToRegion") + require.Nil(failpoint.Enable("tikvclient/invalidCacheAndRetry", "1*off->pause")) + defer failpoint.Disable("tikvclient/invalidCacheAndRetry") + require.Nil(failpoint.Enable("tikvclient/afterSendReqToRegion", "return")) + defer failpoint.Disable("tikvclient/afterSendReqToRegion") + ctx := context.WithValue(context.Background(), "sendReqToRegionFinishHook", hook) + + tx, err := store.Begin() + require.Nil(err) + txn := transaction.TxnProbe{KVTxn: tx} + + vars := *kv.DefaultVars + vars.TxnFileMinMutationSize = 1 + txn.SetVars(&vars) + + for _, key := range [][]byte{[]byte("a"), []byte("z")} { + val := make([]byte, 64) + require.Nil(txn.Set(key, val)) + } + + done := make(chan struct{}) + go func() { + _ = txn.Commit(ctx) + close(done) + }() + + time.Sleep(3 * time.Second) + cluster.Split(regionID, cluster.AllocID(), []byte("h"), []uint64{peerID}, peerID) + require.Nil(failpoint.Disable("tikvclient/invalidCacheAndRetry")) + <-done + + mu.Lock() + defer mu.Unlock() + require.GreaterOrEqual(len(captured), 4, "expected initial send, stale-region retry, and regrouped region requests") + regionErrRetries := 0 + successfulRetryPrewrites := 0 + for _, c := range captured { + require.NotEmpty(c.txnFileChunks) + require.Equal(uint64(2), c.txnSize) + if c.isRetry && c.regionErr { + regionErrRetries++ + } + if c.isRetry && !c.regionErr { + successfulRetryPrewrites++ + } + } + require.GreaterOrEqual(regionErrRetries, 1, "expected a retry-marked txn-file prewrite to hit a region error after the split") + require.GreaterOrEqual(successfulRetryPrewrites, 2, "expected regrouped retry prewrites to reach both post-split regions") +} diff --git a/kv/variables.go b/kv/variables.go index cae78c9c59..7646d0e4a1 100644 --- a/kv/variables.go +++ b/kv/variables.go @@ -49,14 +49,23 @@ type Variables struct { // When its value is 0, it's not killed // When its value is not 0, it's killed, the value indicates concrete reason. Killed *uint32 + + // DisableTxnFile specifies whether file-based txn is disabled. + DisableTxnFile bool + + // TxnFileMinMutationSize is the minimum size of mutations to use file-based txn. + // When its value is 0, use the config of "txn-file-min-mutation-size". + TxnFileMinMutationSize uint64 } // NewVariables create a new Variables instance with default values. func NewVariables(killed *uint32) *Variables { return &Variables{ - BackoffLockFast: DefBackoffLockFast, - BackOffWeight: DefBackOffWeight, - Killed: killed, + BackoffLockFast: DefBackoffLockFast, + BackOffWeight: DefBackOffWeight, + Killed: killed, + DisableTxnFile: false, + TxnFileMinMutationSize: 0, } } diff --git a/metrics/metrics.go b/metrics/metrics.go index 63d10a161f..8eefdb4cbd 100644 --- a/metrics/metrics.go +++ b/metrics/metrics.go @@ -137,6 +137,12 @@ var ( TiKVReadRequestBytes *prometheus.SummaryVec TiKVTxnLagCommitTSWaitHistogram *prometheus.HistogramVec TiKVTxnLagCommitTSAttemptHistogram *prometheus.HistogramVec + + + TiKVTxnFileRequestCounter *prometheus.CounterVec + TiKVTxnFileWriteBytes *prometheus.CounterVec + TiKVTxnFileMutationSizeHistogram *prometheus.HistogramVec + TiKVTxnFileDuration *prometheus.HistogramVec ) // Label constants. @@ -1055,6 +1061,44 @@ func initMetrics(namespace, subsystem string, constLabels prometheus.Labels) { ConstLabels: constLabels, }, []string{LblResult}) + TiKVTxnFileRequestCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Subsystem: subsystem, + Name: "txn_file_requests", + Help: "Counter of file-based transactions requests.", + ConstLabels: constLabels, + }, []string{LblType}) + + TiKVTxnFileWriteBytes = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Subsystem: subsystem, + Name: "txn_file_write_bytes", + Help: "Counter of file-based transactions write bytes.", + ConstLabels: constLabels, + }, []string{LblScope}) + + TiKVTxnFileMutationSizeHistogram = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: namespace, + Subsystem: subsystem, + Name: "txn_file_mutation_size", + Buckets: prometheus.ExponentialBuckets(1<<20, 2, 17), // 1MB ~ 64GB + Help: "Histogram of file-based transactions mutation bytes.", + ConstLabels: constLabels, + }, []string{LblScope}) + + TiKVTxnFileDuration = prometheus.NewHistogramVec( + prometheus.HistogramOpts{ + Namespace: namespace, + Subsystem: subsystem, + Name: "txn_file_duration", + Buckets: prometheus.ExponentialBuckets(0.001, 2, 20), // 1ms ~ 524s + Help: "Duration of executing file-based transactions.", + ConstLabels: constLabels, + }, []string{LblScope}) + initShortcuts() storeMetricVecList.Store(&storeMetrics) } diff --git a/metrics/shortcuts.go b/metrics/shortcuts.go index f2e795f69c..8144cc9c6d 100644 --- a/metrics/shortcuts.go +++ b/metrics/shortcuts.go @@ -204,6 +204,15 @@ var ( LagCommitTSWaitHistogramWithError prometheus.Observer LagCommitTSAttemptHistogramWithOK prometheus.Observer LagCommitTSAttemptHistogramWithError prometheus.Observer + + TxnFileRequestsOk prometheus.Counter + TxnFileRequestsError prometheus.Counter + TxnFileWriteBytesInternal prometheus.Counter + TxnFileWriteBytesGeneral prometheus.Counter + TxnFileMutationSizeInternal prometheus.Observer + TxnFileMutationSizeGeneral prometheus.Observer + TxnFileDurationInternal prometheus.Observer + TxnFileDurationGeneral prometheus.Observer ) func initShortcuts() { @@ -377,4 +386,13 @@ func initShortcuts() { LagCommitTSWaitHistogramWithError = TiKVTxnLagCommitTSWaitHistogram.WithLabelValues("err") LagCommitTSAttemptHistogramWithOK = TiKVTxnLagCommitTSAttemptHistogram.WithLabelValues("ok") LagCommitTSAttemptHistogramWithError = TiKVTxnLagCommitTSAttemptHistogram.WithLabelValues("err") + + TxnFileRequestsOk = TiKVTxnFileRequestCounter.WithLabelValues("ok") + TxnFileRequestsError = TiKVTxnFileRequestCounter.WithLabelValues("err") + TxnFileWriteBytesInternal = TiKVTxnFileWriteBytes.WithLabelValues(LblInternal) + TxnFileWriteBytesGeneral = TiKVTxnFileWriteBytes.WithLabelValues(LblGeneral) + TxnFileMutationSizeInternal = TiKVTxnFileMutationSizeHistogram.WithLabelValues(LblInternal) + TxnFileMutationSizeGeneral = TiKVTxnFileMutationSizeHistogram.WithLabelValues(LblGeneral) + TxnFileDurationInternal = TiKVTxnFileDuration.WithLabelValues(LblInternal) + TxnFileDurationGeneral = TiKVTxnFileDuration.WithLabelValues(LblGeneral) } diff --git a/tikv/split_region.go b/tikv/split_region.go index 7807094a66..7a74a15ef2 100644 --- a/tikv/split_region.go +++ b/tikv/split_region.go @@ -53,6 +53,7 @@ import ( "github.com/tikv/client-go/v2/internal/logutil" "github.com/tikv/client-go/v2/tikvrpc" "github.com/tikv/client-go/v2/txnkv/rangetask" + "github.com/tikv/client-go/v2/txnkv/txnlock" "github.com/tikv/client-go/v2/util" "github.com/tikv/client-go/v2/util/redact" "github.com/tikv/pd/client/opt" @@ -179,6 +180,20 @@ func (s *KVStore) batchSendSingleRegion(bo *Backoffer, batch kvrpc.Batch, scatte } spResp := resp.Resp.(*kvrpcpb.SplitRegionResponse) + + keyErrs := spResp.GetErrors() + if len(keyErrs) > 0 { + err := s.handleSplitRegionKeyErrors(bo, keyErrs) + if err != nil { + batchResp.Error = err + return batchResp + } + resp, err = s.splitBatchRegionsReq(bo, batch.Keys, scatter, tableID) + batchResp.Response = resp + batchResp.Error = err + return batchResp + } + regions := spResp.GetRegions() if len(regions) > 0 { // Divide a region into n, one of them may not need to be scattered, @@ -223,6 +238,46 @@ func (s *KVStore) batchSendSingleRegion(bo *Backoffer, batch kvrpc.Batch, scatte return batchResp } +func (s *KVStore) handleSplitRegionKeyErrors(bo *Backoffer, keyErrs []*kvrpcpb.KeyError) error { + var ( + locks []*txnlock.Lock + startTS uint64 = math.MaxUint64 // Set as MaxUint64 and check txn status will not push the minCommiTS. + ) + for _, keyErr := range keyErrs { + lock, err1 := txnlock.ExtractLockFromKeyErr(keyErr) + if err1 != nil { + // Split region should return key error of locked only. + return err1 + } + logutil.Logger(bo.GetCtx()).Info("split region encounters lock", zap.Stringer("lock", lock)) + locks = append(locks, lock) + } + + token := s.GetLockResolver().RecordResolvingLocks(locks, startTS) + defer s.GetLockResolver().ResolveLocksDone(startTS, token) + + resolveLockOpts := txnlock.ResolveLocksOptions{ + CallerStartTS: startTS, + Locks: locks, + } + resolveLockRes, err := s.GetLockResolver().ResolveLocksWithOpts(bo, resolveLockOpts) + if err != nil { + return errors.WithStack(err) + } + msBeforeExpired := resolveLockRes.TTL + if msBeforeExpired > 0 { + err = bo.BackoffWithCfgAndMaxSleep( + retry.BoTxnLock, + int(msBeforeExpired), + errors.Errorf("split region lockedKeys: %d", len(locks)), + ) + if err != nil { + return errors.WithStack(err) + } + } + return nil +} + const ( splitRegionBackoff = 20000 maxSplitRegionsBackoff = 120000 diff --git a/txnkv/transaction/2pc.go b/txnkv/transaction/2pc.go index f92051ae31..151598e0ec 100644 --- a/txnkv/transaction/2pc.go +++ b/txnkv/transaction/2pc.go @@ -40,6 +40,7 @@ import ( errors2 "errors" "math" "math/rand" + "sort" "strconv" "strings" "sync" @@ -203,6 +204,8 @@ type twoPhaseCommitter struct { primaryOp kvrpcpb.Op pipelinedStart, pipelinedEnd []byte } + + txnFileCtx txnFileCtx } type memBufferMutations struct { @@ -337,6 +340,38 @@ type CommitterMutations interface { NeedConstraintCheckInPrewrite(i int) bool } +func MutationsHasDataInRange(mutations CommitterMutations, start []byte, end []byte) ([]byte /* firstDataKey */, bool) { + isInRange := func(pos int) bool { + return pos < mutations.Len() && (len(end) == 0 || bytes.Compare(mutations.GetKey(pos), end) < 0) + } + isOpForWrite := func(op kvrpcpb.Op) bool { + return op != kvrpcpb.Op_CheckNotExists && + op != kvrpcpb.Op_Lock && + op != kvrpcpb.Op_PessimisticLock + } + + pos := sort.Search(mutations.Len(), func(i int) bool { + return bytes.Compare(mutations.GetKey(i), start) >= 0 + }) + if isInRange(pos) { + var firstDataKey []byte + for { + // Always return primary key if it's in the range. + if pos == 0 || isOpForWrite(mutations.GetOp(pos)) { + firstDataKey = mutations.GetKey(pos) + break + } + + pos++ + if !isInRange(pos) { + break + } + } + return firstDataKey, true + } + return nil, false +} + // PlainMutations contains transaction operations. type PlainMutations struct { ops []kvrpcpb.Op @@ -781,6 +816,9 @@ func (c *twoPhaseCommitter) primary() []byte { if c.mutations != nil { return c.mutations.GetKey(0) } + if c.txnFileCtx.slice.Len() > 0 { + return c.txnFileCtx.slice.chunkRanges[0].smallest + } return nil } return c.primaryKey @@ -1370,7 +1408,7 @@ func keepAlive( ) startTime := time.Now() _, stopHeartBeat, err := sendTxnHeartBeat( - bo, c.store, primaryKey, c.startTS, newTTL, c.minCommitTSMgr.get(), + bo, c.store, primaryKey, c.startTS, newTTL, c.minCommitTSMgr.get(), c.txnFileCtx.slice.Len() > 0, ) if err != nil { keepFail++ @@ -1509,12 +1547,14 @@ func sendTxnHeartBeat( primary []byte, startTS, ttl uint64, minCommitTS uint64, + isTxnFile bool, ) (newTTL uint64, stopHeartBeat bool, err error) { req := tikvrpc.NewRequest(tikvrpc.CmdTxnHeartBeat, &kvrpcpb.TxnHeartBeatRequest{ PrimaryLock: primary, StartVersion: startTS, AdviseLockTtl: ttl, MinCommitTs: minCommitTS, + IsTxnFile: isTxnFile, }) for { loc, err := store.GetRegionCache().LocateKey(bo, primary) diff --git a/txnkv/transaction/2pc_test.go b/txnkv/transaction/2pc_test.go index 3e1a80d1cf..6f0760404e 100644 --- a/txnkv/transaction/2pc_test.go +++ b/txnkv/transaction/2pc_test.go @@ -35,8 +35,10 @@ package transaction import ( + "fmt" "testing" + "github.com/pingcap/kvproto/pkg/kvrpcpb" "github.com/stretchr/testify/assert" ) @@ -122,3 +124,56 @@ func TestMinCommitTsManager(t *testing.T) { }, ) } + +func TestMutationsHasDataInRange(t *testing.T) { + assert := assert.New(t) + + iToKey := func(i int) []byte { + if i < 0 { + return nil + } + return []byte(fmt.Sprintf("%04d", i)) + } + + muts := NewPlainMutations(10) + for i := 10; i < 20; i += 2 { + key := iToKey(i) + var op kvrpcpb.Op + if i%4 == 0 { + op = kvrpcpb.Op_CheckNotExists + } else { + op = kvrpcpb.Op_Put + } + muts.Push(op, key, key, false, false, false, false) + } + + type Case struct { + start int + end int + expectd bool + firstKey int + } + cases := []Case{ + {-1, -1, true, 10}, + {-1, 5, false, -1}, + {0, 10, false, -1}, + {0, 11, true, 10}, + {0, 30, true, 10}, + {0, -1, true, 10}, + {10, 20, true, 10}, + {15, 16, false, -1}, + {15, 17, true, -1}, + {15, -1, true, 18}, + {20, 30, false, -1}, + {21, 30, false, -1}, + {21, -1, false, -1}, + } + + for _, c := range cases { + firstKey, got := MutationsHasDataInRange(&muts, iToKey(c.start), iToKey(c.end)) + assert.Equal(c.expectd, got) + if got { + assert.Equal(iToKey(c.firstKey), firstKey) + } + } +} diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index 52bb657fba..f693b8edf3 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -227,6 +227,9 @@ type KVTxn struct { firstAttemptTS uint64 backoffCnt int } + + // disableTxnFile indicates this transaction should NOT use file-based txn. + disableTxnFile bool } // NewTiKVTxn creates a new KVTxn. @@ -787,6 +790,10 @@ func (txn *KVTxn) GetScope() string { return txn.scope } +func (txn *KVTxn) DisableTxnFile() { + txn.disableTxnFile = true +} + // Commit commits the transaction operations to KV store. func (txn *KVTxn) Commit(ctx context.Context) error { if span := opentracing.SpanFromContext(ctx); span != nil && span.Tracer() != nil { @@ -891,6 +898,23 @@ func (txn *KVTxn) Commit(ctx context.Context) error { } } }() + + useTxnFile := !txn.disableTxnFile + if useTxnFile { + useTxnFile, err = committer.useTxnFile(ctx) + if err != nil { + return err + } + } + if useTxnFile { + err = committer.executeTxnFile(ctx) + // TODO: fall back to normal 2PC when tikv-worker is unavailable. + if val == nil || sessionID > 0 { + txn.onCommitted(err) + } + return err + } + // latches disabled // pessimistic transaction should also bypass latch. // transaction with pipelined memdb should also bypass latch. diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go new file mode 100644 index 0000000000..c620edca8f --- /dev/null +++ b/txnkv/transaction/txn_file.go @@ -0,0 +1,1312 @@ +// 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. + +package transaction + +import ( + "bytes" + "context" + "crypto/tls" + "encoding/binary" + "encoding/json" + "fmt" + "hash/crc32" + "io" + "net/http" + "slices" + "sort" + "strings" + "sync" + "time" + + "github.com/pingcap/kvproto/pkg/errorpb" + "github.com/pingcap/kvproto/pkg/kvrpcpb" + "github.com/pingcap/kvproto/pkg/metapb" + "github.com/pkg/errors" + "github.com/tikv/client-go/v2/config" + "github.com/tikv/client-go/v2/config/retry" + tikverr "github.com/tikv/client-go/v2/error" + "github.com/tikv/client-go/v2/internal/apicodec" + "github.com/tikv/client-go/v2/internal/client" + "github.com/tikv/client-go/v2/internal/locate" + "github.com/tikv/client-go/v2/internal/logutil" + "github.com/tikv/client-go/v2/internal/resourcecontrol" + "github.com/tikv/client-go/v2/kv" + "github.com/tikv/client-go/v2/metrics" + "github.com/tikv/client-go/v2/tikvrpc" + "github.com/tikv/client-go/v2/txnkv/txnlock" + "github.com/tikv/client-go/v2/util" + "github.com/tikv/client-go/v2/util/redact" + resourceControlClient "github.com/tikv/pd/client/resource_group/controller" + atomicutil "go.uber.org/atomic" + "go.uber.org/zap" +) + +var ( + // BuildTxnFileMaxBackoff is max sleep time (in millisecond) to build TxnFile. + BuildTxnFileMaxBackoff = atomicutil.NewUint64(60000) + + buildChunkErrMsg string = "txn file: build chunk failed" +) + +const ( + PreSplitRegionChunks = 4 + + // MaxTxnChunkSizeInParallel is the max parallel size when prewrite/commit txn chunks. + MaxTxnChunkSizeInParallel uint64 = 4 << 30 // 4GB +) + +type txnFileCtx struct { + slice txnChunkSlice +} + +type chunkBatch struct { + txnChunkSlice + region *locate.KeyLocation + sampleKeys [][]byte + isPrimary bool +} + +func (b chunkBatch) String() string { + return fmt.Sprintf("chunkBatch{region: %d, isPrimary: %t, txnChunkSlice: %v}", + b.region.Region.GetID(), b.isPrimary, b.txnChunkSlice.chunkIDs) +} + +func (b *chunkBatch) getSampleKeys() [][]byte { + return b.sampleKeys +} + +func (b *chunkBatch) getBatchTxnSize() uint64 { + var batchTxnSize uint64 + // Chunk range is region-grouped by reference, so if one chunk spans multiple + // regions, each overlapping batch conservatively reuses the full chunk entry + // count rather than computing an exact per-region intersection. + for _, ran := range b.chunkRanges { + batchTxnSize += ran.entries + } + return batchTxnSize +} + +// txnChunkSlice should be sorted by txnChunkRange.smallest and no overlapping. +type txnChunkSlice struct { + chunkIDs []uint64 + chunkRanges []txnChunkRange +} + +func (s txnChunkSlice) String() string { + slice := make([]string, len(s.chunkRanges)) + for i, ran := range s.chunkRanges { + slice[i] = fmt.Sprintf("txnChunkSlice{%v: [%s, %s]}", s.chunkIDs[i], redact.Key(ran.smallest), redact.Key(ran.biggest)) + } + return fmt.Sprintf("[%s]", strings.Join(slice, ", ")) +} + +func (s *txnChunkSlice) Smallest() []byte { + if len(s.chunkRanges) == 0 { + return nil + } + return s.chunkRanges[0].smallest +} + +func (s *txnChunkSlice) Biggest() []byte { + if len(s.chunkRanges) == 0 { + return nil + } + return s.chunkRanges[len(s.chunkRanges)-1].biggest +} + +func (cs *txnChunkSlice) appendSlice(other *txnChunkSlice) { + cs.chunkIDs = append(cs.chunkIDs, other.chunkIDs...) + cs.chunkRanges = append(cs.chunkRanges, other.chunkRanges...) +} + +func (cs *txnChunkSlice) append(chunkID uint64, chunkRange txnChunkRange) { + cs.chunkIDs = append(cs.chunkIDs, chunkID) + cs.chunkRanges = append(cs.chunkRanges, chunkRange) +} + +func (cs *txnChunkSlice) Len() int { + return len(cs.chunkIDs) +} + +func (cs *txnChunkSlice) Swap(i, j int) { + cs.chunkIDs[i], cs.chunkIDs[j] = cs.chunkIDs[j], cs.chunkIDs[i] + cs.chunkRanges[i], cs.chunkRanges[j] = cs.chunkRanges[j], cs.chunkRanges[i] +} + +func (cs *txnChunkSlice) Less(i, j int) bool { + return bytes.Compare(cs.chunkRanges[i].smallest, cs.chunkRanges[j].smallest) < 0 +} + +func (cs *txnChunkSlice) sortAndDedup() { + if len(cs.chunkIDs) <= 1 { + return + } + + sort.Sort(cs) + + newIDs := cs.chunkIDs[:1] + newRanges := cs.chunkRanges[:1] + for i := 1; i < len(cs.chunkIDs); i++ { + if cs.chunkIDs[i] != newIDs[len(newIDs)-1] { + newIDs = append(newIDs, cs.chunkIDs[i]) + newRanges = append(newRanges, cs.chunkRanges[i]) + } + } + cs.chunkIDs = newIDs + cs.chunkRanges = newRanges +} + +// []chunkBatch is sorted by region.StartKey. +// Note: regions may be overlapping. +func (cs *txnChunkSlice) groupToBatches(c *locate.RegionCache, bo *retry.Backoffer, mutations CommitterMutations) ([]chunkBatch, error) { + // Do not use `locate.RegionVerID` as map key to avoid grouping chunks to different batches when `confVer` changes. + type batchMapKey struct { + regionID uint64 + regionVer uint64 + } + batchMap := make(map[batchMapKey]*chunkBatch) + for i, chunkRange := range cs.chunkRanges { + chunkID := cs.chunkIDs[i] + + regions, firstKeys, err := chunkRange.getOverlapRegions(c, bo, mutations) + if err != nil { + return nil, errors.WithStack(err) + } + + for j, r := range regions { + firstKey := firstKeys[j] + + bk := batchMapKey{regionID: r.Region.GetID(), regionVer: r.Region.GetVer()} + if batchMap[bk] == nil { + batchMap[bk] = &chunkBatch{ + region: r, + sampleKeys: make([][]byte, 0, 1), + } + } + + batch := batchMap[bk] + batch.append(chunkID, chunkRange) + if len(firstKey) > 0 { + batch.sampleKeys = append(batch.sampleKeys, firstKey) + } + } + } + + batches := make([]chunkBatch, 0, len(batchMap)) + for _, batch := range batchMap { + batches = append(batches, *batch) + } + sort.Slice(batches, func(i, j int) bool { + // Sort by chunks first, and then by region, to make sure that primary key is in the first batch: + // 1. Different batches may contain the same chunks. + // 2. Different batches may have regions with same start key (if region merge happens during grouping). + // 3. Bigger batches may have regions with smaller start key (if region merge happens during grouping). + cmp := bytes.Compare(batches[i].Smallest(), batches[j].Smallest()) + if cmp == 0 { + return bytes.Compare(batches[i].region.StartKey, batches[j].region.StartKey) < 0 + } + return cmp < 0 + }) + + logutil.Logger(bo.GetCtx()).Debug("txn file group to batches", zap.Stringers("batches", batches)) + return batches, nil +} + +type txnChunkRange struct { + smallest []byte + biggest []byte + // entries tracks the number of serialized mutations in the chunk. When a chunk + // overlaps multiple regions, each overlapping region batch reuses this full + // count as a conservative upper bound for PrewriteRequest.TxnSize. + entries uint64 +} + +func (r txnChunkRange) String() string { + return fmt.Sprintf("txnChunkRange[%s,%s](entries=%d)", redact.Key(r.smallest), redact.Key(r.biggest), r.entries) +} + +func newTxnChunkRange(smallest []byte, biggest []byte, entries uint64) txnChunkRange { + return txnChunkRange{ + smallest: smallest, + biggest: biggest, + entries: entries, + } +} + +func (r *txnChunkRange) getOverlapRegions(c *locate.RegionCache, bo *retry.Backoffer, mutations CommitterMutations) ([]*locate.KeyLocation, [][]byte, error) { + regions := make([]*locate.KeyLocation, 0) + firstKeys := make([][]byte, 0) + startKey := r.smallest + exclusiveBiggest := kv.NextKey(r.biggest) + for bytes.Compare(startKey, r.biggest) <= 0 { + loc, err := c.LocateKey(bo, startKey) + if err != nil { + logutil.Logger(bo.GetCtx()).Error("locate key failed", zap.Error(err), zap.String("startKey", redact.Key(startKey))) + return nil, nil, errors.Wrap(err, "locate key failed") + } + firstKey, ok := MutationsHasDataInRange( + mutations, + util.GetMaxStartKey(r.smallest, loc.StartKey), + util.GetMinEndKey(exclusiveBiggest, loc.EndKey), + ) + if ok { + regions = append(regions, loc) + firstKeys = append(firstKeys, firstKey) + } + if len(loc.EndKey) == 0 { + break + } + startKey = loc.EndKey + } + return regions, firstKeys, nil +} + +type txnFileAction interface { + executeBatch(c *twoPhaseCommitter, bo *retry.Backoffer, batch chunkBatch) (*tikvrpc.Response, error) + onPrimarySuccess(c *twoPhaseCommitter) + extractKeyError(resp *tikvrpc.Response) *kvrpcpb.KeyError + asyncExecuteSecondaries() bool + String() string +} + +type txnFilePrewriteAction struct{} + +var _ txnFileAction = (*txnFilePrewriteAction)(nil) + +func (a txnFilePrewriteAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backoffer, batch chunkBatch) (*tikvrpc.Response, error) { + primaryLock := c.txnFileCtx.slice.chunkRanges[0].smallest + req := tikvrpc.NewRequest(tikvrpc.CmdPrewrite, &kvrpcpb.PrewriteRequest{ + StartVersion: c.startTS, + PrimaryLock: primaryLock, + LockTtl: c.lockTTL, + MaxCommitTs: c.maxCommitTS, + AssertionLevel: kvrpcpb.AssertionLevel_Off, + TxnFileChunks: batch.chunkIDs, + TxnSize: batch.getBatchTxnSize(), + }, kvrpcpb.Context{ + Priority: c.priority, + SyncLog: c.syncLog, + ResourceGroupTag: c.resourceGroupTag, + DiskFullOpt: c.diskFullOpt, + TxnSource: c.txnSource, + MaxExecutionDurationMs: uint64(client.ReadTimeoutMedium.Milliseconds()), + RequestSource: c.txn.GetRequestSource(), + ResourceControlContext: &kvrpcpb.ResourceControlContext{ + ResourceGroupName: c.resourceGroupName, + }, + }) + sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.GetOracle()) + var resolvingRecordToken *int + + for { + if batch.isPrimary { + // Refresh the primary txn-file TTL to observe the latest elapsed time. + c.lockTTL = txnLockTTL(c.txn.startTime, c.txnSize) + req.Prewrite().LockTtl = c.lockTTL + } + resp, _, err := sender.SendReq(bo, req, batch.region.Region, client.ReadTimeoutMedium) + if err != nil { + return nil, err + } + if resp.Resp == nil { + return nil, errors.WithStack(tikverr.ErrBodyMissing) + } + prewriteResp := resp.Resp.(*kvrpcpb.PrewriteResponse) + regionErr := prewriteResp.GetRegionError() + if regionErr != nil { + // For other region error and the fake region error, backoff because + // there's something wrong. + // For the real EpochNotMatch error, don't backoff. + if regionErr.GetEpochNotMatch() == nil || locate.IsFakeRegionError(regionErr) { + err = bo.Backoff(retry.BoRegionMiss, errors.New(regionErr.String())) + if err != nil { + return resp, err + } + } + if regionErr.GetDiskFull() != nil { + return resp, errors.New(regionErr.String()) + } + if len(batch.sampleKeys) > 0 { + loc, err := c.store.GetRegionCache().LocateKey(bo, batch.sampleKeys[0]) + if err != nil { + return nil, err + } + if loc.Region == batch.region.Region { + continue + } + } + + return resp, err + } + + keyErrs := prewriteResp.GetErrors() + if len(keyErrs) == 0 { + return resp, nil + } + var locks []*txnlock.Lock + for _, keyErr := range keyErrs { + // Check already exists error + if alreadyExist := keyErr.GetAlreadyExist(); alreadyExist != nil { + e := &tikverr.ErrKeyExist{AlreadyExist: alreadyExist} + return nil, c.extractKeyExistsErr(e) + } + + // Extract lock from key error + lock, err1 := txnlock.ExtractLockFromKeyErr(keyErr) + if err1 != nil { + return nil, err1 + } + logutil.Logger(bo.GetCtx()).Info( + "prewrite txn file encounters lock", + zap.Uint64("session", c.sessionID), + zap.Uint64("txnID", c.startTS), + zap.Stringer("lock", lock), + ) + // If an optimistic transaction encounters a lock with larger TS, this transaction will certainly + // fail due to a WriteConflict error. So we can construct and return an error here early. + // Pessimistic transactions don't need such an optimization. If this key needs a pessimistic lock, + // TiKV will return a PessimisticLockNotFound error directly if it encounters a different lock. Otherwise, + // TiKV returns lock.TTL = 0, and we still need to resolve the lock. + if lock.TxnID > c.startTS { + return nil, tikverr.NewErrWriteConflictWithArgs( + c.startTS, + lock.TxnID, + 0, + lock.Key, + kvrpcpb.WriteConflict_Optimistic, + ) + } + locks = append(locks, lock) + } + if resolvingRecordToken == nil { + token := c.store.GetLockResolver().RecordResolvingLocks(locks, c.startTS) + resolvingRecordToken = &token + defer c.store.GetLockResolver().ResolveLocksDone(c.startTS, *resolvingRecordToken) + } else { + c.store.GetLockResolver().UpdateResolvingLocks(locks, c.startTS, *resolvingRecordToken) + } + resolveLockOpts := txnlock.ResolveLocksOptions{ + CallerStartTS: c.startTS, + Locks: locks, + Detail: &c.getDetail().ResolveLock, + } + resolveLockRes, err := c.store.GetLockResolver().ResolveLocksWithOpts(bo, resolveLockOpts) + if err != nil { + return nil, err + } + msBeforeExpired := resolveLockRes.TTL + if msBeforeExpired > 0 { + err = bo.BackoffWithCfgAndMaxSleep( + retry.BoTxnLock, + int(msBeforeExpired), + errors.Errorf("2PC txn file prewrite lockedKeys: %d", len(locks)), + ) + if err != nil { + return nil, err + } + } + } +} + +func (a txnFilePrewriteAction) onPrimarySuccess(c *twoPhaseCommitter) { + c.run(c, nil, false) +} + +func (a txnFilePrewriteAction) extractKeyError(resp *tikvrpc.Response) *kvrpcpb.KeyError { + prewriteResp, _ := resp.Resp.(*kvrpcpb.PrewriteResponse) + errs := prewriteResp.GetErrors() + if len(errs) > 0 { + return errs[0] + } + return nil +} + +func (a txnFilePrewriteAction) asyncExecuteSecondaries() bool { + return false +} + +func (a txnFilePrewriteAction) String() string { + return "txnFilePrewrite" +} + +type txnFileCommitAction struct{} + +var _ txnFileAction = (*txnFileCommitAction)(nil) + +func (a txnFileCommitAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backoffer, batch chunkBatch) (*tikvrpc.Response, error) { + req := tikvrpc.NewRequest(tikvrpc.CmdCommit, &kvrpcpb.CommitRequest{ + Keys: batch.getSampleKeys(), // To help detect duplicated request. + StartVersion: c.startTS, + CommitVersion: c.commitTS, + IsTxnFile: true, + }, kvrpcpb.Context{ + Priority: c.priority, + SyncLog: c.syncLog, + ResourceGroupTag: c.resourceGroupTag, + DiskFullOpt: c.diskFullOpt, + TxnSource: c.txnSource, + MaxExecutionDurationMs: uint64(client.ReadTimeoutMedium.Milliseconds()), + RequestSource: c.txn.GetRequestSource(), + ResourceControlContext: &kvrpcpb.ResourceControlContext{ + ResourceGroupName: c.resourceGroupName, + }, + }) + sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.GetOracle()) + for { + resp, _, err := sender.SendReq(bo, req, batch.region.Region, client.ReadTimeoutMedium) + if batch.isPrimary && sender.GetRPCError() != nil { + c.setUndeterminedErr(errors.WithStack(sender.GetRPCError())) + } + // Unexpected error occurs, return it. + if err != nil { + return nil, err + } + if resp.Resp == nil { + return nil, errors.WithStack(tikverr.ErrBodyMissing) + } + commitResp := resp.Resp.(*kvrpcpb.CommitResponse) + if keyErr := commitResp.GetError(); keyErr != nil { + if rejected := keyErr.GetCommitTsExpired(); rejected != nil { + logutil.Logger(bo.GetCtx()).Info("2PC commitTS rejected by TiKV, retry with a newer commitTS", + zap.Uint64("txnStartTS", c.startTS), + zap.Stringer("info", logutil.Hex(rejected))) + + // Do not retry for a txn which has a too large MinCommitTs + // 3600000 << 18 = 943718400000 + if rejected.MinCommitTs-rejected.AttemptedCommitTs > 943718400000 { + return nil, errors.Errorf("2PC MinCommitTS is too large, we got MinCommitTS: %d, and AttemptedCommitTS: %d", + rejected.MinCommitTs, rejected.AttemptedCommitTs) + } + + // Update commit ts and retry. + commitTS, err1 := c.store.GetTimestampWithRetry(bo, c.txn.GetScope()) + if err1 != nil { + logutil.Logger(bo.GetCtx()).Warn("2PC get commitTS failed", + zap.Error(err1), + zap.Uint64("txnStartTS", c.startTS)) + return nil, err1 + } + + c.mu.Lock() + c.commitTS = commitTS + c.mu.Unlock() + // Update the commitTS of the request and retry. + req.Commit().CommitVersion = commitTS + continue + } + return nil, tikverr.ExtractKeyErr(keyErr) + } + return resp, nil + } +} + +func (a txnFileCommitAction) onPrimarySuccess(c *twoPhaseCommitter) { + c.mu.Lock() + c.mu.committed = true + c.mu.Unlock() +} + +func (a txnFileCommitAction) extractKeyError(resp *tikvrpc.Response) *kvrpcpb.KeyError { + commitResp, _ := resp.Resp.(*kvrpcpb.CommitResponse) + return commitResp.GetError() +} + +func (a txnFileCommitAction) asyncExecuteSecondaries() bool { + return true +} + +func (a txnFileCommitAction) String() string { + return "txnFileCommit" +} + +type txnFileRollbackAction struct{} + +var _ txnFileAction = (*txnFileRollbackAction)(nil) + +func (a txnFileRollbackAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backoffer, batch chunkBatch) (*tikvrpc.Response, error) { + req := tikvrpc.NewRequest(tikvrpc.CmdBatchRollback, &kvrpcpb.BatchRollbackRequest{ + Keys: batch.getSampleKeys(), // To help detect duplicated request. + StartVersion: c.startTS, + IsTxnFile: true, + }, kvrpcpb.Context{ + Priority: c.priority, + SyncLog: c.syncLog, + ResourceGroupTag: c.resourceGroupTag, + DiskFullOpt: c.diskFullOpt, + TxnSource: c.txnSource, + MaxExecutionDurationMs: uint64(client.ReadTimeoutShort.Milliseconds()), + RequestSource: c.txn.GetRequestSource(), + ResourceControlContext: &kvrpcpb.ResourceControlContext{ + ResourceGroupName: c.resourceGroupName, + }, + }) + sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.GetOracle()) + resp, _, err1 := sender.SendReq(bo, req, batch.region.Region, client.ReadTimeoutShort) + if err1 != nil { + return nil, err1 + } + return resp, nil +} + +func (a txnFileRollbackAction) onPrimarySuccess(_ *twoPhaseCommitter) { +} + +func (a txnFileRollbackAction) extractKeyError(resp *tikvrpc.Response) *kvrpcpb.KeyError { + rollbackResp, _ := resp.Resp.(*kvrpcpb.BatchRollbackResponse) + return rollbackResp.GetError() +} + +func (a txnFileRollbackAction) asyncExecuteSecondaries() bool { + return true +} + +func (a txnFileRollbackAction) String() string { + return "txnFileRollback" +} + +type step struct { + name string + dur time.Duration +} + +func (s step) String() string { + return fmt.Sprintf("%s:%s", s.name, s.dur.String()) +} + +func (c *twoPhaseCommitter) executeTxnFile(ctx context.Context) (err error) { + if val, err := util.EvalFailpoint("injectErrorOnExecTxnFile"); err == nil { + errVal := val.(string) + if errVal == "writeConflict" { + err = tikverr.NewErrWriteConflictWithArgs( + c.startTS, + c.startTS+1, + 0, + c.primaryKey, + kvrpcpb.WriteConflict_Optimistic, + ) + } else { + err = errors.New("injected error in executeTxnFile: " + errVal) + } + return err + } + + start := time.Now() + steps := make([]step, 0) + stepDone := func(name string) { + now := time.Now() + s := step{name: name, dur: now.Sub(start)} + steps = append(steps, s) + start = now + } + + defer func() { + // Always clean up all written keys if the txn does not commit. + c.mu.RLock() + committed := c.mu.committed + undetermined := c.mu.undeterminedErr != nil + c.mu.RUnlock() + if !committed && !undetermined { + if c.txnFileCtx.slice.Len() > 0 { + err1 := c.executeTxnFileAction(retry.NewBackofferWithVars(ctx, int(CommitMaxBackoff), c.txn.vars), c.txnFileCtx.slice, txnFileRollbackAction{}) + if err1 != nil { + logutil.Logger(ctx).Error("txn file: rollback on error failed", zap.Error(err1)) + } + } + c.reportFailureMetrics() + } else { + c.reportSuccessMetrics(steps) + } + c.txn.commitTS = c.commitTS + + logutil.Logger(ctx).Info("execute txn file finished", + zap.Uint64("startTS", c.startTS), + zap.Uint64("commitTS", c.commitTS), + zap.Error(err), + zap.String("requestSource", c.txn.GetRequestSource()), + zap.Stringers("steps", steps)) + }() + + logutil.Logger(ctx).Debug("execute txn file", zap.Uint64("startTS", c.startTS)) + + buildBo := retry.NewBackofferWithVars(ctx, int(BuildTxnFileMaxBackoff.Load()), c.txn.vars) + + rcInterceptor := client.ResourceControlInterceptor.Load() + var ruDetails *util.RUDetails + if detail := ctx.Value(util.RUDetailsCtxKey); detail != nil { + ruDetails = detail.(*util.RUDetails) + } + reqInfo, err := c.beforeExecuteTxnFile(buildBo, rcInterceptor, ruDetails) + if err != nil { + return + } + + err = c.buildTxnFiles(buildBo, c.mutations) + stepDone("build") + if err != nil { + return + } + + err = c.preSplitTxnFileRegions(buildBo) + stepDone("pre-split") + if err != nil { + return + } + + prewriteBo := retry.NewBackofferWithVars(ctx, int(PrewriteMaxBackoff.Load()), c.txn.vars) + err = c.executeTxnFileAction(prewriteBo, c.txnFileCtx.slice, txnFilePrewriteAction{}) + stepDone("prewrite") + if err != nil { + return + } + + commitBo := retry.NewBackofferWithVars(ctx, int(CommitMaxBackoff), c.txn.vars) + c.commitTS, err = c.store.GetTimestampWithRetry(commitBo, c.txn.GetScope()) + if err != nil { + return + } + err = c.executeTxnFileAction(commitBo, c.txnFileCtx.slice, txnFileCommitAction{}) + stepDone("commit") + if err != nil { + return + } + + err = c.afterExecuteTxnFile(rcInterceptor, reqInfo, ruDetails) + return +} + +func (c *twoPhaseCommitter) executeTxnFileSlice(bo *retry.Backoffer, chunkSlice txnChunkSlice, batches []chunkBatch, action txnFileAction) (txnChunkSlice, error) { + var err error + var regionErrChunks txnChunkSlice + + chunksCount := chunkSlice.Len() + if batches == nil { + batches, err = chunkSlice.groupToBatches(c.store.GetRegionCache(), bo, c.mutations) + if err != nil { + return regionErrChunks, errors.Wrap(err, "txn file: group to batches failed") + } + } + + if len(batches) == 1 { + regionErrSlice, err := c.executeTxnFileSliceSingleBatch(bo, batches[0], action) + if err != nil { + return regionErrChunks, err + } else if regionErrSlice != nil { + regionErrChunks.appendSlice(regionErrSlice) + } + return regionErrChunks, nil + } + + type result struct { + regionErrSlice *txnChunkSlice + err error + } + ch := make(chan result, len(batches)) + + exitCh := make(chan struct{}) + defer close(exitCh) + + bo, cancel := bo.Fork() + defer cancel() + // we do not try to return early in the commit/rollback phase so + // we can commit/rollback as many regions as possible. + _, returnEarly := action.(txnFilePrewriteAction) + + rateLim := len(batches) + cnf := config.GetGlobalConfig() + if rateLim > cnf.CommitterConcurrency { + rateLim = cnf.CommitterConcurrency + } + maxChunksInParallel := int(MaxTxnChunkSizeInParallel / cnf.TiKVClient.TxnChunkMaxSize) // 32 by default + if chunksCount > maxChunksInParallel { + rateLim = maxChunksInParallel + } + rateLimiter := util.NewRateLimit(rateLim) + go func() { + for _, batch := range batches { + batch := batch + bo := bo.Clone() + if exit := rateLimiter.GetToken(exitCh); !exit { + go func() { + defer rateLimiter.PutToken() + regionErrSlice, err := c.executeTxnFileSliceSingleBatch(bo, batch, action) + ch <- result{regionErrSlice, err} + }() + } + } + }() + + err = nil + for i := 0; i < len(batches); i++ { + r := <-ch + if r.err != nil { + if returnEarly { + return regionErrChunks, r.err + } else if err == nil { + err = r.err + } + } else if r.regionErrSlice != nil { + regionErrChunks.appendSlice(r.regionErrSlice) + } + } + regionErrChunks.sortAndDedup() + return regionErrChunks, err +} + +func (c *twoPhaseCommitter) executeTxnFileSliceSingleBatch(bo *retry.Backoffer, batch chunkBatch, action txnFileAction) (*txnChunkSlice, error) { + resp, err1 := action.executeBatch(c, bo, batch) + logutil.Logger(bo.GetCtx()).Debug("txn file: execute batch finished", + zap.Uint64("startTS", c.startTS), + zap.Any("batch", batch), + zap.Stringer("action", action), + zap.Error(err1)) + if err1 != nil { + return nil, err1 + } + if keyErr := action.extractKeyError(resp); keyErr != nil { + if alreadyExist := keyErr.GetAlreadyExist(); alreadyExist != nil { + e := &tikverr.ErrKeyExist{AlreadyExist: alreadyExist} + return nil, c.extractKeyExistsErr(e) + } + lock, err2 := txnlock.ExtractLockFromKeyErr(keyErr) + if err2 != nil { + return nil, err2 + } + if lock.TxnID > c.startTS { + return nil, tikverr.NewErrWriteConflictWithArgs( + c.startTS, + lock.TxnID, + 0, + lock.Key, + kvrpcpb.WriteConflict_Optimistic, + ) + } + } + regionErr, err1 := resp.GetRegionError() + if err1 != nil { + return nil, err1 + } + if regionErr != nil { + logutil.Logger(bo.GetCtx()).Debug("txn file: execute batch failed, region error", + zap.Uint64("startTS", c.startTS), + zap.Stringer("action", action), + zap.Any("batch", batch), + zap.Stringer("regionErr", regionErr)) + return &batch.txnChunkSlice, nil + } + return nil, nil +} + +func (c *twoPhaseCommitter) executeTxnFileSliceWithRetry(bo *retry.Backoffer, chunkSlice txnChunkSlice, batches []chunkBatch, action txnFileAction) error { + currentChunks := chunkSlice + currentBatches := batches + for { + var regionErrChunks txnChunkSlice + regionErrChunks, err := c.executeTxnFileSlice(bo, currentChunks, currentBatches, action) + if err != nil { + return errors.WithStack(err) + } + if regionErrChunks.Len() == 0 { + return nil + } + logutil.Logger(bo.GetCtx()).Debug("txn file meet region errors", zap.Stringer("regionErrChunks", regionErrChunks)) + currentChunks = regionErrChunks + currentBatches = nil + err = bo.Backoff(retry.BoRegionMiss, errors.Errorf("txn file: execute failed, region miss")) + if err != nil { + return errors.WithStack(err) + } + } +} + +func (c *twoPhaseCommitter) executeTxnFilePrimaryBatch(bo *retry.Backoffer, firstBatch chunkBatch, action txnFileAction) (regionErr *errorpb.Error, err error) { + if !firstBatch.region.Contains(c.primary()) { + logutil.Logger(bo.GetCtx()).Error("txn file: primary out of first batch", + zap.Uint64("startTS", c.startTS), + zap.String("primary", redact.Key(c.primary())), + zap.Stringer("action", action), + zap.Stringer("first batch", firstBatch)) + return nil, fmt.Errorf("txn file: primary out of first batch") + } + + firstBatch.isPrimary = true + resp, err := action.executeBatch(c, bo, firstBatch) + logutil.Logger(bo.GetCtx()).Debug("txn file: execute primary batch finished", + zap.Uint64("startTS", c.startTS), + zap.String("primary", redact.Key(c.primary())), + zap.Stringer("action", action), + zap.Stringer("batch", firstBatch), + zap.Error(err)) + if err != nil { + return nil, errors.WithStack(err) + } + regionErr, err = resp.GetRegionError() + if err != nil { + return nil, errors.WithStack(err) + } + if regionErr != nil { + return regionErr, nil + } + action.onPrimarySuccess(c) + return nil, nil +} + +func (c *twoPhaseCommitter) executeTxnFileAction(bo *retry.Backoffer, chunkSlice txnChunkSlice, action txnFileAction) error { + for { + batches, err := chunkSlice.groupToBatches(c.store.GetRegionCache(), bo, c.mutations) + if err != nil { + return errors.Wrap(err, "txn file: group to batches failed") + } + + regionErr, err := c.executeTxnFilePrimaryBatch(bo, batches[0], action) + if err != nil { + return errors.WithStack(err) + } + if regionErr != nil { + errBo := bo.Backoff(retry.BoRegionMiss, errors.Wrap(errors.New(regionErr.String()), "txn file: execute primary batch failed")) + if errBo != nil { + return errors.WithStack(errBo) + } + continue + } + + secondaries := batches[1:] + if len(secondaries) == 0 { + return nil + } + if !action.asyncExecuteSecondaries() { + return c.executeTxnFileSliceWithRetry(bo, chunkSlice, secondaries, action) + } + + c.store.WaitGroup().Add(1) + errGo := c.store.Go(func() { + defer c.store.WaitGroup().Done() + err := c.executeTxnFileSliceWithRetry(bo, chunkSlice, secondaries, action) + logutil.Logger(bo.GetCtx()).Debug("txn file: async execute secondaries finished", + zap.Uint64("startTS", c.startTS), + zap.Stringer("action", action), + zap.Error(err)) + if err != nil { + logutil.Logger(bo.GetCtx()).Warn("txn file: async execute secondaries failed", + zap.Uint64("startTS", c.startTS), zap.Stringer("action", action), zap.Error(err)) + } + }) + if errGo != nil { + c.store.WaitGroup().Done() + logutil.Logger(bo.GetCtx()).Warn("txn file: create goroutine failed", + zap.Uint64("startTS", c.startTS), + zap.Stringer("action", action), + zap.Error(errGo)) + } + return nil + } +} + +func (c *twoPhaseCommitter) buildTxnFiles(bo *retry.Backoffer, mutations CommitterMutations) error { + bo, cancel := bo.Fork() + defer cancel() + + cfg := config.GetGlobalConfig() + maxTxnChunkSize := int(cfg.TiKVClient.TxnChunkMaxSize) + capacity := c.txn.Size() + c.txn.Len()*7 + 4 + totalChunks := (capacity + maxTxnChunkSize - 1) / maxTxnChunkSize + if capacity > maxTxnChunkSize { + capacity = maxTxnChunkSize + } + concurrency := int(cfg.TiKVClient.TxnChunkWriterConcurrency) + if concurrency > totalChunks { + concurrency = totalChunks + } + + writer, err := newChunkWriterClient(c.getKeyspaceID()) + if err != nil { + return errors.Wrap(err, "new chunk writer client failed") + } + + results := make([]buildChunkResult, 0, totalChunks) + resultCh := make(chan buildChunkResult, concurrency) + + totalSize := 0 + inflightChunks := 0 + buf := make([]byte, 0, capacity) + chunkSmallest := mutations.GetKey(0) + chunkEntries := uint64(0) + for i := 0; i < mutations.Len(); i++ { + key := mutations.GetKey(i) + op := mutations.GetOp(i) + val := mutations.GetValue(i) + entrySize := 2 + len(key) + 1 + 4 + len(val) + if len(buf) > 0 && len(buf)+entrySize+4 > cap(buf) { + totalSize += len(buf) + inflightChunks += 1 + ran := newTxnChunkRange(chunkSmallest, mutations.GetKey(i-1), chunkEntries) + writer.asyncBuildChunk(bo, buf, ran, resultCh) + chunkSmallest = key + chunkEntries = 0 + buf = make([]byte, 0, capacity) + + if inflightChunks >= concurrency { + r := <-resultCh + if r.err != nil { + logutil.Logger(bo.GetCtx()).Error(buildChunkErrMsg, zap.Error(err)) + return errors.Wrap(r.err, buildChunkErrMsg) + } + results = append(results, r) + inflightChunks -= 1 + } + } + buf = binary.LittleEndian.AppendUint16(buf, uint16(len(key))) + buf = append(buf, key...) + buf = append(buf, byte(op)) + buf = binary.LittleEndian.AppendUint32(buf, uint32(len(val))) + buf = append(buf, val...) + chunkEntries++ + } + if len(buf) > 0 { + totalSize += len(buf) + inflightChunks += 1 + ran := newTxnChunkRange(chunkSmallest, mutations.GetKey(mutations.Len()-1), chunkEntries) + writer.asyncBuildChunk(bo, buf, ran, resultCh) + } + + for i := 0; i < inflightChunks; i++ { + r := <-resultCh + if r.err != nil { + logutil.Logger(bo.GetCtx()).Error(buildChunkErrMsg, zap.Error(err)) + return errors.Wrap(r.err, buildChunkErrMsg) + } + results = append(results, r) + } + sort.Slice(results, func(i, j int) bool { + return bytes.Compare(results[i].chunkRange.smallest, results[j].chunkRange.smallest) < 0 + }) + for _, r := range results { + c.txnFileCtx.slice.append(r.chunkId, r.chunkRange) + } + + logutil.Logger(bo.GetCtx()).Info("build txn files", + zap.Uint64("startTS", c.startTS), + zap.Int("mutationsLen", mutations.Len()), + zap.Int("totalChunksSize", totalSize), + zap.Any("chunkIDs", c.txnFileCtx.slice.chunkIDs)) + return nil +} + +func (c *twoPhaseCommitter) getKeyspaceID() apicodec.KeyspaceID { + return c.store.GetRegionCache().Codec().GetKeyspaceID() +} + +func (c *twoPhaseCommitter) useTxnFile(ctx context.Context) (bool, error) { + if c.txn == nil || c.txn.vars.DisableTxnFile { + return false, nil + } + conf := config.GetGlobalConfig() + minMutationSize := c.txn.vars.TxnFileMinMutationSize + if minMutationSize == 0 { + minMutationSize = conf.TiKVClient.TxnFileMinMutationSize + } + if c.txn.isInternal() { + // Relax the requirement for internal requests. + minMutationSize = minMutationSize / 2 + } + + if c.txn.isPessimistic || + len(conf.TiKVClient.TxnChunkWriterAddr) == 0 || + uint64(c.txn.GetMemBuffer().Size()) < minMutationSize || + !IsRequestSourceUseTxnFile(c.txn.RequestSource, conf) { + return false, nil + } + + logutil.Logger(ctx).Debug("transaction use txn file", + zap.Uint64("startTS", c.startTS), + zap.Int("size", c.txn.GetMemBuffer().Size()), + zap.Int("len", c.mutations.Len()), + zap.String("reqSource", c.txn.GetRequestSource()), + ) + return true, nil +} + +// IsRequestSourceUseTxnFile checks if the request source is allowed to use file-based txn on the configuration. +func IsRequestSourceUseTxnFile(reqSource *util.RequestSource, conf *config.Config) bool { + if reqSource.RequestSourceInternal { + return slices.Contains( + conf.TiKVClient.TxnFileRequestSourceWhitelist, + reqSource.RequestSourceType, + ) + } + return true +} + +func (c *twoPhaseCommitter) preSplitTxnFileRegions(bo *retry.Backoffer) error { + batches, err := c.txnFileCtx.slice.groupToBatches(c.store.GetRegionCache(), bo, c.mutations) + if err != nil { + return errors.Wrap(err, "group to batches failed") + } + var splitKeys [][]byte + for _, batch := range batches { + if batch.Len() > PreSplitRegionChunks { + for i := PreSplitRegionChunks; i < batch.Len(); i += PreSplitRegionChunks { + splitKeys = append(splitKeys, batch.chunkRanges[i].smallest) + } + } + } + if len(splitKeys) == 0 { + return nil + } + _, err = c.store.SplitRegions(bo.GetCtx(), splitKeys, false, nil) + return errors.Wrap(err, "pre split regions failed") +} + +func (c *twoPhaseCommitter) beforeExecuteTxnFile( + bo *retry.Backoffer, + rcInterceptor *resourceControlClient.ResourceGroupKVInterceptor, + ruDetails *util.RUDetails, +) (*resourcecontrol.RequestInfo, error) { + if rcInterceptor == nil { + return nil, nil + } + + ctx := bo.GetCtx() + regionCache := c.store.GetRegionCache() + + var region *locate.Region + for { + loc, err := regionCache.LocateKey(bo, c.primary()) + if err != nil { + return nil, errors.Wrap(err, "failed to locate primary") + } + + region = regionCache.GetCachedRegionWithRLock(loc.Region) + if region != nil { + break + } + + err = bo.Backoff(retry.BoRegionMiss, errors.New("cached region not found")) + if err != nil { + logutil.Logger(ctx).Error("txn file: cached region not found", + zap.String("key", redact.Key(c.primary())), + zap.Stringer("loc", loc)) + return nil, errors.WithStack(err) + } + } + + var replicaNumber int64 = 0 + for _, peer := range region.GetMeta().GetPeers() { + if peer.GetRole() == metapb.PeerRole_Voter { + replicaNumber++ + } + } + + writeBytes := int64(c.txn.Size()) + discountRatio := config.GetGlobalConfig().TiKVClient.TxnFileRUDiscountRatio + if discountRatio > 0.0 && discountRatio < 1.0 { + writeBytes = int64(float64(writeBytes) * discountRatio) + } + + reqInfo := resourcecontrol.NewRequestInfo( + writeBytes, + region.GetLeaderStoreID(), + replicaNumber, + false, + ) + + consumption, _ /* penalty */, waitDuration, _ /* priority */, err := (*rcInterceptor).OnRequestWait(ctx, c.resourceGroupName, reqInfo) + if err != nil { + return nil, errors.WithStack(err) + } + + if ruDetails != nil { + ruDetails.Update(consumption, waitDuration) + } + + return reqInfo, nil +} + +func (c *twoPhaseCommitter) afterExecuteTxnFile(rcInterceptor *resourceControlClient.ResourceGroupKVInterceptor, reqInfo *resourcecontrol.RequestInfo, ruDetails *util.RUDetails) error { + if rcInterceptor == nil { + return nil + } + + respInfo := &resourcecontrol.ResponseInfo{} + consumption, err := (*rcInterceptor).OnResponse(c.resourceGroupName, reqInfo, respInfo) + if err != nil { + return errors.WithStack(err) + } + if ruDetails != nil { + ruDetails.Update(consumption, time.Duration(0)) + } + + return nil +} + +func (c *twoPhaseCommitter) reportSuccessMetrics(steps []step) { + metrics.TwoPCTxnCounterOk.Inc() + metrics.TxnFileRequestsOk.Inc() + + mutationBytes := c.txn.GetMemBuffer().Size() + var dur time.Duration + for _, step := range steps { + dur += step.dur + } + + if c.txn.isInternal() { + metrics.TxnFileWriteBytesInternal.Add(float64(mutationBytes)) + metrics.TxnFileMutationSizeInternal.Observe(float64(mutationBytes)) + metrics.TxnFileDurationInternal.Observe(dur.Seconds()) + } else { + metrics.TxnFileWriteBytesGeneral.Add(float64(mutationBytes)) + metrics.TxnFileMutationSizeGeneral.Observe(float64(mutationBytes)) + metrics.TxnFileDurationGeneral.Observe(dur.Seconds()) + } +} + +func (c *twoPhaseCommitter) reportFailureMetrics() { + metrics.TwoPCTxnCounterError.Inc() + metrics.TxnFileRequestsError.Inc() +} + +var ( + once sync.Once + scheme string + cli *http.Client + errCli error +) + +func getHTTPClient() (*http.Client, error) { + once.Do(func() { + var ( + cfg = config.GetGlobalConfig() + timeout = time.Duration(BuildTxnFileMaxBackoff.Load()) * time.Millisecond + tlsConfig *tls.Config + ) + + scheme = "http://" + transport := &http.Transport{ + MaxIdleConns: 100, + MaxIdleConnsPerHost: 20, + } + if len(cfg.Security.ClusterSSLCA) != 0 { + scheme = "https://" + + tlsConfig, errCli = cfg.Security.ToTLSConfig() + if errCli != nil { + return + } + transport.TLSClientConfig = tlsConfig + transport.ForceAttemptHTTP2 = true + } + + cli = &http.Client{ + Timeout: timeout, + Transport: transport, + } + }) + return cli, errCli +} + +type chunkWriterClient struct { + cli *http.Client + serviceAddr string +} + +func newChunkWriterClient(keyspaceID apicodec.KeyspaceID) (*chunkWriterClient, error) { + client, err := getHTTPClient() + if err != nil { + return nil, errors.WithStack(err) + } + + cfg := config.GetGlobalConfig() + serviceAddr := fmt.Sprintf("%s%s/txn_chunk?keyspace_id=%v", scheme, cfg.TiKVClient.TxnChunkWriterAddr, keyspaceID) + return &chunkWriterClient{client, serviceAddr}, nil +} + +func (w *chunkWriterClient) buildChunk(bo *retry.Backoffer, buf []byte) (uint64, error) { + hash := crc32.New(crc32.MakeTable(crc32.IEEE)) + hash.Write(buf) + crc := hash.Sum32() + buf = binary.LittleEndian.AppendUint32(buf, crc) + + data, err := w.request(bo, "POST", buf) + if err != nil { + return 0, errors.WithStack(err) + } + v := struct { + ChunkId uint64 `json:"chunk_id"` + }{} + if err = json.Unmarshal(data, &v); err != nil { + return 0, errors.Wrapf(err, "unmarshal response %s", string(data)) + } + logutil.Logger(bo.GetCtx()).Debug("build txn file", zap.Int("size", len(buf)), zap.Uint64("chunkId", v.ChunkId)) + return v.ChunkId, nil +} + +func (w *chunkWriterClient) request(bo *retry.Backoffer, method string, data []byte) ([]byte, error) { + ctx := bo.GetCtx() + for { + if ctx.Err() != nil { + return nil, errors.WithStack(ctx.Err()) + } + + var body io.Reader + if len(data) > 0 { + body = bytes.NewReader(data) + } + + req, err := http.NewRequestWithContext(ctx, method, w.serviceAddr, body) + if err != nil { + return nil, errors.WithStack(err) + } + if body != nil { + req.Header.Set("Content-Type", "application/octet-stream") + } + + resp, err := w.cli.Do(req) + if err != nil { + logutil.Logger(ctx).Warn("request failed", zap.Error(err), zap.String("addr", w.serviceAddr)) + err = bo.Backoff(retry.BoTiKVRPC, errors.WithMessage(err, "request failed")) + if err != nil { + return nil, errors.WithStack(err) + } + continue + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + var bodyStr string + if data, err := io.ReadAll(resp.Body); err == nil { + bodyStr = string(data) + } + logutil.Logger(ctx).Warn("service error", zap.String("http status", resp.Status), zap.String("body", bodyStr)) + err = bo.Backoff(retry.BoTiKVServerBusy, fmt.Errorf("service error, http status %s", resp.Status)) + if err != nil { + return nil, errors.WithStack(err) + } + continue + } + data, err := io.ReadAll(resp.Body) + return data, errors.WithStack(err) + } +} + +type buildChunkResult struct { + chunkId uint64 + chunkRange txnChunkRange + err error +} + +func (w *chunkWriterClient) asyncBuildChunk(bo *retry.Backoffer, buf []byte, chunkRange txnChunkRange, resultCh chan<- buildChunkResult) { + go func() { + chunkId, err := w.buildChunk(bo.Clone(), buf) + resultCh <- buildChunkResult{chunkId, chunkRange, err} + }() +} diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go new file mode 100644 index 0000000000..9ae28a35eb --- /dev/null +++ b/txnkv/transaction/txn_file_test.go @@ -0,0 +1,295 @@ +// 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. + +package transaction + +import ( + "context" + "encoding/json" + "fmt" + "math/rand" + "net/http" + "net/http/httptest" + "slices" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/pingcap/kvproto/pkg/kvrpcpb" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tikv/client-go/v2/config" + "github.com/tikv/client-go/v2/config/retry" + "github.com/tikv/client-go/v2/internal/apicodec" + "github.com/tikv/client-go/v2/internal/client" + "github.com/tikv/client-go/v2/internal/latch" + "github.com/tikv/client-go/v2/internal/locate" + "github.com/tikv/client-go/v2/internal/unionstore" + tikv "github.com/tikv/client-go/v2/kv" + "github.com/tikv/client-go/v2/oracle" + "github.com/tikv/client-go/v2/testutils" + "github.com/tikv/client-go/v2/tikvrpc" + "github.com/tikv/client-go/v2/txnkv/txnlock" + "github.com/tikv/client-go/v2/util" +) + +func TestChunkSliceSortAndDedup(t *testing.T) { + assert := assert.New(t) + + genRndChunkIDs := func() []uint64 { + n := rand.Intn(10) + ids := make([]uint64, 0, n) + for i := 0; i < n; i++ { + ids = append(ids, uint64(rand.Intn(n+n/2+1))) + } + return ids + } + + for i := 0; i < 100; i++ { + ids := genRndChunkIDs() + t.Logf("ids: %v\n", ids) + + expected := make([]uint64, len(ids)) + copy(expected, ids) + slices.Sort(expected) + expected = slices.Compact(expected) + + chunkSlice := txnChunkSlice{ + chunkIDs: make([]uint64, 0, len(ids)), + chunkRanges: make([]txnChunkRange, 0, len(ids)), + } + for _, id := range ids { + chunkSlice.chunkIDs = append(chunkSlice.chunkIDs, id) + chunkSlice.chunkRanges = append(chunkSlice.chunkRanges, txnChunkRange{ + smallest: []byte(fmt.Sprintf("k%04d", id)), + biggest: []byte(fmt.Sprintf("k%04d_end", id)), + entries: id + 1, + }) + } + chunkSlice.sortAndDedup() + + assert.Equal(expected, chunkSlice.chunkIDs) + for j, id := range expected { + assert.Equal(fmt.Sprintf("k%04d", id), string(chunkSlice.chunkRanges[j].smallest), + "smallest mismatch at index %d", j) + assert.Equal(fmt.Sprintf("k%04d_end", id), string(chunkSlice.chunkRanges[j].biggest), + "biggest mismatch at index %d", j) + assert.Equal(id+1, chunkSlice.chunkRanges[j].entries, + "entries mismatch at index %d", j) + } + } +} + +func TestIsRequestSourceUseTxnFile(t *testing.T) { + assert := assert.New(t) + + cases := []struct { + reqSource *util.RequestSource + whitelist []string + expected bool + }{ + { + reqSource: &util.RequestSource{RequestSourceInternal: false}, + whitelist: []string{}, + expected: true, + }, + { + reqSource: &util.RequestSource{RequestSourceType: "ddl_modify_column", RequestSourceInternal: true}, + whitelist: []string{"ddl_modify_column"}, + expected: true, + }, + { + reqSource: &util.RequestSource{RequestSourceType: "ddl_modify_column", RequestSourceInternal: true}, + whitelist: []string{"ddl_alter_partition", "ddl_modify_column"}, + expected: true, + }, + { + reqSource: &util.RequestSource{RequestSourceType: "ddl_modify_column", RequestSourceInternal: true}, + whitelist: []string{}, + expected: false, + }, + { + reqSource: &util.RequestSource{RequestSourceType: "ddl_modify_column", RequestSourceInternal: true}, + whitelist: []string{"ddl_alter_partition"}, + expected: false, + }, + } + + for _, c := range cases { + conf := &config.Config{ + TiKVClient: config.TiKVClient{ + TxnFileRequestSourceWhitelist: c.whitelist, + }, + } + result := IsRequestSourceUseTxnFile(c.reqSource, conf) + assert.Equal(c.expected, result, "Expected %v for request source %v with whitelist %v", c.expected, c.reqSource.RequestSourceType, c.whitelist) + } +} + +// stubKVStore implements kvstore with only GetRegionCache returning a real +// RegionCache backed by the mock PD client. All other methods panic because +// buildTxnFiles does not call them. +type stubKVStore struct { + regionCache *locate.RegionCache +} + +func (s *stubKVStore) GetRegionCache() *locate.RegionCache { return s.regionCache } +func (s *stubKVStore) SplitRegions(_ context.Context, _ [][]byte, _ bool, _ *int64) ([]uint64, error) { + panic("not implemented") +} +func (s *stubKVStore) WaitScatterRegionFinish(_ context.Context, _ uint64, _ int) error { + panic("not implemented") +} +func (s *stubKVStore) GetTimestampWithRetry(_ *retry.Backoffer, _ string) (uint64, error) { + panic("not implemented") +} +func (s *stubKVStore) GetOracle() oracle.Oracle { panic("not implemented") } +func (s *stubKVStore) CurrentTimestamp(_ string) (uint64, error) { panic("not implemented") } +func (s *stubKVStore) SendReq(_ *retry.Backoffer, _ *tikvrpc.Request, _ locate.RegionVerID, _ time.Duration) (*tikvrpc.Response, error) { + panic("not implemented") +} +func (s *stubKVStore) GetTiKVClient() client.Client { panic("not implemented") } +func (s *stubKVStore) GetLockResolver() *txnlock.LockResolver { panic("not implemented") } +func (s *stubKVStore) Ctx() context.Context { panic("not implemented") } +func (s *stubKVStore) WaitGroup() *sync.WaitGroup { panic("not implemented") } +func (s *stubKVStore) TxnLatches() *latch.LatchesScheduler { panic("not implemented") } +func (s *stubKVStore) GetClusterID() uint64 { return 0 } +func (s *stubKVStore) IsClose() bool { return false } +func (s *stubKVStore) Go(_ func()) error { panic("not implemented") } + +func TestBuildTxnFilesEntryCounting(t *testing.T) { + require := require.New(t) + + var chunkIDCounter atomic.Uint64 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Method != http.MethodPost { + http.Error(w, "method not allowed", http.StatusMethodNotAllowed) + return + } + id := chunkIDCounter.Add(1) + resp, _ := json.Marshal(map[string]uint64{"chunk_id": id}) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + _, _ = w.Write(resp) + })) + defer srv.Close() + + // Wire size per entry: 2 (key-len) + 3 (key "kXX") + 1 (op) + 4 (val-len) + 3 (val "vXX") = 13 bytes. + // Flush condition: len(buf)+entrySize+4 > cap(buf), where +4 is the CRC trailer reserved space. + // maxChunkSize=50: after 3 entries (39 bytes), 39+13+4=56 > 50 → flush; 2 entries: 26+13+4=43 ≤ 50 → fits. + const maxChunkSize = 50 + + origCfg := config.GetGlobalConfig() + newCfg := *origCfg + newCfg.TiKVClient.TxnChunkWriterAddr = srv.Listener.Addr().String() + newCfg.TiKVClient.TxnChunkMaxSize = maxChunkSize + newCfg.TiKVClient.TxnChunkWriterConcurrency = 4 + config.StoreGlobalConfig(&newCfg) + defer func() { + config.StoreGlobalConfig(origCfg) + once = sync.Once{} + cli = nil + errCli = nil + scheme = "" + }() + + once = sync.Once{} + cli = srv.Client() + errCli = nil + scheme = "http://" + + _, _, pdClient, err := testutils.NewMockTiKV("", nil) + require.NoError(err) + regionCache := locate.NewRegionCache(pdClient) + defer regionCache.Close() + + store := &stubKVStore{regionCache: regionCache} + + memDB := unionstore.NewMemDB() + ops := []kvrpcpb.Op{ + kvrpcpb.Op_Put, + kvrpcpb.Op_Del, + kvrpcpb.Op_Insert, + kvrpcpb.Op_Lock, + kvrpcpb.Op_CheckNotExists, + kvrpcpb.Op_Put, + kvrpcpb.Op_Del, + kvrpcpb.Op_Insert, + kvrpcpb.Op_Lock, + } + for i, op := range ops { + key := []byte(fmt.Sprintf("k%02d", i)) + val := []byte(fmt.Sprintf("v%02d", i)) + flags := tikv.KeyFlags(0) + _ = flags + _ = op + require.NoError(memDB.Set(key, val)) + } + + txn := &KVTxn{ + store: store, + startTS: 1, + valid: true, + vars: tikv.DefaultVars, + us: unionstore.NewUnionStore(memDB, nil), + } + + muts := NewPlainMutations(len(ops)) + for i, op := range ops { + key := []byte(fmt.Sprintf("k%02d", i)) + val := []byte(fmt.Sprintf("v%02d", i)) + muts.Push(op, key, val, false, false, false, false) + } + + c := &twoPhaseCommitter{ + store: store, + txn: txn, + startTS: 1, + regionTxnSize: map[uint64]int{}, + } + + bo := retry.NewBackofferWithVars(context.Background(), 60000, nil) + require.NoError(c.buildTxnFiles(bo, &muts)) + + slice := c.txnFileCtx.slice + require.Equal(3, slice.Len(), "expected 3 chunks") + + totalEntries := uint64(0) + for i := 0; i < slice.Len(); i++ { + totalEntries += slice.chunkRanges[i].entries + } + require.Equal(uint64(len(ops)), totalEntries, "total entries must equal mutation count") + + for i := 0; i < slice.Len(); i++ { + require.Equal(uint64(3), slice.chunkRanges[i].entries, + "chunk %d should have 3 entries", i) + } + + opsSeen := make(map[kvrpcpb.Op]bool) + for i := 0; i < muts.Len(); i++ { + opsSeen[muts.GetOp(i)] = true + } + require.True(opsSeen[kvrpcpb.Op_Put]) + require.True(opsSeen[kvrpcpb.Op_Del]) + require.True(opsSeen[kvrpcpb.Op_Insert]) + require.True(opsSeen[kvrpcpb.Op_Lock]) + require.True(opsSeen[kvrpcpb.Op_CheckNotExists]) +} + +// Ensure stubKVStore satisfies the kvstore interface at compile time. +var _ kvstore = (*stubKVStore)(nil) + +// Ensure the codec used by the test RegionCache returns keyspace ID 0 (codecV1). +var _ apicodec.KeyspaceID = apicodec.NullspaceID diff --git a/txnkv/txnlock/lock_resolver.go b/txnkv/txnlock/lock_resolver.go index 15c3d74e46..2e45ea4c71 100644 --- a/txnkv/txnlock/lock_resolver.go +++ b/txnkv/txnlock/lock_resolver.go @@ -208,6 +208,7 @@ type Lock struct { UseAsyncCommit bool LockForUpdateTS uint64 MinCommitTS uint64 + IsTxnFile bool } func (l *Lock) IsPessimistic() bool { @@ -242,6 +243,7 @@ func NewLock(l *kvrpcpb.LockInfo) *Lock { UseAsyncCommit: l.UseAsyncCommit, LockForUpdateTS: l.LockForUpdateTs, MinCommitTS: l.MinCommitTs, + IsTxnFile: l.IsTxnFile, } } @@ -296,6 +298,7 @@ func (lr *LockResolver) BatchResolveLocks(bo *retry.Backoffer, locks []*Lock, lo expiredLocks := locks txnInfos := make(map[uint64]uint64) + txnFileIDs := make(map[uint64]bool) startTime := time.Now() for _, l := range expiredLocks { logutil.Logger(bo.GetCtx()).Debug("BatchResolveLocks handling lock", zap.Stringer("lock", l)) @@ -357,6 +360,9 @@ func (lr *LockResolver) BatchResolveLocks(bo *retry.Backoffer, locks []*Lock, lo } txnInfos[l.TxnID] = status.commitTS + if l.IsTxnFile { + txnFileIDs[l.TxnID] = true + } } logutil.BgLogger().Info("BatchResolveLocks: lookup txn status", zap.Duration("cost time", time.Since(startTime)), @@ -365,8 +371,9 @@ func (lr *LockResolver) BatchResolveLocks(bo *retry.Backoffer, locks []*Lock, lo listTxnInfos := make([]*kvrpcpb.TxnInfo, 0, len(txnInfos)) for txnID, status := range txnInfos { listTxnInfos = append(listTxnInfos, &kvrpcpb.TxnInfo{ - Txn: txnID, - Status: status, + Txn: txnID, + Status: status, + IsTxnFile: txnFileIDs[txnID], }) } @@ -1008,6 +1015,7 @@ func (lr *LockResolver) getTxnStatus(bo *retry.Backoffer, txnID uint64, primary var status TxnStatus resolvingPessimisticLock := lockInfo != nil && lockInfo.IsPessimistic() + isTxnFile := lockInfo != nil && lockInfo.IsTxnFile req := tikvrpc.NewRequest(tikvrpc.CmdCheckTxnStatus, &kvrpcpb.CheckTxnStatusRequest{ PrimaryKey: primary, LockTs: txnID, @@ -1017,6 +1025,7 @@ func (lr *LockResolver) getTxnStatus(bo *retry.Backoffer, txnID uint64, primary ForceSyncCommit: forceSyncCommit, ResolvingPessimisticLock: resolvingPessimisticLock, VerifyIsPrimary: true, + IsTxnFile: isTxnFile, }, kvrpcpb.Context{ RequestSource: util.RequestSourceFromCtx(bo.GetCtx()), ResourceControlContext: &kvrpcpb.ResourceControlContext{ @@ -1460,6 +1469,7 @@ func (lr *LockResolver) batchLiteResolveLocks(bo *retry.Backoffer, l *Lock, keys func (lr *LockResolver) resolveRegionLocks(bo *retry.Backoffer, l *Lock, region locate.RegionVerID, keys [][]byte, status TxnStatus) error { lreq := &kvrpcpb.ResolveLockRequest{ StartVersion: l.TxnID, + IsTxnFile: l.IsTxnFile, } if status.IsCommitted() { lreq.CommitVersion = status.CommitTS() @@ -1560,6 +1570,7 @@ func (lr *LockResolver) resolveLock(bo *retry.Backoffer, l *Lock, status TxnStat } lreq := &kvrpcpb.ResolveLockRequest{ StartVersion: l.TxnID, + IsTxnFile: l.IsTxnFile, } if status.IsCommitted() { lreq.CommitVersion = status.CommitTS() From 5560ce5308c8caebc6be3dbc6d41a785f1829057 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Tue, 9 Jun 2026 23:33:02 +0800 Subject: [PATCH 02/33] fix build error Signed-off-by: Ping Yu --- internal/locate/region_cache.go | 5 +++ internal/resourcecontrol/resource_control.go | 10 ++++++ txnkv/transaction/test_probe.go | 2 +- txnkv/transaction/txn_file.go | 2 +- util/misc.go | 21 +++++++++++ util/misc_test.go | 37 +++++++++++++++++++- 6 files changed, 74 insertions(+), 3 deletions(-) diff --git a/internal/locate/region_cache.go b/internal/locate/region_cache.go index e2b69bde6c..ac9c9d9fbb 100644 --- a/internal/locate/region_cache.go +++ b/internal/locate/region_cache.go @@ -3056,6 +3056,11 @@ func (c *RegionCache) UpdateBucketsIfNeeded(regionID RegionVerID, requestBuckets } } +// Codec returns the API codec used by this region cache. +func (c *RegionCache) Codec() apicodec.Codec { + return c.codec +} + const cleanCacheInterval = time.Second const cleanRegionNumPerRound = 50 const refreshStoreListInterval = 10 * time.Second diff --git a/internal/resourcecontrol/resource_control.go b/internal/resourcecontrol/resource_control.go index ec851322f2..ae62aae481 100644 --- a/internal/resourcecontrol/resource_control.go +++ b/internal/resourcecontrol/resource_control.go @@ -43,6 +43,16 @@ type RequestInfo struct { bypass bool } +// NewRequestInfo builds request information from precomputed resource-control fields. +func NewRequestInfo(writeBytes int64, storeID uint64, replicaNumber int64, bypass bool) *RequestInfo { + return &RequestInfo{ + writeBytes: writeBytes, + storeID: storeID, + replicaNumber: replicaNumber, + bypass: bypass, + } +} + func toPDAccessLocationType(accessType kv.AccessLocationType) controller.AccessLocationType { switch accessType { case kv.AccessLocalZone: diff --git a/txnkv/transaction/test_probe.go b/txnkv/transaction/test_probe.go index dd2b429b5e..a7567d6ef2 100644 --- a/txnkv/transaction/test_probe.go +++ b/txnkv/transaction/test_probe.go @@ -396,7 +396,7 @@ func (c CommitterProbe) ResolveFlushedLocks(bo *retry.Backoffer, start, end []by // SendTxnHeartBeat renews a txn's ttl. func SendTxnHeartBeat(bo *retry.Backoffer, store kvstore, primary []byte, startTS, ttl uint64) (newTTL uint64, stopHeartBeat bool, err error) { - return sendTxnHeartBeat(bo, store, primary, startTS, ttl, 0) + return sendTxnHeartBeat(bo, store, primary, startTS, ttl, 0, false) } // ConfigProbe exposes configurations and global variables for testing purpose. diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index c620edca8f..d068085eff 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -329,7 +329,7 @@ func (a txnFilePrewriteAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back // For other region error and the fake region error, backoff because // there's something wrong. // For the real EpochNotMatch error, don't backoff. - if regionErr.GetEpochNotMatch() == nil || locate.IsFakeRegionError(regionErr) { + if regionErr.GetEpochNotMatch() == nil || retry.IsFakeRegionError(regionErr) { err = bo.Backoff(retry.BoRegionMiss, errors.New(regionErr.String())) if err != nil { return resp, err diff --git a/util/misc.go b/util/misc.go index 87f9c85276..9b65163952 100644 --- a/util/misc.go +++ b/util/misc.go @@ -35,6 +35,7 @@ package util import ( + "bytes" "context" "fmt" "strconv" @@ -175,3 +176,23 @@ func None[T interface{}]() Option[T] { func (o Option[T]) Inner() *T { return o.inner } + +func GetMaxStartKey(lhs []byte, rhs []byte) []byte { + if bytes.Compare(lhs, rhs) > 0 { + return lhs + } + return rhs +} + +func GetMinEndKey(lhs []byte, rhs []byte) []byte { + if len(rhs) == 0 { + return lhs + } + if len(lhs) == 0 { + return rhs + } + if bytes.Compare(lhs, rhs) < 0 { + return lhs + } + return rhs +} diff --git a/util/misc_test.go b/util/misc_test.go index fd863759d3..fbeeb1058b 100644 --- a/util/misc_test.go +++ b/util/misc_test.go @@ -101,5 +101,40 @@ func TestTimeDetail(t *testing.T) { KvGrpcWaitTime: time.Millisecond * 7, TotalRPCWallTime: time.Millisecond * 8, } - assert.Equal(t, "time_detail: {total_process_time: 2ms, total_suspend_time: 3ms, total_wait_time: 4ms, total_kv_read_wall_time: 5ms, tikv_grpc_process_time: 6ms, tikv_grpc_wait_time: 7ms, tikv_wall_time: 8ms}", detail.String()) + assert.Equal(t, "time_detail: {total_process_time: 2ms, total_suspend_time: 3ms, total_wait_time: 4ms, total_kv_read_wall_time: 5ms, tikv_wall_time: 6ms}", detail.String()) +} + +func TestGetMaxStartKey(t *testing.T) { + assert := assert.New(t) + + cases := []struct { + lhs, rhs, expected string + }{ + {"", "", ""}, + {"", "a", "a"}, + {"a", "a", "a"}, + } + + for _, c := range cases { + assert.Equal([]byte(c.expected), GetMaxStartKey([]byte(c.lhs), []byte(c.rhs))) + assert.Equal([]byte(c.expected), GetMaxStartKey([]byte(c.rhs), []byte(c.lhs))) + } +} + +func TestGetMinEndKey(t *testing.T) { + assert := assert.New(t) + + cases := []struct { + lhs, rhs, expected string + }{ + {"", "", ""}, + {"a", "", "a"}, + {"a", "a", "a"}, + {"a", "b", "a"}, + } + + for _, c := range cases { + assert.Equal([]byte(c.expected), GetMinEndKey([]byte(c.lhs), []byte(c.rhs))) + assert.Equal([]byte(c.expected), GetMinEndKey([]byte(c.rhs), []byte(c.lhs))) + } } From 8db13623604949a51fff97e610844532584e691b Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Wed, 10 Jun 2026 08:48:39 +0800 Subject: [PATCH 03/33] fix CI errors Signed-off-by: Ping Yu --- integration_tests/txn_file_test.go | 17 +++++++++-------- kv/variables.go | 10 +++++----- metrics/metrics.go | 3 +-- tikv/kv_test.go | 4 ++++ txnkv/transaction/txn_file.go | 11 ++++++++++- util/misc_test.go | 2 +- 6 files changed, 30 insertions(+), 17 deletions(-) diff --git a/integration_tests/txn_file_test.go b/integration_tests/txn_file_test.go index e25a4b7a10..eb82e0d6a5 100644 --- a/integration_tests/txn_file_test.go +++ b/integration_tests/txn_file_test.go @@ -206,16 +206,16 @@ func TestTxnFilePrewriteTxnSizeAfterRegionRegroup(t *testing.T) { var mu sync.Mutex var captured []capturedPrewrite - hook := func(req *tikvrpc.Request, resp *tikvrpc.Response, sendErr error) { + hook := func(req *tikvrpc.Request, resp *tikvrpc.Response, sendErr error) (*tikvrpc.Response, error) { if req.Type != tikvrpc.CmdPrewrite { - return + return resp, sendErr } inner, ok := req.Req.(*kvrpcpb.PrewriteRequest) if !ok || len(inner.TxnFileChunks) == 0 { - return + return resp, sendErr } if sendErr != nil { - return + return resp, sendErr } chunks := make([]uint64, len(inner.TxnFileChunks)) copy(chunks, inner.TxnFileChunks) @@ -229,19 +229,20 @@ func TestTxnFilePrewriteTxnSizeAfterRegionRegroup(t *testing.T) { captured = append(captured, capturedPrewrite{ txnFileChunks: chunks, txnSize: inner.TxnSize, - isRetry: req.Context.IsRetryRequest, + isRetry: req.IsRetryRequest || req.Context.IsRetryRequest, regionErr: regionErr, }) mu.Unlock() + return resp, sendErr } require.Nil(failpoint.Enable("tikvclient/mockRetrySendReqToRegion", "1*return(true)->return(false)")) defer failpoint.Disable("tikvclient/mockRetrySendReqToRegion") require.Nil(failpoint.Enable("tikvclient/invalidCacheAndRetry", "1*off->pause")) defer failpoint.Disable("tikvclient/invalidCacheAndRetry") - require.Nil(failpoint.Enable("tikvclient/afterSendReqToRegion", "return")) - defer failpoint.Disable("tikvclient/afterSendReqToRegion") - ctx := context.WithValue(context.Background(), "sendReqToRegionFinishHook", hook) + require.Nil(failpoint.Enable("tikvclient/onRPCFinishedHook", "return")) + defer failpoint.Disable("tikvclient/onRPCFinishedHook") + ctx := context.WithValue(context.Background(), "onRPCFinishedHook", hook) tx, err := store.Begin() require.Nil(err) diff --git a/kv/variables.go b/kv/variables.go index 7646d0e4a1..90e137aaa3 100644 --- a/kv/variables.go +++ b/kv/variables.go @@ -61,11 +61,11 @@ type Variables struct { // NewVariables create a new Variables instance with default values. func NewVariables(killed *uint32) *Variables { return &Variables{ - BackoffLockFast: DefBackoffLockFast, - BackOffWeight: DefBackOffWeight, - Killed: killed, - DisableTxnFile: false, - TxnFileMinMutationSize: 0, + BackoffLockFast: DefBackoffLockFast, + BackOffWeight: DefBackOffWeight, + Killed: killed, + DisableTxnFile: false, + TxnFileMinMutationSize: 0, } } diff --git a/metrics/metrics.go b/metrics/metrics.go index 8eefdb4cbd..fbb89c75c1 100644 --- a/metrics/metrics.go +++ b/metrics/metrics.go @@ -138,7 +138,6 @@ var ( TiKVTxnLagCommitTSWaitHistogram *prometheus.HistogramVec TiKVTxnLagCommitTSAttemptHistogram *prometheus.HistogramVec - TiKVTxnFileRequestCounter *prometheus.CounterVec TiKVTxnFileWriteBytes *prometheus.CounterVec TiKVTxnFileMutationSizeHistogram *prometheus.HistogramVec @@ -1061,7 +1060,7 @@ func initMetrics(namespace, subsystem string, constLabels prometheus.Labels) { ConstLabels: constLabels, }, []string{LblResult}) - TiKVTxnFileRequestCounter = prometheus.NewCounterVec( + TiKVTxnFileRequestCounter = prometheus.NewCounterVec( prometheus.CounterOpts{ Namespace: namespace, Subsystem: subsystem, diff --git a/tikv/kv_test.go b/tikv/kv_test.go index 8a7042819d..8854369a49 100644 --- a/tikv/kv_test.go +++ b/tikv/kv_test.go @@ -32,6 +32,7 @@ import ( "github.com/tikv/client-go/v2/testutils" "github.com/tikv/client-go/v2/tikvrpc" "github.com/tikv/client-go/v2/util" + "github.com/tikv/client-go/v2/util/intest" pdhttp "github.com/tikv/pd/client/http" ) @@ -277,6 +278,9 @@ func (s *testKVSuite) TestMinSafeTsFromMixed2() { } func (s *testKVSuite) TestErrorHalfwayInNewKVStore() { + if !intest.InTest { + s.T().Skip("requires intest error injection") + } // this is a leak test, TestMain will check goroutine leak _, err := NewKVStore("TestErrorHalfwayInNewKVStore", s.store.pdClient, NewMockSafePointKV(), &mocktikv.RPCClient{}) require.Error(s.T(), err) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index d068085eff..fefb41013c 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -80,7 +80,7 @@ type chunkBatch struct { func (b chunkBatch) String() string { return fmt.Sprintf("chunkBatch{region: %d, isPrimary: %t, txnChunkSlice: %v}", - b.region.Region.GetID(), b.isPrimary, b.txnChunkSlice.chunkIDs) + b.region.Region.GetID(), b.isPrimary, b.chunkIDs) } func (b *chunkBatch) getSampleKeys() [][]byte { @@ -98,6 +98,12 @@ func (b *chunkBatch) getBatchTxnSize() uint64 { return batchTxnSize } +func markTxnFileRetryRequest(req *tikvrpc.Request, bo *retry.Backoffer) { + if bo.GetTotalSleep() > 0 { + req.IsRetryRequest = true + } +} + // txnChunkSlice should be sorted by txnChunkRange.smallest and no overlapping. type txnChunkSlice struct { chunkIDs []uint64 @@ -307,6 +313,7 @@ func (a txnFilePrewriteAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back ResourceGroupName: c.resourceGroupName, }, }) + markTxnFileRetryRequest(req, bo) sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.GetOracle()) var resolvingRecordToken *int @@ -463,6 +470,7 @@ func (a txnFileCommitAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backof ResourceGroupName: c.resourceGroupName, }, }) + markTxnFileRetryRequest(req, bo) sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.GetOracle()) for { resp, _, err := sender.SendReq(bo, req, batch.region.Region, client.ReadTimeoutMedium) @@ -552,6 +560,7 @@ func (a txnFileRollbackAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back ResourceGroupName: c.resourceGroupName, }, }) + markTxnFileRetryRequest(req, bo) sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.GetOracle()) resp, _, err1 := sender.SendReq(bo, req, batch.region.Region, client.ReadTimeoutShort) if err1 != nil { diff --git a/util/misc_test.go b/util/misc_test.go index fbeeb1058b..0da09a885c 100644 --- a/util/misc_test.go +++ b/util/misc_test.go @@ -101,7 +101,7 @@ func TestTimeDetail(t *testing.T) { KvGrpcWaitTime: time.Millisecond * 7, TotalRPCWallTime: time.Millisecond * 8, } - assert.Equal(t, "time_detail: {total_process_time: 2ms, total_suspend_time: 3ms, total_wait_time: 4ms, total_kv_read_wall_time: 5ms, tikv_wall_time: 6ms}", detail.String()) + assert.Equal(t, "time_detail: {total_process_time: 2ms, total_suspend_time: 3ms, total_wait_time: 4ms, total_kv_read_wall_time: 5ms, tikv_grpc_process_time: 6ms, tikv_grpc_wait_time: 7ms, tikv_wall_time: 8ms}", detail.String()) } func TestGetMaxStartKey(t *testing.T) { From f558577f948c02723bc54b511a63502976fb132d Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Mon, 6 Jul 2026 22:37:26 +0800 Subject: [PATCH 04/33] address comments Signed-off-by: Ping Yu --- txnkv/transaction/2pc.go | 3 +++ txnkv/transaction/txn.go | 1 + txnkv/transaction/txn_file.go | 21 +++++++++++++-------- 3 files changed, 17 insertions(+), 8 deletions(-) diff --git a/txnkv/transaction/2pc.go b/txnkv/transaction/2pc.go index 151598e0ec..88a3d45c79 100644 --- a/txnkv/transaction/2pc.go +++ b/txnkv/transaction/2pc.go @@ -340,6 +340,9 @@ type CommitterMutations interface { NeedConstraintCheckInPrewrite(i int) bool } +// MutationsHasDataInRange returns whether mutations has data in the range [start, end). +// If it has, it returns the primary or first write key in the range. +// Note that the firstDataKey can be empty when the range contains only non-write ops (and not the primary at pos 0). func MutationsHasDataInRange(mutations CommitterMutations, start []byte, end []byte) ([]byte /* firstDataKey */, bool) { isInRange := func(pos int) bool { return pos < mutations.Len() && (len(end) == 0 || bytes.Compare(mutations.GetKey(pos), end) < 0) diff --git a/txnkv/transaction/txn.go b/txnkv/transaction/txn.go index f693b8edf3..cf2847fd76 100644 --- a/txnkv/transaction/txn.go +++ b/txnkv/transaction/txn.go @@ -790,6 +790,7 @@ func (txn *KVTxn) GetScope() string { return txn.scope } +// DisableTxnFile disables the file-based transaction path for this transaction. func (txn *KVTxn) DisableTxnFile() { txn.disableTxnFile = true } diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index fefb41013c..92e11d6ec3 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -898,10 +898,15 @@ func (c *twoPhaseCommitter) executeTxnFileAction(bo *retry.Backoffer, chunkSlice return c.executeTxnFileSliceWithRetry(bo, chunkSlice, secondaries, action) } - c.store.WaitGroup().Add(1) - errGo := c.store.Go(func() { - defer c.store.WaitGroup().Done() - err := c.executeTxnFileSliceWithRetry(bo, chunkSlice, secondaries, action) + secondaryBo := retry.NewBackofferWithVars(c.store.Ctx(), CommitSecondaryMaxBackoff, c.txn.vars) + if c.store.IsClose() { + logutil.Logger(bo.GetCtx()).Warn("the store is closed", + zap.Uint64("startTS", c.startTS), zap.Uint64("commitTS", c.commitTS), + zap.Stringer("action", action)) + return nil + } + errGo := c.txn.spawnWithStorePool(func() { + err := c.executeTxnFileSliceWithRetry(secondaryBo, chunkSlice, secondaries, action) logutil.Logger(bo.GetCtx()).Debug("txn file: async execute secondaries finished", zap.Uint64("startTS", c.startTS), zap.Stringer("action", action), @@ -912,7 +917,6 @@ func (c *twoPhaseCommitter) executeTxnFileAction(bo *retry.Backoffer, chunkSlice } }) if errGo != nil { - c.store.WaitGroup().Done() logutil.Logger(bo.GetCtx()).Warn("txn file: create goroutine failed", zap.Uint64("startTS", c.startTS), zap.Stringer("action", action), @@ -968,7 +972,7 @@ func (c *twoPhaseCommitter) buildTxnFiles(bo *retry.Backoffer, mutations Committ if inflightChunks >= concurrency { r := <-resultCh if r.err != nil { - logutil.Logger(bo.GetCtx()).Error(buildChunkErrMsg, zap.Error(err)) + logutil.Logger(bo.GetCtx()).Error(buildChunkErrMsg, zap.Error(r.err)) return errors.Wrap(r.err, buildChunkErrMsg) } results = append(results, r) @@ -992,7 +996,7 @@ func (c *twoPhaseCommitter) buildTxnFiles(bo *retry.Backoffer, mutations Committ for i := 0; i < inflightChunks; i++ { r := <-resultCh if r.err != nil { - logutil.Logger(bo.GetCtx()).Error(buildChunkErrMsg, zap.Error(err)) + logutil.Logger(bo.GetCtx()).Error(buildChunkErrMsg, zap.Error(r.err)) return errors.Wrap(r.err, buildChunkErrMsg) } results = append(results, r) @@ -1289,12 +1293,12 @@ func (w *chunkWriterClient) request(bo *retry.Backoffer, method string, data []b } continue } - defer resp.Body.Close() if resp.StatusCode != http.StatusOK { var bodyStr string if data, err := io.ReadAll(resp.Body); err == nil { bodyStr = string(data) } + resp.Body.Close() logutil.Logger(ctx).Warn("service error", zap.String("http status", resp.Status), zap.String("body", bodyStr)) err = bo.Backoff(retry.BoTiKVServerBusy, fmt.Errorf("service error, http status %s", resp.Status)) if err != nil { @@ -1303,6 +1307,7 @@ func (w *chunkWriterClient) request(bo *retry.Backoffer, method string, data []b continue } data, err := io.ReadAll(resp.Body) + resp.Body.Close() return data, errors.WithStack(err) } } From 44fa21205d27ee447dc95e6ba9626897762c9854 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Tue, 14 Jul 2026 00:24:13 +0800 Subject: [PATCH 05/33] prepareTxnFileCommitTS Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 40 +++- txnkv/transaction/txn_file_test.go | 284 +++++++++++++++++++++++++++++ 2 files changed, 322 insertions(+), 2 deletions(-) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 92e11d6ec3..bed7769dc4 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -44,6 +44,7 @@ import ( "github.com/tikv/client-go/v2/internal/resourcecontrol" "github.com/tikv/client-go/v2/kv" "github.com/tikv/client-go/v2/metrics" + "github.com/tikv/client-go/v2/oracle" "github.com/tikv/client-go/v2/tikvrpc" "github.com/tikv/client-go/v2/txnkv/txnlock" "github.com/tikv/client-go/v2/util" @@ -452,6 +453,38 @@ type txnFileCommitAction struct{} var _ txnFileAction = (*txnFileCommitAction)(nil) +func (c *twoPhaseCommitter) prepareTxnFileCommitTS(ctx context.Context) (uint64, error) { + start := time.Now() + logutil.Event(ctx, "start get commit ts") + commitTS, err := c.txn.GetTimestampForCommit(retry.NewBackofferWithVars(ctx, TsoMaxBackoff, c.txn.vars), c.txn.GetScope()) + if err != nil { + logutil.Logger(ctx).Warn("txn file get commitTS failed", + zap.Error(err), + zap.Uint64("txnStartTS", c.startTS)) + return 0, errors.WithStack(err) + } + commitDetail := c.getDetail() + commitDetail.GetCommitTsTime = time.Since(start) + c.txn.fillCommitTSLagDetails(&commitDetail.LagDetails) + logutil.Event(ctx, "finish get commit ts") + logutil.SetTag(ctx, "commitTs", commitTS) + + if err = c.checkSchemaValid(ctx, commitTS, c.txn.schemaVer); err != nil { + return 0, errors.WithStack(err) + } + + if c.store.GetOracle().IsExpired(c.startTS, MaxTxnTimeUse, &oracle.Option{TxnScope: oracle.GlobalTxnScope}) { + return 0, errors.Errorf("session %d txn takes too much time, txnStartTS: %d, comm: %d", + c.sessionID, c.startTS, commitTS) + } + + if c.txn.commitTSUpperBoundCheck != nil && !c.txn.commitTSUpperBoundCheck(commitTS) { + return 0, errors.Errorf("session %d check commit ts upper bound fail, txnStartTS: %d, comm: %d", + c.sessionID, c.startTS, commitTS) + } + return commitTS, nil +} + func (a txnFileCommitAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backoffer, batch chunkBatch) (*tikvrpc.Response, error) { req := tikvrpc.NewRequest(tikvrpc.CmdCommit, &kvrpcpb.CommitRequest{ Keys: batch.getSampleKeys(), // To help detect duplicated request. @@ -490,6 +523,9 @@ func (a txnFileCommitAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backof logutil.Logger(bo.GetCtx()).Info("2PC commitTS rejected by TiKV, retry with a newer commitTS", zap.Uint64("txnStartTS", c.startTS), zap.Stringer("info", logutil.Hex(rejected))) + if !batch.isPrimary { + return nil, errors.New("2PC commitTS rejected by TiKV, but the txn-file batch is not the primary batch") + } // Do not retry for a txn which has a too large MinCommitTs // 3600000 << 18 = 943718400000 @@ -499,7 +535,7 @@ func (a txnFileCommitAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backof } // Update commit ts and retry. - commitTS, err1 := c.store.GetTimestampWithRetry(bo, c.txn.GetScope()) + commitTS, err1 := c.prepareTxnFileCommitTS(bo.GetCtx()) if err1 != nil { logutil.Logger(bo.GetCtx()).Warn("2PC get commitTS failed", zap.Error(err1), @@ -681,7 +717,7 @@ func (c *twoPhaseCommitter) executeTxnFile(ctx context.Context) (err error) { } commitBo := retry.NewBackofferWithVars(ctx, int(CommitMaxBackoff), c.txn.vars) - c.commitTS, err = c.store.GetTimestampWithRetry(commitBo, c.txn.GetScope()) + c.commitTS, err = c.prepareTxnFileCommitTS(ctx) if err != nil { return } diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 9ae28a35eb..a1c3026179 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -17,6 +17,7 @@ package transaction import ( "context" "encoding/json" + "errors" "fmt" "math/rand" "net/http" @@ -45,6 +46,289 @@ import ( "github.com/tikv/client-go/v2/util" ) +type txnFileCommitTSOracle struct { + unimplementedOracle + + expired bool + calls int + startTS uint64 + ttl uint64 + option *oracle.Option +} + +func (o *txnFileCommitTSOracle) IsExpired(startTS uint64, ttl uint64, option *oracle.Option) bool { + o.calls++ + o.startTS = startTS + o.ttl = ttl + o.option = option + return o.expired +} + +type txnFileCommitTSStore struct { + unimplementedKVStore + + timestamps []uint64 + timestampCalls int + timestampErr error + oracle *txnFileCommitTSOracle + regionCache *locate.RegionCache + client client.Client +} + +func (s *txnFileCommitTSStore) GetTimestampWithRetry(_ *retry.Backoffer, _ string) (uint64, error) { + s.timestampCalls++ + if s.timestampErr != nil { + return 0, s.timestampErr + } + if len(s.timestamps) == 0 { + return 0, errors.New("no timestamp configured") + } + ts := s.timestamps[0] + s.timestamps = s.timestamps[1:] + return ts, nil +} + +func (s *txnFileCommitTSStore) GetOracle() oracle.Oracle { + return s.oracle +} + +func (s *txnFileCommitTSStore) GetRegionCache() *locate.RegionCache { + return s.regionCache +} + +func (s *txnFileCommitTSStore) GetTiKVClient() client.Client { + return s.client +} + +type txnFileSchemaVer int64 + +func (v txnFileSchemaVer) SchemaMetaVersion() int64 { + return int64(v) +} + +type txnFileSchemaLeaseChecker struct { + err error + calls int + checkTS uint64 + schemaVer SchemaVer +} + +func (c *txnFileSchemaLeaseChecker) CheckBySchemaVer(checkTS uint64, schemaVer SchemaVer) (*RelatedSchemaChange, error) { + c.calls++ + c.checkTS = checkTS + c.schemaVer = schemaVer + return nil, c.err +} + +func newTxnFileCommitTSTestCommitter( + store *txnFileCommitTSStore, + checker SchemaLeaseChecker, + upperBoundCheck func(uint64) bool, +) *twoPhaseCommitter { + txn := &KVTxn{ + store: store, + startTS: 1, + schemaVer: txnFileSchemaVer(10), + schemaLeaseChecker: checker, + scope: oracle.GlobalTxnScope, + commitTSUpperBoundCheck: upperBoundCheck, + } + committer := &twoPhaseCommitter{ + store: store, + txn: txn, + startTS: txn.startTS, + sessionID: 7, + } + committer.setDetail(&util.CommitDetails{}) + return committer +} + +func TestPrepareTxnFileCommitTS(t *testing.T) { + t.Run("success", func(t *testing.T) { + commitOracle := &txnFileCommitTSOracle{} + store := &txnFileCommitTSStore{ + timestamps: []uint64{100, 102}, + oracle: commitOracle, + } + checker := &txnFileSchemaLeaseChecker{} + upperBoundCalls := 0 + committer := newTxnFileCommitTSTestCommitter(store, checker, func(commitTS uint64) bool { + upperBoundCalls++ + return commitTS == 102 + }) + committer.txn.SetCommitWaitUntilTSO(101) + committer.txn.SetCommitWaitUntilTSOTimeout(time.Second) + + commitTS, err := committer.prepareTxnFileCommitTS(context.Background()) + + require.NoError(t, err) + require.Equal(t, uint64(102), commitTS) + require.Equal(t, 2, store.timestampCalls) + require.Equal(t, 1, checker.calls) + require.Equal(t, commitTS, checker.checkTS) + require.Equal(t, txnFileSchemaVer(10), checker.schemaVer) + require.Equal(t, 1, commitOracle.calls) + require.Equal(t, uint64(1), commitOracle.startTS) + require.Equal(t, uint64(MaxTxnTimeUse), commitOracle.ttl) + require.Equal(t, oracle.GlobalTxnScope, commitOracle.option.TxnScope) + require.Equal(t, 1, upperBoundCalls) + require.Equal(t, uint64(100), committer.getDetail().LagDetails.FirstLagTS) + require.Equal(t, uint64(101), committer.getDetail().LagDetails.WaitUntilTS) + require.Equal(t, 1, committer.getDetail().LagDetails.BackoffCnt) + }) + + t.Run("schema invalid", func(t *testing.T) { + schemaErr := errors.New("schema changed") + commitOracle := &txnFileCommitTSOracle{} + store := &txnFileCommitTSStore{timestamps: []uint64{100}, oracle: commitOracle} + checker := &txnFileSchemaLeaseChecker{err: schemaErr} + upperBoundCalls := 0 + committer := newTxnFileCommitTSTestCommitter(store, checker, func(uint64) bool { + upperBoundCalls++ + return true + }) + + commitTS, err := committer.prepareTxnFileCommitTS(context.Background()) + + require.Zero(t, commitTS) + require.ErrorIs(t, err, schemaErr) + require.Equal(t, 1, checker.calls) + require.Zero(t, commitOracle.calls) + require.Zero(t, upperBoundCalls) + }) + + t.Run("transaction expired", func(t *testing.T) { + commitOracle := &txnFileCommitTSOracle{expired: true} + store := &txnFileCommitTSStore{timestamps: []uint64{100}, oracle: commitOracle} + checker := &txnFileSchemaLeaseChecker{} + upperBoundCalls := 0 + committer := newTxnFileCommitTSTestCommitter(store, checker, func(uint64) bool { + upperBoundCalls++ + return true + }) + + commitTS, err := committer.prepareTxnFileCommitTS(context.Background()) + + require.Zero(t, commitTS) + require.ErrorContains(t, err, "txn takes too much time") + require.Equal(t, 1, checker.calls) + require.Equal(t, 1, commitOracle.calls) + require.Zero(t, upperBoundCalls) + }) + + t.Run("commit timestamp exceeds upper bound", func(t *testing.T) { + commitOracle := &txnFileCommitTSOracle{} + store := &txnFileCommitTSStore{timestamps: []uint64{100}, oracle: commitOracle} + checker := &txnFileSchemaLeaseChecker{} + upperBoundCalls := 0 + committer := newTxnFileCommitTSTestCommitter(store, checker, func(uint64) bool { + upperBoundCalls++ + return false + }) + + commitTS, err := committer.prepareTxnFileCommitTS(context.Background()) + + require.Zero(t, commitTS) + require.ErrorContains(t, err, "check commit ts upper bound fail") + require.Equal(t, 1, checker.calls) + require.Equal(t, 1, commitOracle.calls) + require.Equal(t, 1, upperBoundCalls) + }) +} + +func TestTxnFileCommitTSExpiredRetryUsesPreparedTimestamp(t *testing.T) { + pd := &mockPDClient{} + regionCache := locate.NewTestRegionCache() + regionCache.SetPDClient(pd) + defer regionCache.Close() + + commitOracle := &txnFileCommitTSOracle{} + checker := &txnFileSchemaLeaseChecker{} + requestCount := 0 + kvClient := &fnClient{} + kvClient.onSend = func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + requestCount++ + if requestCount == 1 { + return &tikvrpc.Response{Resp: &kvrpcpb.CommitResponse{Error: &kvrpcpb.KeyError{ + CommitTsExpired: &kvrpcpb.CommitTsExpired{ + StartTs: 1, + AttemptedCommitTs: req.Commit().CommitVersion, + MinCommitTs: 100, + }, + }}}, nil + } + require.Equal(t, uint64(102), req.Commit().CommitVersion) + return &tikvrpc.Response{Resp: &kvrpcpb.CommitResponse{}}, nil + } + store := &txnFileCommitTSStore{ + timestamps: []uint64{102}, + oracle: commitOracle, + regionCache: regionCache, + client: kvClient, + } + upperBoundCalls := 0 + committer := newTxnFileCommitTSTestCommitter(store, checker, func(commitTS uint64) bool { + upperBoundCalls++ + return commitTS == 102 + }) + committer.commitTS = 2 + + bo := retry.NewBackoffer(context.Background(), 1000) + location, err := regionCache.LocateKey(bo, []byte("k")) + require.NoError(t, err) + batch := chunkBatch{ + txnChunkSlice: txnChunkSlice{ + chunkIDs: []uint64{1}, + chunkRanges: []txnChunkRange{{ + smallest: []byte("k"), + biggest: []byte("k"), + }}, + }, + region: location, + sampleKeys: [][]byte{[]byte("k")}, + isPrimary: true, + } + + _, err = (txnFileCommitAction{}).executeBatch(committer, bo, batch) + + require.NoError(t, err) + require.Equal(t, 2, requestCount) + require.Equal(t, uint64(102), committer.commitTS) + require.Equal(t, 1, store.timestampCalls) + require.Equal(t, 1, checker.calls) + require.Equal(t, 1, commitOracle.calls) + require.Equal(t, 1, upperBoundCalls) + + requestCount = 0 + store.timestamps = []uint64{104} + store.timestampCalls = 0 + checker.calls = 0 + commitOracle.calls = 0 + upperBoundCalls = 0 + committer.commitTS = 2 + batch.isPrimary = false + kvClient.onSend = func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + requestCount++ + return &tikvrpc.Response{Resp: &kvrpcpb.CommitResponse{Error: &kvrpcpb.KeyError{ + CommitTsExpired: &kvrpcpb.CommitTsExpired{ + StartTs: 1, + AttemptedCommitTs: req.Commit().CommitVersion, + MinCommitTs: 100, + }, + }}}, nil + } + + _, err = (txnFileCommitAction{}).executeBatch(committer, bo, batch) + + require.ErrorContains(t, err, "txn-file batch is not the primary batch") + require.Equal(t, 1, requestCount) + require.Equal(t, uint64(2), committer.commitTS) + require.Zero(t, store.timestampCalls) + require.Zero(t, checker.calls) + require.Zero(t, commitOracle.calls) + require.Zero(t, upperBoundCalls) +} + func TestChunkSliceSortAndDedup(t *testing.T) { assert := assert.New(t) From a209cbe8ac21b9c1f3a6b1da91d0ea40f092e8da Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Tue, 14 Jul 2026 17:55:03 +0800 Subject: [PATCH 06/33] bo Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 9 +++++---- txnkv/transaction/txn_file_test.go | 13 ++++++++----- 2 files changed, 13 insertions(+), 9 deletions(-) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index bed7769dc4..258b4e3772 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -453,10 +453,11 @@ type txnFileCommitAction struct{} var _ txnFileAction = (*txnFileCommitAction)(nil) -func (c *twoPhaseCommitter) prepareTxnFileCommitTS(ctx context.Context) (uint64, error) { +func (c *twoPhaseCommitter) prepareTxnFileCommitTS(bo *retry.Backoffer) (uint64, error) { + ctx := bo.GetCtx() start := time.Now() logutil.Event(ctx, "start get commit ts") - commitTS, err := c.txn.GetTimestampForCommit(retry.NewBackofferWithVars(ctx, TsoMaxBackoff, c.txn.vars), c.txn.GetScope()) + commitTS, err := c.txn.GetTimestampForCommit(bo, c.txn.GetScope()) if err != nil { logutil.Logger(ctx).Warn("txn file get commitTS failed", zap.Error(err), @@ -535,7 +536,7 @@ func (a txnFileCommitAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backof } // Update commit ts and retry. - commitTS, err1 := c.prepareTxnFileCommitTS(bo.GetCtx()) + commitTS, err1 := c.prepareTxnFileCommitTS(bo) if err1 != nil { logutil.Logger(bo.GetCtx()).Warn("2PC get commitTS failed", zap.Error(err1), @@ -717,7 +718,7 @@ func (c *twoPhaseCommitter) executeTxnFile(ctx context.Context) (err error) { } commitBo := retry.NewBackofferWithVars(ctx, int(CommitMaxBackoff), c.txn.vars) - c.commitTS, err = c.prepareTxnFileCommitTS(ctx) + c.commitTS, err = c.prepareTxnFileCommitTS(retry.NewBackofferWithVars(ctx, TsoMaxBackoff, c.txn.vars)) if err != nil { return } diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index a1c3026179..1b1a75b8b9 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -70,13 +70,15 @@ type txnFileCommitTSStore struct { timestamps []uint64 timestampCalls int timestampErr error + backoffer *retry.Backoffer oracle *txnFileCommitTSOracle regionCache *locate.RegionCache client client.Client } -func (s *txnFileCommitTSStore) GetTimestampWithRetry(_ *retry.Backoffer, _ string) (uint64, error) { +func (s *txnFileCommitTSStore) GetTimestampWithRetry(bo *retry.Backoffer, _ string) (uint64, error) { s.timestampCalls++ + s.backoffer = bo if s.timestampErr != nil { return 0, s.timestampErr } @@ -159,7 +161,7 @@ func TestPrepareTxnFileCommitTS(t *testing.T) { committer.txn.SetCommitWaitUntilTSO(101) committer.txn.SetCommitWaitUntilTSOTimeout(time.Second) - commitTS, err := committer.prepareTxnFileCommitTS(context.Background()) + commitTS, err := committer.prepareTxnFileCommitTS(retry.NewBackoffer(context.Background(), TsoMaxBackoff)) require.NoError(t, err) require.Equal(t, uint64(102), commitTS) @@ -188,7 +190,7 @@ func TestPrepareTxnFileCommitTS(t *testing.T) { return true }) - commitTS, err := committer.prepareTxnFileCommitTS(context.Background()) + commitTS, err := committer.prepareTxnFileCommitTS(retry.NewBackoffer(context.Background(), TsoMaxBackoff)) require.Zero(t, commitTS) require.ErrorIs(t, err, schemaErr) @@ -207,7 +209,7 @@ func TestPrepareTxnFileCommitTS(t *testing.T) { return true }) - commitTS, err := committer.prepareTxnFileCommitTS(context.Background()) + commitTS, err := committer.prepareTxnFileCommitTS(retry.NewBackoffer(context.Background(), TsoMaxBackoff)) require.Zero(t, commitTS) require.ErrorContains(t, err, "txn takes too much time") @@ -226,7 +228,7 @@ func TestPrepareTxnFileCommitTS(t *testing.T) { return false }) - commitTS, err := committer.prepareTxnFileCommitTS(context.Background()) + commitTS, err := committer.prepareTxnFileCommitTS(retry.NewBackoffer(context.Background(), TsoMaxBackoff)) require.Zero(t, commitTS) require.ErrorContains(t, err, "check commit ts upper bound fail") @@ -295,6 +297,7 @@ func TestTxnFileCommitTSExpiredRetryUsesPreparedTimestamp(t *testing.T) { require.Equal(t, 2, requestCount) require.Equal(t, uint64(102), committer.commitTS) require.Equal(t, 1, store.timestampCalls) + require.Same(t, bo, store.backoffer) require.Equal(t, 1, checker.calls) require.Equal(t, 1, commitOracle.calls) require.Equal(t, 1, upperBoundCalls) From d50071dc597f8d61582ec4da8206fefbee418576 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Tue, 14 Jul 2026 20:14:55 +0800 Subject: [PATCH 07/33] handle binlog Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 258b4e3772..efec9cfcb8 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -1058,7 +1058,9 @@ func (c *twoPhaseCommitter) getKeyspaceID() apicodec.KeyspaceID { } func (c *twoPhaseCommitter) useTxnFile(ctx context.Context) (bool, error) { - if c.txn == nil || c.txn.vars.DisableTxnFile { + // shouldWriteBinlog(): Disable txn file if should write binlog for safe. + // Txn file is enabled on TiDBCloud only, while binlog is just for on-premise. + if c.txn == nil || c.txn.vars.DisableTxnFile || c.shouldWriteBinlog() { return false, nil } conf := config.GetGlobalConfig() From 26bd55b00e3ba8920ef6857ba1d587b290c5570e Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Tue, 14 Jul 2026 23:40:05 +0800 Subject: [PATCH 08/33] undetermined Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 30 +++- txnkv/transaction/txn_file_test.go | 227 ++++++++++++++++++++++++++++- 2 files changed, 253 insertions(+), 4 deletions(-) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index efec9cfcb8..1d74e53d09 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -519,13 +519,32 @@ func (a txnFileCommitAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backof return nil, errors.WithStack(tikverr.ErrBodyMissing) } commitResp := resp.Resp.(*kvrpcpb.CommitResponse) + if regionErr := commitResp.GetRegionError(); regionErr != nil { + if batch.isPrimary && regionErr.GetUndeterminedResult() != nil { + // Keep the RPC error, if any, as the cause of the ambiguity. + if c.getUndeterminedErr() == nil { + c.setUndeterminedErr(errors.New(regionErr.String())) + } + return nil, errors.WithStack(tikverr.ErrResultUndetermined) + } + return resp, nil + } + if batch.isPrimary { + // TiKV has definitively processed the primary commit request. + c.setUndeterminedErr(nil) + } if keyErr := commitResp.GetError(); keyErr != nil { if rejected := keyErr.GetCommitTsExpired(); rejected != nil { logutil.Logger(bo.GetCtx()).Info("2PC commitTS rejected by TiKV, retry with a newer commitTS", zap.Uint64("txnStartTS", c.startTS), zap.Stringer("info", logutil.Hex(rejected))) - if !batch.isPrimary { - return nil, errors.New("2PC commitTS rejected by TiKV, but the txn-file batch is not the primary batch") + if !batch.isPrimary || !bytes.Equal(rejected.Key, c.primary()) { + logutil.Logger(bo.GetCtx()).Error("2PC commitTS rejected by TiKV, but the key is not the primary key", + zap.Uint64("txnStartTS", c.startTS), + zap.String("key", redact.Key(rejected.Key)), + zap.String("primary", redact.Key(c.primary())), + zap.Bool("batchIsPrimary", batch.isPrimary)) + return nil, errors.New("2PC commitTS rejected by TiKV, but the key is not the primary key") } // Do not retry for a txn which has a too large MinCommitTs @@ -725,6 +744,13 @@ func (c *twoPhaseCommitter) executeTxnFile(ctx context.Context) (err error) { err = c.executeTxnFileAction(commitBo, c.txnFileCtx.slice, txnFileCommitAction{}) stepDone("commit") if err != nil { + if undeterminedErr := c.getUndeterminedErr(); undeterminedErr != nil { + logutil.Logger(ctx).Warn("txn file commit result undetermined", + zap.Error(err), + zap.NamedError("rpcErr", undeterminedErr), + zap.Uint64("txnStartTS", c.startTS)) + err = errors.WithStack(tikverr.ErrResultUndetermined) + } return } diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 1b1a75b8b9..820857d44c 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -17,7 +17,6 @@ package transaction import ( "context" "encoding/json" - "errors" "fmt" "math/rand" "net/http" @@ -28,11 +27,14 @@ import ( "testing" "time" + "github.com/pingcap/kvproto/pkg/errorpb" "github.com/pingcap/kvproto/pkg/kvrpcpb" + "github.com/pkg/errors" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/tikv/client-go/v2/config" "github.com/tikv/client-go/v2/config/retry" + tikverr "github.com/tikv/client-go/v2/error" "github.com/tikv/client-go/v2/internal/apicodec" "github.com/tikv/client-go/v2/internal/client" "github.com/tikv/client-go/v2/internal/latch" @@ -256,6 +258,7 @@ func TestTxnFileCommitTSExpiredRetryUsesPreparedTimestamp(t *testing.T) { StartTs: 1, AttemptedCommitTs: req.Commit().CommitVersion, MinCommitTs: 100, + Key: []byte("k"), }, }}}, nil } @@ -274,6 +277,7 @@ func TestTxnFileCommitTSExpiredRetryUsesPreparedTimestamp(t *testing.T) { return commitTS == 102 }) committer.commitTS = 2 + committer.primaryKey = []byte("k") bo := retry.NewBackoffer(context.Background(), 1000) location, err := regionCache.LocateKey(bo, []byte("k")) @@ -323,13 +327,232 @@ func TestTxnFileCommitTSExpiredRetryUsesPreparedTimestamp(t *testing.T) { _, err = (txnFileCommitAction{}).executeBatch(committer, bo, batch) - require.ErrorContains(t, err, "txn-file batch is not the primary batch") + require.ErrorContains(t, err, "key is not the primary key") require.Equal(t, 1, requestCount) require.Equal(t, uint64(2), committer.commitTS) require.Zero(t, store.timestampCalls) require.Zero(t, checker.calls) require.Zero(t, commitOracle.calls) require.Zero(t, upperBoundCalls) + + requestCount = 0 + store.timestamps = []uint64{106} + store.timestampCalls = 0 + checker.calls = 0 + commitOracle.calls = 0 + upperBoundCalls = 0 + committer.commitTS = 2 + batch.isPrimary = true + kvClient.onSend = func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + requestCount++ + return &tikvrpc.Response{Resp: &kvrpcpb.CommitResponse{Error: &kvrpcpb.KeyError{ + CommitTsExpired: &kvrpcpb.CommitTsExpired{ + StartTs: 1, + AttemptedCommitTs: req.Commit().CommitVersion, + MinCommitTs: 100, + Key: []byte("not-primary"), + }, + }}}, nil + } + + _, err = (txnFileCommitAction{}).executeBatch(committer, bo, batch) + + require.ErrorContains(t, err, "key is not the primary key") + require.Equal(t, 1, requestCount) + require.Equal(t, uint64(2), committer.commitTS) + require.Zero(t, store.timestampCalls) + require.Zero(t, checker.calls) + require.Zero(t, commitOracle.calls) + require.Zero(t, upperBoundCalls) +} + +func newTxnFileCommitTestBatch( + t *testing.T, + onSend func(context.Context, string, *tikvrpc.Request, time.Duration) (*tikvrpc.Response, error), +) (*twoPhaseCommitter, *retry.Backoffer, chunkBatch) { + t.Helper() + + pd := &mockPDClient{} + regionCache := locate.NewTestRegionCache() + regionCache.SetPDClient(pd) + t.Cleanup(regionCache.Close) + + store := &txnFileCommitTSStore{ + oracle: &txnFileCommitTSOracle{}, + regionCache: regionCache, + client: &fnClient{onSend: onSend}, + } + committer := newTxnFileCommitTSTestCommitter(store, &txnFileSchemaLeaseChecker{}, nil) + committer.commitTS = 2 + + bo := retry.NewBackoffer(context.Background(), 1000) + location, err := regionCache.LocateKey(bo, []byte("k")) + require.NoError(t, err) + batch := chunkBatch{ + txnChunkSlice: txnChunkSlice{ + chunkIDs: []uint64{1}, + chunkRanges: []txnChunkRange{{ + smallest: []byte("k"), + biggest: []byte("k"), + }}, + }, + region: location, + sampleKeys: [][]byte{[]byte("k")}, + isPrimary: true, + } + return committer, bo, batch +} + +func TestTxnFileCommitPrimaryRPCErrorMarksResultUndetermined(t *testing.T) { + committer, bo, batch := newTxnFileCommitTestBatch(t, func(context.Context, string, *tikvrpc.Request, time.Duration) (*tikvrpc.Response, error) { + return nil, context.Canceled + }) + + _, err := (txnFileCommitAction{}).executeBatch(committer, bo, batch) + + require.ErrorIs(t, err, context.Canceled) + require.Equal(t, context.Canceled, errors.Cause(committer.getUndeterminedErr())) +} + +func TestTxnFileCommitSecondaryRPCErrorIsNotResultUndetermined(t *testing.T) { + committer, bo, batch := newTxnFileCommitTestBatch(t, func(context.Context, string, *tikvrpc.Request, time.Duration) (*tikvrpc.Response, error) { + return nil, context.Canceled + }) + batch.isPrimary = false + + _, err := (txnFileCommitAction{}).executeBatch(committer, bo, batch) + + require.ErrorIs(t, err, context.Canceled) + require.Nil(t, committer.getUndeterminedErr()) +} + +func TestTxnFileCommitClearsUndeterminedErrOnDefinitivePrimaryResponse(t *testing.T) { + tests := []struct { + name string + resp *kvrpcpb.CommitResponse + }{ + { + name: "success", + resp: &kvrpcpb.CommitResponse{}, + }, + { + name: "key error", + resp: &kvrpcpb.CommitResponse{Error: &kvrpcpb.KeyError{Abort: "aborted"}}, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + committer, bo, batch := newTxnFileCommitTestBatch(t, func(context.Context, string, *tikvrpc.Request, time.Duration) (*tikvrpc.Response, error) { + return &tikvrpc.Response{Resp: tt.resp}, nil + }) + committer.setUndeterminedErr(errors.New("stale RPC error")) + + _, err := (txnFileCommitAction{}).executeBatch(committer, bo, batch) + + if tt.resp.GetError() == nil { + require.NoError(t, err) + } else { + require.Error(t, err) + } + require.Nil(t, committer.getUndeterminedErr()) + }) + } +} + +func TestTxnFileCommitPrimaryUndeterminedRegionError(t *testing.T) { + regionErr := &errorpb.Error{UndeterminedResult: &errorpb.UndeterminedResult{}} + requestCount := 0 + committer, bo, batch := newTxnFileCommitTestBatch(t, func(context.Context, string, *tikvrpc.Request, time.Duration) (*tikvrpc.Response, error) { + requestCount++ + return &tikvrpc.Response{Resp: &kvrpcpb.CommitResponse{RegionError: regionErr}}, nil + }) + + _, err := (txnFileCommitAction{}).executeBatch(committer, bo, batch) + + require.ErrorIs(t, err, tikverr.ErrResultUndetermined) + require.Equal(t, regionErr.String(), errors.Cause(committer.getUndeterminedErr()).Error()) + require.Equal(t, 1, requestCount) +} + +func TestTxnFileCommitPrimaryRPCErrorIsNormalized(t *testing.T) { + pd := &mockPDClient{} + regionCache := locate.NewTestRegionCache() + regionCache.SetPDClient(pd) + defer regionCache.Close() + + chunkWriter := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, http.MethodPost, r.Method) + _, err := w.Write([]byte(`{"chunk_id":1}`)) + require.NoError(t, err) + })) + defer chunkWriter.Close() + + origCfg := config.GetGlobalConfig() + newCfg := *origCfg + newCfg.TiKVClient.TxnChunkWriterAddr = chunkWriter.Listener.Addr().String() + config.StoreGlobalConfig(&newCfg) + defer func() { + config.StoreGlobalConfig(origCfg) + once = sync.Once{} + cli = nil + errCli = nil + scheme = "" + }() + + once = sync.Once{} + cli = nil + errCli = nil + scheme = "" + + var commitRequestCount atomic.Int64 + var rollbackRequestCount atomic.Int64 + store := &txnFileCommitTSStore{ + timestamps: []uint64{2}, + oracle: &txnFileCommitTSOracle{}, + regionCache: regionCache, + client: &fnClient{onSend: func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + switch req.Type { + case tikvrpc.CmdPrewrite: + return &tikvrpc.Response{Resp: &kvrpcpb.PrewriteResponse{}}, nil + case tikvrpc.CmdCommit: + commitRequestCount.Add(1) + return nil, context.Canceled + case tikvrpc.CmdBatchRollback: + rollbackRequestCount.Add(1) + return &tikvrpc.Response{Resp: &kvrpcpb.BatchRollbackResponse{}}, nil + default: + return nil, errors.Errorf("unexpected request type %s", req.Type) + } + }}, + } + memDB := unionstore.NewMemDB() + require.NoError(t, memDB.Set([]byte("k"), []byte("v"))) + txn := &KVTxn{ + store: store, + startTS: 1, + startTime: time.Now(), + schemaVer: txnFileSchemaVer(10), + schemaLeaseChecker: &txnFileSchemaLeaseChecker{}, + scope: oracle.GlobalTxnScope, + vars: tikv.DefaultVars, + us: unionstore.NewUnionStore(memDB, nil), + } + committer := &twoPhaseCommitter{ + store: store, + txn: txn, + startTS: txn.startTS, + regionTxnSize: map[uint64]int{}, + } + require.NoError(t, committer.initKeysAndMutations(context.Background())) + committer.ttlManager.state = stateRunning + + err := committer.executeTxnFile(context.Background()) + + require.ErrorIs(t, err, tikverr.ErrResultUndetermined) + require.Equal(t, context.Canceled, errors.Cause(committer.getUndeterminedErr())) + require.Equal(t, int64(1), commitRequestCount.Load()) + require.Zero(t, rollbackRequestCount.Load()) } func TestChunkSliceSortAndDedup(t *testing.T) { From 70cee6ea7398774946af3af0d836ddc8d209a7ea Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Wed, 15 Jul 2026 09:21:44 +0800 Subject: [PATCH 09/33] resource group tag Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 30 ++++++ txnkv/transaction/txn_file_test.go | 162 +++++++++++++++++++++++++++++ 2 files changed, 192 insertions(+) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 1d74e53d09..4a814a53f8 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -30,6 +30,7 @@ import ( "sync" "time" + "github.com/golang/protobuf/proto" //nolint:staticcheck "github.com/pingcap/kvproto/pkg/errorpb" "github.com/pingcap/kvproto/pkg/kvrpcpb" "github.com/pingcap/kvproto/pkg/metapb" @@ -105,6 +106,32 @@ func markTxnFileRetryRequest(req *tikvrpc.Request, bo *retry.Backoffer) { } } +func (c *twoPhaseCommitter) applyTxnFileResourceGroupTagger(req *tikvrpc.Request) { + if c.resourceGroupTag == nil && c.resourceGroupTagger != nil { + c.resourceGroupTagger(req) + } +} + +func (c *twoPhaseCommitter) applyTxnFilePrewriteResourceGroupTagger(req *tikvrpc.Request, sampleKeys [][]byte) { + if len(sampleKeys) == 0 || c.resourceGroupTag != nil || c.resourceGroupTagger == nil { + return + } + + prewrite := req.Prewrite() + tagReq := tikvrpc.NewRequest(tikvrpc.CmdPrewrite, &kvrpcpb.PrewriteRequest{ + Mutations: []*kvrpcpb.Mutation{{Key: slices.Clone(sampleKeys[0])}}, + PrimaryLock: slices.Clone(prewrite.PrimaryLock), + StartVersion: prewrite.StartVersion, + LockTtl: prewrite.LockTtl, + MaxCommitTs: prewrite.MaxCommitTs, + AssertionLevel: prewrite.AssertionLevel, + TxnFileChunks: slices.Clone(prewrite.TxnFileChunks), + TxnSize: prewrite.TxnSize, + }, *proto.Clone(&req.Context).(*kvrpcpb.Context)) + c.applyTxnFileResourceGroupTagger(tagReq) + req.ResourceGroupTag = tagReq.ResourceGroupTag +} + // txnChunkSlice should be sorted by txnChunkRange.smallest and no overlapping. type txnChunkSlice struct { chunkIDs []uint64 @@ -314,6 +341,7 @@ func (a txnFilePrewriteAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back ResourceGroupName: c.resourceGroupName, }, }) + c.applyTxnFilePrewriteResourceGroupTagger(req, batch.sampleKeys) markTxnFileRetryRequest(req, bo) sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.GetOracle()) var resolvingRecordToken *int @@ -504,6 +532,7 @@ func (a txnFileCommitAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backof ResourceGroupName: c.resourceGroupName, }, }) + c.applyTxnFileResourceGroupTagger(req) markTxnFileRetryRequest(req, bo) sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.GetOracle()) for { @@ -616,6 +645,7 @@ func (a txnFileRollbackAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back ResourceGroupName: c.resourceGroupName, }, }) + c.applyTxnFileResourceGroupTagger(req) markTxnFileRetryRequest(req, bo) sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.GetOracle()) resp, _, err1 := sender.SendReq(bo, req, batch.region.Region, client.ReadTimeoutShort) diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 820857d44c..314e263161 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -278,6 +278,10 @@ func TestTxnFileCommitTSExpiredRetryUsesPreparedTimestamp(t *testing.T) { }) committer.commitTS = 2 committer.primaryKey = []byte("k") + taggerCalls := 0 + committer.resourceGroupTagger = func(*tikvrpc.Request) { + taggerCalls++ + } bo := retry.NewBackoffer(context.Background(), 1000) location, err := regionCache.LocateKey(bo, []byte("k")) @@ -305,6 +309,7 @@ func TestTxnFileCommitTSExpiredRetryUsesPreparedTimestamp(t *testing.T) { require.Equal(t, 1, checker.calls) require.Equal(t, 1, commitOracle.calls) require.Equal(t, 1, upperBoundCalls) + require.Equal(t, 1, taggerCalls) requestCount = 0 store.timestamps = []uint64{104} @@ -400,9 +405,166 @@ func newTxnFileCommitTestBatch( sampleKeys: [][]byte{[]byte("k")}, isPrimary: true, } + committer.txnFileCtx = txnFileCtx{slice: batch.txnChunkSlice} return committer, bo, batch } +func TestTxnFileActionsApplyResourceGroupTagger(t *testing.T) { + tests := []struct { + name string + action txnFileAction + requestType tikvrpc.CmdType + newResponse func() *tikvrpc.Response + }{ + { + name: "prewrite", + action: txnFilePrewriteAction{}, + requestType: tikvrpc.CmdPrewrite, + newResponse: func() *tikvrpc.Response { + return &tikvrpc.Response{Resp: &kvrpcpb.PrewriteResponse{}} + }, + }, + { + name: "commit", + action: txnFileCommitAction{}, + requestType: tikvrpc.CmdCommit, + newResponse: func() *tikvrpc.Response { + return &tikvrpc.Response{Resp: &kvrpcpb.CommitResponse{}} + }, + }, + { + name: "rollback", + action: txnFileRollbackAction{}, + requestType: tikvrpc.CmdBatchRollback, + newResponse: func() *tikvrpc.Response { + return &tikvrpc.Response{Resp: &kvrpcpb.BatchRollbackResponse{}} + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + taggerCalls := 0 + committer, bo, batch := newTxnFileCommitTestBatch(t, func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + require.Equal(t, tt.requestType, req.Type) + require.Equal(t, []byte("dynamic-tag"), req.ResourceGroupTag) + switch req.Type { + case tikvrpc.CmdPrewrite: + prewrite := req.Prewrite() + require.Empty(t, prewrite.Mutations) + require.Equal(t, []uint64{1}, prewrite.TxnFileChunks) + require.Equal(t, []byte("k"), prewrite.PrimaryLock) + require.NotNil(t, req.ResourceControlContext) + require.Equal(t, "txn-file-test", req.ResourceControlContext.ResourceGroupName) + case tikvrpc.CmdCommit: + require.Equal(t, [][]byte{[]byte("k")}, req.Commit().Keys) + case tikvrpc.CmdBatchRollback: + require.Equal(t, [][]byte{[]byte("k")}, req.BatchRollback().Keys) + } + return tt.newResponse(), nil + }) + committer.resourceGroupName = "txn-file-test" + committer.resourceGroupTagger = func(req *tikvrpc.Request) { + taggerCalls++ + require.Equal(t, tt.requestType, req.Type) + require.NotNil(t, req.ResourceControlContext) + require.Equal(t, "txn-file-test", req.ResourceControlContext.ResourceGroupName) + switch req.Type { + case tikvrpc.CmdPrewrite: + prewrite := req.Prewrite() + require.Len(t, prewrite.Mutations, 1) + require.Equal(t, batch.sampleKeys[0], prewrite.Mutations[0].Key) + prewrite.PrimaryLock[0] = 'x' + prewrite.TxnFileChunks[0] = 99 + req.ResourceControlContext.ResourceGroupName = "tagger-mutated" + case tikvrpc.CmdCommit: + require.Equal(t, batch.sampleKeys, req.Commit().Keys) + case tikvrpc.CmdBatchRollback: + require.Equal(t, batch.sampleKeys, req.BatchRollback().Keys) + } + req.ResourceGroupTag = []byte("dynamic-tag") + } + + _, err := tt.action.executeBatch(committer, bo, batch) + + require.NoError(t, err) + require.Equal(t, 1, taggerCalls) + }) + } +} + +func TestTxnFileActionsPreserveStaticResourceGroupTag(t *testing.T) { + tests := []struct { + name string + action txnFileAction + requestType tikvrpc.CmdType + newResponse func() *tikvrpc.Response + }{ + { + name: "prewrite", + action: txnFilePrewriteAction{}, + requestType: tikvrpc.CmdPrewrite, + newResponse: func() *tikvrpc.Response { + return &tikvrpc.Response{Resp: &kvrpcpb.PrewriteResponse{}} + }, + }, + { + name: "commit", + action: txnFileCommitAction{}, + requestType: tikvrpc.CmdCommit, + newResponse: func() *tikvrpc.Response { + return &tikvrpc.Response{Resp: &kvrpcpb.CommitResponse{}} + }, + }, + { + name: "rollback", + action: txnFileRollbackAction{}, + requestType: tikvrpc.CmdBatchRollback, + newResponse: func() *tikvrpc.Response { + return &tikvrpc.Response{Resp: &kvrpcpb.BatchRollbackResponse{}} + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + taggerCalls := 0 + committer, bo, batch := newTxnFileCommitTestBatch(t, func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + require.Equal(t, tt.requestType, req.Type) + require.Equal(t, []byte("static-tag"), req.ResourceGroupTag) + return tt.newResponse(), nil + }) + committer.resourceGroupTag = []byte("static-tag") + committer.resourceGroupTagger = func(*tikvrpc.Request) { + taggerCalls++ + } + + _, err := tt.action.executeBatch(committer, bo, batch) + + require.NoError(t, err) + require.Zero(t, taggerCalls) + }) + } +} + +func TestTxnFilePrewriteTaggerSkipsBatchWithoutSampleKeys(t *testing.T) { + taggerCalls := 0 + committer, bo, batch := newTxnFileCommitTestBatch(t, func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + require.Empty(t, req.ResourceGroupTag) + require.Empty(t, req.Prewrite().Mutations) + return &tikvrpc.Response{Resp: &kvrpcpb.PrewriteResponse{}}, nil + }) + batch.sampleKeys = nil + committer.resourceGroupTagger = func(*tikvrpc.Request) { + taggerCalls++ + } + + _, err := (txnFilePrewriteAction{}).executeBatch(committer, bo, batch) + + require.NoError(t, err) + require.Zero(t, taggerCalls) +} + func TestTxnFileCommitPrimaryRPCErrorMarksResultUndetermined(t *testing.T) { committer, bo, batch := newTxnFileCommitTestBatch(t, func(context.Context, string, *tikvrpc.Request, time.Duration) (*tikvrpc.Response, error) { return nil, context.Canceled From b15c470e282c77f9184463744642936e671fb0d1 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Wed, 15 Jul 2026 15:16:16 +0800 Subject: [PATCH 10/33] note for skip case Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 1 + 1 file changed, 1 insertion(+) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 4a814a53f8..7c540f7aaa 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -113,6 +113,7 @@ func (c *twoPhaseCommitter) applyTxnFileResourceGroupTagger(req *tikvrpc.Request } func (c *twoPhaseCommitter) applyTxnFilePrewriteResourceGroupTagger(req *tikvrpc.Request, sampleKeys [][]byte) { + // Note: Batches containing only non-write operations have no sample key, so dynamic tagging is skipped. if len(sampleKeys) == 0 || c.resourceGroupTag != nil || c.resourceGroupTagger == nil { return } From dd46c190e4cdd1f0b19f11955c35ac05edf0010a Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Wed, 15 Jul 2026 15:37:47 +0800 Subject: [PATCH 11/33] register metrics Signed-off-by: Ping Yu --- metrics/metrics.go | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/metrics/metrics.go b/metrics/metrics.go index fbb89c75c1..69c74a13fa 100644 --- a/metrics/metrics.go +++ b/metrics/metrics.go @@ -1211,6 +1211,10 @@ func RegisterMetrics() { prometheus.MustRegister(TiKVTxnLagCommitTSWaitHistogram) prometheus.MustRegister(TiKVTxnLagCommitTSAttemptHistogram) prometheus.MustRegister(TiKVStaleBucketFromPDCounter) + prometheus.MustRegister(TiKVTxnFileRequestCounter) + prometheus.MustRegister(TiKVTxnFileWriteBytes) + prometheus.MustRegister(TiKVTxnFileMutationSizeHistogram) + prometheus.MustRegister(TiKVTxnFileDuration) } // readCounter reads the value of a prometheus.Counter. From 06541e7b8cacc3aea9352e170b42fe0f34d5d29e Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Wed, 15 Jul 2026 18:17:38 +0800 Subject: [PATCH 12/33] validate config Signed-off-by: Ping Yu --- config/client.go | 16 ++++++++++ config/config_test.go | 69 +++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+) diff --git a/config/client.go b/config/client.go index 429ba21628..f6de4bfb12 100644 --- a/config/client.go +++ b/config/client.go @@ -267,6 +267,22 @@ func (config *TiKVClient) Valid() error { if config.GetGrpcKeepAliveTimeout() < time.Millisecond*50 { return fmt.Errorf("grpc-keepalive-timeout should be at least 0.05, but got %f", config.GrpcKeepAliveTimeout) } + return validateTxnFileConfig(config) +} + +func validateTxnFileConfig(config *TiKVClient) error { + if config.TxnChunkMaxSize == 0 { + return fmt.Errorf("txn-chunk-max-size should be greater than 0") + } + if config.TxnChunkMaxSize > math.MaxInt { + return fmt.Errorf("txn-chunk-max-size should not exceed %d, but got %d", math.MaxInt, config.TxnChunkMaxSize) + } + if config.TxnChunkWriterConcurrency == 0 { + return fmt.Errorf("txn-chunk-writer-concurrency should be greater than 0") + } + if config.TxnChunkWriterConcurrency > math.MaxInt { + return fmt.Errorf("txn-chunk-writer-concurrency should not exceed %d, but got %d", math.MaxInt, config.TxnChunkWriterConcurrency) + } return nil } diff --git a/config/config_test.go b/config/config_test.go index a9baace9c9..6857d591dc 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -35,6 +35,8 @@ package config import ( + "fmt" + "math" "testing" "time" @@ -89,3 +91,70 @@ func TestValidateGRPCKeepAliveTimeout(t *testing.T) { assert.NotNil(t, cfg.Valid()) assert.Equal(t, "grpc-keepalive-timeout should be at least 0.05, but got 0.040000", cfg.Valid().Error()) } + +func TestValidateTxnFileConfig(t *testing.T) { + maxInt := uint64(math.MaxInt) + tests := []struct { + name string + configure func(*TiKVClient) + err string + }{ + { + name: "default", + }, + { + name: "zero chunk size", + configure: func(cfg *TiKVClient) { + cfg.TxnChunkMaxSize = 0 + }, + err: "txn-chunk-max-size should be greater than 0", + }, + { + name: "maximum chunk size", + configure: func(cfg *TiKVClient) { + cfg.TxnChunkMaxSize = maxInt + }, + }, + { + name: "chunk size exceeds int", + configure: func(cfg *TiKVClient) { + cfg.TxnChunkMaxSize = maxInt + 1 + }, + err: fmt.Sprintf("txn-chunk-max-size should not exceed %d, but got %d", maxInt, maxInt+1), + }, + { + name: "zero writer concurrency", + configure: func(cfg *TiKVClient) { + cfg.TxnChunkWriterConcurrency = 0 + }, + err: "txn-chunk-writer-concurrency should be greater than 0", + }, + { + name: "maximum writer concurrency", + configure: func(cfg *TiKVClient) { + cfg.TxnChunkWriterConcurrency = uint(maxInt) + }, + }, + { + name: "writer concurrency exceeds int", + configure: func(cfg *TiKVClient) { + cfg.TxnChunkWriterConcurrency = uint(maxInt) + 1 + }, + err: fmt.Sprintf("txn-chunk-writer-concurrency should not exceed %d, but got %d", maxInt, uint(maxInt)+1), + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + cfg := DefaultTiKVClient() + if test.configure != nil { + test.configure(&cfg) + } + if test.err == "" { + assert.NoError(t, cfg.Valid()) + return + } + assert.EqualError(t, cfg.Valid(), test.err) + }) + } +} From ace38d1b7b008ba0065aaf432148315b2e5227b5 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Wed, 15 Jul 2026 18:35:47 +0800 Subject: [PATCH 13/33] comment for GetMaxStartKey/GetMinEndKey Signed-off-by: Ping Yu --- util/misc.go | 2 ++ 1 file changed, 2 insertions(+) diff --git a/util/misc.go b/util/misc.go index 9b65163952..dad9f727a3 100644 --- a/util/misc.go +++ b/util/misc.go @@ -177,6 +177,7 @@ func (o Option[T]) Inner() *T { return o.inner } +// GetMaxStartKey returns the lexicographically larger start key, treating an empty key as unbounded below. func GetMaxStartKey(lhs []byte, rhs []byte) []byte { if bytes.Compare(lhs, rhs) > 0 { return lhs @@ -184,6 +185,7 @@ func GetMaxStartKey(lhs []byte, rhs []byte) []byte { return rhs } +// GetMinEndKey returns the lexicographically smaller end key, treating an empty key as unbounded above. func GetMinEndKey(lhs []byte, rhs []byte) []byte { if len(rhs) == 0 { return lhs From 4f653f2847e64c0b7b1036d1496f1ac5a5204544 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Wed, 15 Jul 2026 18:48:52 +0800 Subject: [PATCH 14/33] require -> assert Signed-off-by: Ping Yu --- txnkv/transaction/txn_file_test.go | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 314e263161..f182d80be4 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -644,9 +644,9 @@ func TestTxnFileCommitPrimaryRPCErrorIsNormalized(t *testing.T) { defer regionCache.Close() chunkWriter := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - require.Equal(t, http.MethodPost, r.Method) + assert.Equal(t, http.MethodPost, r.Method) _, err := w.Write([]byte(`{"chunk_id":1}`)) - require.NoError(t, err) + assert.NoError(t, err) })) defer chunkWriter.Close() From 707c2cf1eb2749cced8d16f859362a07397e9eab Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Tue, 28 Jul 2026 23:23:17 +0800 Subject: [PATCH 15/33] fix CI Signed-off-by: Ping Yu --- txnkv/transaction/txn_file_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index f182d80be4..d8eeea60b4 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -707,7 +707,7 @@ func TestTxnFileCommitPrimaryRPCErrorIsNormalized(t *testing.T) { regionTxnSize: map[uint64]int{}, } require.NoError(t, committer.initKeysAndMutations(context.Background())) - committer.ttlManager.state = stateRunning + committer.state = stateRunning err := committer.executeTxnFile(context.Background()) From 8b878c4428756fa26dd618d3bd16796973e1c0fc Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Thu, 6 Aug 2026 17:42:04 +0800 Subject: [PATCH 16/33] always get resource group tag Signed-off-by: Ping Yu --- txnkv/transaction/2pc.go | 32 ++++++++--------- txnkv/transaction/2pc_test.go | 38 ++++++++++---------- txnkv/transaction/txn_file.go | 57 ++++++++++++++++++------------ txnkv/transaction/txn_file_test.go | 51 +++++++++++++++++++------- 4 files changed, 106 insertions(+), 72 deletions(-) diff --git a/txnkv/transaction/2pc.go b/txnkv/transaction/2pc.go index 88a3d45c79..067c662fe1 100644 --- a/txnkv/transaction/2pc.go +++ b/txnkv/transaction/2pc.go @@ -340,10 +340,9 @@ type CommitterMutations interface { NeedConstraintCheckInPrewrite(i int) bool } -// MutationsHasDataInRange returns whether mutations has data in the range [start, end). -// If it has, it returns the primary or first write key in the range. -// Note that the firstDataKey can be empty when the range contains only non-write ops (and not the primary at pos 0). -func MutationsHasDataInRange(mutations CommitterMutations, start []byte, end []byte) ([]byte /* firstDataKey */, bool) { +// MutationsHasDataInRange returns the first mutation key and first data key in [start, end). +// The first data key is the primary or first write key, and can be nil when the range only contains non-write operations. +func MutationsHasDataInRange(mutations CommitterMutations, start []byte, end []byte) (firstKey, firstDataKey []byte, ok bool) { isInRange := func(pos int) bool { return pos < mutations.Len() && (len(end) == 0 || bytes.Compare(mutations.GetKey(pos), end) < 0) } @@ -356,23 +355,20 @@ func MutationsHasDataInRange(mutations CommitterMutations, start []byte, end []b pos := sort.Search(mutations.Len(), func(i int) bool { return bytes.Compare(mutations.GetKey(i), start) >= 0 }) - if isInRange(pos) { - var firstDataKey []byte - for { - // Always return primary key if it's in the range. - if pos == 0 || isOpForWrite(mutations.GetOp(pos)) { - firstDataKey = mutations.GetKey(pos) - break - } + if !isInRange(pos) { + return nil, nil, false + } - pos++ - if !isInRange(pos) { - break - } + firstKey = mutations.GetKey(pos) + for isInRange(pos) { + // Always return primary key if it's in the range. + if pos == 0 || isOpForWrite(mutations.GetOp(pos)) { + firstDataKey = mutations.GetKey(pos) + break } - return firstDataKey, true + pos++ } - return nil, false + return firstKey, firstDataKey, true } // PlainMutations contains transaction operations. diff --git a/txnkv/transaction/2pc_test.go b/txnkv/transaction/2pc_test.go index 6f0760404e..e31cee105b 100644 --- a/txnkv/transaction/2pc_test.go +++ b/txnkv/transaction/2pc_test.go @@ -148,32 +148,34 @@ func TestMutationsHasDataInRange(t *testing.T) { } type Case struct { - start int - end int - expectd bool - firstKey int + start int + end int + expectd bool + firstKey int + firstDataKey int } cases := []Case{ - {-1, -1, true, 10}, - {-1, 5, false, -1}, - {0, 10, false, -1}, - {0, 11, true, 10}, - {0, 30, true, 10}, - {0, -1, true, 10}, - {10, 20, true, 10}, - {15, 16, false, -1}, - {15, 17, true, -1}, - {15, -1, true, 18}, - {20, 30, false, -1}, - {21, 30, false, -1}, - {21, -1, false, -1}, + {-1, -1, true, 10, 10}, + {-1, 5, false, -1, -1}, + {0, 10, false, -1, -1}, + {0, 11, true, 10, 10}, + {0, 30, true, 10, 10}, + {0, -1, true, 10, 10}, + {10, 20, true, 10, 10}, + {15, 16, false, -1, -1}, + {15, 17, true, 16, -1}, + {15, -1, true, 16, 18}, + {20, 30, false, -1, -1}, + {21, 30, false, -1, -1}, + {21, -1, false, -1, -1}, } for _, c := range cases { - firstKey, got := MutationsHasDataInRange(&muts, iToKey(c.start), iToKey(c.end)) + firstKey, firstDataKey, got := MutationsHasDataInRange(&muts, iToKey(c.start), iToKey(c.end)) assert.Equal(c.expectd, got) if got { assert.Equal(iToKey(c.firstKey), firstKey) + assert.Equal(iToKey(c.firstDataKey), firstDataKey) } } } diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 7c540f7aaa..7b0b75ef9e 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -75,9 +75,10 @@ type txnFileCtx struct { type chunkBatch struct { txnChunkSlice - region *locate.KeyLocation - sampleKeys [][]byte - isPrimary bool + region *locate.KeyLocation + firstKey []byte + sampleDataKeys [][]byte + isPrimary bool } func (b chunkBatch) String() string { @@ -85,8 +86,8 @@ func (b chunkBatch) String() string { b.region.Region.GetID(), b.isPrimary, b.chunkIDs) } -func (b *chunkBatch) getSampleKeys() [][]byte { - return b.sampleKeys +func (b *chunkBatch) getSampleDataKeys() [][]byte { + return b.sampleDataKeys } func (b *chunkBatch) getBatchTxnSize() uint64 { @@ -112,15 +113,18 @@ func (c *twoPhaseCommitter) applyTxnFileResourceGroupTagger(req *tikvrpc.Request } } -func (c *twoPhaseCommitter) applyTxnFilePrewriteResourceGroupTagger(req *tikvrpc.Request, sampleKeys [][]byte) { - // Note: Batches containing only non-write operations have no sample key, so dynamic tagging is skipped. - if len(sampleKeys) == 0 || c.resourceGroupTag != nil || c.resourceGroupTagger == nil { +func (c *twoPhaseCommitter) applyTxnFilePrewriteResourceGroupTagger(req *tikvrpc.Request, firstKey []byte) { + if c.resourceGroupTag != nil || c.resourceGroupTagger == nil { return } prewrite := req.Prewrite() + mutations := make([]*kvrpcpb.Mutation, 0, 1) + if len(firstKey) > 0 { + mutations = append(mutations, &kvrpcpb.Mutation{Key: slices.Clone(firstKey)}) + } tagReq := tikvrpc.NewRequest(tikvrpc.CmdPrewrite, &kvrpcpb.PrewriteRequest{ - Mutations: []*kvrpcpb.Mutation{{Key: slices.Clone(sampleKeys[0])}}, + Mutations: mutations, PrimaryLock: slices.Clone(prewrite.PrimaryLock), StartVersion: prewrite.StartVersion, LockTtl: prewrite.LockTtl, @@ -215,26 +219,31 @@ func (cs *txnChunkSlice) groupToBatches(c *locate.RegionCache, bo *retry.Backoff for i, chunkRange := range cs.chunkRanges { chunkID := cs.chunkIDs[i] - regions, firstKeys, err := chunkRange.getOverlapRegions(c, bo, mutations) + regions, firstKeys, firstDataKeys, err := chunkRange.getOverlapRegions(c, bo, mutations) if err != nil { return nil, errors.WithStack(err) } for j, r := range regions { firstKey := firstKeys[j] + firstDataKey := firstDataKeys[j] bk := batchMapKey{regionID: r.Region.GetID(), regionVer: r.Region.GetVer()} if batchMap[bk] == nil { batchMap[bk] = &chunkBatch{ - region: r, - sampleKeys: make([][]byte, 0, 1), + region: r, + firstKey: firstKey, + sampleDataKeys: make([][]byte, 0, 1), } } batch := batchMap[bk] + if len(batch.firstKey) == 0 { + batch.firstKey = firstKey + } batch.append(chunkID, chunkRange) - if len(firstKey) > 0 { - batch.sampleKeys = append(batch.sampleKeys, firstKey) + if len(firstDataKey) > 0 { + batch.sampleDataKeys = append(batch.sampleDataKeys, firstDataKey) } } } @@ -280,18 +289,19 @@ func newTxnChunkRange(smallest []byte, biggest []byte, entries uint64) txnChunkR } } -func (r *txnChunkRange) getOverlapRegions(c *locate.RegionCache, bo *retry.Backoffer, mutations CommitterMutations) ([]*locate.KeyLocation, [][]byte, error) { +func (r *txnChunkRange) getOverlapRegions(c *locate.RegionCache, bo *retry.Backoffer, mutations CommitterMutations) ([]*locate.KeyLocation, [][]byte, [][]byte, error) { regions := make([]*locate.KeyLocation, 0) firstKeys := make([][]byte, 0) + firstDataKeys := make([][]byte, 0) startKey := r.smallest exclusiveBiggest := kv.NextKey(r.biggest) for bytes.Compare(startKey, r.biggest) <= 0 { loc, err := c.LocateKey(bo, startKey) if err != nil { logutil.Logger(bo.GetCtx()).Error("locate key failed", zap.Error(err), zap.String("startKey", redact.Key(startKey))) - return nil, nil, errors.Wrap(err, "locate key failed") + return nil, nil, nil, errors.Wrap(err, "locate key failed") } - firstKey, ok := MutationsHasDataInRange( + firstKey, firstDataKey, ok := MutationsHasDataInRange( mutations, util.GetMaxStartKey(r.smallest, loc.StartKey), util.GetMinEndKey(exclusiveBiggest, loc.EndKey), @@ -299,13 +309,14 @@ func (r *txnChunkRange) getOverlapRegions(c *locate.RegionCache, bo *retry.Backo if ok { regions = append(regions, loc) firstKeys = append(firstKeys, firstKey) + firstDataKeys = append(firstDataKeys, firstDataKey) } if len(loc.EndKey) == 0 { break } startKey = loc.EndKey } - return regions, firstKeys, nil + return regions, firstKeys, firstDataKeys, nil } type txnFileAction interface { @@ -342,7 +353,7 @@ func (a txnFilePrewriteAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back ResourceGroupName: c.resourceGroupName, }, }) - c.applyTxnFilePrewriteResourceGroupTagger(req, batch.sampleKeys) + c.applyTxnFilePrewriteResourceGroupTagger(req, batch.firstKey) markTxnFileRetryRequest(req, bo) sender := locate.NewRegionRequestSender(c.store.GetRegionCache(), c.store.GetTiKVClient(), c.store.GetOracle()) var resolvingRecordToken *int @@ -375,8 +386,8 @@ func (a txnFilePrewriteAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back if regionErr.GetDiskFull() != nil { return resp, errors.New(regionErr.String()) } - if len(batch.sampleKeys) > 0 { - loc, err := c.store.GetRegionCache().LocateKey(bo, batch.sampleKeys[0]) + if len(batch.sampleDataKeys) > 0 { + loc, err := c.store.GetRegionCache().LocateKey(bo, batch.sampleDataKeys[0]) if err != nil { return nil, err } @@ -517,7 +528,7 @@ func (c *twoPhaseCommitter) prepareTxnFileCommitTS(bo *retry.Backoffer) (uint64, func (a txnFileCommitAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backoffer, batch chunkBatch) (*tikvrpc.Response, error) { req := tikvrpc.NewRequest(tikvrpc.CmdCommit, &kvrpcpb.CommitRequest{ - Keys: batch.getSampleKeys(), // To help detect duplicated request. + Keys: batch.getSampleDataKeys(), // To help detect duplicated request. StartVersion: c.startTS, CommitVersion: c.commitTS, IsTxnFile: true, @@ -631,7 +642,7 @@ var _ txnFileAction = (*txnFileRollbackAction)(nil) func (a txnFileRollbackAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backoffer, batch chunkBatch) (*tikvrpc.Response, error) { req := tikvrpc.NewRequest(tikvrpc.CmdBatchRollback, &kvrpcpb.BatchRollbackRequest{ - Keys: batch.getSampleKeys(), // To help detect duplicated request. + Keys: batch.getSampleDataKeys(), // To help detect duplicated request. StartVersion: c.startTS, IsTxnFile: true, }, kvrpcpb.Context{ diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index d8eeea60b4..34f1680953 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -294,9 +294,10 @@ func TestTxnFileCommitTSExpiredRetryUsesPreparedTimestamp(t *testing.T) { biggest: []byte("k"), }}, }, - region: location, - sampleKeys: [][]byte{[]byte("k")}, - isPrimary: true, + region: location, + sampleDataKeys: [][]byte{[]byte("k")}, + firstKey: []byte("k"), + isPrimary: true, } _, err = (txnFileCommitAction{}).executeBatch(committer, bo, batch) @@ -401,9 +402,10 @@ func newTxnFileCommitTestBatch( biggest: []byte("k"), }}, }, - region: location, - sampleKeys: [][]byte{[]byte("k")}, - isPrimary: true, + region: location, + sampleDataKeys: [][]byte{[]byte("k")}, + firstKey: []byte("k"), + isPrimary: true, } committer.txnFileCtx = txnFileCtx{slice: batch.txnChunkSlice} return committer, bo, batch @@ -473,14 +475,14 @@ func TestTxnFileActionsApplyResourceGroupTagger(t *testing.T) { case tikvrpc.CmdPrewrite: prewrite := req.Prewrite() require.Len(t, prewrite.Mutations, 1) - require.Equal(t, batch.sampleKeys[0], prewrite.Mutations[0].Key) + require.Equal(t, batch.firstKey, prewrite.Mutations[0].Key) prewrite.PrimaryLock[0] = 'x' prewrite.TxnFileChunks[0] = 99 req.ResourceControlContext.ResourceGroupName = "tagger-mutated" case tikvrpc.CmdCommit: - require.Equal(t, batch.sampleKeys, req.Commit().Keys) + require.Equal(t, batch.sampleDataKeys, req.Commit().Keys) case tikvrpc.CmdBatchRollback: - require.Equal(t, batch.sampleKeys, req.BatchRollback().Keys) + require.Equal(t, batch.sampleDataKeys, req.BatchRollback().Keys) } req.ResourceGroupTag = []byte("dynamic-tag") } @@ -547,22 +549,45 @@ func TestTxnFileActionsPreserveStaticResourceGroupTag(t *testing.T) { } } -func TestTxnFilePrewriteTaggerSkipsBatchWithoutSampleKeys(t *testing.T) { +func TestTxnFilePrewriteTaggerUsesFirstKeyWithoutSampleDataKeys(t *testing.T) { taggerCalls := 0 committer, bo, batch := newTxnFileCommitTestBatch(t, func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { require.Empty(t, req.ResourceGroupTag) require.Empty(t, req.Prewrite().Mutations) return &tikvrpc.Response{Resp: &kvrpcpb.PrewriteResponse{}}, nil }) - batch.sampleKeys = nil - committer.resourceGroupTagger = func(*tikvrpc.Request) { + batch.sampleDataKeys = nil + committer.resourceGroupTagger = func(req *tikvrpc.Request) { taggerCalls++ + require.Len(t, req.Prewrite().Mutations, 1) + require.Equal(t, batch.firstKey, req.Prewrite().Mutations[0].Key) } _, err := (txnFilePrewriteAction{}).executeBatch(committer, bo, batch) require.NoError(t, err) - require.Zero(t, taggerCalls) + require.Equal(t, 1, taggerCalls) +} + +func TestTxnFilePrewriteTaggerAppliesWithoutFirstKey(t *testing.T) { + taggerCalls := 0 + committer, bo, batch := newTxnFileCommitTestBatch(t, func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + require.Equal(t, []byte("metadata-tag"), req.ResourceGroupTag) + require.Empty(t, req.Prewrite().Mutations) + return &tikvrpc.Response{Resp: &kvrpcpb.PrewriteResponse{}}, nil + }) + batch.firstKey = nil + batch.sampleDataKeys = nil + committer.resourceGroupTagger = func(req *tikvrpc.Request) { + taggerCalls++ + require.Empty(t, req.Prewrite().Mutations) + req.ResourceGroupTag = []byte("metadata-tag") + } + + _, err := (txnFilePrewriteAction{}).executeBatch(committer, bo, batch) + + require.NoError(t, err) + require.Equal(t, 1, taggerCalls) } func TestTxnFileCommitPrimaryRPCErrorMarksResultUndetermined(t *testing.T) { From b30dcc8b5417dbcc22650f6dcc015db03110cd0e Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Thu, 6 Aug 2026 20:25:40 +0800 Subject: [PATCH 17/33] handle primary not first Signed-off-by: Ping Yu --- txnkv/transaction/2pc.go | 7 +++---- txnkv/transaction/2pc_test.go | 16 ++++++++-------- txnkv/transaction/txn_file.go | 18 +++++++++++++++++- txnkv/transaction/txn_file_test.go | 25 +++++++++++++++++++++++++ 4 files changed, 53 insertions(+), 13 deletions(-) diff --git a/txnkv/transaction/2pc.go b/txnkv/transaction/2pc.go index 067c662fe1..b9fed7be2e 100644 --- a/txnkv/transaction/2pc.go +++ b/txnkv/transaction/2pc.go @@ -340,8 +340,8 @@ type CommitterMutations interface { NeedConstraintCheckInPrewrite(i int) bool } -// MutationsHasDataInRange returns the first mutation key and first data key in [start, end). -// The first data key is the primary or first write key, and can be nil when the range only contains non-write operations. +// MutationsHasDataInRange returns the first mutation key and first write key in [start, end). +// The first write key can be nil when the range only contains non-write operations. func MutationsHasDataInRange(mutations CommitterMutations, start []byte, end []byte) (firstKey, firstDataKey []byte, ok bool) { isInRange := func(pos int) bool { return pos < mutations.Len() && (len(end) == 0 || bytes.Compare(mutations.GetKey(pos), end) < 0) @@ -361,8 +361,7 @@ func MutationsHasDataInRange(mutations CommitterMutations, start []byte, end []b firstKey = mutations.GetKey(pos) for isInRange(pos) { - // Always return primary key if it's in the range. - if pos == 0 || isOpForWrite(mutations.GetOp(pos)) { + if isOpForWrite(mutations.GetOp(pos)) { firstDataKey = mutations.GetKey(pos) break } diff --git a/txnkv/transaction/2pc_test.go b/txnkv/transaction/2pc_test.go index e31cee105b..084faa8ff9 100644 --- a/txnkv/transaction/2pc_test.go +++ b/txnkv/transaction/2pc_test.go @@ -139,7 +139,7 @@ func TestMutationsHasDataInRange(t *testing.T) { for i := 10; i < 20; i += 2 { key := iToKey(i) var op kvrpcpb.Op - if i%4 == 0 { + if i%4 == 2 { op = kvrpcpb.Op_CheckNotExists } else { op = kvrpcpb.Op_Put @@ -155,16 +155,16 @@ func TestMutationsHasDataInRange(t *testing.T) { firstDataKey int } cases := []Case{ - {-1, -1, true, 10, 10}, + {-1, -1, true, 10, 12}, {-1, 5, false, -1, -1}, {0, 10, false, -1, -1}, - {0, 11, true, 10, 10}, - {0, 30, true, 10, 10}, - {0, -1, true, 10, 10}, - {10, 20, true, 10, 10}, + {0, 11, true, 10, -1}, + {0, 30, true, 10, 12}, + {0, -1, true, 10, 12}, + {10, 20, true, 10, 12}, {15, 16, false, -1, -1}, - {15, 17, true, 16, -1}, - {15, -1, true, 16, 18}, + {15, 17, true, 16, 16}, + {15, -1, true, 16, 16}, {20, 30, false, -1, -1}, {21, 30, false, -1, -1}, {21, -1, false, -1, -1}, diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 7b0b75ef9e..fd3517504a 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -332,7 +332,7 @@ type txnFilePrewriteAction struct{} var _ txnFileAction = (*txnFilePrewriteAction)(nil) func (a txnFilePrewriteAction) executeBatch(c *twoPhaseCommitter, bo *retry.Backoffer, batch chunkBatch) (*tikvrpc.Response, error) { - primaryLock := c.txnFileCtx.slice.chunkRanges[0].smallest + primaryLock := c.primary() req := tikvrpc.NewRequest(tikvrpc.CmdPrewrite, &kvrpcpb.PrewriteRequest{ StartVersion: c.startTS, PrimaryLock: primaryLock, @@ -976,6 +976,16 @@ func (c *twoPhaseCommitter) executeTxnFilePrimaryBatch(bo *retry.Backoffer, firs return nil, nil } +func (c *twoPhaseCommitter) txnFilePrimaryBatchIndex(batches []chunkBatch) (int, error) { + primary := c.primary() + for i := range batches { + if batches[i].region.Contains(primary) { + return i, nil + } + } + return -1, fmt.Errorf("txn file: primary out of batches") +} + func (c *twoPhaseCommitter) executeTxnFileAction(bo *retry.Backoffer, chunkSlice txnChunkSlice, action txnFileAction) error { for { batches, err := chunkSlice.groupToBatches(c.store.GetRegionCache(), bo, c.mutations) @@ -983,6 +993,12 @@ func (c *twoPhaseCommitter) executeTxnFileAction(bo *retry.Backoffer, chunkSlice return errors.Wrap(err, "txn file: group to batches failed") } + primaryBatchIndex, err := c.txnFilePrimaryBatchIndex(batches) + if err != nil { + return errors.WithStack(err) + } + batches[0], batches[primaryBatchIndex] = batches[primaryBatchIndex], batches[0] + regionErr, err := c.executeTxnFilePrimaryBatch(bo, batches[0], action) if err != nil { return errors.WithStack(err) diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 34f1680953..209cbbeda7 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -411,6 +411,31 @@ func newTxnFileCommitTestBatch( return committer, bo, batch } +func TestTxnFilePrewriteUsesPrimaryKey(t *testing.T) { + committer, bo, batch := newTxnFileCommitTestBatch(t, func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + require.Equal(t, []byte("primary"), req.Prewrite().PrimaryLock) + return &tikvrpc.Response{Resp: &kvrpcpb.PrewriteResponse{}}, nil + }) + committer.primaryKey = []byte("primary") + + _, err := (txnFilePrewriteAction{}).executeBatch(committer, bo, batch) + + require.NoError(t, err) +} + +func TestTxnFilePrimaryBatchIndexFindsPrimaryRegion(t *testing.T) { + committer := &twoPhaseCommitter{primaryKey: []byte("primary")} + batches := []chunkBatch{ + {region: &locate.KeyLocation{EndKey: []byte("primary")}}, + {region: &locate.KeyLocation{StartKey: []byte("primary")}}, + } + + index, err := committer.txnFilePrimaryBatchIndex(batches) + + require.NoError(t, err) + require.Equal(t, 1, index) +} + func TestTxnFileActionsApplyResourceGroupTagger(t *testing.T) { tests := []struct { name string From 83ca9736ed8c549f5eb600f2a9de122ab5f4c627 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Thu, 6 Aug 2026 20:40:07 +0800 Subject: [PATCH 18/33] no pipeline txn Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 1 + txnkv/transaction/txn_file_test.go | 17 +++++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index fd3517504a..33a6b827c4 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -1158,6 +1158,7 @@ func (c *twoPhaseCommitter) useTxnFile(ctx context.Context) (bool, error) { } if c.txn.isPessimistic || + c.txn.isPipelined || len(conf.TiKVClient.TxnChunkWriterAddr) == 0 || uint64(c.txn.GetMemBuffer().Size()) < minMutationSize || !IsRequestSourceUseTxnFile(c.txn.RequestSource, conf) { diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 209cbbeda7..440113f406 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -860,6 +860,23 @@ func TestIsRequestSourceUseTxnFile(t *testing.T) { } } +func TestUseTxnFileExcludesPipelinedTxn(t *testing.T) { + restore := config.UpdateGlobal(func(conf *config.Config) { + conf.TiKVClient.TxnChunkWriterAddr = "127.0.0.1" + conf.TiKVClient.TxnFileMinMutationSize = 0 + }) + t.Cleanup(restore) + + txn := newTestTxn(t, 1) + txn.isPipelined = true + committer := &twoPhaseCommitter{txn: txn.KVTxn} + + useTxnFile, err := committer.useTxnFile(context.Background()) + + require.NoError(t, err) + require.False(t, useTxnFile) +} + // stubKVStore implements kvstore with only GetRegionCache returning a real // RegionCache backed by the mock PD client. All other methods panic because // buildTxnFiles does not call them. From 88763ffa550bcd353e9746996a16edded0185134 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Thu, 6 Aug 2026 21:41:24 +0800 Subject: [PATCH 19/33] handle shared lock Signed-off-by: Ping Yu --- integration_tests/shared_lock_test.go | 55 ++++++++++++++++++++++ tikv/kv_test.go | 28 ++++++++++++ tikv/split_region.go | 13 ++++-- txnkv/transaction/txn_file.go | 66 +++++++++++++-------------- txnkv/transaction/txn_file_test.go | 35 ++++++++++++++ txnkv/txnlock/lock.go | 15 ++++++ 6 files changed, 174 insertions(+), 38 deletions(-) diff --git a/integration_tests/shared_lock_test.go b/integration_tests/shared_lock_test.go index 9d08ff6b3e..d3bbae6f6e 100644 --- a/integration_tests/shared_lock_test.go +++ b/integration_tests/shared_lock_test.go @@ -458,6 +458,61 @@ func (s *testSharedLockSuite) TestPrewriteResolveExpiredSharedLock() { s.Nil(txn1.Rollback()) } +func (s *testSharedLockSuite) TestPrewriteResolveExpiredSharedLockWithActiveHolder() { + originManagedLockTTL := atomic.LoadUint64(&transaction.ManagedLockTTL) + atomic.StoreUint64(&transaction.ManagedLockTTL, 500) + defer atomic.StoreUint64(&transaction.ManagedLockTTL, originManagedLockTTL) + + expiredTxn := s.begin() + activeTxn := s.begin() + sharedKey := s.key("TestPrewriteResolveExpiredSharedLockWithActiveHolder_key") + primaryKeys := [][]byte{ + s.key("TestPrewriteResolveExpiredSharedLockWithActiveHolder_expired_pk"), + s.key("TestPrewriteResolveExpiredSharedLockWithActiveHolder_active_pk"), + } + for i, txn := range []transaction.TxnProbe{expiredTxn, activeTxn} { + s.Nil(txn.LockKeys(context.Background(), kv.NewLockCtx(s.getTS(), 1000, time.Now()), primaryKeys[i])) + lockCtx := kv.NewLockCtx(s.getTS(), 1000, time.Now()) + lockCtx.InShareMode = true + s.Nil(txn.LockKeys(context.Background(), lockCtx, sharedKey)) + } + s.waitLocks(sharedKey, s.getTS(), 2, "expect two shared lock holders") + expiredTxn.GetCommitter().CloseTTLManager() + time.Sleep(time.Duration(atomic.LoadUint64(&transaction.ManagedLockTTL))*time.Millisecond + 200*time.Millisecond) + s.True(activeTxn.GetCommitter().IsTTLRunning()) + + contender, err := s.store.Begin() + s.Nil(err) + value := []byte("contender-value") + s.Nil(contender.Set(sharedKey, value)) + commitDone := make(chan error, 1) + go func() { + commitDone <- contender.Commit(context.Background()) + }() + + locks := s.waitLocks(sharedKey, s.getTS(), 1, "expired holder should be removed while active holder remains") + s.Equal(activeTxn.StartTS(), locks[0].TxnID) + select { + case err := <-commitDone: + s.FailNow("prewrite returned while the active shared lock remained", err) + case <-time.After(200 * time.Millisecond): + } + + s.Nil(activeTxn.Rollback()) + select { + case err := <-commitDone: + s.Nil(err) + case <-time.After(5 * time.Second): + s.FailNow("prewrite did not finish after the active holder released") + } + + snapshot := s.store.GetSnapshot(contender.CommitTS()) + got, err := snapshot.Get(context.Background(), sharedKey) + s.Nil(err) + s.Equal(value, got.Value) + s.Nil(expiredTxn.Rollback()) +} + func (s *testSharedLockSuite) TestForceLockRetryOnSharedLock() { if config.NextGen { s.T().Skip("NextGen does not support allow_lock_with_conflict / ForceLock yet") diff --git a/tikv/kv_test.go b/tikv/kv_test.go index 8854369a49..8e33517170 100644 --- a/tikv/kv_test.go +++ b/tikv/kv_test.go @@ -27,10 +27,12 @@ import ( "github.com/pingcap/kvproto/pkg/metapb" "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" + "github.com/tikv/client-go/v2/config/retry" "github.com/tikv/client-go/v2/internal/mockstore/mocktikv" "github.com/tikv/client-go/v2/oracle" "github.com/tikv/client-go/v2/testutils" "github.com/tikv/client-go/v2/tikvrpc" + "github.com/tikv/client-go/v2/txnkv/txnlock" "github.com/tikv/client-go/v2/util" "github.com/tikv/client-go/v2/util/intest" pdhttp "github.com/tikv/pd/client/http" @@ -166,6 +168,32 @@ func (s *testKVSuite) TestMinSafeTsFromStores() { s.Require().Equal(mockClient.tikvSafeTs, ts) } +func (s *testKVSuite) TestHandleSplitRegionKeyErrorsExpandsSharedLockHolders() { + var observed []*txnlock.Lock + resolver := txnlock.LockResolverProbe{LockResolver: s.store.GetLockResolver()} + resolver.SetMeetLockCallback(func(locks []*txnlock.Lock) { + observed = locks + panic("captured shared locks") + }) + keyErrs := []*kvrpcpb.KeyError{{ + Locked: &kvrpcpb.LockInfo{ + Key: []byte("k"), + LockType: kvrpcpb.Op_SharedLock, + SharedLockInfos: []*kvrpcpb.LockInfo{ + {Key: []byte("k"), LockVersion: 1, LockType: kvrpcpb.Op_PessimisticLock}, + {Key: []byte("k"), LockVersion: 1, LockType: kvrpcpb.Op_Lock}, + }, + }, + }} + + require.PanicsWithValue(s.T(), "captured shared locks", func() { + _ = s.store.handleSplitRegionKeyErrors(retry.NewBackoffer(context.Background(), 1000), keyErrs) + }) + require.Len(s.T(), observed, 2) + require.Equal(s.T(), kvrpcpb.Op_PessimisticLock, observed[0].LockType) + require.Equal(s.T(), kvrpcpb.Op_Lock, observed[1].LockType) +} + func (s *testKVSuite) TestMinSafeTsFromStoresWithAllZeros() { // ref https://github.com/tikv/client-go/issues/1276 mockClient := newStoreSafeTsMockClient(s) diff --git a/tikv/split_region.go b/tikv/split_region.go index 7a74a15ef2..b435bec827 100644 --- a/tikv/split_region.go +++ b/tikv/split_region.go @@ -244,21 +244,24 @@ func (s *KVStore) handleSplitRegionKeyErrors(bo *Backoffer, keyErrs []*kvrpcpb.K startTS uint64 = math.MaxUint64 // Set as MaxUint64 and check txn status will not push the minCommiTS. ) for _, keyErr := range keyErrs { - lock, err1 := txnlock.ExtractLockFromKeyErr(keyErr) + locksFromKeyErr, err1 := txnlock.ExtractLocksFromKeyErr(keyErr) if err1 != nil { // Split region should return key error of locked only. return err1 } - logutil.Logger(bo.GetCtx()).Info("split region encounters lock", zap.Stringer("lock", lock)) - locks = append(locks, lock) + for _, lock := range locksFromKeyErr { + logutil.Logger(bo.GetCtx()).Info("split region encounters lock", zap.Stringer("lock", lock)) + locks = append(locks, lock) + } } token := s.GetLockResolver().RecordResolvingLocks(locks, startTS) defer s.GetLockResolver().ResolveLocksDone(startTS, token) resolveLockOpts := txnlock.ResolveLocksOptions{ - CallerStartTS: startTS, - Locks: locks, + CallerStartTS: startTS, + Locks: locks, + PessimisticRegionResolve: true, } resolveLockRes, err := s.GetLockResolver().ResolveLocksWithOpts(bo, resolveLockOpts) if err != nil { diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 33a6b827c4..3f0ef09d48 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -411,32 +411,29 @@ func (a txnFilePrewriteAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back return nil, c.extractKeyExistsErr(e) } - // Extract lock from key error - lock, err1 := txnlock.ExtractLockFromKeyErr(keyErr) + locksFromKeyErr, err1 := txnlock.ExtractLocksFromKeyErr(keyErr) if err1 != nil { return nil, err1 } - logutil.Logger(bo.GetCtx()).Info( - "prewrite txn file encounters lock", - zap.Uint64("session", c.sessionID), - zap.Uint64("txnID", c.startTS), - zap.Stringer("lock", lock), - ) - // If an optimistic transaction encounters a lock with larger TS, this transaction will certainly - // fail due to a WriteConflict error. So we can construct and return an error here early. - // Pessimistic transactions don't need such an optimization. If this key needs a pessimistic lock, - // TiKV will return a PessimisticLockNotFound error directly if it encounters a different lock. Otherwise, - // TiKV returns lock.TTL = 0, and we still need to resolve the lock. - if lock.TxnID > c.startTS { - return nil, tikverr.NewErrWriteConflictWithArgs( - c.startTS, - lock.TxnID, - 0, - lock.Key, - kvrpcpb.WriteConflict_Optimistic, + for _, lock := range locksFromKeyErr { + logutil.Logger(bo.GetCtx()).Info( + "prewrite txn file encounters lock", + zap.Uint64("session", c.sessionID), + zap.Uint64("txnID", c.startTS), + zap.Stringer("lock", lock), ) + if (lock.TxnID > c.startTS && !c.isPessimistic) || + c.txn.prewriteEncounterLockPolicy == NoResolvePolicy { + return nil, tikverr.NewErrWriteConflictWithArgs( + c.startTS, + lock.TxnID, + 0, + lock.Key, + kvrpcpb.WriteConflict_Optimistic, + ) + } + locks = append(locks, lock) } - locks = append(locks, lock) } if resolvingRecordToken == nil { token := c.store.GetLockResolver().RecordResolvingLocks(locks, c.startTS) @@ -446,9 +443,10 @@ func (a txnFilePrewriteAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back c.store.GetLockResolver().UpdateResolvingLocks(locks, c.startTS, *resolvingRecordToken) } resolveLockOpts := txnlock.ResolveLocksOptions{ - CallerStartTS: c.startTS, - Locks: locks, - Detail: &c.getDetail().ResolveLock, + CallerStartTS: c.startTS, + Locks: locks, + Detail: &c.getDetail().ResolveLock, + PessimisticRegionResolve: true, } resolveLockRes, err := c.store.GetLockResolver().ResolveLocksWithOpts(bo, resolveLockOpts) if err != nil { @@ -893,18 +891,20 @@ func (c *twoPhaseCommitter) executeTxnFileSliceSingleBatch(bo *retry.Backoffer, e := &tikverr.ErrKeyExist{AlreadyExist: alreadyExist} return nil, c.extractKeyExistsErr(e) } - lock, err2 := txnlock.ExtractLockFromKeyErr(keyErr) + locks, err2 := txnlock.ExtractLocksFromKeyErr(keyErr) if err2 != nil { return nil, err2 } - if lock.TxnID > c.startTS { - return nil, tikverr.NewErrWriteConflictWithArgs( - c.startTS, - lock.TxnID, - 0, - lock.Key, - kvrpcpb.WriteConflict_Optimistic, - ) + for _, lock := range locks { + if lock.TxnID > c.startTS { + return nil, tikverr.NewErrWriteConflictWithArgs( + c.startTS, + lock.TxnID, + 0, + lock.Key, + kvrpcpb.WriteConflict_Optimistic, + ) + } } } regionErr, err1 := resp.GetRegionError() diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 440113f406..15c77d2db8 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -76,6 +76,7 @@ type txnFileCommitTSStore struct { oracle *txnFileCommitTSOracle regionCache *locate.RegionCache client client.Client + lockResolver *txnlock.LockResolver } func (s *txnFileCommitTSStore) GetTimestampWithRetry(bo *retry.Backoffer, _ string) (uint64, error) { @@ -104,6 +105,10 @@ func (s *txnFileCommitTSStore) GetTiKVClient() client.Client { return s.client } +func (s *txnFileCommitTSStore) GetLockResolver() *txnlock.LockResolver { + return s.lockResolver +} + type txnFileSchemaVer int64 func (v txnFileSchemaVer) SchemaMetaVersion() int64 { @@ -388,6 +393,8 @@ func newTxnFileCommitTestBatch( regionCache: regionCache, client: &fnClient{onSend: onSend}, } + store.lockResolver = txnlock.NewLockResolver(store) + t.Cleanup(store.lockResolver.Close) committer := newTxnFileCommitTSTestCommitter(store, &txnFileSchemaLeaseChecker{}, nil) committer.commitTS = 2 @@ -423,6 +430,34 @@ func TestTxnFilePrewriteUsesPrimaryKey(t *testing.T) { require.NoError(t, err) } +func TestTxnFilePrewriteExpandsSharedLockHolders(t *testing.T) { + committer, bo, batch := newTxnFileCommitTestBatch(t, func(_ context.Context, _ string, _ *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + return &tikvrpc.Response{Resp: &kvrpcpb.PrewriteResponse{Errors: []*kvrpcpb.KeyError{{ + Locked: &kvrpcpb.LockInfo{ + Key: []byte("k"), + LockType: kvrpcpb.Op_SharedLock, + SharedLockInfos: []*kvrpcpb.LockInfo{ + {Key: []byte("k"), LockVersion: 1, LockType: kvrpcpb.Op_PessimisticLock}, + {Key: []byte("k"), LockVersion: 1, LockType: kvrpcpb.Op_Lock}, + }, + }, + }}}}, nil + }) + var observed []*txnlock.Lock + resolver := txnlock.LockResolverProbe{LockResolver: committer.store.GetLockResolver()} + resolver.SetMeetLockCallback(func(locks []*txnlock.Lock) { + observed = locks + panic("captured shared locks") + }) + + require.PanicsWithValue(t, "captured shared locks", func() { + _, _ = (txnFilePrewriteAction{}).executeBatch(committer, bo, batch) + }) + require.Len(t, observed, 2) + require.Equal(t, kvrpcpb.Op_PessimisticLock, observed[0].LockType) + require.Equal(t, kvrpcpb.Op_Lock, observed[1].LockType) +} + func TestTxnFilePrimaryBatchIndexFindsPrimaryRegion(t *testing.T) { committer := &twoPhaseCommitter{primaryKey: []byte("primary")} batches := []chunkBatch{ diff --git a/txnkv/txnlock/lock.go b/txnkv/txnlock/lock.go index 1777e45da6..478a0fc738 100644 --- a/txnkv/txnlock/lock.go +++ b/txnkv/txnlock/lock.go @@ -26,3 +26,18 @@ func ExtractLockFromKeyErr(keyErr *kvrpcpb.KeyError) (*Lock, error) { } return nil, tikverr.ExtractKeyErr(keyErr) } + +// ExtractLocksFromKeyErr extracts all locks represented by a KeyError. +func ExtractLocksFromKeyErr(keyErr *kvrpcpb.KeyError) ([]*Lock, error) { + if locked := keyErr.GetLocked(); locked != nil { + if sharedLockInfos := locked.GetSharedLockInfos(); len(sharedLockInfos) > 0 { + locks := make([]*Lock, 0, len(sharedLockInfos)) + for _, sharedLockInfo := range sharedLockInfos { + locks = append(locks, NewLock(sharedLockInfo)) + } + return locks, nil + } + return []*Lock{NewLock(locked)}, nil + } + return nil, tikverr.ExtractKeyErr(keyErr) +} From 46861abf0d5ce59bc9997de4b0365ddb8503ba35 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Thu, 6 Aug 2026 22:01:16 +0800 Subject: [PATCH 20/33] add lock_test Signed-off-by: Ping Yu --- txnkv/txnlock/lock_test.go | 63 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 63 insertions(+) create mode 100644 txnkv/txnlock/lock_test.go diff --git a/txnkv/txnlock/lock_test.go b/txnkv/txnlock/lock_test.go new file mode 100644 index 0000000000..8d9a80c55c --- /dev/null +++ b/txnkv/txnlock/lock_test.go @@ -0,0 +1,63 @@ +// Copyright 2026 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. + +package txnlock + +import ( + "testing" + + "github.com/pingcap/kvproto/pkg/kvrpcpb" + "github.com/stretchr/testify/require" +) + +func TestExtractLocksFromKeyErrExpandsSharedLockHolders(t *testing.T) { + keyErr := &kvrpcpb.KeyError{ + Locked: &kvrpcpb.LockInfo{ + Key: []byte("shared-key"), + LockType: kvrpcpb.Op_SharedLock, + LockVersion: 100, + SharedLockInfos: []*kvrpcpb.LockInfo{ + {Key: []byte("shared-key"), LockVersion: 101, LockType: kvrpcpb.Op_PessimisticLock}, + {Key: []byte("shared-key"), LockVersion: 102, LockType: kvrpcpb.Op_Lock}, + }, + }, + } + + locks, err := ExtractLocksFromKeyErr(keyErr) + + require.NoError(t, err) + require.Len(t, locks, 2) + require.Equal(t, uint64(101), locks[0].TxnID) + require.Equal(t, kvrpcpb.Op_PessimisticLock, locks[0].LockType) + require.Equal(t, uint64(102), locks[1].TxnID) + require.Equal(t, kvrpcpb.Op_Lock, locks[1].LockType) +} + +func TestExtractLocksFromKeyErrPreservesExclusiveLock(t *testing.T) { + locks, err := ExtractLocksFromKeyErr(&kvrpcpb.KeyError{ + Locked: &kvrpcpb.LockInfo{Key: []byte("key"), LockVersion: 7, LockType: kvrpcpb.Op_Lock}, + }) + + require.NoError(t, err) + require.Len(t, locks, 1) + require.Equal(t, uint64(7), locks[0].TxnID) +} + +func TestExtractLocksFromKeyErrReturnsKeyError(t *testing.T) { + _, err := ExtractLocksFromKeyErr(&kvrpcpb.KeyError{ + AlreadyExist: &kvrpcpb.AlreadyExist{Key: []byte("key")}, + }) + + require.Error(t, err) +} From d82de8faab1b256caa8ce541470ac5f59b6dd4b9 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Thu, 6 Aug 2026 22:02:17 +0800 Subject: [PATCH 21/33] cleanup ctx Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 7 ++++++- txnkv/transaction/txn_file_test.go | 10 ++++++++++ 2 files changed, 16 insertions(+), 1 deletion(-) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 3f0ef09d48..ba69abfd5c 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -690,6 +690,10 @@ func (s step) String() string { return fmt.Sprintf("%s:%s", s.name, s.dur.String()) } +func txnFileCleanupContext(storeCtx, txnCtx context.Context) context.Context { + return context.WithValue(storeCtx, retry.TxnStartKey, txnCtx.Value(retry.TxnStartKey)) +} + func (c *twoPhaseCommitter) executeTxnFile(ctx context.Context) (err error) { if val, err := util.EvalFailpoint("injectErrorOnExecTxnFile"); err == nil { errVal := val.(string) @@ -724,7 +728,8 @@ func (c *twoPhaseCommitter) executeTxnFile(ctx context.Context) (err error) { c.mu.RUnlock() if !committed && !undetermined { if c.txnFileCtx.slice.Len() > 0 { - err1 := c.executeTxnFileAction(retry.NewBackofferWithVars(ctx, int(CommitMaxBackoff), c.txn.vars), c.txnFileCtx.slice, txnFileRollbackAction{}) + cleanupCtx := txnFileCleanupContext(c.store.Ctx(), ctx) + err1 := c.executeTxnFileAction(retry.NewBackofferWithVars(cleanupCtx, int(CommitMaxBackoff), c.txn.vars), c.txnFileCtx.slice, txnFileRollbackAction{}) if err1 != nil { logutil.Logger(ctx).Error("txn file: rollback on error failed", zap.Error(err1)) } diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 15c77d2db8..6e8aef2dfc 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -48,6 +48,16 @@ import ( "github.com/tikv/client-go/v2/util" ) +func TestTxnFileCleanupContextUsesStoreContext(t *testing.T) { + transactionCtx := context.WithValue(context.Background(), retry.TxnStartKey, uint64(42)) + cancelledCtx, cancel := context.WithCancel(transactionCtx) + cleanupCtx := txnFileCleanupContext(context.Background(), cancelledCtx) + cancel() + + require.NoError(t, cleanupCtx.Err()) + require.Equal(t, uint64(42), cleanupCtx.Value(retry.TxnStartKey)) +} + type txnFileCommitTSOracle struct { unimplementedOracle From ebab9a900a9a0b5090bcf662c46d75c410600ef0 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Thu, 6 Aug 2026 22:18:17 +0800 Subject: [PATCH 22/33] MaxTxnChunkSizeInParallel Signed-off-by: Ping Yu --- config/client.go | 5 +++++ config/config_test.go | 9 ++++++++- txnkv/transaction/txn_file.go | 13 +++++++++---- txnkv/transaction/txn_file_test.go | 30 ++++++++++++++++++++++++++++++ 4 files changed, 52 insertions(+), 5 deletions(-) diff --git a/config/client.go b/config/client.go index f6de4bfb12..c64b6997b2 100644 --- a/config/client.go +++ b/config/client.go @@ -49,6 +49,8 @@ const ( DefGrpcInitialConnWindowSize = 1 << 27 // 128MiB DefMaxConcurrencyRequestLimit = math.MaxInt64 DefBatchPolicy = BatchPolicyStandard + // MaxTxnChunkSizeInParallel is the maximum total size of transaction chunks processed in parallel. + MaxTxnChunkSizeInParallel uint64 = 4 << 30 // 4GB ) const ( @@ -277,6 +279,9 @@ func validateTxnFileConfig(config *TiKVClient) error { if config.TxnChunkMaxSize > math.MaxInt { return fmt.Errorf("txn-chunk-max-size should not exceed %d, but got %d", math.MaxInt, config.TxnChunkMaxSize) } + if config.TxnChunkMaxSize > MaxTxnChunkSizeInParallel { + return fmt.Errorf("txn-chunk-max-size should not exceed %d, but got %d", MaxTxnChunkSizeInParallel, config.TxnChunkMaxSize) + } if config.TxnChunkWriterConcurrency == 0 { return fmt.Errorf("txn-chunk-writer-concurrency should be greater than 0") } diff --git a/config/config_test.go b/config/config_test.go index 6857d591dc..8df6df001a 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -112,9 +112,16 @@ func TestValidateTxnFileConfig(t *testing.T) { { name: "maximum chunk size", configure: func(cfg *TiKVClient) { - cfg.TxnChunkMaxSize = maxInt + cfg.TxnChunkMaxSize = MaxTxnChunkSizeInParallel }, }, + { + name: "chunk size exceeds parallel budget", + configure: func(cfg *TiKVClient) { + cfg.TxnChunkMaxSize = 4<<30 + 1 + }, + err: fmt.Sprintf("txn-chunk-max-size should not exceed %d, but got %d", uint64(4<<30), uint64(4<<30)+1), + }, { name: "chunk size exceeds int", configure: func(cfg *TiKVClient) { diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index ba69abfd5c..443c3e4b0d 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -64,9 +64,6 @@ var ( const ( PreSplitRegionChunks = 4 - - // MaxTxnChunkSizeInParallel is the max parallel size when prewrite/commit txn chunks. - MaxTxnChunkSizeInParallel uint64 = 4 << 30 // 4GB ) type txnFileCtx struct { @@ -845,7 +842,7 @@ func (c *twoPhaseCommitter) executeTxnFileSlice(bo *retry.Backoffer, chunkSlice if rateLim > cnf.CommitterConcurrency { rateLim = cnf.CommitterConcurrency } - maxChunksInParallel := int(MaxTxnChunkSizeInParallel / cnf.TiKVClient.TxnChunkMaxSize) // 32 by default + maxChunksInParallel := txnFileMaxChunksInParallel(cnf.TiKVClient.TxnChunkMaxSize) // 32 by default if chunksCount > maxChunksInParallel { rateLim = maxChunksInParallel } @@ -881,6 +878,14 @@ func (c *twoPhaseCommitter) executeTxnFileSlice(bo *retry.Backoffer, chunkSlice return regionErrChunks, err } +func txnFileMaxChunksInParallel(txnChunkMaxSize uint64) int { + maxChunksInParallel := int(config.MaxTxnChunkSizeInParallel / txnChunkMaxSize) + if maxChunksInParallel < 1 { + return 1 + } + return maxChunksInParallel +} + func (c *twoPhaseCommitter) executeTxnFileSliceSingleBatch(bo *retry.Backoffer, batch chunkBatch, action txnFileAction) (*txnChunkSlice, error) { resp, err1 := action.executeBatch(c, bo, batch) logutil.Logger(bo.GetCtx()).Debug("txn file: execute batch finished", diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 6e8aef2dfc..8625cd15ac 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -58,6 +58,36 @@ func TestTxnFileCleanupContextUsesStoreContext(t *testing.T) { require.Equal(t, uint64(42), cleanupCtx.Value(retry.TxnStartKey)) } +func TestTxnFileMaxChunksInParallel(t *testing.T) { + tests := []struct { + name string + chunkMaxSize uint64 + expectedResult int + }{ + { + name: "default chunk size", + chunkMaxSize: 128 * 1024 * 1024, + expectedResult: 32, + }, + { + name: "parallel budget boundary", + chunkMaxSize: config.MaxTxnChunkSizeInParallel, + expectedResult: 1, + }, + { + name: "chunk size exceeds parallel budget", + chunkMaxSize: config.MaxTxnChunkSizeInParallel + 1, + expectedResult: 1, + }, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + assert.Equal(t, test.expectedResult, txnFileMaxChunksInParallel(test.chunkMaxSize)) + }) + } +} + type txnFileCommitTSOracle struct { unimplementedOracle From fdca77b0d23eb70895e1ec6a4fd0e0c1c3f4fbe2 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Thu, 6 Aug 2026 23:01:08 +0800 Subject: [PATCH 23/33] resource control Signed-off-by: Ping Yu --- internal/client/client_interceptor.go | 9 ++++ internal/client/client_interceptor_test.go | 51 ++++++++++++++++++++-- txnkv/transaction/txn_file.go | 19 +++++--- txnkv/transaction/txn_file_test.go | 5 +++ 4 files changed, 76 insertions(+), 8 deletions(-) diff --git a/internal/client/client_interceptor.go b/internal/client/client_interceptor.go index a6a67c9ad2..ae1d8f96f6 100644 --- a/internal/client/client_interceptor.go +++ b/internal/client/client_interceptor.go @@ -174,6 +174,15 @@ func getResourceControlInfo(ctx context.Context, req *tikvrpc.Request) ( return resourceGroupName, resourceControlInterceptor, reqInfo } +// GetResourceControlInfo applies the normal resource-control selection policy to req. +func GetResourceControlInfo(ctx context.Context, req *tikvrpc.Request) ( + string, + resourceControlClient.ResourceGroupKVInterceptor, + *resourcecontrol.RequestInfo, +) { + return getResourceControlInfo(ctx, req) +} + // buildResourceControlInterceptor builds a resource control interceptor with // the given resource group name. // diff --git a/internal/client/client_interceptor_test.go b/internal/client/client_interceptor_test.go index df730abc66..98282b84fa 100644 --- a/internal/client/client_interceptor_test.go +++ b/internal/client/client_interceptor_test.go @@ -112,8 +112,9 @@ func TestAppendChainedInterceptor(t *testing.T) { // benign zero values so the interceptor wiring can be exercised end-to-end // without touching real PD state. type recordingInterceptor struct { - waitCalls int - respCalls int + waitCalls int + respCalls int + background bool } var recordingRequestConsumption = &rmpb.Consumption{RRU: 11, WRU: 3} @@ -139,7 +140,7 @@ func (r *recordingInterceptor) OnResponseWait( } func (r *recordingInterceptor) IsBackgroundRequest(context.Context, string, string) bool { - return false + return r.background } func (r *recordingInterceptor) GetRUVersion() resourceControlClient.RUVersion { @@ -181,6 +182,50 @@ func newRGRequest() *tikvrpc.Request { return req } +func TestGetResourceControlInfoHonorsSelectionPolicy(t *testing.T) { + tests := []struct { + name string + enabled bool + resourceGroup string + background bool + requestSource string + wantSelected bool + }{ + {name: "enabled request", enabled: true, resourceGroup: "test-rg", wantSelected: true}, + {name: "disabled switch", enabled: false, resourceGroup: "test-rg"}, + {name: "empty group", enabled: true}, + {name: "background request", enabled: true, resourceGroup: "test-rg", background: true}, + {name: "bypassed request", enabled: true, resourceGroup: "test-rg", requestSource: util.InternalRequestPrefix + util.InternalTxnOthers}, + } + + for _, test := range tests { + t.Run(test.name, func(t *testing.T) { + rec := &recordingInterceptor{background: test.background} + var iface resourceControlClient.ResourceGroupKVInterceptor = rec + ResourceControlSwitch.Store(test.enabled) + ResourceControlInterceptor.Store(&iface) + t.Cleanup(func() { + ResourceControlSwitch.Store(false) + ResourceControlInterceptor.Store(nil) + }) + + req := newRGRequest() + req.ResourceControlContext.ResourceGroupName = test.resourceGroup + req.RequestSource = test.requestSource + group, interceptor, reqInfo := GetResourceControlInfo(context.Background(), req) + if test.wantSelected { + assert.Equal(t, test.resourceGroup, group) + assert.NotNil(t, interceptor) + assert.NotNil(t, reqInfo) + return + } + assert.Empty(t, group) + assert.Nil(t, interceptor) + assert.Nil(t, reqInfo) + }) + } +} + func TestSendRequestDoesNotSettleAndKeepsRUDetailsOnTransportFailure(t *testing.T) { rec := withRecordingInterceptor(t) client := NewInterceptedClient(failingClient{}) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 443c3e4b0d..beca05479d 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -749,7 +749,13 @@ func (c *twoPhaseCommitter) executeTxnFile(ctx context.Context) (err error) { buildBo := retry.NewBackofferWithVars(ctx, int(BuildTxnFileMaxBackoff.Load()), c.txn.vars) - rcInterceptor := client.ResourceControlInterceptor.Load() + rcReq := tikvrpc.NewRequest(tikvrpc.CmdPrewrite, &kvrpcpb.PrewriteRequest{}, kvrpcpb.Context{ + RequestSource: c.txn.GetRequestSource(), + ResourceControlContext: &kvrpcpb.ResourceControlContext{ + ResourceGroupName: c.resourceGroupName, + }, + }) + _, rcInterceptor, _ := client.GetResourceControlInfo(buildBo.GetCtx(), rcReq) var ruDetails *util.RUDetails if detail := ctx.Value(util.RUDetailsCtxKey); detail != nil { ruDetails = detail.(*util.RUDetails) @@ -879,6 +885,9 @@ func (c *twoPhaseCommitter) executeTxnFileSlice(bo *retry.Backoffer, chunkSlice } func txnFileMaxChunksInParallel(txnChunkMaxSize uint64) int { + if txnChunkMaxSize == 0 { + return 1 + } maxChunksInParallel := int(config.MaxTxnChunkSizeInParallel / txnChunkMaxSize) if maxChunksInParallel < 1 { return 1 @@ -1217,7 +1226,7 @@ func (c *twoPhaseCommitter) preSplitTxnFileRegions(bo *retry.Backoffer) error { func (c *twoPhaseCommitter) beforeExecuteTxnFile( bo *retry.Backoffer, - rcInterceptor *resourceControlClient.ResourceGroupKVInterceptor, + rcInterceptor resourceControlClient.ResourceGroupKVInterceptor, ruDetails *util.RUDetails, ) (*resourcecontrol.RequestInfo, error) { if rcInterceptor == nil { @@ -1268,7 +1277,7 @@ func (c *twoPhaseCommitter) beforeExecuteTxnFile( false, ) - consumption, _ /* penalty */, waitDuration, _ /* priority */, err := (*rcInterceptor).OnRequestWait(ctx, c.resourceGroupName, reqInfo) + consumption, _ /* penalty */, waitDuration, _ /* priority */, err := rcInterceptor.OnRequestWait(ctx, c.resourceGroupName, reqInfo) if err != nil { return nil, errors.WithStack(err) } @@ -1280,13 +1289,13 @@ func (c *twoPhaseCommitter) beforeExecuteTxnFile( return reqInfo, nil } -func (c *twoPhaseCommitter) afterExecuteTxnFile(rcInterceptor *resourceControlClient.ResourceGroupKVInterceptor, reqInfo *resourcecontrol.RequestInfo, ruDetails *util.RUDetails) error { +func (c *twoPhaseCommitter) afterExecuteTxnFile(rcInterceptor resourceControlClient.ResourceGroupKVInterceptor, reqInfo *resourcecontrol.RequestInfo, ruDetails *util.RUDetails) error { if rcInterceptor == nil { return nil } respInfo := &resourcecontrol.ResponseInfo{} - consumption, err := (*rcInterceptor).OnResponse(c.resourceGroupName, reqInfo, respInfo) + consumption, err := rcInterceptor.OnResponse(c.resourceGroupName, reqInfo, respInfo) if err != nil { return errors.WithStack(err) } diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 8625cd15ac..10e867bea5 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -79,6 +79,11 @@ func TestTxnFileMaxChunksInParallel(t *testing.T) { chunkMaxSize: config.MaxTxnChunkSizeInParallel + 1, expectedResult: 1, }, + { + name: "zero chunk size", + chunkMaxSize: 0, + expectedResult: 1, + }, } for _, test := range tests { From b8b8131ab918c766e5c1ec5fd34442e550e48b52 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Thu, 6 Aug 2026 23:23:09 +0800 Subject: [PATCH 24/33] skip valid config Signed-off-by: Ping Yu --- config/client.go | 4 ++++ config/config_test.go | 5 +++++ 2 files changed, 9 insertions(+) diff --git a/config/client.go b/config/client.go index c64b6997b2..065ac49076 100644 --- a/config/client.go +++ b/config/client.go @@ -273,6 +273,10 @@ func (config *TiKVClient) Valid() error { } func validateTxnFileConfig(config *TiKVClient) error { + if config.TxnChunkWriterAddr == "" { + // Skip validation when txn file is not enabled. + return nil + } if config.TxnChunkMaxSize == 0 { return fmt.Errorf("txn-chunk-max-size should be greater than 0") } diff --git a/config/config_test.go b/config/config_test.go index 8df6df001a..30749530a3 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -93,6 +93,10 @@ func TestValidateGRPCKeepAliveTimeout(t *testing.T) { } func TestValidateTxnFileConfig(t *testing.T) { + configWithoutTxnFile := DefaultTiKVClient() + configWithoutTxnFile.TxnChunkMaxSize = 0 + assert.NoError(t, configWithoutTxnFile.Valid()) + maxInt := uint64(math.MaxInt) tests := []struct { name string @@ -161,6 +165,7 @@ func TestValidateTxnFileConfig(t *testing.T) { assert.NoError(t, cfg.Valid()) return } + cfg.TxnChunkWriterAddr = "127.0.0.1" assert.EqualError(t, cfg.Valid(), test.err) }) } From 0d123892765b316a5d7a6fc60f26335b43ae9d8c Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Fri, 7 Aug 2026 01:03:38 +0800 Subject: [PATCH 25/33] txn file split region Signed-off-by: Ping Yu --- tikv/kv_test.go | 86 ++++++++++++++++++++++++++++++ tikv/split_region.go | 44 ++++++++++----- txnkv/transaction/2pc.go | 2 + txnkv/transaction/test_util.go | 4 ++ txnkv/transaction/txn_file.go | 2 +- txnkv/transaction/txn_file_test.go | 50 +++++++++++++++-- 6 files changed, 169 insertions(+), 19 deletions(-) diff --git a/tikv/kv_test.go b/tikv/kv_test.go index 8e33517170..cc6a0fa161 100644 --- a/tikv/kv_test.go +++ b/tikv/kv_test.go @@ -109,6 +109,22 @@ type storeSafeTsMockClient struct { tiflashSafeTs uint64 } +type splitRegionKeyErrorMockClient struct { + Client + responses []*tikvrpc.Response + calls atomic.Int32 +} + +func (c *splitRegionKeyErrorMockClient) SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) { + if req.Type == tikvrpc.CmdSplitRegion { + call := int(c.calls.Add(1)) - 1 + if call < len(c.responses) { + return c.responses[call], nil + } + } + return c.Client.SendRequest(ctx, addr, req, timeout) +} + func newStoreSafeTsMockClient(s *testKVSuite) *storeSafeTsMockClient { return &storeSafeTsMockClient{ Client: s.store.GetTiKVClient(), @@ -118,6 +134,76 @@ func newStoreSafeTsMockClient(s *testKVSuite) *storeSafeTsMockClient { } } +func (s *testKVSuite) TestSplitRegionsPreservesLegacyKeyErrorBehavior() { + // Given + client := &splitRegionKeyErrorMockClient{ + Client: s.store.GetTiKVClient(), + responses: []*tikvrpc.Response{{Resp: &kvrpcpb.SplitRegionResponse{ + Errors: []*kvrpcpb.KeyError{{Locked: &kvrpcpb.LockInfo{ + Key: []byte("k"), LockVersion: 1, LockTtl: 1, + }}}, + }}}, + } + s.store.SetTiKVClient(client) + resolver := txnlock.LockResolverProbe{LockResolver: s.store.GetLockResolver()} + resolver.SetMeetLockCallback(func([]*txnlock.Lock) { + panic("generic split must not resolve locks") + }) + + // When + var regionIDs []uint64 + var err error + require.NotPanics(s.T(), func() { + regionIDs, err = s.store.SplitRegions(context.Background(), [][]byte{[]byte("k")}, false, nil) + }) + + // Then + s.Require().NoError(err) + s.Require().Empty(regionIDs) + s.Require().Equal(int32(1), client.calls.Load()) +} + +func (s *testKVSuite) TestSplitTxnFileRegionsResolvesLockAndRetries() { + // Given + client := &splitRegionKeyErrorMockClient{ + Client: s.store.GetTiKVClient(), + responses: []*tikvrpc.Response{{Resp: &kvrpcpb.SplitRegionResponse{ + Errors: []*kvrpcpb.KeyError{{Locked: &kvrpcpb.LockInfo{ + Key: []byte("k"), PrimaryLock: []byte("k"), LockVersion: 1, + }}}, + }}}, + } + s.store.SetTiKVClient(client) + var observed atomic.Int32 + resolver := txnlock.LockResolverProbe{LockResolver: s.store.GetLockResolver()} + resolver.SetMeetLockCallback(func([]*txnlock.Lock) { + observed.Add(1) + }) + + // When + err := s.store.SplitTxnFileRegions(context.Background(), [][]byte{[]byte("k")}) + + // Then + s.Require().NoError(err) + s.Require().GreaterOrEqual(observed.Load(), int32(1)) + s.Require().GreaterOrEqual(client.calls.Load(), int32(2)) +} + +func (s *testKVSuite) TestSplitTxnFileRegionsSplitsWithoutScattering() { + // Given + require.NoError(s.T(), failpoint.Enable("tikvclient/mockScatterRegionTimeout", `return(true)`)) + s.T().Cleanup(func() { + require.NoError(s.T(), failpoint.Disable("tikvclient/mockScatterRegionTimeout")) + }) + splitKey := []byte("txn-file-split") + + // When + err := s.store.SplitTxnFileRegions(context.Background(), [][]byte{splitKey}) + + // Then + s.Require().NoError(err) +} + func (c *storeSafeTsMockClient) SendRequest(ctx context.Context, addr string, req *tikvrpc.Request, timeout time.Duration) (*tikvrpc.Response, error) { if req.Type != tikvrpc.CmdStoreSafeTS { return c.Client.SendRequest(ctx, addr, req, timeout) diff --git a/tikv/split_region.go b/tikv/split_region.go index b435bec827..94acceec86 100644 --- a/tikv/split_region.go +++ b/tikv/split_region.go @@ -62,11 +62,18 @@ import ( const splitBatchRegionLimit = 2048 +type splitRegionMode uint8 + +const ( + splitRegionLegacy splitRegionMode = iota + splitRegionResolveLocks +) + func equalRegionStartKey(key, regionStartKey []byte) bool { return bytes.Equal(key, regionStartKey) } -func (s *KVStore) splitBatchRegionsReq(bo *Backoffer, keys [][]byte, scatter bool, tableID *int64) (*tikvrpc.Response, error) { +func (s *KVStore) splitBatchRegionsReq(bo *Backoffer, keys [][]byte, scatter bool, tableID *int64, mode splitRegionMode) (*tikvrpc.Response, error) { // equalRegionStartKey is used to filter split keys. // If the split key is equal to the start key of the region, then the key has been split, we need to skip the split key. groups, _, err := s.regionCache.GroupKeysByRegion(bo, keys, equalRegionStartKey) @@ -91,7 +98,7 @@ func (s *KVStore) splitBatchRegionsReq(bo *Backoffer, keys [][]byte, scatter boo zap.String("first split key", redact.Key(batches[0].Keys[0]))) } if len(batches) == 1 { - resp := s.batchSendSingleRegion(bo, batches[0], scatter, tableID) + resp := s.batchSendSingleRegion(bo, batches[0], scatter, tableID, mode) return resp.Response, resp.Error } ch := make(chan kvrpc.BatchResult, len(batches)) @@ -102,7 +109,7 @@ func (s *KVStore) splitBatchRegionsReq(bo *Backoffer, keys [][]byte, scatter boo defer cancel() util.WithRecovery(func() { - batchResult := s.batchSendSingleRegion(backoffer, b, scatter, tableID) + batchResult := s.batchSendSingleRegion(backoffer, b, scatter, tableID, mode) lastForkedBo.Store(backoffer) select { case ch <- batchResult: @@ -138,7 +145,7 @@ func (s *KVStore) splitBatchRegionsReq(bo *Backoffer, keys [][]byte, scatter boo return &tikvrpc.Response{Resp: srResp}, err } -func (s *KVStore) batchSendSingleRegion(bo *Backoffer, batch kvrpc.Batch, scatter bool, tableID *int64) kvrpc.BatchResult { +func (s *KVStore) batchSendSingleRegion(bo *Backoffer, batch kvrpc.Batch, scatter bool, tableID *int64, mode splitRegionMode) kvrpc.BatchResult { if val, err := util.EvalFailpoint("mockSplitRegionTimeout"); err == nil { if val.(bool) { if _, ok := bo.GetCtx().Deadline(); ok { @@ -173,7 +180,7 @@ func (s *KVStore) batchSendSingleRegion(bo *Backoffer, batch kvrpc.Batch, scatte batchResp.Error = err return batchResp } - resp, err = s.splitBatchRegionsReq(bo, batch.Keys, scatter, tableID) + resp, err = s.splitBatchRegionsReq(bo, batch.Keys, scatter, tableID, mode) batchResp.Response = resp batchResp.Error = err return batchResp @@ -181,17 +188,19 @@ func (s *KVStore) batchSendSingleRegion(bo *Backoffer, batch kvrpc.Batch, scatte spResp := resp.Resp.(*kvrpcpb.SplitRegionResponse) - keyErrs := spResp.GetErrors() - if len(keyErrs) > 0 { - err := s.handleSplitRegionKeyErrors(bo, keyErrs) - if err != nil { + if mode == splitRegionResolveLocks { + keyErrs := spResp.GetErrors() + if len(keyErrs) > 0 { + err := s.handleSplitRegionKeyErrors(bo, keyErrs) + if err != nil { + batchResp.Error = err + return batchResp + } + resp, err = s.splitBatchRegionsReq(bo, batch.Keys, scatter, tableID, mode) + batchResp.Response = resp batchResp.Error = err return batchResp } - resp, err = s.splitBatchRegionsReq(bo, batch.Keys, scatter, tableID) - batchResp.Response = resp - batchResp.Error = err - return batchResp } regions := spResp.GetRegions() @@ -289,7 +298,7 @@ const ( // SplitRegions splits regions by splitKeys. func (s *KVStore) SplitRegions(ctx context.Context, splitKeys [][]byte, scatter bool, tableID *int64) (regionIDs []uint64, err error) { bo := retry.NewBackofferWithVars(ctx, int(math.Min(float64(len(splitKeys))*splitRegionBackoff, maxSplitRegionsBackoff)), nil) - resp, err := s.splitBatchRegionsReq(bo, splitKeys, scatter, tableID) + resp, err := s.splitBatchRegionsReq(bo, splitKeys, scatter, tableID, splitRegionLegacy) regionIDs = make([]uint64, 0, len(splitKeys)) if resp != nil && resp.Resp != nil { spResp := resp.Resp.(*kvrpcpb.SplitRegionResponse) @@ -301,6 +310,13 @@ func (s *KVStore) SplitRegions(ctx context.Context, splitKeys [][]byte, scatter return regionIDs, err } +// SplitTxnFileRegions splits regions for file-based transactions without scattering and resolves locks encountered by TiKV. +func (s *KVStore) SplitTxnFileRegions(ctx context.Context, splitKeys [][]byte) error { + bo := retry.NewBackofferWithVars(ctx, int(math.Min(float64(len(splitKeys))*splitRegionBackoff, maxSplitRegionsBackoff)), nil) + _, err := s.splitBatchRegionsReq(bo, splitKeys, false, nil, splitRegionResolveLocks) + return err +} + func (s *KVStore) scatterRegion(bo *Backoffer, regionID uint64, tableID *int64) error { logutil.BgLogger().Info("start scatter region", zap.Uint64("regionID", regionID)) diff --git a/txnkv/transaction/2pc.go b/txnkv/transaction/2pc.go index b9fed7be2e..79a79876e3 100644 --- a/txnkv/transaction/2pc.go +++ b/txnkv/transaction/2pc.go @@ -98,6 +98,8 @@ type kvstore interface { GetRegionCache() *locate.RegionCache // SplitRegions splits regions by splitKeys. SplitRegions(ctx context.Context, splitKeys [][]byte, scatter bool, tableID *int64) (regionIDs []uint64, err error) + // SplitTxnFileRegions splits regions for file-based transactions and resolves locks encountered by TiKV. + SplitTxnFileRegions(ctx context.Context, splitKeys [][]byte) error // WaitScatterRegionFinish implements SplittableStore interface. // backOff is the back off time of the wait scatter region.(Milliseconds) // if backOff <= 0, the default wait scatter back off time will be used. diff --git a/txnkv/transaction/test_util.go b/txnkv/transaction/test_util.go index 5c1b6b14bf..f350eedd48 100644 --- a/txnkv/transaction/test_util.go +++ b/txnkv/transaction/test_util.go @@ -464,6 +464,10 @@ func (u *unimplementedKVStore) SplitRegions(ctx context.Context, splitKeys [][]b panic("unimplemented") } +func (u *unimplementedKVStore) SplitTxnFileRegions(ctx context.Context, splitKeys [][]byte) error { + panic("unimplemented") +} + // TxnLatches implements kvstore. func (u *unimplementedKVStore) TxnLatches() *latch.LatchesScheduler { panic("unimplemented") diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index beca05479d..1289a3b073 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -1220,7 +1220,7 @@ func (c *twoPhaseCommitter) preSplitTxnFileRegions(bo *retry.Backoffer) error { if len(splitKeys) == 0 { return nil } - _, err = c.store.SplitRegions(bo.GetCtx(), splitKeys, false, nil) + err = c.store.SplitTxnFileRegions(bo.GetCtx(), splitKeys) return errors.Wrap(err, "pre split regions failed") } diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 10e867bea5..5f51cea774 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -957,16 +957,23 @@ func TestUseTxnFileExcludesPipelinedTxn(t *testing.T) { require.False(t, useTxnFile) } -// stubKVStore implements kvstore with only GetRegionCache returning a real -// RegionCache backed by the mock PD client. All other methods panic because +// stubKVStore implements kvstore with only GetRegionCache and split-call +// recording backed by the mock PD client. All other methods panic because // buildTxnFiles does not call them. type stubKVStore struct { - regionCache *locate.RegionCache + regionCache *locate.RegionCache + splitRegionsCalls atomic.Uint32 + splitTxnFileRegionsCalls atomic.Uint32 } func (s *stubKVStore) GetRegionCache() *locate.RegionCache { return s.regionCache } func (s *stubKVStore) SplitRegions(_ context.Context, _ [][]byte, _ bool, _ *int64) ([]uint64, error) { - panic("not implemented") + s.splitRegionsCalls.Add(1) + panic("unexpected generic split path") +} +func (s *stubKVStore) SplitTxnFileRegions(_ context.Context, _ [][]byte) error { + s.splitTxnFileRegionsCalls.Add(1) + return nil } func (s *stubKVStore) WaitScatterRegionFinish(_ context.Context, _ uint64, _ int) error { panic("not implemented") @@ -988,6 +995,41 @@ func (s *stubKVStore) GetClusterID() uint64 { return 0 } func (s *stubKVStore) IsClose() bool { return false } func (s *stubKVStore) Go(_ func()) error { panic("not implemented") } +func TestPreSplitTxnFileRegionsUsesDedicatedSplitPath(t *testing.T) { + // Given + txn := newTestTxn(t, 1) + for i := 1; i <= 5; i++ { + require.NoError(t, txn.Set([]byte(fmt.Sprintf("k%d", i)), []byte("v"))) + } + committer, err := newTwoPhaseCommitter(txn.KVTxn, 1) + require.NoError(t, err) + require.NoError(t, committer.initKeysAndMutations(context.Background())) + + store := &stubKVStore{regionCache: txn.store.cache} + committer.store = store + slice := txnChunkSlice{ + chunkIDs: []uint64{1, 2, 3, 4, 5}, + chunkRanges: []txnChunkRange{ + newTxnChunkRange([]byte("k1"), []byte("k1"), 1), + newTxnChunkRange([]byte("k2"), []byte("k2"), 1), + newTxnChunkRange([]byte("k3"), []byte("k3"), 1), + newTxnChunkRange([]byte("k4"), []byte("k4"), 1), + newTxnChunkRange([]byte("k5"), []byte("k5"), 1), + }, + } + committer.txnFileCtx = txnFileCtx{slice: slice} + + // When + require.NotPanics(t, func() { + err = committer.preSplitTxnFileRegions(retry.NewBackoffer(context.Background(), 1000)) + }) + + // Then + require.NoError(t, err) + require.Equal(t, uint32(1), store.splitTxnFileRegionsCalls.Load()) + require.Equal(t, uint32(0), store.splitRegionsCalls.Load()) +} + func TestBuildTxnFilesEntryCounting(t *testing.T) { require := require.New(t) From 9b6b05ce6247cedfc358902ac4b566f61ab3c195 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Fri, 7 Aug 2026 09:10:14 +0800 Subject: [PATCH 26/33] no txn file for shared lock Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 1 + txnkv/transaction/txn_file_test.go | 23 +++++++++++++++++++++++ 2 files changed, 24 insertions(+) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 1289a3b073..5d426327ab 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -1178,6 +1178,7 @@ func (c *twoPhaseCommitter) useTxnFile(ctx context.Context) (bool, error) { if c.txn.isPessimistic || c.txn.isPipelined || + c.hasSharedLocks || len(conf.TiKVClient.TxnChunkWriterAddr) == 0 || uint64(c.txn.GetMemBuffer().Size()) < minMutationSize || !IsRequestSourceUseTxnFile(c.txn.RequestSource, conf) { diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 5f51cea774..641b15cf9d 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -957,6 +957,29 @@ func TestUseTxnFileExcludesPipelinedTxn(t *testing.T) { require.False(t, useTxnFile) } +func TestUseTxnFileExcludesSharedLockTxn(t *testing.T) { + // Given + restore := config.UpdateGlobal(func(conf *config.Config) { + conf.TiKVClient.TxnChunkWriterAddr = "127.0.0.1" + conf.TiKVClient.TxnFileMinMutationSize = 0 + }) + t.Cleanup(restore) + + txn := newTestTxn(t, 1) + require.NoError(t, txn.Set([]byte("key"), []byte("value"))) + committer, err := newTwoPhaseCommitter(txn.KVTxn, 1) + require.NoError(t, err) + require.NoError(t, committer.initKeysAndMutations(context.Background())) + committer.hasSharedLocks = true + + // When + useTxnFile, err := committer.useTxnFile(context.Background()) + + // Then + require.NoError(t, err) + require.False(t, useTxnFile) +} + // stubKVStore implements kvstore with only GetRegionCache and split-call // recording backed by the mock PD client. All other methods panic because // buildTxnFiles does not call them. From f0a60e91690de616fce7b5808a6985642c693b52 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Fri, 7 Aug 2026 10:53:14 +0800 Subject: [PATCH 27/33] handle assertion level Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 9 ++++- txnkv/transaction/txn_file_test.go | 64 ++++++++++++++++++++++++++++++ 2 files changed, 72 insertions(+), 1 deletion(-) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 5d426327ab..440db340fc 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -1175,7 +1175,6 @@ func (c *twoPhaseCommitter) useTxnFile(ctx context.Context) (bool, error) { // Relax the requirement for internal requests. minMutationSize = minMutationSize / 2 } - if c.txn.isPessimistic || c.txn.isPipelined || c.hasSharedLocks || @@ -1184,6 +1183,14 @@ func (c *twoPhaseCommitter) useTxnFile(ctx context.Context) (bool, error) { !IsRequestSourceUseTxnFile(c.txn.RequestSource, conf) { return false, nil } + if c.txn.assertionLevel != kvrpcpb.AssertionLevel_Off { + // Txn-file chunks do not preserve per-mutation assertions. + for i := 0; i < c.mutations.Len(); i++ { + if c.mutations.IsAssertExists(i) || c.mutations.IsAssertNotExist(i) { + return false, nil + } + } + } logutil.Logger(ctx).Debug("transaction use txn file", zap.Uint64("startTS", c.startTS), diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 641b15cf9d..cffc7c1a13 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -949,6 +949,7 @@ func TestUseTxnFileExcludesPipelinedTxn(t *testing.T) { txn := newTestTxn(t, 1) txn.isPipelined = true + txn.SetAssertionLevel(kvrpcpb.AssertionLevel_Strict) committer := &twoPhaseCommitter{txn: txn.KVTxn} useTxnFile, err := committer.useTxnFile(context.Background()) @@ -980,6 +981,69 @@ func TestUseTxnFileExcludesSharedLockTxn(t *testing.T) { require.False(t, useTxnFile) } +func TestUseTxnFileExcludesMutationAssertions(t *testing.T) { + // Given + restore := config.UpdateGlobal(func(conf *config.Config) { + conf.TiKVClient.TxnChunkWriterAddr = "127.0.0.1" + conf.TiKVClient.TxnFileMinMutationSize = 0 + }) + t.Cleanup(restore) + + tests := []struct { + name string + assertionLevel kvrpcpb.AssertionLevel + flag tikv.FlagsOp + want bool + }{ + { + name: "strict assert exists", + assertionLevel: kvrpcpb.AssertionLevel_Strict, + flag: tikv.SetAssertExist, + want: false, + }, + { + name: "strict assert not exists", + assertionLevel: kvrpcpb.AssertionLevel_Strict, + flag: tikv.SetAssertNotExist, + want: false, + }, + { + name: "strict without mutation assertion", + assertionLevel: kvrpcpb.AssertionLevel_Strict, + want: true, + }, + { + name: "assertion off", + assertionLevel: kvrpcpb.AssertionLevel_Off, + flag: tikv.SetAssertExist, + want: true, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + // Given + txn := newTestTxn(t, 1) + txn.SetAssertionLevel(tt.assertionLevel) + key := []byte("key") + require.NoError(t, txn.Set(key, []byte("value"))) + if tt.flag != 0 { + txn.GetMemBuffer().UpdateFlags(key, tt.flag) + } + committer, err := newTwoPhaseCommitter(txn.KVTxn, 1) + require.NoError(t, err) + require.NoError(t, committer.initKeysAndMutations(context.Background())) + + // When + useTxnFile, err := committer.useTxnFile(context.Background()) + + // Then + require.NoError(t, err) + require.Equal(t, tt.want, useTxnFile) + }) + } +} + // stubKVStore implements kvstore with only GetRegionCache and split-call // recording backed by the mock PD client. All other methods panic because // buildTxnFiles does not call them. From fd466357bd10185a4f749fef2379cdd04fd3dc3f Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Fri, 7 Aug 2026 17:56:07 +0800 Subject: [PATCH 28/33] txn file assertion Signed-off-by: Ping Yu --- integration_tests/assertion_test.go | 29 +++++++++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/integration_tests/assertion_test.go b/integration_tests/assertion_test.go index 49f9a9b41d..2313a59f0b 100644 --- a/integration_tests/assertion_test.go +++ b/integration_tests/assertion_test.go @@ -25,6 +25,7 @@ import ( "github.com/pingcap/failpoint" "github.com/pingcap/kvproto/pkg/kvrpcpb" "github.com/stretchr/testify/suite" + "github.com/tikv/client-go/v2/config" tikverr "github.com/tikv/client-go/v2/error" "github.com/tikv/client-go/v2/kv" "github.com/tikv/client-go/v2/oracle" @@ -195,6 +196,34 @@ func (s *testAssertionSuite) TestPrewriteAssertion() { s.testAssertionImpl(prefix+"c", true, true, kvrpcpb.AssertionLevel_Strict) } +func (s *testAssertionSuite) TestPrewriteAssertionWithTxnFileEnabled() { + ts, err := s.store.CurrentTimestamp(oracle.GlobalTxnScope) + s.Require().NoError(err) + key := encodeKey("~assertion", fmt.Sprintf("test-txn-file-assertion-%d", ts)) + + prepareTxn, err := s.store.Begin() + s.Require().NoError(err) + s.Require().NoError(prepareTxn.Set(key, []byte("existing"))) + s.Require().NoError(prepareTxn.Commit(context.Background())) + + restore := config.UpdateGlobal(func(conf *config.Config) { + conf.TiKVClient.TxnChunkWriterAddr = "127.0.0.1" + conf.TiKVClient.TxnFileMinMutationSize = 1 + }) + s.T().Cleanup(restore) + + txn, err := s.store.Begin() + s.Require().NoError(err) + txn.SetAssertionLevel(kvrpcpb.AssertionLevel_Strict) + s.Require().NoError(txn.GetMemBuffer().SetWithFlags(key, []byte("updated"), kv.SetAssertNotExist)) + + err = txn.Commit(context.Background()) + assertionFailed, ok := errors.Cause(err).(*tikverr.ErrAssertionFailed) + s.Require().True(ok, "expected ErrAssertionFailed, got %v", err) + s.Equal(key, assertionFailed.Key) + s.Equal(kvrpcpb.Assertion_NotExist, assertionFailed.Assertion) +} + func (s *testAssertionSuite) TestFastAssertion() { // When the test cases runs with TiKV, the TiKV cluster can be reused, thus there may be deleted versions caused by // previous tests. This test case may meet different behavior if there are deleted versions. To avoid it, compose a From 0bb1ee0485b42c900ce54d65285336ee6c4591ca Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Fri, 7 Aug 2026 19:02:41 +0800 Subject: [PATCH 29/33] accouting error Signed-off-by: Ping Yu --- metrics/metrics.go | 11 +++ metrics/shortcuts.go | 2 + txnkv/transaction/txn_file.go | 9 ++- txnkv/transaction/txn_file_test.go | 122 +++++++++++++++++++++++++++++ 4 files changed, 142 insertions(+), 2 deletions(-) diff --git a/metrics/metrics.go b/metrics/metrics.go index 69c74a13fa..a68e05c585 100644 --- a/metrics/metrics.go +++ b/metrics/metrics.go @@ -139,6 +139,7 @@ var ( TiKVTxnLagCommitTSAttemptHistogram *prometheus.HistogramVec TiKVTxnFileRequestCounter *prometheus.CounterVec + TiKVTxnFileErrorCounter *prometheus.CounterVec TiKVTxnFileWriteBytes *prometheus.CounterVec TiKVTxnFileMutationSizeHistogram *prometheus.HistogramVec TiKVTxnFileDuration *prometheus.HistogramVec @@ -1069,6 +1070,15 @@ func initMetrics(namespace, subsystem string, constLabels prometheus.Labels) { ConstLabels: constLabels, }, []string{LblType}) + TiKVTxnFileErrorCounter = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Namespace: namespace, + Subsystem: subsystem, + Name: "txn_file_errors", + Help: "Counter of file-based transaction errors.", + ConstLabels: constLabels, + }, []string{LblType}) + TiKVTxnFileWriteBytes = prometheus.NewCounterVec( prometheus.CounterOpts{ Namespace: namespace, @@ -1212,6 +1222,7 @@ func RegisterMetrics() { prometheus.MustRegister(TiKVTxnLagCommitTSAttemptHistogram) prometheus.MustRegister(TiKVStaleBucketFromPDCounter) prometheus.MustRegister(TiKVTxnFileRequestCounter) + prometheus.MustRegister(TiKVTxnFileErrorCounter) prometheus.MustRegister(TiKVTxnFileWriteBytes) prometheus.MustRegister(TiKVTxnFileMutationSizeHistogram) prometheus.MustRegister(TiKVTxnFileDuration) diff --git a/metrics/shortcuts.go b/metrics/shortcuts.go index 8144cc9c6d..11c886067b 100644 --- a/metrics/shortcuts.go +++ b/metrics/shortcuts.go @@ -207,6 +207,7 @@ var ( TxnFileRequestsOk prometheus.Counter TxnFileRequestsError prometheus.Counter + TxnFileErrorAccounting prometheus.Counter TxnFileWriteBytesInternal prometheus.Counter TxnFileWriteBytesGeneral prometheus.Counter TxnFileMutationSizeInternal prometheus.Observer @@ -389,6 +390,7 @@ func initShortcuts() { TxnFileRequestsOk = TiKVTxnFileRequestCounter.WithLabelValues("ok") TxnFileRequestsError = TiKVTxnFileRequestCounter.WithLabelValues("err") + TxnFileErrorAccounting = TiKVTxnFileErrorCounter.WithLabelValues("accounting") TxnFileWriteBytesInternal = TiKVTxnFileWriteBytes.WithLabelValues(LblInternal) TxnFileWriteBytesGeneral = TiKVTxnFileWriteBytes.WithLabelValues(LblGeneral) TxnFileMutationSizeInternal = TiKVTxnFileMutationSizeHistogram.WithLabelValues(LblInternal) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 440db340fc..6c9bbd6b75 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -802,8 +802,13 @@ func (c *twoPhaseCommitter) executeTxnFile(ctx context.Context) (err error) { return } - err = c.afterExecuteTxnFile(rcInterceptor, reqInfo, ruDetails) - return + if accountingErr := c.afterExecuteTxnFile(rcInterceptor, reqInfo, ruDetails); accountingErr != nil { + metrics.TxnFileErrorAccounting.Inc() + logutil.Logger(ctx).Warn("txn file: resource control accounting failed after commit", + zap.Uint64("txnStartTS", c.startTS), + zap.Error(accountingErr)) + } + return nil } func (c *twoPhaseCommitter) executeTxnFileSlice(bo *retry.Backoffer, chunkSlice txnChunkSlice, batches []chunkBatch, action txnFileAction) (txnChunkSlice, error) { diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index cffc7c1a13..354f3b2f67 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -29,6 +29,7 @@ import ( "github.com/pingcap/kvproto/pkg/errorpb" "github.com/pingcap/kvproto/pkg/kvrpcpb" + rmpb "github.com/pingcap/kvproto/pkg/resource_manager" "github.com/pkg/errors" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" @@ -46,6 +47,7 @@ import ( "github.com/tikv/client-go/v2/tikvrpc" "github.com/tikv/client-go/v2/txnkv/txnlock" "github.com/tikv/client-go/v2/util" + resourceControlClient "github.com/tikv/pd/client/resource_group/controller" ) func TestTxnFileCleanupContextUsesStoreContext(t *testing.T) { @@ -847,6 +849,126 @@ func TestTxnFileCommitPrimaryRPCErrorIsNormalized(t *testing.T) { require.Zero(t, rollbackRequestCount.Load()) } +type txnFileResponseErrorInterceptor struct { + responseCalls int +} + +func (i *txnFileResponseErrorInterceptor) OnRequestWait( + context.Context, string, resourceControlClient.RequestInfo, +) (*rmpb.Consumption, *rmpb.Consumption, time.Duration, uint32, error) { + return &rmpb.Consumption{}, &rmpb.Consumption{}, 0, 0, nil +} + +func (i *txnFileResponseErrorInterceptor) OnResponse( + string, resourceControlClient.RequestInfo, resourceControlClient.ResponseInfo, +) (*rmpb.Consumption, error) { + i.responseCalls++ + return nil, errors.New("post-commit accounting failed") +} + +func (i *txnFileResponseErrorInterceptor) OnResponseWait( + context.Context, string, resourceControlClient.RequestInfo, resourceControlClient.ResponseInfo, +) (*rmpb.Consumption, time.Duration, error) { + return &rmpb.Consumption{}, 0, nil +} + +func (i *txnFileResponseErrorInterceptor) IsBackgroundRequest(context.Context, string, string) bool { + return false +} + +func (i *txnFileResponseErrorInterceptor) GetRUVersion() resourceControlClient.RUVersion { + return resourceControlClient.DefaultRUVersion +} + +func TestTxnFileCommitPreservesCommitOnResourceControlResponseError(t *testing.T) { + pd := &mockPDClient{} + regionCache := locate.NewTestRegionCache() + regionCache.SetPDClient(pd) + defer regionCache.Close() + + chunkWriter := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + assert.Equal(t, http.MethodPost, r.Method) + _, err := w.Write([]byte(`{"chunk_id":1}`)) + assert.NoError(t, err) + })) + defer chunkWriter.Close() + + origCfg := config.GetGlobalConfig() + newCfg := *origCfg + newCfg.TiKVClient.TxnChunkWriterAddr = chunkWriter.Listener.Addr().String() + newCfg.TiKVClient.TxnFileMinMutationSize = 1 + config.StoreGlobalConfig(&newCfg) + defer func() { + config.StoreGlobalConfig(origCfg) + once = sync.Once{} + cli = nil + errCli = nil + scheme = "" + }() + + once = sync.Once{} + cli = nil + errCli = nil + scheme = "" + + interceptor := &txnFileResponseErrorInterceptor{} + var rcInterceptor resourceControlClient.ResourceGroupKVInterceptor = interceptor + client.ResourceControlSwitch.Store(true) + client.ResourceControlInterceptor.Store(&rcInterceptor) + t.Cleanup(func() { + client.ResourceControlSwitch.Store(false) + client.ResourceControlInterceptor.Store(nil) + }) + + var commitRequestCount atomic.Int64 + var rollbackRequestCount atomic.Int64 + store := &txnFileCommitTSStore{ + timestamps: []uint64{2}, + oracle: &txnFileCommitTSOracle{}, + regionCache: regionCache, + client: &fnClient{onSend: func(_ context.Context, _ string, req *tikvrpc.Request, _ time.Duration) (*tikvrpc.Response, error) { + switch req.Type { + case tikvrpc.CmdPrewrite: + return &tikvrpc.Response{Resp: &kvrpcpb.PrewriteResponse{}}, nil + case tikvrpc.CmdCommit: + commitRequestCount.Add(1) + return &tikvrpc.Response{Resp: &kvrpcpb.CommitResponse{}}, nil + case tikvrpc.CmdBatchRollback: + rollbackRequestCount.Add(1) + return &tikvrpc.Response{Resp: &kvrpcpb.BatchRollbackResponse{}}, nil + default: + return nil, errors.Errorf("unexpected request type %s", req.Type) + } + }}, + } + memDB := unionstore.NewMemDB() + require.NoError(t, memDB.Set([]byte("k"), []byte("v"))) + txn := &KVTxn{ + store: store, + startTS: 1, + startTime: time.Now(), + valid: true, + schemaVer: txnFileSchemaVer(10), + schemaLeaseChecker: &txnFileSchemaLeaseChecker{}, + scope: oracle.GlobalTxnScope, + vars: tikv.DefaultVars, + us: unionstore.NewUnionStore(memDB, nil), + resourceGroupName: "txn-file-test", + RequestSource: &util.RequestSource{}, + } + committer, err := newTwoPhaseCommitter(txn, 0) + require.NoError(t, err) + txn.committer = committer + + err = txn.Commit(context.Background()) + + require.NoError(t, err) + require.Equal(t, int64(1), commitRequestCount.Load()) + require.Zero(t, rollbackRequestCount.Load()) + require.Equal(t, 1, interceptor.responseCalls) + require.True(t, committer.mu.committed) +} + func TestChunkSliceSortAndDedup(t *testing.T) { assert := assert.New(t) From 8cf5a023b53d759615261cbb8c396b9de2204d14 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Fri, 7 Aug 2026 21:44:06 +0800 Subject: [PATCH 30/33] rollback key error Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 7 +++++++ txnkv/transaction/txn_file_test.go | 13 +++++++++++++ 2 files changed, 20 insertions(+) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 6c9bbd6b75..0439fbc433 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -659,6 +659,13 @@ func (a txnFileRollbackAction) executeBatch(c *twoPhaseCommitter, bo *retry.Back if err1 != nil { return nil, err1 } + if keyErr := resp.Resp.(*kvrpcpb.BatchRollbackResponse).GetError(); keyErr != nil { + err := errors.Errorf("session %d txn file cleanup failed: %s", c.sessionID, keyErr) + logutil.BgLogger().Debug("txn file failed cleanup key", + zap.Error(err), + zap.Uint64("txnStartTS", c.startTS)) + return nil, err + } return resp, nil } diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 354f3b2f67..ad851c5caf 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -518,6 +518,19 @@ func TestTxnFilePrimaryBatchIndexFindsPrimaryRegion(t *testing.T) { require.Equal(t, 1, index) } +func TestTxnFilePrimaryRollbackPropagatesKeyError(t *testing.T) { + committer, bo, batch := newTxnFileCommitTestBatch(t, func(context.Context, string, *tikvrpc.Request, time.Duration) (*tikvrpc.Response, error) { + return &tikvrpc.Response{Resp: &kvrpcpb.BatchRollbackResponse{Error: &kvrpcpb.KeyError{ + Abort: "primary rollback failed", + }}}, nil + }) + + regionErr, err := committer.executeTxnFilePrimaryBatch(bo, batch, txnFileRollbackAction{}) + + require.Nil(t, regionErr) + require.ErrorContains(t, err, "session 7 txn file cleanup failed") +} + func TestTxnFileActionsApplyResourceGroupTagger(t *testing.T) { tests := []struct { name string From 98432af4047e114f62bf7e735e6a6d6000f8a789 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Fri, 7 Aug 2026 23:20:28 +0800 Subject: [PATCH 31/33] http close Signed-off-by: Ping Yu --- txnkv/client.go | 8 +++ txnkv/client_test.go | 26 +++++++++ txnkv/transaction/txn_file.go | 19 ++++-- txnkv/transaction/txn_file_test.go | 92 ++++++++++++++++++++++++++++-- 4 files changed, 135 insertions(+), 10 deletions(-) create mode 100644 txnkv/client_test.go diff --git a/txnkv/client.go b/txnkv/client.go index bfcaad47a7..2930b5c0db 100644 --- a/txnkv/client.go +++ b/txnkv/client.go @@ -119,6 +119,14 @@ func NewClient(pdAddrs []string, opts ...ClientOpt) (*Client, error) { return &Client{KVStore: s}, nil } +// Close releases resources owned by the transactional client, including shared +// idle HTTP connections opened for txn-file chunk uploads. +func (c Client) Close() error { + err := c.KVStore.Close() + transaction.CloseTxnFileIdleConnections() + return err +} + // GetTimestamp returns the current global timestamp. func (c *Client) GetTimestamp(ctx context.Context) (uint64, error) { bo := retry.NewBackofferWithVars(ctx, transaction.TsoMaxBackoff, nil) diff --git a/txnkv/client_test.go b/txnkv/client_test.go new file mode 100644 index 0000000000..2bf1c260af --- /dev/null +++ b/txnkv/client_test.go @@ -0,0 +1,26 @@ +// Copyright 2026 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. + +package txnkv_test + +import ( + "io" + + "github.com/tikv/client-go/v2/txnkv" +) + +var ( + _ io.Closer = txnkv.Client{} + _ io.Closer = (*txnkv.Client)(nil) +) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 0439fbc433..8c143cb6e0 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -28,6 +28,7 @@ import ( "sort" "strings" "sync" + "sync/atomic" "time" "github.com/golang/protobuf/proto" //nolint:staticcheck @@ -1355,7 +1356,7 @@ func (c *twoPhaseCommitter) reportFailureMetrics() { var ( once sync.Once scheme string - cli *http.Client + cli atomic.Pointer[http.Client] errCli error ) @@ -1371,6 +1372,7 @@ func getHTTPClient() (*http.Client, error) { transport := &http.Transport{ MaxIdleConns: 100, MaxIdleConnsPerHost: 20, + IdleConnTimeout: 90 * time.Second, } if len(cfg.Security.ClusterSSLCA) != 0 { scheme = "https://" @@ -1383,12 +1385,21 @@ func getHTTPClient() (*http.Client, error) { transport.ForceAttemptHTTP2 = true } - cli = &http.Client{ + cli.Store(&http.Client{ Timeout: timeout, Transport: transport, - } + }) }) - return cli, errCli + return cli.Load(), errCli +} + +// CloseTxnFileIdleConnections closes idle HTTP connections opened by txn-file +// chunk uploads. It does not interrupt active requests and is safe to call +// while another client is using the shared uploader. +func CloseTxnFileIdleConnections() { + if client := cli.Load(); client != nil { + client.CloseIdleConnections() + } } type chunkWriterClient struct { diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index ad851c5caf..b23c906062 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -18,7 +18,9 @@ import ( "context" "encoding/json" "fmt" + "io" "math/rand" + "net" "net/http" "net/http/httptest" "slices" @@ -105,6 +107,84 @@ type txnFileCommitTSOracle struct { option *oracle.Option } +func TestCloseTxnFileIdleConnections(t *testing.T) { + original := cli.Load() + t.Cleanup(func() { + cli.Store(original) + }) + + idle := make(chan struct{}, 1) + closed := make(chan struct{}, 1) + server := httptest.NewUnstartedServer(http.HandlerFunc(func(http.ResponseWriter, *http.Request) {})) + server.Config.ConnState = func(_ net.Conn, state http.ConnState) { + switch state { + case http.StateIdle: + select { + case idle <- struct{}{}: + default: + } + case http.StateClosed: + select { + case closed <- struct{}{}: + default: + } + } + } + server.Start() + t.Cleanup(server.Close) + client := server.Client() + cli.Store(client) + + resp, err := client.Get(server.URL) + require.NoError(t, err) + _, err = io.Copy(io.Discard, resp.Body) + require.NoError(t, err) + require.NoError(t, resp.Body.Close()) + select { + case <-idle: + case <-time.After(5 * time.Second): + require.FailNow(t, "HTTP connection did not become idle") + } + + CloseTxnFileIdleConnections() + + select { + case <-closed: + case <-time.After(5 * time.Second): + require.FailNow(t, "idle HTTP connection was not closed") + } +} + +func TestCloseTxnFileIdleConnectionsBeforeInitialization(t *testing.T) { + original := cli.Load() + t.Cleanup(func() { + cli.Store(original) + }) + cli.Store(nil) + + require.NotPanics(t, CloseTxnFileIdleConnections) +} + +func TestTxnFileHTTPClientHasIdleConnectionTimeout(t *testing.T) { + t.Cleanup(func() { + once = sync.Once{} + cli.Store(nil) + errCli = nil + scheme = "" + }) + once = sync.Once{} + cli.Store(nil) + errCli = nil + scheme = "" + + client, err := getHTTPClient() + + require.NoError(t, err) + transport, ok := client.Transport.(*http.Transport) + require.True(t, ok) + require.Equal(t, 90*time.Second, transport.IdleConnTimeout) +} + func (o *txnFileCommitTSOracle) IsExpired(startTS uint64, ttl uint64, option *oracle.Option) bool { o.calls++ o.startTS = startTS @@ -802,13 +882,13 @@ func TestTxnFileCommitPrimaryRPCErrorIsNormalized(t *testing.T) { defer func() { config.StoreGlobalConfig(origCfg) once = sync.Once{} - cli = nil + cli.Store(nil) errCli = nil scheme = "" }() once = sync.Once{} - cli = nil + cli.Store(nil) errCli = nil scheme = "" @@ -914,13 +994,13 @@ func TestTxnFileCommitPreservesCommitOnResourceControlResponseError(t *testing.T defer func() { config.StoreGlobalConfig(origCfg) once = sync.Once{} - cli = nil + cli.Store(nil) errCli = nil scheme = "" }() once = sync.Once{} - cli = nil + cli.Store(nil) errCli = nil scheme = "" @@ -1283,13 +1363,13 @@ func TestBuildTxnFilesEntryCounting(t *testing.T) { defer func() { config.StoreGlobalConfig(origCfg) once = sync.Once{} - cli = nil + cli.Store(nil) errCli = nil scheme = "" }() once = sync.Once{} - cli = srv.Client() + cli.Store(srv.Client()) errCli = nil scheme = "http://" From 18258918aba9f5fea72d995db360f806a4ff7635 Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Sat, 8 Aug 2026 00:21:00 +0800 Subject: [PATCH 32/33] discard value Signed-off-by: Ping Yu --- txnkv/transaction/txn_file.go | 1 + txnkv/transaction/txn_file_test.go | 23 ++++++++++++++++++++--- 2 files changed, 21 insertions(+), 3 deletions(-) diff --git a/txnkv/transaction/txn_file.go b/txnkv/transaction/txn_file.go index 8c143cb6e0..227d93c0dd 100644 --- a/txnkv/transaction/txn_file.go +++ b/txnkv/transaction/txn_file.go @@ -797,6 +797,7 @@ func (c *twoPhaseCommitter) executeTxnFile(ctx context.Context) (err error) { if err != nil { return } + c.txn.GetMemBuffer().GetMemDB().DiscardValues() err = c.executeTxnFileAction(commitBo, c.txnFileCtx.slice, txnFileCommitAction{}) stepDone("commit") if err != nil { diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index b23c906062..7be5ae9995 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -16,8 +16,10 @@ package transaction import ( "context" + "encoding/binary" "encoding/json" "fmt" + "hash/crc32" "io" "math/rand" "net" @@ -979,9 +981,15 @@ func TestTxnFileCommitPreservesCommitOnResourceControlResponseError(t *testing.T regionCache.SetPDClient(pd) defer regionCache.Close() + memDB := unionstore.NewMemDB() + require.NoError(t, memDB.Set([]byte("k"), []byte("v"))) + uploadedChunks := make(chan []byte, 1) chunkWriter := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { assert.Equal(t, http.MethodPost, r.Method) - _, err := w.Write([]byte(`{"chunk_id":1}`)) + chunk, err := io.ReadAll(r.Body) + assert.NoError(t, err) + uploadedChunks <- chunk + _, err = w.Write([]byte(`{"chunk_id":1}`)) assert.NoError(t, err) })) defer chunkWriter.Close() @@ -1034,8 +1042,6 @@ func TestTxnFileCommitPreservesCommitOnResourceControlResponseError(t *testing.T } }}, } - memDB := unionstore.NewMemDB() - require.NoError(t, memDB.Set([]byte("k"), []byte("v"))) txn := &KVTxn{ store: store, startTS: 1, @@ -1056,10 +1062,21 @@ func TestTxnFileCommitPreservesCommitOnResourceControlResponseError(t *testing.T err = txn.Commit(context.Background()) require.NoError(t, err) + expectedChunk := binary.LittleEndian.AppendUint16(nil, uint16(len("k"))) + expectedChunk = append(expectedChunk, "k"...) + expectedChunk = append(expectedChunk, byte(kvrpcpb.Op_Put)) + expectedChunk = binary.LittleEndian.AppendUint32(expectedChunk, uint32(len("v"))) + expectedChunk = append(expectedChunk, "v"...) + expectedChunk = binary.LittleEndian.AppendUint32(expectedChunk, crc32.ChecksumIEEE(expectedChunk)) + require.Equal(t, expectedChunk, <-uploadedChunks) require.Equal(t, int64(1), commitRequestCount.Load()) require.Zero(t, rollbackRequestCount.Load()) require.Equal(t, 1, interceptor.responseCalls) require.True(t, committer.mu.committed) + // DiscardValues invalidates the MemDB value log, so reading a value after commit panics by contract. + require.Panics(t, func() { + _, _ = memDB.Get(context.Background(), []byte("k")) + }) } func TestChunkSliceSortAndDedup(t *testing.T) { From 77facf49667ad666cfc91927acf4097172e6d79e Mon Sep 17 00:00:00 2001 From: Ping Yu Date: Sat, 8 Aug 2026 09:08:43 +0800 Subject: [PATCH 33/33] fix ci Signed-off-by: Ping Yu --- txnkv/transaction/txn_file_test.go | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/txnkv/transaction/txn_file_test.go b/txnkv/transaction/txn_file_test.go index 7be5ae9995..055d5cbd58 100644 --- a/txnkv/transaction/txn_file_test.go +++ b/txnkv/transaction/txn_file_test.go @@ -137,7 +137,9 @@ func TestCloseTxnFileIdleConnections(t *testing.T) { client := server.Client() cli.Store(client) - resp, err := client.Get(server.URL) + req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, server.URL, nil) + require.NoError(t, err) + resp, err := client.Do(req) require.NoError(t, err) _, err = io.Copy(io.Discard, resp.Body) require.NoError(t, err)