transaction: Support file based transaction

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{},
 	}
 }
 

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")
+}

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