diff --git a/client/clients/tso/dispatcher.go b/client/clients/tso/dispatcher.go index 5802343b306..949fe626b41 100644 --- a/client/clients/tso/dispatcher.go +++ b/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())) diff --git a/client/clients/tso/dispatcher_test.go b/client/clients/tso/dispatcher_test.go index a2b9384fb76..1220602b17f 100644 --- a/client/clients/tso/dispatcher_test.go +++ b/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. diff --git a/client/opt/option.go b/client/opt/option.go index c46edb8f924..61f80559d13 100644 --- a/client/opt/option.go +++ b/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) { diff --git a/client/opt/option_test.go b/client/opt/option_test.go index 08caea59a69..d9074629b16 100644 --- a/client/opt/option_test.go +++ b/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) +} diff --git a/client/pkg/retry/backoff.go b/client/pkg/retry/backoff.go index 7030433f16d..8f1f26fe360 100644 --- a/client/pkg/retry/backoff.go +++ b/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 diff --git a/pkg/storage/endpoint/tso.go b/pkg/storage/endpoint/tso.go index 66a11daeb3b..3ce9a60b510 100644 --- a/pkg/storage/endpoint/tso.go +++ b/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 +} diff --git a/pkg/tso/tso.go b/pkg/tso/tso.go index 1254bed8701..23cc6847a75 100644 --- a/pkg/tso/tso.go +++ b/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) diff --git a/pkg/utils/keypath/absolute_key_path.go b/pkg/utils/keypath/absolute_key_path.go index a6c789b6160..ba9fa2f6338 100644 --- a/pkg/utils/keypath/absolute_key_path.go +++ b/pkg/utils/keypath/absolute_key_path.go @@ -22,6 +22,7 @@ import ( "strings" "github.com/tikv/pd/pkg/keyspace/constant" + mcsconstant "github.com/tikv/pd/pkg/mcs/utils/constant" ) const ( @@ -140,6 +141,32 @@ type MsParam struct { GroupID uint32 // only used for tso keyspace group } +// MsTimestampPrefix returns the prefix for all TSO timestamps in microservice. +func MsTimestampPrefix() string { + return fmt.Sprintf("/ms/%d/tso/", ClusterID()) +} + +// MsTimestampPrefixRangeEnd returns the range end for TSO timestamps in microservice. +func MsTimestampPrefixRangeEnd() string { + return fmt.Sprintf("/ms/%d/tso_%05d", ClusterID(), mcsconstant.MaxKeyspaceGroupCount) +} + +// ExtractKeyspaceGroupIDFromMsTimestamp extracts the keyspace group ID from microservice TSO timestamp path. +// Expected path format: /ms/{cluster_id}/tso/{group_id}/gta/timestamp +func ExtractKeyspaceGroupIDFromMsTimestamp(path string) uint32 { + parts := strings.Split(path, "/") + if len(parts) < 5 { + return 0 + } + // parts should be: ["", "ms", "{cluster_id}", "tso", "{group_id}", "gta", "timestamp"] + groupIDStr := parts[4] + groupID, err := strconv.ParseUint(groupIDStr, 10, 32) + if err != nil { + return 0 + } + return uint32(groupID) +} + // Prefix returns the parent directory of the given path. func Prefix(str string) string { return path.Dir(str) diff --git a/tests/cluster.go b/tests/cluster.go index c3371fb9417..a26d270a8d8 100644 --- a/tests/cluster.go +++ b/tests/cluster.go @@ -80,6 +80,37 @@ var ( defaultMaxRetryTimes = 5 ) +type startServersRetryAction int + +const ( + startServersNoRetry startServersRetryAction = iota + startServersRetryCurrent + startServersRetryRecreate +) + +func shouldRetryCurrentServers(err error) bool { + if err == nil { + return false + } + errMsg := err.Error() + return strings.Contains(errMsg, "ErrCancelStartEtcd") || strings.Contains(errMsg, "ErrStartEtcd") +} + +func classifyInitialServersError(err error) startServersRetryAction { + if err == nil { + return startServersNoRetry + } + errMsg := err.Error() + switch { + case strings.Contains(errMsg, "address already in use") || strings.Contains(errMsg, "Etcd cluster ID mismatch"): + return startServersRetryRecreate + case shouldRetryCurrentServers(err): + return startServersRetryCurrent + default: + return startServersNoRetry + } +} + // TestServer is only for test. type TestServer struct { syncutil.RWMutex @@ -689,9 +720,8 @@ func (c *TestCluster) runInitialServersWithRetry(maxRetries int) error { return nil } - errMsg := lastErr.Error() - switch { - case strings.Contains(errMsg, "address already in use") || strings.Contains(errMsg, "Etcd cluster ID mismatch"): + switch classifyInitialServersError(lastErr) { + case startServersRetryRecreate: // `Etcd cluster ID mismatch` can happen when the allocated peer URL happens to // connect to another test's etcd cluster (port reuse across concurrent `go test` // processes). Treat it as a port conflict and recreate servers with new ports. @@ -738,7 +768,7 @@ func (c *TestCluster) runInitialServersWithRetry(maxRetries int) error { } time.Sleep(backoff) continue - case strings.Contains(errMsg, "ErrStartEtcd"): + case startServersRetryCurrent: log.Warn("etcd start failed, will retry", zap.Error(lastErr)) for _, s := range servers { if s.State() == Running { @@ -764,9 +794,7 @@ func RunServersWithRetry(servers []*TestServer, maxRetries int) error { return nil } - // If the error is related to etcd start cancellation, we should retry - if strings.Contains(lastErr.Error(), "ErrCancelStartEtcd") || - strings.Contains(lastErr.Error(), "ErrStartEtcd") { + if shouldRetryCurrentServers(lastErr) { log.Warn("etcd start failed, will retry", zap.Error(lastErr)) // Stop any partially started servers before retrying for _, s := range servers { diff --git a/tests/cluster_test.go b/tests/cluster_test.go new file mode 100644 index 00000000000..f012b1ce979 --- /dev/null +++ b/tests/cluster_test.go @@ -0,0 +1,56 @@ +// Copyright 2026 TiKV Project 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 tests + +import ( + "testing" + + "github.com/stretchr/testify/require" + "go.uber.org/goleak" + + "github.com/pingcap/errors" + + "github.com/tikv/pd/pkg/utils/testutil" +) + +func TestMain(m *testing.M) { + goleak.VerifyTestMain(m, testutil.LeakOptions...) +} + +func TestShouldRetryCurrentServers(t *testing.T) { + t.Parallel() + + re := require.New(t) + re.True(shouldRetryCurrentServers(errors.New("[PD:server:ErrCancelStartEtcd]etcd start canceled"))) + re.True(shouldRetryCurrentServers(errors.New("[PD:etcd:ErrStartEtcd]start etcd failed"))) + re.True(shouldRetryCurrentServers(errors.New("[PD:etcd:ErrStartEtcd]start etcd failed: listen tcp 127.0.0.1:2379: bind: address already in use"))) + re.False(shouldRetryCurrentServers(errors.New("listen tcp 127.0.0.1:2379: bind: address already in use"))) + re.False(shouldRetryCurrentServers(errors.New("Etcd cluster ID mismatch"))) + re.False(shouldRetryCurrentServers(errors.New("some other error"))) + re.False(shouldRetryCurrentServers(nil)) +} + +func TestClassifyInitialServersError(t *testing.T) { + t.Parallel() + + re := require.New(t) + re.Equal(startServersRetryCurrent, classifyInitialServersError(errors.New("[PD:server:ErrCancelStartEtcd]etcd start canceled"))) + re.Equal(startServersRetryCurrent, classifyInitialServersError(errors.New("[PD:etcd:ErrStartEtcd]start etcd failed"))) + re.Equal(startServersRetryRecreate, classifyInitialServersError(errors.New("[PD:etcd:ErrStartEtcd]start etcd failed: listen tcp 127.0.0.1:2379: bind: address already in use"))) + re.Equal(startServersRetryRecreate, classifyInitialServersError(errors.New("listen tcp 127.0.0.1:2379: bind: address already in use"))) + re.Equal(startServersRetryRecreate, classifyInitialServersError(errors.New("Etcd cluster ID mismatch"))) + re.Equal(startServersNoRetry, classifyInitialServersError(errors.New("some other error"))) + re.Equal(startServersNoRetry, classifyInitialServersError(nil)) +}