From 27bda8a6ed01cfdf8f6e728518e496525fb8d418 Mon Sep 17 00:00:00 2001 From: Jiale Lin <63439129+ljluestc@users.noreply.github.com> Date: Mon, 23 Mar 2026 22:27:25 -0700 Subject: [PATCH 1/3] tso, keypath: enable non-default keyspace group to read all timestamps Non-default keyspace group TSO services now read all keyspace group timestamps at startup and use the maximum value, ensuring TSO monotonicity without requiring metadata cleanup during TSO node start/stop operations. Signed-off-by: Jiale Lin <63439129+ljluestc@users.noreply.github.com> --- PR_DESCRIPTION_issue6972.md | 78 ++++++++++++++++++++++++++ pkg/storage/endpoint/tso.go | 52 +++++++++++++++++ pkg/tso/tso.go | 18 +++++- pkg/utils/keypath/absolute_key_path.go | 27 +++++++++ tests/cluster.go | 42 +++++++++++--- tests/cluster_test.go | 56 ++++++++++++++++++ 6 files changed, 265 insertions(+), 8 deletions(-) create mode 100644 PR_DESCRIPTION_issue6972.md create mode 100644 tests/cluster_test.go diff --git a/PR_DESCRIPTION_issue6972.md b/PR_DESCRIPTION_issue6972.md new file mode 100644 index 00000000000..c12bfad42bd --- /dev/null +++ b/PR_DESCRIPTION_issue6972.md @@ -0,0 +1,78 @@ +# Issue #6972: Enable Non-Default Keyspace Group TSO to Read All Timestamp Keys + +## 问题概述 + +目前，只有 PD leader TSO 和默认 keyspace group TSO 服务会在启动时读取所有前缀的 timestamp key 并获取最大的最新值。非默认 keyspace group TSO 服务在启动同步时没有相同的行为。 + +通过本修复，所有 keyspace group 的 TSO 服务（包括非默认的 keyspace group）在启动时都将读取所有 keyspace group 的 timestamp 并获取最大值，以确保 TSO 的单调性。 + +## 背景 + +在 keyspace group 分片场景下： +- 每个 keyspace group 维护自己的 timestamp，存储在 etcd 路径：`/ms/{cluster_id}/tso/{group_id}/gta/timestamp` +- 当非默认 keyspace group 重启时，如果只读取自己的 timestamp key，可能会遇到 TSO 回退的问题 +- 这需要运维人员在启动/停止 TSO 节点时额外考虑元数据清理 + +## 解决方案 + +### 核心变更 + +1. **扩展 TSOStorage 接口**（`pkg/storage/endpoint/tso.go`） + - 添加新方法：`LoadMaxTimestampAllGroups() (time.Time, error)` + - 该方法读取所有 keyspace group 的 timestamp key 并返回最大值 + - 使用 etcd 范围查询读取前缀：`/ms/{cluster_id}/tso/` 下的所有数据 + - 解析每个 timestamp 并找到最大值 + - 处理错误情况（无效值、解析错误等） + +2. **修改 TSO 同步逻辑**（`pkg/tso/tso.go`） + - 在 `timestampOracle.syncTimestamp()` 中添加条件判断 + - 对于非默认 keyspace group（`keyspaceGroupID > 0` 且不是默认组 ID），调用 `LoadMaxTimestampAllGroups()` + - 对于默认 keyspace group，使用原有的 `LoadTimestamp()` 行为 + - 确保 TSO 单调性 + +3. **添加辅助函数**（`pkg/utils/keypath/absolute_key_path.go`） + - `MsTimestampPrefix()`: 返回微服务 TSO timestamp 键的前缀 + - `MsTimestampPrefixRangeEnd()`: 返回前缀范围的结束键 + - `ExtractKeyspaceGroupIDFromMsTimestamp()`: 从路径中提取 keyspace group ID + +## 实现优势 + +1. **简化运维**：允许 TSO 节点更安全地启动/停止，无需考虑元数据清理 +2. **保证单调性**：非默认 keyspace group 初始化时总是从最新的 TSO 开始 +3. **保持向后兼容**：默认 keyspace group 行为不变 +4. **最小化变更**：只需修改两个核心文件和添加辅助函数 + +## 测试验证 + +- **编译验证**：`make build` 成功 +- **基本测试**：`make basic-test` 通过（除一个环境相关的 cgroup 测试） +- **代码审查**：遵循 PD 项目规范 + +## 影响范围 + +### 修改的文件 +1. `pkg/storage/endpoint/tso.go` - 添加接口方法 +2. `pkg/tso/tso.go` - 修改同步逻辑 +3. `pkg/utils/keypath/absolute_key_path.go` - 添加辅助函数 + +### 行为变更 +- **非默认 keyspace group**：启动时读取所有 keyspace group 的最大 timestamp +- **默认 keyspace group**：保持原有行为，只读取自己的 timestamp + +## 关键细节 + +- `LoadMaxTimestampAllGroups()` 使用 etcd 范围查询读取所有 keyspace group 的 timestamp +- 使用前缀 `/ms/{cluster_id}/tso/` 和范围结束键 `/ms/{cluster_id}/tso_99999` +- 跳过空值或无效值，记录警告日志 +- 返回的 timestamp 为所有有效 values 中的最大值 +- 如果没有任何有效 timestamp，返回 `typeutil.ZeroTime` + +## 相关设计考虑 + +这个修复借鉴了现有的 `mergingChecker` 的设计模式（`pkg/tso/keyspace_group_manager.go:1396-1414`），该模式已经是读取多个 keyspace group 的 timestamp 并找到最大值的实现。 + +## 向后兼容性 + +- 默认 keyspace group 的行为完全不变 +- 现有的 API 接口保持不变 +- 只添加了新的接口方法，不修改现有方法签名 \ No newline at end of file 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)) +} From 583bc7c2a49386bcfa419feec762eb16f543d06a Mon Sep 17 00:00:00 2001 From: Jiale Lin <63439129+ljluestc@users.noreply.github.com> Date: Mon, 23 Mar 2026 23:39:46 -0700 Subject: [PATCH 2/3] fix CI --- PR_DESCRIPTION_issue6972.md | 78 ------------------------------------- 1 file changed, 78 deletions(-) delete mode 100644 PR_DESCRIPTION_issue6972.md diff --git a/PR_DESCRIPTION_issue6972.md b/PR_DESCRIPTION_issue6972.md deleted file mode 100644 index c12bfad42bd..00000000000 --- a/PR_DESCRIPTION_issue6972.md +++ /dev/null @@ -1,78 +0,0 @@ -# Issue #6972: Enable Non-Default Keyspace Group TSO to Read All Timestamp Keys - -## 问题概述 - -目前，只有 PD leader TSO 和默认 keyspace group TSO 服务会在启动时读取所有前缀的 timestamp key 并获取最大的最新值。非默认 keyspace group TSO 服务在启动同步时没有相同的行为。 - -通过本修复，所有 keyspace group 的 TSO 服务（包括非默认的 keyspace group）在启动时都将读取所有 keyspace group 的 timestamp 并获取最大值，以确保 TSO 的单调性。 - -## 背景 - -在 keyspace group 分片场景下： -- 每个 keyspace group 维护自己的 timestamp，存储在 etcd 路径：`/ms/{cluster_id}/tso/{group_id}/gta/timestamp` -- 当非默认 keyspace group 重启时，如果只读取自己的 timestamp key，可能会遇到 TSO 回退的问题 -- 这需要运维人员在启动/停止 TSO 节点时额外考虑元数据清理 - -## 解决方案 - -### 核心变更 - -1. **扩展 TSOStorage 接口**（`pkg/storage/endpoint/tso.go`） - - 添加新方法：`LoadMaxTimestampAllGroups() (time.Time, error)` - - 该方法读取所有 keyspace group 的 timestamp key 并返回最大值 - - 使用 etcd 范围查询读取前缀：`/ms/{cluster_id}/tso/` 下的所有数据 - - 解析每个 timestamp 并找到最大值 - - 处理错误情况（无效值、解析错误等） - -2. **修改 TSO 同步逻辑**（`pkg/tso/tso.go`） - - 在 `timestampOracle.syncTimestamp()` 中添加条件判断 - - 对于非默认 keyspace group（`keyspaceGroupID > 0` 且不是默认组 ID），调用 `LoadMaxTimestampAllGroups()` - - 对于默认 keyspace group，使用原有的 `LoadTimestamp()` 行为 - - 确保 TSO 单调性 - -3. **添加辅助函数**（`pkg/utils/keypath/absolute_key_path.go`） - - `MsTimestampPrefix()`: 返回微服务 TSO timestamp 键的前缀 - - `MsTimestampPrefixRangeEnd()`: 返回前缀范围的结束键 - - `ExtractKeyspaceGroupIDFromMsTimestamp()`: 从路径中提取 keyspace group ID - -## 实现优势 - -1. **简化运维**：允许 TSO 节点更安全地启动/停止，无需考虑元数据清理 -2. **保证单调性**：非默认 keyspace group 初始化时总是从最新的 TSO 开始 -3. **保持向后兼容**：默认 keyspace group 行为不变 -4. **最小化变更**：只需修改两个核心文件和添加辅助函数 - -## 测试验证 - -- **编译验证**：`make build` 成功 -- **基本测试**：`make basic-test` 通过（除一个环境相关的 cgroup 测试） -- **代码审查**：遵循 PD 项目规范 - -## 影响范围 - -### 修改的文件 -1. `pkg/storage/endpoint/tso.go` - 添加接口方法 -2. `pkg/tso/tso.go` - 修改同步逻辑 -3. `pkg/utils/keypath/absolute_key_path.go` - 添加辅助函数 - -### 行为变更 -- **非默认 keyspace group**：启动时读取所有 keyspace group 的最大 timestamp -- **默认 keyspace group**：保持原有行为，只读取自己的 timestamp - -## 关键细节 - -- `LoadMaxTimestampAllGroups()` 使用 etcd 范围查询读取所有 keyspace group 的 timestamp -- 使用前缀 `/ms/{cluster_id}/tso/` 和范围结束键 `/ms/{cluster_id}/tso_99999` -- 跳过空值或无效值，记录警告日志 -- 返回的 timestamp 为所有有效 values 中的最大值 -- 如果没有任何有效 timestamp，返回 `typeutil.ZeroTime` - -## 相关设计考虑 - -这个修复借鉴了现有的 `mergingChecker` 的设计模式（`pkg/tso/keyspace_group_manager.go:1396-1414`），该模式已经是读取多个 keyspace group 的 timestamp 并找到最大值的实现。 - -## 向后兼容性 - -- 默认 keyspace group 的行为完全不变 -- 现有的 API 接口保持不变 -- 只添加了新的接口方法，不修改现有方法签名 \ No newline at end of file From d61bff7c63f47c393ae654b391f54c26a7388f45 Mon Sep 17 00:00:00 2001 From: Jiale Lin <63439129+ljluestc@users.noreply.github.com> Date: Sat, 4 Apr 2026 15:46:32 -0700 Subject: [PATCH 3/3] client: fix TSO stream-loop timeout reset Use the effective TSO timeout when resetting the dispatcher stream loop timer on each batch, so timeout behavior no longer falls back to pd-server-timeout. Add a dispatcher unit test that verifies stream-loop retries honor TSO timeout instead of the PD timeout, and align suite setup timeout defaults with the new semantics. Remove local PR description artifact from tracked files. Closes: #9469 Signed-off-by: Jiale Lin <63439129+ljluestc@users.noreply.github.com> --- client/clients/tso/dispatcher.go | 9 ++-- client/clients/tso/dispatcher_test.go | 49 +++++++++++++++++ client/opt/option.go | 34 +++++++++++- client/opt/option_test.go | 78 +++++++++++++++++++++++++++ client/pkg/retry/backoff.go | 6 +++ 5 files changed, 171 insertions(+), 5 deletions(-) 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