Private/issue 2677 decision statistics pr desc clean

client/clients/tso/dispatcher.go

@@ -167,7 +167,10 @@ func (td *tsoDispatcher) handleDispatcher(wg *sync.WaitGroup) {
 		stream    *tsoStream
 	)
 	// Loop through each batch of TSO requests and send them for processing.
-	streamLoopTimer := time.NewTimer(option.Timeout)
+	// Use TSO-specific timeout which can be configured via WithCustomTSOTimeoutOption.
+	// GetTSOTimeout also considers the backoffer's total time limit if set.
+	tsoTimeout := option.GetTSOTimeout()
+	streamLoopTimer := time.NewTimer(tsoTimeout)
 	defer streamLoopTimer.Stop()
 
 	// Create a not-started-timer to be used for collecting batches for concurrent RPC.
@@ -215,7 +218,7 @@ tsoBatchLoop:
 		if maxBatchWaitInterval >= 0 {
 			tsoBatchController.AdjustBestBatchSize()
 		}
-		streamLoopTimer.Reset(option.Timeout)
+		streamLoopTimer.Reset(tsoTimeout)
 		// Choose a stream to send the TSO gRPC request.
 	streamChoosingLoop:
 		for {
@@ -323,7 +326,7 @@ tsoBatchLoop:
 			}
 		}
 
