transaction: Support file based transaction

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 (
@@ -121,6 +123,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 +249,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{},
 	}
 }
 
@@ -246,6 +269,29 @@ 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.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")
+	}
+	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")
+	}
+	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
 }
 

config/config_test.go

@@ -35,6 +35,8 @@
 package config
 
 import (
+	"fmt"
+	"math"
 	"testing"
 	"time"
 
@@ -89,3 +91,82 @@ 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) {
+	configWithoutTxnFile := DefaultTiKVClient()
+	configWithoutTxnFile.TxnChunkMaxSize = 0
+	assert.NoError(t, configWithoutTxnFile.Valid())
+
+	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 = 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) {
+				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
+			}
+			cfg.TxnChunkWriterAddr = "127.0.0.1"
+			assert.EqualError(t, cfg.Valid(), test.err)
+		})
+	}
+}

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

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