Allow setting concurrency for pipelined flush and resolveLocks

examples/gcworker/go.mod

@@ -1,6 +1,6 @@
 module gcworker
 
-go 1.23
+go 1.23.2
 
 require github.com/tikv/client-go/v2 v2.0.0
 

examples/rawkv/go.mod

@@ -1,6 +1,6 @@
 module rawkv
 
-go 1.23
+go 1.23.2
 
 require github.com/tikv/client-go/v2 v2.0.0
 

examples/txnkv/1pc_txn/go.mod

@@ -1,6 +1,6 @@
 module 1pc_txn
 
-go 1.23
+go 1.23.2
 
 require github.com/tikv/client-go/v2 v2.0.0
 

examples/txnkv/async_commit/go.mod

@@ -1,6 +1,6 @@
 module async_commit
 
-go 1.23
+go 1.23.2
 
 require github.com/tikv/client-go/v2 v2.0.0
 

examples/txnkv/delete_range/go.mod

@@ -1,6 +1,6 @@
 module delete_range
 
-go 1.23
+go 1.23.2
 
 require github.com/tikv/client-go/v2 v2.0.0
 

examples/txnkv/go.mod

@@ -1,6 +1,6 @@
 module txnkv
 
-go 1.23
+go 1.23.2
 
 require github.com/tikv/client-go/v2 v2.0.0
 

examples/txnkv/pessimistic_txn/go.mod

@@ -1,6 +1,6 @@
 module pessimistic_txn
 
-go 1.23
+go 1.23.2
 
 require github.com/tikv/client-go/v2 v2.0.0
 

examples/txnkv/unsafedestoryrange/go.mod

@@ -1,6 +1,6 @@
 module unsafedestoryrange
 
-go 1.23
+go 1.23.2
 
 require github.com/tikv/client-go/v2 v2.0.0
 

go.mod