-		done := td.deadlineWatcher.Start(ctx, option.Timeout, cancel)
+		done := td.deadlineWatcher.Start(ctx, tsoTimeout, cancel)
 		if done == nil {
 			// Finish the collected requests if the context is canceled.
 			td.cancelCollectedRequests(tsoBatchController, invalidStreamID, errors.WithStack(ctx.Err()))

client/clients/tso/dispatcher_test.go

@@ -43,6 +43,7 @@ func TestMain(m *testing.M) {
 type mockTSOServiceProvider struct {
 	option       *opt.Option
 	createStream func(ctx context.Context) *tsoStream
+	updateConCtx func(ctx context.Context) bool
 	conCtxMgr    *cctx.Manager[*tsoStream]
 }
 
@@ -67,6 +68,9 @@ func (m *mockTSOServiceProvider) getConnectionCtxMgr() *cctx.Manager[*tsoStream]
 }
 
 func (m *mockTSOServiceProvider) updateConnectionCtxs(ctx context.Context) bool {
+	if m.updateConCtx != nil {
+		return m.updateConCtx(ctx)
+	}
 	if m.conCtxMgr.Exist(mockStreamURL) {
 		return true
 	}
@@ -81,6 +85,50 @@ func (m *mockTSOServiceProvider) updateConnectionCtxs(ctx context.Context) bool
 	return true
 }
 
+func TestTSODispatcherStreamLoopUsesTSOTimeout(t *testing.T) {
+	re := require.New(t)
+	option := opt.NewOption()
+	option.Timeout = 200 * time.Millisecond
+	option.TSOTimeout = 700 * time.Millisecond
+
+	provider := newMockTSOServiceProvider(option, nil)
+	provider.updateConCtx = func(context.Context) bool {
+		return false
+	}
+
+	dispatcher := newTSODispatcher(context.Background(), defaultMaxTSOBatchSize, provider)
+	var dispatcherWg sync.WaitGroup
+	dispatcherWg.Add(1)
+	go dispatcher.handleDispatcher(&dispatcherWg)
+	defer func() {
+		dispatcher.close()
+		dispatcherWg.Wait()
+	}()
+
+	reqPool := &sync.Pool{
+		New: func() any {
+			return &Request{
+				done: make(chan error, 1),
+			}
+		},
+	}
+	req := reqPool.Get().(*Request)
+	req.clientCtx = context.Background()
+	req.requestCtx = context.Background()
+	req.start = time.Now()
+	req.pool = reqPool
+	req.physical = 0
+	req.logical = 0
+
+	start := time.Now()
+	dispatcher.push(req)
+	_, _, err := req.waitTimeout(3 * time.Second)
+	elapsed := time.Since(start)
+
+	re.Error(err)
+	re.Greater(elapsed, option.Timeout+200*time.Millisecond)
+}
+
 type testTSODispatcherSuite struct {
 	suite.Suite
 	re *require.Assertions
@@ -98,6 +146,7 @@ func (s *testTSODispatcherSuite) SetupTest() {
 	s.re = require.New(s.T())
 	s.option = opt.NewOption()
 	s.option.Timeout = time.Hour
+	s.option.TSOTimeout = time.Hour
 	// As the internal logic of the tsoDispatcher allows it to create streams multiple times, but our tests needs
 	// single stable access to the inner stream, we do not allow it to create it more than once in these tests.
 	// To avoid data race on reading `stream` and `streamInner` fields.

client/opt/option.go

@@ -34,6 +34,7 @@ const (
 	defaultEnableFollowerHandle                  = false
 	defaultTSOClientRPCConcurrency               = 1
 	defaultEnableRouterClient                    = true
+	defaultTSOTimeout                            = 15 * time.Second
 )
 
 // DynamicOption is used to distinguish the dynamic option type.
@@ -63,6 +64,7 @@ type Option struct {
 	// Static options.
 	GRPCDialOptions   []grpc.DialOption
 	Timeout           time.Duration
+	TSOTimeout         time.Duration
 	MaxRetryTimes     int
 	EnableForwarding  bool
 	UseTSOServerProxy bool
@@ -82,11 +84,12 @@ type Option struct {
 func NewOption() *Option {
 	co := &Option{
 		Timeout:                  defaultPDTimeout,
+		TSOTimeout:               defaultTSOTimeout,
 		MaxRetryTimes:            maxInitClusterRetries,
 		EnableTSOFollowerProxyCh: make(chan struct{}, 1),
 		EnableFollowerHandleCh:   make(chan struct{}, 1),
-		EnableRouterClientCh:     make(chan struct{}, 1),
-		InitMetrics:              true,
+		EnableRouterClientCh:      make(chan struct{}, 1),
+		InitMetrics:               true,
 	}
 
 	co.dynamicOptions[MaxTSOBatchWaitInterval].Store(defaultMaxTSOBatchWaitInterval)
@@ -184,6 +187,33 @@ func WithCustomTimeoutOption(timeout time.Duration) ClientOption {
 	}
 }
 
+// WithCustomTSOTimeoutOption configures the client with TSO timeout option.
+// This timeout affects how long TSO requests will wait before timing out.
+// The TSO timeout can also be affected by the backoffer if set via WithBackoffer.
+func WithCustomTSOTimeoutOption(timeout time.Duration) ClientOption {
+	return func(op *Option) {
+		op.TSOTimeout = timeout
+	}
+}
+
+// GetTSOTimeout returns the effective TSO timeout considering the backoffer's total time limit.
+// If a backoffer is set with a non-zero total time, the TSO timeout will be capped by that limit.
+// This allows tidb_backoff_weight to indirectly affect TSO timeout via the backoffer.
+func (o *Option) GetTSOTimeout() time.Duration {
+	tsoTimeout := o.TSOTimeout
+	if tsoTimeout <= 0 {
+		tsoTimeout = o.Timeout
+	}
+	// If backoffer has a total time limit, cap the TSO timeout accordingly
+	if o.Backoffer != nil {
+		backofferTotal := o.Backoffer.TotalTime()
+		if backofferTotal > 0 && backofferTotal < tsoTimeout {
+			return backofferTotal
+		}
+	}
+	return tsoTimeout
+}
+
 // WithForwardingOption configures the client with forwarding option.
 func WithForwardingOption(enableForwarding bool) ClientOption {
 	return func(op *Option) {

client/opt/option_test.go

@@ -21,6 +21,7 @@ import (
 	"github.com/stretchr/testify/require"
 	"go.uber.org/goleak"
 
+	"github.com/tikv/pd/client/pkg/retry"
 	"github.com/tikv/pd/client/pkg/utils/testutil"
 )
 
@@ -149,3 +150,80 @@ func ensureNoNotification(t *testing.T, ch chan struct{}) {
 		// No notification received as expected.
 	}
 }
+
+func TestGetTSOTimeout(t *testing.T) {
+	re := require.New(t)
+
+	tests := []struct {
+		name           string
+		option         *Option
+		expectedResult time.Duration
+	}{
+		{
+			name:           "default timeout",
+			option:         NewOption(),
+			expectedResult: defaultTSOTimeout,
+		},
+		{
+			name: "custom TSO timeout",
+			option: &Option{
+				Timeout:    defaultPDTimeout,
+				TSOTimeout:  30 * time.Second,
+			},
+			expectedResult: 30 * time.Second,
+		},
+		{
+			name: "backoffer with shorter total time",
+			option: &Option{
+				Timeout:    defaultPDTimeout,
+				TSOTimeout: 15 * time.Second,
+				Backoffer:  retry.InitialBackoffer(100*time.Millisecond, 1*time.Second, 5*time.Second),
+			},
+			expectedResult: 5 * time.Second,
+		},
+		{
+			name: "backoffer with longer total time",
+			option: &Option{
+				Timeout:    defaultPDTimeout,
+				TSOTimeout: 10 * time.Second,
+				Backoffer:  retry.InitialBackoffer(100*time.Millisecond, 1*time.Second, 30*time.Second),
+			},
+			expectedResult: 10 * time.Second,
+		},
+		{
+			name: "backoffer with infinite retry",
+			option: &Option{
+				Timeout:    defaultPDTimeout,
+				TSOTimeout: 15 * time.Second,
+				Backoffer:  retry.InitialBackoffer(100*time.Millisecond, 1*time.Second, 0),
+			},
+			expectedResult: 15 * time.Second,
+		},
+		{
+			name: "fallback to Timeout when TSOTimeout is zero",
+			option: &Option{
+				Timeout:    5 * time.Second,
+				TSOTimeout: 0,
+			},
+			expectedResult: 5 * time.Second,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			result := tt.option.GetTSOTimeout()
+			re.Equal(tt.expectedResult, result)
+		})
+	}
+}
+
+func TestWithCustomTSOTimeoutOption(t *testing.T) {
+	re := require.New(t)
+
+	option := NewOption()
+	re.Equal(defaultTSOTimeout, option.TSOTimeout)
+
+	customTimeout := 20 * time.Second
+	WithCustomTSOTimeoutOption(customTimeout)(option)
+	re.Equal(customTimeout, option.TSOTimeout)
+}

client/pkg/retry/backoff.go

@@ -181,6 +181,12 @@ func (bo *Backoffer) resetBackoff() {
 	bo.nextLogTime = 0
 }
 
+// TotalTime returns the maximum total time duration for the backoff.
+// Returns 0 if infinite retry is enabled.
+func (bo *Backoffer) TotalTime() time.Duration {
+	return bo.total
+}
+
 // Only used for test.
 var testBackOffExecuteFlag atomic.Bool
 

pkg/storage/endpoint/tso.go

@@ -33,6 +33,8 @@ import (
 // TSOStorage is the interface for timestamp storage.
 type TSOStorage interface {
 	LoadTimestamp(groupID uint32) (time.Time, error)
+	// LoadMaxTimestampAllGroups retrieves the maximum TSO timestamp from all keyspace groups
+	LoadMaxTimestampAllGroups() (time.Time, error)
 	SaveTimestamp(ctx context.Context, groupID uint32, ts time.Time, leadership *election.Leadership) error
 	DeleteTimestamp(ctx context.Context, groupID uint32) error
 }
@@ -119,3 +121,53 @@ func (se *StorageEndpoint) DeleteTimestamp(ctx context.Context, groupID uint32)
 		return txn.Remove(keypath.TimestampPath(groupID))
 	})
 }
+
+// LoadMaxTimestampAllGroups retrieves the maximum TSO timestamp from all keyspace groups.
+// It reads all timestamp keys and returns the maximum value.
+// This is used for non-default keyspace groups to ensure TSO monotonicity.
+func (se *StorageEndpoint) LoadMaxTimestampAllGroups() (time.Time, error) {
+	prefixStart := keypath.MsTimestampPrefix()
+	prefixEnd := keypath.MsTimestampPrefixRangeEnd()
+	keys, values, err := se.LoadRange(prefixStart, prefixEnd, 0)
+	if err != nil {
+		return typeutil.ZeroTime, err
+	}
+	if len(keys) == 0 {
+		log.Info("no timestamp keys found in tso nodes")
+		return typeutil.ZeroTime, nil
+	}
+	tsByKeyspaceGroup := make(map[uint32]time.Time)
+	for i, key := range keys {
+		value := values[i]
+		if len(value) == 0 {
+			continue
+		}
+		tsWindow, err := typeutil.ParseTimestamp([]byte(value))
+		if err != nil {
+			log.Warn("parse timestamp window failed",
+				zap.String("ts-window-key", key),
+				zap.Error(err))
+			continue
+		}
+		groupID := keypath.ExtractKeyspaceGroupIDFromMsTimestamp(key)
+		if groupID == 0 {
+			log.Warn("failed to extract keyspace group id from timestamp key",
+				zap.String("ts-window-key", key))
+			continue
+		}
+		tsByKeyspaceGroup[groupID] = tsWindow
+	}
+	if len(tsByKeyspaceGroup) == 0 {
+		return typeutil.ZeroTime, nil
+	}
+	var maxTS time.Time
+	for _, ts := range tsByKeyspaceGroup {
+		if ts.After(maxTS) {
+			maxTS = ts
+		}
+	}
+	log.Info("loaded max timestamp from all keyspace groups",
+		zap.Uint32("keyspace-group-count", uint32(len(tsByKeyspaceGroup))),
+		zap.Time("max-timestamp", maxTS))
+	return maxTS, nil
+}

pkg/tso/tso.go

@@ -28,6 +28,7 @@ import (
 	"github.com/pingcap/log"
 
 	"github.com/tikv/pd/pkg/errs"
+	"github.com/tikv/pd/pkg/keyspace/constant"
 	"github.com/tikv/pd/pkg/member"
 	"github.com/tikv/pd/pkg/storage/endpoint"
 	"github.com/tikv/pd/pkg/utils/logutil"
@@ -170,11 +171,22 @@ func (t *timestampOracle) syncTimestamp() error {
 		time.Sleep(time.Second)
 	})
 
-	last, err := t.storage.LoadTimestamp(t.keyspaceGroupID)
+	var last time.Time
+	var err error
+	// For non-default keyspace groups, read all timestamps to get the maximum
+	// to ensure TSO monotonicity.
+	if constant.DefaultKeyspaceGroupID != t.keyspaceGroupID && t.keyspaceGroupID > 0 {
+		log.Info("loading max timestamp from all keyspace groups for non-default group",
+			zap.Uint32("keyspace-group-id", t.keyspaceGroupID))
+		last, err = t.storage.LoadMaxTimestampAllGroups()
+	} else {
+		last, err = t.storage.LoadTimestamp(t.keyspaceGroupID)
+	}
 	if err != nil {
 		return err
 	}
 	lastSavedTime := t.getLastSavedTime()
+
 	// We could skip the synchronization if the following conditions are met:
 	//   1. The timestamp in memory has been initialized.
 	//   2. The last saved timestamp in etcd is not zero.
@@ -193,6 +205,10 @@ func (t *timestampOracle) syncTimestamp() error {
 		return nil
 	}
 
+	log.Info("loaded timestamp for synchronization",
+		logutil.CondUint32("keyspace-group-id", t.keyspaceGroupID, t.keyspaceGroupID > 0),
+		zap.Time("last-timestamp", last))
+
 	next := time.Now()
 	failpoint.Inject("fallBackSync", func() {
 		next = next.Add(time.Hour)