@@ -3,6 +3,7 @@ module github.com/tikv/client-go/v2
 go 1.23
 
 require (
+	github.com/VividCortex/ewma v1.2.0
 	github.com/cznic/mathutil v0.0.0-20181122101859-297441e03548
 	github.com/dgryski/go-farm v0.0.0-20240924180020-3414d57e47da
 	github.com/docker/go-units v0.5.0

go.sum

@@ -1,6 +1,8 @@
 cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
 github.com/BurntSushi/toml v0.3.1 h1:WXkYYl6Yr3qBf1K79EBnL4mak0OimBfB0XUf9Vl28OQ=
 github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
+github.com/VividCortex/ewma v1.2.0 h1:f58SaIzcDXrSy3kWaHNvuJgJ3Nmz59Zji6XoJR/q1ow=
+github.com/VividCortex/ewma v1.2.0/go.mod h1:nz4BbCtbLyFDeC9SUHbtcT5644juEuWfUAUnGx7j5l4=
 github.com/benbjohnson/clock v1.1.0/go.mod h1:J11/hYXuz8f4ySSvYwY0FKfm+ezbsZBKZxNJlLklBHA=
 github.com/beorn7/perks v1.0.1 h1:VlbKKnNfV8bJzeqoa4cOKqO6bYr3WgKZxO8Z16+hsOM=
 github.com/beorn7/perks v1.0.1/go.mod h1:G2ZrVWU2WbWT9wwq4/hrbKbnv/1ERSJQ0ibhJ6rlkpw=

integration_tests/go.mod

@@ -1,6 +1,6 @@
 module integration_tests
 
-go 1.23
+go 1.23.2
 
 require (
 	github.com/ninedraft/israce v0.0.3

integration_tests/pipelined_memdb_test.go

@@ -97,7 +97,7 @@ func (s *testPipelinedMemDBSuite) mustGetLock(key []byte) *txnkv.Lock {
 
 func (s *testPipelinedMemDBSuite) TestPipelinedAndFlush() {
 	ctx := context.Background()
-	txn, err := s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err := s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 	txn1, err := s.store.Begin()
 	s.Nil(err)
@@ -139,7 +139,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedAndFlush() {
 
 func (s *testPipelinedMemDBSuite) TestPipelinedMemDBBufferGet() {
 	ctx := context.Background()
-	txn, err := s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err := s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 	for i := 0; i < 100; i++ {
 		key := []byte(strconv.Itoa(i))
@@ -161,7 +161,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedMemDBBufferGet() {
 }
 
 func (s *testPipelinedMemDBSuite) TestPipelinedFlushBlock() {
-	txn, err := s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err := s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 	txn.Set([]byte("key1"), []byte("value1"))
 
@@ -192,7 +192,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedFlushBlock() {
 }
 
 func (s *testPipelinedMemDBSuite) TestPipelinedSkipFlushedLock() {
-	txn, err := s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err := s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 	txn.Set([]byte("key1"), []byte("value1"))
 	flushed, err := txn.GetMemBuffer().Flush(true)
@@ -206,7 +206,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedSkipFlushedLock() {
 	s.Nil(txn.Commit(context.Background()))
 
 	// can see it after commit
-	txn, err = s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err = s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 	defer txn.Rollback()
 	val, err := txn.Get(context.Background(), []byte("key1"))
@@ -222,7 +222,7 @@ func (s *testPipelinedMemDBSuite) TestResolveLockRace() {
 		failpoint.Disable("tikvclient/pipelinedCommitFail")
 	}()
 	for i := 0; i < 100; i++ {
-		txn, err := s.store.Begin(tikv.WithPipelinedMemDB())
+		txn, err := s.store.Begin(tikv.WithDefaultPipelinedTxn())
 		startTS := txn.StartTS()
 		s.Nil(err)
 		for j := 0; j < 100; j++ {
@@ -246,7 +246,7 @@ func (s *testPipelinedMemDBSuite) TestResolveLockRace() {
 }
 
 func (s *testPipelinedMemDBSuite) TestPipelinedCommit() {
-	txn, err := s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err := s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 	for i := 0; i < 100; i++ {
 		key := []byte(strconv.Itoa(i))
@@ -279,7 +279,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedCommit() {
 	}
 
 	// check the result
-	txn, err = s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err = s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 	for i := 0; i < 100; i++ {
 		key := []byte(strconv.Itoa(i))
@@ -290,7 +290,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedCommit() {
 }
 
 func (s *testPipelinedMemDBSuite) TestPipelinedRollback() {
-	txn, err := s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err := s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	startTS := txn.StartTS()
 	s.Nil(err)
 	keys := make([][]byte, 0, 100)
@@ -304,7 +304,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedRollback() {
 	s.Nil(txn.GetMemBuffer().FlushWait())
 	s.Nil(txn.Rollback())
 	s.Eventuallyf(func() bool {
-		txn, err := s.store.Begin(tikv.WithStartTS(startTS), tikv.WithPipelinedMemDB())
+		txn, err := s.store.Begin(tikv.WithStartTS(startTS), tikv.WithDefaultPipelinedTxn())
 		s.Nil(err)
 		defer func() { s.Nil(txn.Rollback()) }()
 		storageBufferedValues, err := txn.GetSnapshot().BatchGetWithTier(context.Background(), keys, txnsnapshot.BatchGetBufferTier)
@@ -322,7 +322,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedRollback() {
 func (s *testPipelinedMemDBSuite) TestPipelinedPrefetch() {
 	failpoint.Enable("tikvclient/beforeSendReqToRegion", "return")
 	defer failpoint.Disable("tikvclient/beforeSendReqToRegion")
-	txn, err := s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err := s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 
 	mustFlush := func(txn *transaction.KVTxn) {
@@ -384,7 +384,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedPrefetch() {
 
 	s.Nil(txn.Commit(context.Background()))
 
-	txn, err = s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err = s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 	txn.Set([]byte("100"), []byte("100")) // snapshot: [0, 1, ..., 99], membuffer: [100]
 	m, err := txn.BatchGet(context.Background(), [][]byte{[]byte("99"), []byte("100"), []byte("101")})
@@ -417,7 +417,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedPrefetch() {
 	txn.Rollback()
 
 	// empty memdb should also cache the not exist result.
-	txn, err = s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err = s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	// batch get cache: [99 -> not exist]
 	m, err = txn.BatchGet(context.Background(), [][]byte{[]byte("99")})
 	s.Nil(err)
@@ -433,7 +433,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedDMLFailedByPKRollback() {
 	atomic.StoreUint64(&transaction.ManagedLockTTL, 100) // set to 100ms
 	defer atomic.StoreUint64(&transaction.ManagedLockTTL, originManageTTLVal)
 
-	txn, err := s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err := s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 	txn.Set([]byte("key1"), []byte("value1"))
 	txnProbe := transaction.TxnProbe{KVTxn: txn}
@@ -482,7 +482,7 @@ func (s *testPipelinedMemDBSuite) TestPipelinedDMLFailedByPKMaxTTLExceeded() {
 		restoreGlobalConfFunc()
 	}()
 
-	txn, err := s.store.Begin(tikv.WithPipelinedMemDB())
+	txn, err := s.store.Begin(tikv.WithDefaultPipelinedTxn())
 	s.Nil(err)
 	txn.Set([]byte("key1"), []byte("value1"))
 	txnProbe := transaction.TxnProbe{KVTxn: txn}

metrics/metrics.go

@@ -117,6 +117,7 @@ var (
 	TiKVPipelinedFlushDuration                     prometheus.Histogram
 	TiKVValidateReadTSFromPDCount                  prometheus.Counter
 	TiKVLowResolutionTSOUpdateIntervalSecondsGauge prometheus.Gauge
+	TiKVPipelinedFlushThrottleSecondsHistogram     prometheus.Histogram
 )
 
 // Label constants.
@@ -852,6 +853,15 @@ func initMetrics(namespace, subsystem string, constLabels prometheus.Labels) {
 			Help:      "The actual working update interval for the low resolution TSO. As there are adaptive mechanism internally, this value may differ from the config.",
 		})
 
+	TiKVPipelinedFlushThrottleSecondsHistogram = prometheus.NewHistogram(
+		prometheus.HistogramOpts{
+			Namespace: namespace,
+			Subsystem: subsystem,
+			Name:      "pipelined_flush_throttle_seconds",
+			Help:      "Throttle durations of pipelined flushes.",
+			Buckets:   prometheus.ExponentialBuckets(0.0005, 2, 28), // 0.5ms ~ 18h
+		})
+
 	initShortcuts()
 }
 
@@ -948,6 +958,7 @@ func RegisterMetrics() {
 	prometheus.MustRegister(TiKVPipelinedFlushDuration)
 	prometheus.MustRegister(TiKVValidateReadTSFromPDCount)
 	prometheus.MustRegister(TiKVLowResolutionTSOUpdateIntervalSecondsGauge)
+	prometheus.MustRegister(TiKVPipelinedFlushThrottleSecondsHistogram)
 }
 
 // readCounter reads the value of a prometheus.Counter.

tikv/kv.go

@@ -80,6 +80,10 @@ const (
 	// DCLabelKey indicates the key of label which represents the dc for Store.
 	DCLabelKey           = "zone"
 	safeTSUpdateInterval = time.Second * 2
+
+	defaultPipelinedFlushConcurrency       = 128
+	defaultPipelinedResolveLockConcurrency = 8
+	defaultPipelinedWriteThrottleRatio     = 0.0
 )
 
 func createEtcdKV(addrs []string, tlsConfig *tls.Config) (*clientv3.Client, error) {
@@ -944,10 +948,31 @@ func WithStartTS(startTS uint64) TxnOption {
 	}
 }
 
-// WithPipelinedMemDB creates transaction with pipelined memdb.
-func WithPipelinedMemDB() TxnOption {
+// WithDefaultPipelinedTxn creates pipelined txn with default parameters
+func WithDefaultPipelinedTxn() TxnOption {
 	return func(st *transaction.TxnOptions) {
-		st.PipelinedMemDB = true
+		st.PipelinedTxn = transaction.PipelinedTxnOptions{
+			Enable:                 true,
+			FlushConcurrency:       defaultPipelinedFlushConcurrency,
+			ResolveLockConcurrency: defaultPipelinedResolveLockConcurrency,
+			WriteThrottleRatio:     defaultPipelinedWriteThrottleRatio,
+		}
+	}
+}
+
+// WithPipelinedTxn creates pipelined txn with specified parameters
+func WithPipelinedTxn(
+	flushConcurrency,
+	resolveLockConcurrency int,
+	writeThrottleRatio float64,
+) TxnOption {
+	return func(st *transaction.TxnOptions) {
+		st.PipelinedTxn = transaction.PipelinedTxnOptions{
+			Enable:                 true,
+			FlushConcurrency:       flushConcurrency,
+			ResolveLockConcurrency: resolveLockConcurrency,
+			WriteThrottleRatio:     writeThrottleRatio,
+		}
 	}
 }
 

txnkv/transaction/2pc.go

@@ -1102,16 +1102,28 @@ func (c *twoPhaseCommitter) doActionOnBatches(
 		}
 		return nil
 	}
-	rateLim := len(batches)
+	rateLim := c.calcActionConcurrency(len(batches), action)
+	batchExecutor := newBatchExecutor(rateLim, c, action, bo)
+	return batchExecutor.process(batches)
+}
+
+func (c *twoPhaseCommitter) calcActionConcurrency(
+	numBatches int, action twoPhaseCommitAction,
+) int {
+	rateLim := numBatches
 	// Set rateLim here for the large transaction.
 	// If the rate limit is too high, tikv will report service is busy.
 	// If the rate limit is too low, we can't full utilize the tikv's throughput.
 	// TODO: Find a self-adaptive way to control the rate limit here.
-	if rateLim > config.GetGlobalConfig().CommitterConcurrency {
-		rateLim = config.GetGlobalConfig().CommitterConcurrency
+	switch action.(type) {
+	case actionPipelinedFlush:
+		rateLim = min(rateLim, max(1, c.txn.pipelinedFlushConcurrency))
+	default:
+		if rateLim > config.GetGlobalConfig().CommitterConcurrency {
+			rateLim = config.GetGlobalConfig().CommitterConcurrency
+		}
 	}
-	batchExecutor := newBatchExecutor(rateLim, c, action, bo)
-	return batchExecutor.process(batches)
+	return rateLim
 }
 
 func (c *twoPhaseCommitter) keyValueSize(key, value []byte) int {

txnkv/transaction/pipelined_flush.go

@@ -454,7 +454,6 @@ func (c *twoPhaseCommitter) buildPipelinedResolveHandler(commit bool, resolved *
 // resolveFlushedLocks resolves all locks in the given range [start, end) with the given status.
 // The resolve process is running in another goroutine so this function won't block.
 func (c *twoPhaseCommitter) resolveFlushedLocks(bo *retry.Backoffer, start, end []byte, commit bool) {
-	const RESOLVE_CONCURRENCY = 8
 	var resolved atomic.Uint64
 	handler, err := c.buildPipelinedResolveHandler(commit, &resolved)
 	commitTs := uint64(0)
@@ -481,7 +480,7 @@ func (c *twoPhaseCommitter) resolveFlushedLocks(bo *retry.Backoffer, start, end
 		fmt.Sprintf("pipelined-dml-%s", status),
 		fmt.Sprintf("pipelined-dml-%s-%d", status, c.startTS),
 		c.store,
-		RESOLVE_CONCURRENCY,
+		c.txn.pipelinedResolveLockConcurrency,
 		handler,
 	)
 	runner.SetStatLogInterval(30 * time.Second)