From 03068161954cd02eadde83cb7ed1ceb5df6bf701 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 11:13:06 +0800 Subject: [PATCH 01/15] tests: avoid TestServer startup deadlock Signed-off-by: Ryan Leung --- tests/cluster.go | 16 +++++++++++++--- 1 file changed, 13 insertions(+), 3 deletions(-) diff --git a/tests/cluster.go b/tests/cluster.go index a26d270a8d8..f487ae64494 100644 --- a/tests/cluster.go +++ b/tests/cluster.go @@ -173,14 +173,24 @@ func NewTestServer(ctx context.Context, cfg *config.Config, services []string, h // Run starts to run a TestServer. func (s *TestServer) Run() error { s.Lock() - defer s.Unlock() if s.state != Initial && s.state != Stop { - return errors.Errorf("server(state%d) cannot run", s.state) + state := s.state + s.Unlock() + return errors.Errorf("server(state%d) cannot run", state) } + prevState := s.state + // Treat startup as running so retry cleanup can close a blocked server.Run. + s.state = Running + s.Unlock() + if err := s.server.Run(); err != nil { + s.Lock() + if s.state == Running { + s.state = prevState + } + s.Unlock() return err } - s.state = Running return nil } From e7edcaede8c7560e127cc98a7b22e32fe204e5c2 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 15:10:45 +0800 Subject: [PATCH 02/15] tso: update keyspace group revision after initial load Signed-off-by: Ryan Leung --- pkg/tso/keyspace_group_manager.go | 5 ++ pkg/tso/keyspace_group_manager_test.go | 36 +++++++++++++ pkg/utils/etcdutil/etcdutil.go | 11 ++++ pkg/utils/etcdutil/etcdutil_test.go | 32 ++++++++++++ .../mcs/keyspace/tso_keyspace_group_test.go | 50 ++++++++++++++----- 5 files changed, 122 insertions(+), 12 deletions(-) diff --git a/pkg/tso/keyspace_group_manager.go b/pkg/tso/keyspace_group_manager.go index 2e810e9e12c..a592df06a64 100644 --- a/pkg/tso/keyspace_group_manager.go +++ b/pkg/tso/keyspace_group_manager.go @@ -597,6 +597,11 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { postEventsFn, true, /* withPrefix */ ) + kgm.groupWatcher.SetPostLoadRevisionHook(func(revision int64) { + if revision > 0 { + kgm.SetModRevision(uint64(revision)) + } + }) if kgm.loadFromEtcdMaxRetryTimes > 0 { kgm.groupWatcher.SetLoadRetryTimes(kgm.loadFromEtcdMaxRetryTimes) } diff --git a/pkg/tso/keyspace_group_manager_test.go b/pkg/tso/keyspace_group_manager_test.go index 12c4b217cad..36d7b15fe29 100644 --- a/pkg/tso/keyspace_group_manager_test.go +++ b/pkg/tso/keyspace_group_manager_test.go @@ -210,6 +210,42 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsAssignment() { suite.runTestLoadKeyspaceGroupsAssignment(re, maxCountInUse+1, 0, 10) } +func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsSetsModRevision() { + re := suite.Require() + + mgr := suite.newUniqueKeyspaceGroupManager(1) + re.NotNil(mgr) + defer mgr.Close() + + const ( + groupID = uint32(1) + keyspaceID = uint32(101) + ) + err := addKeyspaceGroupAssignment( + suite.ctx, + suite.etcdClient, + groupID, + []string{mgr.tsoServiceID.ServiceAddr}, + []int{mcs.DefaultKeyspaceGroupReplicaPriority}, + []uint32{keyspaceID}, + ) + re.NoError(err) + + resp, err := suite.etcdClient.Get(suite.ctx, keypath.KeyspaceGroupIDPath(groupID)) + re.NoError(err) + re.Len(resp.Kvs, 1) + targetRevision := uint64(resp.Kvs[0].ModRevision) + + err = mgr.Initialize() + re.NoError(err) + + _, kg, loadedGroupID, loadedRevision, err := mgr.FindGroupByKeyspaceID(keyspaceID) + re.NoError(err) + re.NotNil(kg) + re.Equal(groupID, loadedGroupID) + re.GreaterOrEqual(loadedRevision, targetRevision) +} + // TestLoadWithDifferentBatchSize tests the loading of the keyspace group assignment with the different batch size. func (suite *keyspaceGroupManagerTestSuite) TestLoadWithDifferentBatchSize() { re := suite.Require() diff --git a/pkg/utils/etcdutil/etcdutil.go b/pkg/utils/etcdutil/etcdutil.go index aa9f729aa7a..7547f245ad3 100644 --- a/pkg/utils/etcdutil/etcdutil.go +++ b/pkg/utils/etcdutil/etcdutil.go @@ -372,6 +372,9 @@ type LoopWatcher struct { deleteFn func(*mvccpb.KeyValue) error // postEventsFn is used to call after handling all events. postEventsFn func([]*clientv3.Event) error + // postLoadRevisionFn is called after the initial load has completed. The + // revision is the etcd snapshot revision that the load reflects. + postLoadRevisionFn func(int64) // preEventsFn is used to call before handling all events. preEventsFn func([]*clientv3.Event) error @@ -490,6 +493,9 @@ func (lw *LoopWatcher) initFromEtcd(ctx context.Context) int64 { if err != nil { log.Warn("meet error when loading in watch loop", zap.String("name", lw.name), zap.String("key", lw.key), zap.Error(err)) } else { + if lw.postLoadRevisionFn != nil && watchStartRevision > 0 { + lw.postLoadRevisionFn(watchStartRevision - 1) + } log.Info("load finished in watch loop", zap.String("name", lw.name), zap.String("key", lw.key)) } lw.isLoadedCh <- err @@ -768,3 +774,8 @@ func (lw *LoopWatcher) SetLoadRetryTimes(times int) { func (lw *LoopWatcher) SetLoadBatchSize(size int64) { lw.loadBatchSize = size } + +// SetPostLoadRevisionHook sets the callback invoked after initial load finishes. +func (lw *LoopWatcher) SetPostLoadRevisionHook(hook func(int64)) { + lw.postLoadRevisionFn = hook +} diff --git a/pkg/utils/etcdutil/etcdutil_test.go b/pkg/utils/etcdutil/etcdutil_test.go index 970c42325b5..57f251d6e8b 100644 --- a/pkg/utils/etcdutil/etcdutil_test.go +++ b/pkg/utils/etcdutil/etcdutil_test.go @@ -467,6 +467,38 @@ func (suite *loopWatcherTestSuite) TestLoadNoExistedKey() { re.Empty(cache) } +func (suite *loopWatcherTestSuite) TestPostLoadRevisionHook() { + re := suite.Require() + ctx, cancel := context.WithCancel(suite.ctx) + defer cancel() + + key := "TestPostLoadRevisionHook" + resp, err := suite.client.Put(ctx, key, "") + re.NoError(err) + targetRevision := resp.Header.Revision + + var loadedRevision atomic.Int64 + watcher := NewLoopWatcher( + ctx, + &suite.wg, + suite.client, + "test", + key, + func([]*clientv3.Event) error { return nil }, + func(*mvccpb.KeyValue) error { return nil }, + func(*mvccpb.KeyValue) error { return nil }, + func([]*clientv3.Event) error { return nil }, + false, /* withPrefix */ + ) + watcher.SetPostLoadRevisionHook(func(revision int64) { + loadedRevision.Store(revision) + }) + watcher.StartWatchLoop() + err = watcher.WaitLoad() + re.NoError(err) + re.GreaterOrEqual(loadedRevision.Load(), targetRevision) +} + func (suite *loopWatcherTestSuite) TestLoadWithLimitChange() { re := suite.Require() re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/utils/etcdutil/meetEtcdError", `return()`)) diff --git a/tests/integrations/mcs/keyspace/tso_keyspace_group_test.go b/tests/integrations/mcs/keyspace/tso_keyspace_group_test.go index 04cf95ad8f7..33acabcce89 100644 --- a/tests/integrations/mcs/keyspace/tso_keyspace_group_test.go +++ b/tests/integrations/mcs/keyspace/tso_keyspace_group_test.go @@ -39,8 +39,10 @@ import ( bs "github.com/tikv/pd/pkg/basicserver" "github.com/tikv/pd/pkg/keyspace" "github.com/tikv/pd/pkg/keyspace/constant" + tsoserver "github.com/tikv/pd/pkg/mcs/tso/server" mcs "github.com/tikv/pd/pkg/mcs/utils/constant" "github.com/tikv/pd/pkg/storage/endpoint" + "github.com/tikv/pd/pkg/utils/keypath" "github.com/tikv/pd/pkg/utils/tempurl" "github.com/tikv/pd/pkg/utils/testutil" "github.com/tikv/pd/pkg/utils/tsoutil" @@ -146,6 +148,27 @@ func (suite *keyspaceGroupTestSuite) closeAllTSONodesAndWait(re *require.Asserti }, testutil.WithWaitFor(10*time.Second), testutil.WithTickInterval(500*time.Millisecond)) } +func (suite *keyspaceGroupTestSuite) waitTSOKeyspaceGroupReady( + re *require.Assertions, + node *tsoserver.Server, + keyspaceID, keyspaceGroupID uint32, +) { + testutil.Eventually(re, func() bool { + resp, err := node.GetClient().Get(suite.ctx, keypath.KeyspaceGroupIDPath(keyspaceGroupID)) + if err != nil || len(resp.Kvs) == 0 { + return false + } + targetRevision := uint64(resp.Kvs[0].ModRevision) + _, kg, loadedGroupID, loadedRevision, err := + node.GetKeyspaceGroupManager().FindGroupByKeyspaceID(keyspaceID) + return err == nil && + kg != nil && + loadedGroupID == keyspaceGroupID && + loadedRevision >= targetRevision && + node.IsKeyspaceServingByGroup(keyspaceID, keyspaceGroupID) + }, testutil.WithWaitFor(30*time.Second), testutil.WithTickInterval(100*time.Millisecond)) +} + func (suite *keyspaceGroupTestSuite) TestAllocNodesUpdate() { re := suite.Require() // add three nodes. @@ -558,12 +581,13 @@ func (suite *keyspaceGroupTestSuite) trySetNodesForKeyspaceGroup(re *require.Ass // tsoTestSetup holds the setup information for TSO tests type tsoTestSetup struct { - nodes map[string]bs.Server - cleanups []func() - client pd.Client - initialTS uint64 - firstNodeAddr string - keyspaceID uint32 + nodes map[string]bs.Server + cleanups []func() + client pd.Client + initialTS uint64 + firstNodeAddr string + keyspaceID uint32 + keyspaceGroupID uint32 } // setupTSONodesAndClient creates TSO nodes, keyspace group, and returns initialized client @@ -649,12 +673,13 @@ func (suite *keyspaceGroupTestSuite) setupTSONodesAndClient(re *require.Assertio break } return &tsoTestSetup{ - nodes: nodes, - cleanups: cleanups, - client: client, - initialTS: initialTS, - firstNodeAddr: firstNodeAddr, - keyspaceID: keyspaceID, + nodes: nodes, + cleanups: cleanups, + client: client, + initialTS: initialTS, + firstNodeAddr: firstNodeAddr, + keyspaceID: keyspaceID, + keyspaceGroupID: keyspaceGroupID, } } @@ -736,6 +761,7 @@ func (suite *keyspaceGroupTestSuite) TestUpdateMemberWhenRecovery() { setup.cleanups = append(setup.cleanups, cleanup) nodes[newNode.GetAddr()] = newNode tests.WaitForPrimaryServing(re, map[string]bs.Server{newNode.GetAddr(): newNode}) + suite.waitTSOKeyspaceGroupReady(re, newNode, setup.keyspaceID, setup.keyspaceGroupID) // Step 7: Verify eventual recovery after node restart. // The in-flight GetTS may stay attached to stale discovery/metadata during From a4e4df77fa8c0fb4a7ecd4b171d8624f5ade7496 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 15:30:11 +0800 Subject: [PATCH 03/15] tso: avoid keyspace group load revision hook Signed-off-by: Ryan Leung --- pkg/tso/keyspace_group_manager.go | 8 +++----- pkg/utils/etcdutil/etcdutil.go | 17 +++++++++-------- pkg/utils/etcdutil/etcdutil_test.go | 10 +++------- 3 files changed, 15 insertions(+), 20 deletions(-) diff --git a/pkg/tso/keyspace_group_manager.go b/pkg/tso/keyspace_group_manager.go index a592df06a64..33651206921 100644 --- a/pkg/tso/keyspace_group_manager.go +++ b/pkg/tso/keyspace_group_manager.go @@ -597,11 +597,6 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { postEventsFn, true, /* withPrefix */ ) - kgm.groupWatcher.SetPostLoadRevisionHook(func(revision int64) { - if revision > 0 { - kgm.SetModRevision(uint64(revision)) - } - }) if kgm.loadFromEtcdMaxRetryTimes > 0 { kgm.groupWatcher.SetLoadRetryTimes(kgm.loadFromEtcdMaxRetryTimes) } @@ -615,6 +610,9 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { kgm.Close() return errs.ErrLoadKeyspaceGroupsTerminated.Wrap(err) } + if loadedRevision := kgm.groupWatcher.GetLoadedRevision(); loadedRevision > 0 { + kgm.SetModRevision(uint64(loadedRevision)) + } if !defaultKGConfigured { log.Info("initializing default keyspace group") diff --git a/pkg/utils/etcdutil/etcdutil.go b/pkg/utils/etcdutil/etcdutil.go index 7547f245ad3..f785db54b2e 100644 --- a/pkg/utils/etcdutil/etcdutil.go +++ b/pkg/utils/etcdutil/etcdutil.go @@ -22,6 +22,7 @@ import ( "os" "strings" "sync" + "sync/atomic" "time" "github.com/gogo/protobuf/proto" @@ -372,12 +373,12 @@ type LoopWatcher struct { deleteFn func(*mvccpb.KeyValue) error // postEventsFn is used to call after handling all events. postEventsFn func([]*clientv3.Event) error - // postLoadRevisionFn is called after the initial load has completed. The - // revision is the etcd snapshot revision that the load reflects. - postLoadRevisionFn func(int64) // preEventsFn is used to call before handling all events. preEventsFn func([]*clientv3.Event) error + // loadedRevision is the etcd snapshot revision loaded during initialization. + loadedRevision atomic.Int64 + // forceLoadMu is used to ensure two force loads have minimal interval. forceLoadMu syncutil.RWMutex // lastTimeForceLoad is used to record the last time force loading data from etcd. @@ -493,8 +494,8 @@ func (lw *LoopWatcher) initFromEtcd(ctx context.Context) int64 { if err != nil { log.Warn("meet error when loading in watch loop", zap.String("name", lw.name), zap.String("key", lw.key), zap.Error(err)) } else { - if lw.postLoadRevisionFn != nil && watchStartRevision > 0 { - lw.postLoadRevisionFn(watchStartRevision - 1) + if watchStartRevision > 0 { + lw.loadedRevision.Store(watchStartRevision - 1) } log.Info("load finished in watch loop", zap.String("name", lw.name), zap.String("key", lw.key)) } @@ -775,7 +776,7 @@ func (lw *LoopWatcher) SetLoadBatchSize(size int64) { lw.loadBatchSize = size } -// SetPostLoadRevisionHook sets the callback invoked after initial load finishes. -func (lw *LoopWatcher) SetPostLoadRevisionHook(hook func(int64)) { - lw.postLoadRevisionFn = hook +// GetLoadedRevision returns the etcd snapshot revision loaded during initialization. +func (lw *LoopWatcher) GetLoadedRevision() int64 { + return lw.loadedRevision.Load() } diff --git a/pkg/utils/etcdutil/etcdutil_test.go b/pkg/utils/etcdutil/etcdutil_test.go index 57f251d6e8b..95031bf16a6 100644 --- a/pkg/utils/etcdutil/etcdutil_test.go +++ b/pkg/utils/etcdutil/etcdutil_test.go @@ -467,17 +467,16 @@ func (suite *loopWatcherTestSuite) TestLoadNoExistedKey() { re.Empty(cache) } -func (suite *loopWatcherTestSuite) TestPostLoadRevisionHook() { +func (suite *loopWatcherTestSuite) TestGetLoadedRevision() { re := suite.Require() ctx, cancel := context.WithCancel(suite.ctx) defer cancel() - key := "TestPostLoadRevisionHook" + key := "TestGetLoadedRevision" resp, err := suite.client.Put(ctx, key, "") re.NoError(err) targetRevision := resp.Header.Revision - var loadedRevision atomic.Int64 watcher := NewLoopWatcher( ctx, &suite.wg, @@ -490,13 +489,10 @@ func (suite *loopWatcherTestSuite) TestPostLoadRevisionHook() { func([]*clientv3.Event) error { return nil }, false, /* withPrefix */ ) - watcher.SetPostLoadRevisionHook(func(revision int64) { - loadedRevision.Store(revision) - }) watcher.StartWatchLoop() err = watcher.WaitLoad() re.NoError(err) - re.GreaterOrEqual(loadedRevision.Load(), targetRevision) + re.GreaterOrEqual(watcher.GetLoadedRevision(), targetRevision) } func (suite *loopWatcherTestSuite) TestLoadWithLimitChange() { From 615b1c439743433b8995fa3584706b4535d7f130 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 16:04:50 +0800 Subject: [PATCH 04/15] etcdutil: return revision after initial load Signed-off-by: Ryan Leung --- pkg/tso/keyspace_group_manager.go | 7 +++--- pkg/utils/etcdutil/etcdutil.go | 37 +++++++++++++++++------------ pkg/utils/etcdutil/etcdutil_test.go | 8 +++---- 3 files changed, 30 insertions(+), 22 deletions(-) diff --git a/pkg/tso/keyspace_group_manager.go b/pkg/tso/keyspace_group_manager.go index 33651206921..2fe40f9e7fd 100644 --- a/pkg/tso/keyspace_group_manager.go +++ b/pkg/tso/keyspace_group_manager.go @@ -604,14 +604,15 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { kgm.groupWatcher.SetLoadBatchSize(kgm.loadKeyspaceGroupsBatchSize) } kgm.groupWatcher.StartWatchLoop() - if err := kgm.groupWatcher.WaitLoad(); err != nil { + snapshotRevision, err := kgm.groupWatcher.WaitLoadRevision() + if err != nil { log.Error("failed to initialize keyspace group manager", errs.ZapError(err)) // We might have partially loaded/initialized the keyspace groups. Close the manager to clean up. kgm.Close() return errs.ErrLoadKeyspaceGroupsTerminated.Wrap(err) } - if loadedRevision := kgm.groupWatcher.GetLoadedRevision(); loadedRevision > 0 { - kgm.SetModRevision(uint64(loadedRevision)) + if snapshotRevision > 0 { + kgm.SetModRevision(uint64(snapshotRevision)) } if !defaultKGConfigured { diff --git a/pkg/utils/etcdutil/etcdutil.go b/pkg/utils/etcdutil/etcdutil.go index f785db54b2e..25c83dd3b70 100644 --- a/pkg/utils/etcdutil/etcdutil.go +++ b/pkg/utils/etcdutil/etcdutil.go @@ -22,7 +22,6 @@ import ( "os" "strings" "sync" - "sync/atomic" "time" "github.com/gogo/protobuf/proto" @@ -365,7 +364,7 @@ type LoopWatcher struct { // forceLoadCh is used to force loading data from etcd. forceLoadCh chan struct{} // isLoadedCh is used to notify that the data has been loaded from etcd first time. - isLoadedCh chan error + isLoadedCh chan loopWatcherLoadResult // putFn is used to handle the put event. putFn func(*mvccpb.KeyValue) error @@ -376,9 +375,6 @@ type LoopWatcher struct { // preEventsFn is used to call before handling all events. preEventsFn func([]*clientv3.Event) error - // loadedRevision is the etcd snapshot revision loaded during initialization. - loadedRevision atomic.Int64 - // forceLoadMu is used to ensure two force loads have minimal interval. forceLoadMu syncutil.RWMutex // lastTimeForceLoad is used to record the last time force loading data from etcd. @@ -397,6 +393,11 @@ type LoopWatcher struct { watchChTimeoutDuration time.Duration } +type loopWatcherLoadResult struct { + revision int64 + err error +} + // NewLoopWatcher creates a new LoopWatcher. func NewLoopWatcher( ctx context.Context, wg *sync.WaitGroup, @@ -414,7 +415,7 @@ func NewLoopWatcher( key: key, wg: wg, forceLoadCh: make(chan struct{}, 1), - isLoadedCh: make(chan error, 1), + isLoadedCh: make(chan loopWatcherLoadResult, 1), updateClientCh: make(chan *clientv3.Client, 1), putFn: putFn, deleteFn: deleteFn, @@ -486,7 +487,7 @@ func (lw *LoopWatcher) initFromEtcd(ctx context.Context) int64 { } select { case <-ctx.Done(): - lw.isLoadedCh <- errors.Errorf("ctx is done before load data from etcd") + lw.isLoadedCh <- loopWatcherLoadResult{err: errors.Errorf("ctx is done before load data from etcd")} return watchStartRevision case <-ticker.C: } @@ -494,12 +495,16 @@ func (lw *LoopWatcher) initFromEtcd(ctx context.Context) int64 { if err != nil { log.Warn("meet error when loading in watch loop", zap.String("name", lw.name), zap.String("key", lw.key), zap.Error(err)) } else { - if watchStartRevision > 0 { - lw.loadedRevision.Store(watchStartRevision - 1) - } log.Info("load finished in watch loop", zap.String("name", lw.name), zap.String("key", lw.key)) } - lw.isLoadedCh <- err + snapshotRevision := int64(0) + if watchStartRevision > 0 { + snapshotRevision = watchStartRevision - 1 + } + lw.isLoadedCh <- loopWatcherLoadResult{ + revision: snapshotRevision, + err: err, + } return watchStartRevision } @@ -763,7 +768,8 @@ func (lw *LoopWatcher) ForceLoad() { // WaitLoad waits for the result to obtain whether data is loaded. func (lw *LoopWatcher) WaitLoad() error { - return <-lw.isLoadedCh + _, err := lw.WaitLoadRevision() + return err } // SetLoadRetryTimes sets the retry times when loading data from etcd. @@ -776,7 +782,8 @@ func (lw *LoopWatcher) SetLoadBatchSize(size int64) { lw.loadBatchSize = size } -// GetLoadedRevision returns the etcd snapshot revision loaded during initialization. -func (lw *LoopWatcher) GetLoadedRevision() int64 { - return lw.loadedRevision.Load() +// WaitLoadRevision waits for the initial load and returns the loaded etcd snapshot revision. +func (lw *LoopWatcher) WaitLoadRevision() (int64, error) { + result := <-lw.isLoadedCh + return result.revision, result.err } diff --git a/pkg/utils/etcdutil/etcdutil_test.go b/pkg/utils/etcdutil/etcdutil_test.go index 95031bf16a6..23429f60e3e 100644 --- a/pkg/utils/etcdutil/etcdutil_test.go +++ b/pkg/utils/etcdutil/etcdutil_test.go @@ -467,12 +467,12 @@ func (suite *loopWatcherTestSuite) TestLoadNoExistedKey() { re.Empty(cache) } -func (suite *loopWatcherTestSuite) TestGetLoadedRevision() { +func (suite *loopWatcherTestSuite) TestWaitLoadRevision() { re := suite.Require() ctx, cancel := context.WithCancel(suite.ctx) defer cancel() - key := "TestGetLoadedRevision" + key := "TestWaitLoadRevision" resp, err := suite.client.Put(ctx, key, "") re.NoError(err) targetRevision := resp.Header.Revision @@ -490,9 +490,9 @@ func (suite *loopWatcherTestSuite) TestGetLoadedRevision() { false, /* withPrefix */ ) watcher.StartWatchLoop() - err = watcher.WaitLoad() + snapshotRevision, err := watcher.WaitLoadRevision() re.NoError(err) - re.GreaterOrEqual(watcher.GetLoadedRevision(), targetRevision) + re.GreaterOrEqual(snapshotRevision, targetRevision) } func (suite *loopWatcherTestSuite) TestLoadWithLimitChange() { From 1b38a893c78168eec73ad6a1b888d096ff03c455 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 16:15:48 +0800 Subject: [PATCH 05/15] tso: keep loaded group revision in manager Signed-off-by: Ryan Leung --- pkg/tso/keyspace_group_manager.go | 18 ++++++++++++------ pkg/utils/etcdutil/etcdutil.go | 29 +++++------------------------ pkg/utils/etcdutil/etcdutil_test.go | 28 ---------------------------- 3 files changed, 17 insertions(+), 58 deletions(-) diff --git a/pkg/tso/keyspace_group_manager.go b/pkg/tso/keyspace_group_manager.go index 2fe40f9e7fd..1f1ced44152 100644 --- a/pkg/tso/keyspace_group_manager.go +++ b/pkg/tso/keyspace_group_manager.go @@ -536,6 +536,11 @@ func (kgm *KeyspaceGroupManager) InitializeTSOServerWatchLoop() error { // Value: endpoint.KeyspaceGroup func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { defaultKGConfigured := false + maxLoadedModRevision := uint64(0) + preEventsFn := func([]*clientv3.Event) error { + maxLoadedModRevision = 0 + return nil + } putFn := func(kv *mvccpb.KeyValue) error { group := &endpoint.KeyspaceGroup{} if err := json.Unmarshal(kv.Value, group); err != nil { @@ -547,6 +552,9 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { failpoint.Return(nil) } }) + if modRevision := uint64(kv.ModRevision); modRevision > maxLoadedModRevision { + maxLoadedModRevision = modRevision + } kgm.updateKeyspaceGroup(group) if group.ID == constant.DefaultKeyspaceGroupID { defaultKGConfigured = true @@ -581,6 +589,8 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { zap.Uint64("new-mod-revision", uint64(last.Kv.ModRevision)), ) } + } else if maxLoadedModRevision > 0 { + kgm.SetModRevision(maxLoadedModRevision) } return nil } @@ -591,7 +601,7 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { "keyspace-watcher", // To keep the consistency with the previous code, we should trim the suffix `/`. strings.TrimSuffix(keypath.KeyspaceGroupIDPrefix(), "/"), - func([]*clientv3.Event) error { return nil }, + preEventsFn, putFn, deleteFn, postEventsFn, @@ -604,16 +614,12 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { kgm.groupWatcher.SetLoadBatchSize(kgm.loadKeyspaceGroupsBatchSize) } kgm.groupWatcher.StartWatchLoop() - snapshotRevision, err := kgm.groupWatcher.WaitLoadRevision() - if err != nil { + if err := kgm.groupWatcher.WaitLoad(); err != nil { log.Error("failed to initialize keyspace group manager", errs.ZapError(err)) // We might have partially loaded/initialized the keyspace groups. Close the manager to clean up. kgm.Close() return errs.ErrLoadKeyspaceGroupsTerminated.Wrap(err) } - if snapshotRevision > 0 { - kgm.SetModRevision(uint64(snapshotRevision)) - } if !defaultKGConfigured { log.Info("initializing default keyspace group") diff --git a/pkg/utils/etcdutil/etcdutil.go b/pkg/utils/etcdutil/etcdutil.go index 25c83dd3b70..aa9f729aa7a 100644 --- a/pkg/utils/etcdutil/etcdutil.go +++ b/pkg/utils/etcdutil/etcdutil.go @@ -364,7 +364,7 @@ type LoopWatcher struct { // forceLoadCh is used to force loading data from etcd. forceLoadCh chan struct{} // isLoadedCh is used to notify that the data has been loaded from etcd first time. - isLoadedCh chan loopWatcherLoadResult + isLoadedCh chan error // putFn is used to handle the put event. putFn func(*mvccpb.KeyValue) error @@ -393,11 +393,6 @@ type LoopWatcher struct { watchChTimeoutDuration time.Duration } -type loopWatcherLoadResult struct { - revision int64 - err error -} - // NewLoopWatcher creates a new LoopWatcher. func NewLoopWatcher( ctx context.Context, wg *sync.WaitGroup, @@ -415,7 +410,7 @@ func NewLoopWatcher( key: key, wg: wg, forceLoadCh: make(chan struct{}, 1), - isLoadedCh: make(chan loopWatcherLoadResult, 1), + isLoadedCh: make(chan error, 1), updateClientCh: make(chan *clientv3.Client, 1), putFn: putFn, deleteFn: deleteFn, @@ -487,7 +482,7 @@ func (lw *LoopWatcher) initFromEtcd(ctx context.Context) int64 { } select { case <-ctx.Done(): - lw.isLoadedCh <- loopWatcherLoadResult{err: errors.Errorf("ctx is done before load data from etcd")} + lw.isLoadedCh <- errors.Errorf("ctx is done before load data from etcd") return watchStartRevision case <-ticker.C: } @@ -497,14 +492,7 @@ func (lw *LoopWatcher) initFromEtcd(ctx context.Context) int64 { } else { log.Info("load finished in watch loop", zap.String("name", lw.name), zap.String("key", lw.key)) } - snapshotRevision := int64(0) - if watchStartRevision > 0 { - snapshotRevision = watchStartRevision - 1 - } - lw.isLoadedCh <- loopWatcherLoadResult{ - revision: snapshotRevision, - err: err, - } + lw.isLoadedCh <- err return watchStartRevision } @@ -768,8 +756,7 @@ func (lw *LoopWatcher) ForceLoad() { // WaitLoad waits for the result to obtain whether data is loaded. func (lw *LoopWatcher) WaitLoad() error { - _, err := lw.WaitLoadRevision() - return err + return <-lw.isLoadedCh } // SetLoadRetryTimes sets the retry times when loading data from etcd. @@ -781,9 +768,3 @@ func (lw *LoopWatcher) SetLoadRetryTimes(times int) { func (lw *LoopWatcher) SetLoadBatchSize(size int64) { lw.loadBatchSize = size } - -// WaitLoadRevision waits for the initial load and returns the loaded etcd snapshot revision. -func (lw *LoopWatcher) WaitLoadRevision() (int64, error) { - result := <-lw.isLoadedCh - return result.revision, result.err -} diff --git a/pkg/utils/etcdutil/etcdutil_test.go b/pkg/utils/etcdutil/etcdutil_test.go index 23429f60e3e..970c42325b5 100644 --- a/pkg/utils/etcdutil/etcdutil_test.go +++ b/pkg/utils/etcdutil/etcdutil_test.go @@ -467,34 +467,6 @@ func (suite *loopWatcherTestSuite) TestLoadNoExistedKey() { re.Empty(cache) } -func (suite *loopWatcherTestSuite) TestWaitLoadRevision() { - re := suite.Require() - ctx, cancel := context.WithCancel(suite.ctx) - defer cancel() - - key := "TestWaitLoadRevision" - resp, err := suite.client.Put(ctx, key, "") - re.NoError(err) - targetRevision := resp.Header.Revision - - watcher := NewLoopWatcher( - ctx, - &suite.wg, - suite.client, - "test", - key, - func([]*clientv3.Event) error { return nil }, - func(*mvccpb.KeyValue) error { return nil }, - func(*mvccpb.KeyValue) error { return nil }, - func([]*clientv3.Event) error { return nil }, - false, /* withPrefix */ - ) - watcher.StartWatchLoop() - snapshotRevision, err := watcher.WaitLoadRevision() - re.NoError(err) - re.GreaterOrEqual(snapshotRevision, targetRevision) -} - func (suite *loopWatcherTestSuite) TestLoadWithLimitChange() { re := suite.Require() re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/utils/etcdutil/meetEtcdError", `return()`)) From 14bff2fd4f9e5f3a09872dd2878a5ef80583bd10 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 6 Jul 2026 16:46:00 +0800 Subject: [PATCH 06/15] tests: avoid RunServer goroutine leak Signed-off-by: Ryan Leung --- tests/cluster.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/cluster.go b/tests/cluster.go index f487ae64494..c109745e322 100644 --- a/tests/cluster.go +++ b/tests/cluster.go @@ -689,7 +689,7 @@ func restartTestCluster( // RunServer starts to run TestServer. func RunServer(server *TestServer) <-chan error { - resC := make(chan error) + resC := make(chan error, 1) go func() { resC <- server.Run() }() return resC } From 4777a5615224310d99952fdb6d61081a4758e1a6 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Fri, 10 Jul 2026 15:30:57 +0800 Subject: [PATCH 07/15] etcdutil: preserve initial load snapshot revision Signed-off-by: Ryan Leung --- pkg/tso/keyspace_group_manager.go | 21 ++++++------- pkg/tso/keyspace_group_manager_test.go | 19 ++++++++++-- pkg/utils/etcdutil/etcdutil.go | 23 +++++++++++++- pkg/utils/etcdutil/etcdutil_test.go | 42 ++++++++++++++++++++++++++ 4 files changed, 91 insertions(+), 14 deletions(-) diff --git a/pkg/tso/keyspace_group_manager.go b/pkg/tso/keyspace_group_manager.go index 1f1ced44152..ea280f97862 100644 --- a/pkg/tso/keyspace_group_manager.go +++ b/pkg/tso/keyspace_group_manager.go @@ -536,11 +536,6 @@ func (kgm *KeyspaceGroupManager) InitializeTSOServerWatchLoop() error { // Value: endpoint.KeyspaceGroup func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { defaultKGConfigured := false - maxLoadedModRevision := uint64(0) - preEventsFn := func([]*clientv3.Event) error { - maxLoadedModRevision = 0 - return nil - } putFn := func(kv *mvccpb.KeyValue) error { group := &endpoint.KeyspaceGroup{} if err := json.Unmarshal(kv.Value, group); err != nil { @@ -552,9 +547,6 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { failpoint.Return(nil) } }) - if modRevision := uint64(kv.ModRevision); modRevision > maxLoadedModRevision { - maxLoadedModRevision = modRevision - } kgm.updateKeyspaceGroup(group) if group.ID == constant.DefaultKeyspaceGroupID { defaultKGConfigured = true @@ -589,8 +581,6 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { zap.Uint64("new-mod-revision", uint64(last.Kv.ModRevision)), ) } - } else if maxLoadedModRevision > 0 { - kgm.SetModRevision(maxLoadedModRevision) } return nil } @@ -601,7 +591,7 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { "keyspace-watcher", // To keep the consistency with the previous code, we should trim the suffix `/`. strings.TrimSuffix(keypath.KeyspaceGroupIDPrefix(), "/"), - preEventsFn, + func([]*clientv3.Event) error { return nil }, putFn, deleteFn, postEventsFn, @@ -620,6 +610,15 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { kgm.Close() return errs.ErrLoadKeyspaceGroupsTerminated.Wrap(err) } + if loadedRevision := kgm.groupWatcher.GetLoadedRevision(); loadedRevision > 0 { + // The watch loop can advance the manager after the initial load finishes + // but before WaitLoad returns. In that case, keeping the newer revision is + // correct and SetModRevision returns false. + if !kgm.SetModRevision(uint64(loadedRevision)) { + log.Debug("keyspace group revision is already at or ahead of the loaded snapshot", + zap.Int64("loaded-revision", loadedRevision)) + } + } if !defaultKGConfigured { log.Info("initializing default keyspace group") diff --git a/pkg/tso/keyspace_group_manager_test.go b/pkg/tso/keyspace_group_manager_test.go index 36d7b15fe29..9d42be9e252 100644 --- a/pkg/tso/keyspace_group_manager_test.go +++ b/pkg/tso/keyspace_group_manager_test.go @@ -234,7 +234,22 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsSetsModRevisio resp, err := suite.etcdClient.Get(suite.ctx, keypath.KeyspaceGroupIDPath(groupID)) re.NoError(err) re.Len(resp.Kvs, 1) - targetRevision := uint64(resp.Kvs[0].ModRevision) + groupRevision := uint64(resp.Kvs[0].ModRevision) + + const deletedGroupID = uint32(2) + err = addKeyspaceGroupAssignment( + suite.ctx, + suite.etcdClient, + deletedGroupID, + []string{mgr.tsoServiceID.ServiceAddr}, + []int{mcs.DefaultKeyspaceGroupReplicaPriority}, + []uint32{keyspaceID + 1}, + ) + re.NoError(err) + deleteResp, err := suite.etcdClient.Delete(suite.ctx, keypath.KeyspaceGroupIDPath(deletedGroupID)) + re.NoError(err) + deletedRevision := uint64(deleteResp.Header.Revision) + re.Greater(deletedRevision, groupRevision) err = mgr.Initialize() re.NoError(err) @@ -243,7 +258,7 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsSetsModRevisio re.NoError(err) re.NotNil(kg) re.Equal(groupID, loadedGroupID) - re.GreaterOrEqual(loadedRevision, targetRevision) + re.GreaterOrEqual(loadedRevision, deletedRevision) } // TestLoadWithDifferentBatchSize tests the loading of the keyspace group assignment with the different batch size. diff --git a/pkg/utils/etcdutil/etcdutil.go b/pkg/utils/etcdutil/etcdutil.go index aa9f729aa7a..bf23492324e 100644 --- a/pkg/utils/etcdutil/etcdutil.go +++ b/pkg/utils/etcdutil/etcdutil.go @@ -22,6 +22,7 @@ import ( "os" "strings" "sync" + "sync/atomic" "time" "github.com/gogo/protobuf/proto" @@ -374,6 +375,8 @@ type LoopWatcher struct { postEventsFn func([]*clientv3.Event) error // preEventsFn is used to call before handling all events. preEventsFn func([]*clientv3.Event) error + // loadedRevision is the etcd snapshot revision loaded during initialization. + loadedRevision atomic.Int64 // forceLoadMu is used to ensure two force loads have minimal interval. forceLoadMu syncutil.RWMutex @@ -490,6 +493,9 @@ func (lw *LoopWatcher) initFromEtcd(ctx context.Context) int64 { if err != nil { log.Warn("meet error when loading in watch loop", zap.String("name", lw.name), zap.String("key", lw.key), zap.Error(err)) } else { + if watchStartRevision > 0 { + lw.loadedRevision.Store(watchStartRevision - 1) + } log.Info("load finished in watch loop", zap.String("name", lw.name), zap.String("key", lw.key)) } lw.isLoadedCh <- err @@ -639,6 +645,7 @@ func (lw *LoopWatcher) watch(ctx context.Context, revision int64) (nextRevision func (lw *LoopWatcher) load(ctx context.Context) (nextRevision int64, err error) { startKey := lw.key limit := lw.loadBatchSize + snapshotRevision := int64(0) opts := lw.buildLoadingOpts(limit) if err := lw.preEventsFn([]*clientv3.Event{}); err != nil { @@ -677,10 +684,19 @@ func (lw *LoopWatcher) load(ctx context.Context) (nextRevision int64, err error) return 0, err } opts = lw.buildLoadingOpts(limit) + if snapshotRevision > 0 { + opts = append(opts, clientv3.WithRev(snapshotRevision)) + } continue } return 0, err } + if snapshotRevision == 0 { + snapshotRevision = resp.Header.Revision + // Keep all remaining pages on the same snapshot. Otherwise a write + // between pages could be skipped when the watch starts. + opts = append(opts, clientv3.WithRev(snapshotRevision)) + } for i, item := range resp.Kvs { if i == len(resp.Kvs)-1 && resp.More { // If there are more keys, we need to load the next batch. @@ -701,7 +717,7 @@ func (lw *LoopWatcher) load(ctx context.Context) (nextRevision int64, err error) } // Note: if there are no keys in etcd, the resp.More is false. It also means the load is finished. if !resp.More { - return resp.Header.Revision + 1, err + return snapshotRevision + 1, err } } } @@ -759,6 +775,11 @@ func (lw *LoopWatcher) WaitLoad() error { return <-lw.isLoadedCh } +// GetLoadedRevision returns the etcd snapshot revision loaded during initialization. +func (lw *LoopWatcher) GetLoadedRevision() int64 { + return lw.loadedRevision.Load() +} + // SetLoadRetryTimes sets the retry times when loading data from etcd. func (lw *LoopWatcher) SetLoadRetryTimes(times int) { lw.loadRetryTimes = times diff --git a/pkg/utils/etcdutil/etcdutil_test.go b/pkg/utils/etcdutil/etcdutil_test.go index 970c42325b5..c59950ff717 100644 --- a/pkg/utils/etcdutil/etcdutil_test.go +++ b/pkg/utils/etcdutil/etcdutil_test.go @@ -467,6 +467,48 @@ func (suite *loopWatcherTestSuite) TestLoadNoExistedKey() { re.Empty(cache) } +func (suite *loopWatcherTestSuite) TestGetLoadedRevision() { + re := suite.Require() + ctx, cancel := context.WithCancel(suite.ctx) + defer cancel() + + prefix := "TestGetLoadedRevision/" + _, err := suite.client.Txn(ctx).Then( + clientv3.OpPut(prefix+"a", ""), + clientv3.OpPut(prefix+"b", ""), + ).Commit() + re.NoError(err) + resp, err := suite.client.Get(ctx, prefix, clientv3.WithPrefix()) + re.NoError(err) + re.Len(resp.Kvs, 2) + expectedRevision := resp.Header.Revision + + var once sync.Once + var putErr error + watcher := NewLoopWatcher( + ctx, + &suite.wg, + suite.client, + "test", + prefix, + func([]*clientv3.Event) error { return nil }, + func(*mvccpb.KeyValue) error { + once.Do(func() { + _, putErr = suite.client.Put(ctx, prefix+"c", "") + }) + return putErr + }, + func(*mvccpb.KeyValue) error { return nil }, + func([]*clientv3.Event) error { return nil }, + true, /* withPrefix */ + ) + watcher.SetLoadBatchSize(1) + watcher.StartWatchLoop() + err = watcher.WaitLoad() + re.NoError(err) + re.Equal(expectedRevision, watcher.GetLoadedRevision()) +} + func (suite *loopWatcherTestSuite) TestLoadWithLimitChange() { re := suite.Require() re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/utils/etcdutil/meetEtcdError", `return()`)) From 18b55daeff68c7c7d1648b8fe9ca6365b1279bdb Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 13 Jul 2026 11:53:27 +0800 Subject: [PATCH 08/15] tso: advance keyspace group revision on watch progress Signed-off-by: Ryan Leung --- pkg/tso/keyspace_group_manager.go | 5 +++++ pkg/utils/etcdutil/etcdutil.go | 15 ++++++++++++++ pkg/utils/etcdutil/etcdutil_test.go | 32 +++++++++++++++++++++++++++++ 3 files changed, 52 insertions(+) diff --git a/pkg/tso/keyspace_group_manager.go b/pkg/tso/keyspace_group_manager.go index ea280f97862..6f3512e0f6a 100644 --- a/pkg/tso/keyspace_group_manager.go +++ b/pkg/tso/keyspace_group_manager.go @@ -597,6 +597,11 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { postEventsFn, true, /* withPrefix */ ) + kgm.groupWatcher.SetRevisionUpdatedCallback(func(revision int64) { + if revision > 0 { + kgm.SetModRevision(uint64(revision)) + } + }) if kgm.loadFromEtcdMaxRetryTimes > 0 { kgm.groupWatcher.SetLoadRetryTimes(kgm.loadFromEtcdMaxRetryTimes) } diff --git a/pkg/utils/etcdutil/etcdutil.go b/pkg/utils/etcdutil/etcdutil.go index bf23492324e..a10792b122f 100644 --- a/pkg/utils/etcdutil/etcdutil.go +++ b/pkg/utils/etcdutil/etcdutil.go @@ -377,6 +377,8 @@ type LoopWatcher struct { preEventsFn func([]*clientv3.Event) error // loadedRevision is the etcd snapshot revision loaded during initialization. loadedRevision atomic.Int64 + // revisionUpdatedFn is called when the watcher has caught up to a newer etcd revision. + revisionUpdatedFn func(int64) // forceLoadMu is used to ensure two force loads have minimal interval. forceLoadMu syncutil.RWMutex @@ -419,6 +421,7 @@ func NewLoopWatcher( deleteFn: deleteFn, postEventsFn: postEventsFn, preEventsFn: preEventsFn, + revisionUpdatedFn: func(int64) {}, isWithPrefix: isWithPrefix, lastTimeForceLoad: time.Now(), loadRetryTimes: defaultLoadFromEtcdRetryTimes, @@ -603,6 +606,7 @@ func (lw *LoopWatcher) watch(ctx context.Context, revision int64) (nextRevision } else if wresp.IsProgressNotify() { log.Debug("watcher receives progress notify in watch loop", zap.Int64("revision", revision), zap.String("name", lw.name), zap.String("key", lw.key)) + lw.revisionUpdatedFn(wresp.Header.Revision) goto watchChanLoop } if err := lw.preEventsFn(wresp.Events); err != nil { @@ -636,6 +640,7 @@ func (lw *LoopWatcher) watch(ctx context.Context, revision int64) (nextRevision log.Error("run post event failed in watch loop", zap.Error(err), zap.Int64("revision", revision), zap.String("name", lw.name), zap.String("key", lw.key)) } + lw.revisionUpdatedFn(wresp.Header.Revision) revision = wresp.Header.Revision + 1 } goto watchChanLoop // Use goto to avoid creating a new watchChan @@ -717,6 +722,7 @@ func (lw *LoopWatcher) load(ctx context.Context) (nextRevision int64, err error) } // Note: if there are no keys in etcd, the resp.More is false. It also means the load is finished. if !resp.More { + lw.revisionUpdatedFn(snapshotRevision) return snapshotRevision + 1, err } } @@ -780,6 +786,15 @@ func (lw *LoopWatcher) GetLoadedRevision() int64 { return lw.loadedRevision.Load() } +// SetRevisionUpdatedCallback sets a callback that is invoked after the watcher catches up to an etcd revision. +func (lw *LoopWatcher) SetRevisionUpdatedCallback(fn func(int64)) { + if fn == nil { + lw.revisionUpdatedFn = func(int64) {} + return + } + lw.revisionUpdatedFn = fn +} + // SetLoadRetryTimes sets the retry times when loading data from etcd. func (lw *LoopWatcher) SetLoadRetryTimes(times int) { lw.loadRetryTimes = times diff --git a/pkg/utils/etcdutil/etcdutil_test.go b/pkg/utils/etcdutil/etcdutil_test.go index c59950ff717..ba46ef8eb14 100644 --- a/pkg/utils/etcdutil/etcdutil_test.go +++ b/pkg/utils/etcdutil/etcdutil_test.go @@ -816,6 +816,38 @@ func (suite *loopWatcherTestSuite) TestWatcherRequestProgress() { checkWatcherRequestProgress(true) } +func (suite *loopWatcherTestSuite) TestWatcherReportsProgressRevision() { + re := suite.Require() + ctx, cancel := context.WithCancel(suite.ctx) + defer cancel() + + watcher := NewLoopWatcher( + ctx, + &suite.wg, + suite.client, + "test", + "TestWatcherReportsProgressRevision", + func([]*clientv3.Event) error { return nil }, + func(*mvccpb.KeyValue) error { return nil }, + func(*mvccpb.KeyValue) error { return nil }, + func([]*clientv3.Event) error { return nil }, + false, /* withPrefix */ + ) + var updatedRevision atomic.Int64 + watcher.SetRevisionUpdatedCallback(updatedRevision.Store) + + suite.wg.Add(1) + go func() { + defer suite.wg.Done() + _, err := watcher.watch(ctx, 0) + re.NoError(err) + }() + + testutil.Eventually(re, func() bool { + return updatedRevision.Load() > 0 + }) +} + func (suite *loopWatcherTestSuite) startEtcd(re *require.Assertions) { etcd1, err := embed.StartEtcd(suite.config) re.NoError(err) From 2662822ebedf11fb5a367abe6787ffe33d2d2dff Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 13 Jul 2026 12:18:39 +0800 Subject: [PATCH 09/15] tso: preserve skipped watch revision semantics Signed-off-by: Ryan Leung --- pkg/tso/keyspace_group_manager.go | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/pkg/tso/keyspace_group_manager.go b/pkg/tso/keyspace_group_manager.go index 6f3512e0f6a..29820e31395 100644 --- a/pkg/tso/keyspace_group_manager.go +++ b/pkg/tso/keyspace_group_manager.go @@ -598,6 +598,12 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { true, /* withPrefix */ ) kgm.groupWatcher.SetRevisionUpdatedCallback(func(revision int64) { + failpoint.Inject("SkipKeyspaceWatch", func(val failpoint.Value) { + addr, ok := val.(string) + if ok && addr == kgm.electionNamePrefix { + failpoint.Return() + } + }) if revision > 0 { kgm.SetModRevision(uint64(revision)) } From 684425303baa2f36e65d6d5b363ec59485f679a4 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 13 Jul 2026 13:45:38 +0800 Subject: [PATCH 10/15] tests: wait for server startup goroutines before cleanup Signed-off-by: Ryan Leung --- tests/cluster.go | 71 ++++++++++++++++++++++++++++++++++++++++--- tests/cluster_test.go | 59 +++++++++++++++++++++++++++++++++++ 2 files changed, 126 insertions(+), 4 deletions(-) diff --git a/tests/cluster.go b/tests/cluster.go index c109745e322..08302af9335 100644 --- a/tests/cluster.go +++ b/tests/cluster.go @@ -172,16 +172,26 @@ func NewTestServer(ctx context.Context, cfg *config.Config, services []string, h // Run starts to run a TestServer. func (s *TestServer) Run() error { + return s.runWithStartSignal(nil) +} + +func (s *TestServer) runWithStartSignal(started chan<- struct{}) error { s.Lock() if s.state != Initial && s.state != Stop { state := s.state s.Unlock() + if started != nil { + close(started) + } return errors.Errorf("server(state%d) cannot run", state) } prevState := s.state // Treat startup as running so retry cleanup can close a blocked server.Run. s.state = Running s.Unlock() + if started != nil { + close(started) + } if err := s.server.Run(); err != nil { s.Lock() @@ -696,12 +706,65 @@ func RunServer(server *TestServer) <-chan error { // RunServers starts to run multiple TestServer. func RunServers(servers []*TestServer) error { - res := make([]<-chan error, len(servers)) + runners := make([]testServerRunner, 0, len(servers)) + for _, server := range servers { + runners = append(runners, server) + } + return runTestServers(runners) +} + +type testServerRunner interface { + runWithStartSignal(chan<- struct{}) error + Stop() error + State() int32 +} + +func runTestServers(servers []testServerRunner) error { + type runResult struct { + index int + err error + } + resC := make(chan runResult, len(servers)) for i, s := range servers { - res[i] = RunServer(s) + index := i + server := s + started := make(chan struct{}) + go func() { + resC <- runResult{index: index, err: server.runWithStartSignal(started)} + }() + // runWithStartSignal changes the state to Running before notifying started. + // Wait for that transition so an early failure cannot race with a server + // that has not entered run yet and therefore cannot be stopped. + <-started + } + + errs := make([]error, len(servers)) + stopping := false + for range servers { + result := <-resC + errs[result.index] = result.err + if result.err != nil && !stopping { + stopping = true + // Another server may still be blocked in Run. Stop all in-flight servers + // to unblock them, then keep draining resC before retry cleanup can destroy + // their data directories. + var wg sync.WaitGroup + for _, s := range servers { + if s.State() != Running { + continue + } + server := s + wg.Add(1) + go func() { + defer wg.Done() + _ = server.Stop() + }() + } + wg.Wait() + } } - for _, c := range res { - if err := <-c; err != nil { + for _, err := range errs { + if err != nil { return errors.WithStack(err) } } diff --git a/tests/cluster_test.go b/tests/cluster_test.go index f012b1ce979..c7c0826d113 100644 --- a/tests/cluster_test.go +++ b/tests/cluster_test.go @@ -15,6 +15,7 @@ package tests import ( + "sync/atomic" "testing" "github.com/stretchr/testify/require" @@ -54,3 +55,61 @@ func TestClassifyInitialServersError(t *testing.T) { re.Equal(startServersNoRetry, classifyInitialServersError(errors.New("some other error"))) re.Equal(startServersNoRetry, classifyInitialServersError(nil)) } + +type stubTestServer struct { + state atomic.Int32 + stopCh chan struct{} + runErr error + stopped atomic.Bool + finished atomic.Bool +} + +func newStubTestServer() *stubTestServer { + return &stubTestServer{ + stopCh: make(chan struct{}), + } +} + +func (s *stubTestServer) runWithStartSignal(started chan<- struct{}) error { + s.state.Store(Running) + if started != nil { + close(started) + } + if s.runErr != nil { + s.state.Store(Initial) + s.finished.Store(true) + return s.runErr + } + <-s.stopCh + s.finished.Store(true) + return nil +} + +func (s *stubTestServer) Stop() error { + if !s.state.CompareAndSwap(Running, Stop) { + return errors.New("server is not running") + } + s.stopped.Store(true) + close(s.stopCh) + return nil +} + +func (s *stubTestServer) State() int32 { + return s.state.Load() +} + +func TestRunServersWaitsForInFlightRuns(t *testing.T) { + t.Parallel() + + re := require.New(t) + failedServer := newStubTestServer() + failedServer.runErr = errors.New("start failed") + blockedServer := newStubTestServer() + + // Start the failed server first to verify that RunServers waits until the + // following server is stoppable before handling the failure. + err := runTestServers([]testServerRunner{failedServer, blockedServer}) + re.EqualError(err, "start failed") + re.True(blockedServer.stopped.Load()) + re.True(blockedServer.finished.Load()) +} From bba572c90998e63d534d41ec9ee1472c88bdd6c0 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 13 Jul 2026 16:35:59 +0800 Subject: [PATCH 11/15] tests: clean up partially started servers Signed-off-by: Ryan Leung --- server/server.go | 67 +++++++++++++++++-- tests/cluster.go | 78 ++++++++++++++++++---- tests/cluster_test.go | 18 +++-- tests/integrations/tso/consistency_test.go | 52 +++++++++++++++ 4 files changed, 190 insertions(+), 25 deletions(-) diff --git a/server/server.go b/server/server.go index fe0a2a37d1c..56acc5a9afd 100644 --- a/server/server.go +++ b/server/server.go @@ -671,28 +671,87 @@ func (s *Server) IsClosed() bool { // Run runs the pd server. func (s *Server) Run() error { + return s.RunWithContext(s.ctx) +} + +// RunWithContext runs the PD server with a context that controls startup and +// the server loops. The caller must keep the context alive while the server is +// running and call Close after canceling it. +func (s *Server) RunWithContext(ctx context.Context) (retErr error) { go systimemon.StartMonitor(s.ctx, time.Now, func() { log.Error("system time jumps backward", errs.ZapError(errs.ErrIncorrectSystemTime)) timeJumpBackCounter.Inc() }) - if err := s.startEtcd(s.ctx); err != nil { + if err := s.startEtcd(ctx); err != nil { return err } + defer func() { + if retErr != nil { + s.cleanupFailedStart() + } + }() + failpoint.Inject("failAfterStartEtcd", func() { + failpoint.Return(errors.New("injected error after etcd startup")) + }) - if err := s.startServer(s.ctx); err != nil { + if err := s.startServer(ctx); err != nil { return err } - s.cgMonitor.StartMonitor(s.ctx) + s.cgMonitor.StartMonitor(ctx) failpoint.Inject("delayStartServerLoop", func() { time.Sleep(2 * time.Second) }) - s.startServerLoop(s.ctx) + s.startServerLoop(ctx) return nil } +func (s *Server) cleanupFailedStart() { + if s.cluster != nil { + s.cluster.Stop() + } + if s.IsKeyspaceGroupEnabled() && s.keyspaceGroupManager != nil { + s.keyspaceGroupManager.Close() + } + if s.tsoAllocator != nil { + s.tsoAllocator.Close() + } + if s.meteringWriter != nil { + s.meteringWriter.Stop() + } + if s.client != nil { + if err := s.client.Close(); err != nil { + log.Error("close etcd client meet error", errs.ZapError(errs.ErrCloseEtcdClient, err)) + } + } + if s.electionClient != nil { + if err := s.electionClient.Close(); err != nil { + log.Error("close election client meet error", errs.ZapError(errs.ErrCloseEtcdClient, err)) + } + } + if s.httpClient != nil { + s.httpClient.CloseIdleConnections() + } + if s.member.Etcd() != nil { + s.member.Close() + } + if s.hbStreams != nil { + s.hbStreams.Close() + } + if s.storage != nil { + if err := s.storage.Close(); err != nil { + log.Error("close storage meet error", errs.ZapError(err)) + } + } + if s.hotRegionStorage != nil { + if err := s.hotRegionStorage.Close(); err != nil { + log.Error("close hot region storage meet error", errs.ZapError(err)) + } + } +} + // SetServiceAuditBackendForHTTP is used to register service audit config for HTTP. func (s *Server) SetServiceAuditBackendForHTTP(route *mux.Route, labels ...string) { if len(route.GetName()) == 0 { diff --git a/tests/cluster.go b/tests/cluster.go index 5a816624bfc..e234eb34625 100644 --- a/tests/cluster.go +++ b/tests/cluster.go @@ -77,6 +77,10 @@ var ( // defaultMaxRetryTimes is the default maximum retry times for starting servers. defaultMaxRetryTimes = 5 + // runServersCleanupGracePeriod allows the remaining servers to finish startup + // naturally after a peer fails. Canceling etcd before it becomes ready can + // leave its embedded listeners blocked in shutdown. + runServersCleanupGracePeriod = 10 * time.Second ) type startServersRetryAction int @@ -116,6 +120,8 @@ type TestServer struct { server *server.Server grpcServer *server.GrpcServer state int32 + runCancel context.CancelFunc + runDone chan struct{} } var zapLogOnce sync.Once @@ -185,17 +191,26 @@ func (s *TestServer) runWithStartSignal(started chan<- struct{}) error { return errors.Errorf("server(state%d) cannot run", state) } prevState := s.state + runCtx, runCancel := context.WithCancel(s.server.Context()) + runDone := make(chan struct{}) // Treat startup as running so retry cleanup can close a blocked server.Run. s.state = Running + s.runCancel = runCancel + s.runDone = runDone s.Unlock() if started != nil { close(started) } - if err := s.server.Run(); err != nil { + err := s.server.RunWithContext(runCtx) + close(runDone) + if err != nil { s.Lock() if s.state == Running { s.state = prevState + runCancel() + s.runCancel = nil + s.runDone = nil } s.Unlock() return err @@ -210,7 +225,15 @@ func (s *TestServer) Stop() error { if s.state != Running { return errors.Errorf("server(state%d) cannot stop", s.state) } + if s.runCancel != nil { + s.runCancel() + } + if s.runDone != nil { + <-s.runDone + } s.server.Close() + s.runCancel = nil + s.runDone = nil s.state = Stop return nil } @@ -220,7 +243,15 @@ func (s *TestServer) Destroy() error { s.Lock() defer s.Unlock() if s.state == Running { + if s.runCancel != nil { + s.runCancel() + } + if s.runDone != nil { + <-s.runDone + } s.server.Close() + s.runCancel = nil + s.runDone = nil } if err := os.RemoveAll(s.server.GetConfig().DataDir); err != nil { return err @@ -737,16 +768,21 @@ func runTestServers(servers []testServerRunner) error { <-started } - errs := make([]error, len(servers)) - stopping := false - for range servers { - result := <-resC - errs[result.index] = result.err - if result.err != nil && !stopping { - stopping = true - // Another server may still be blocked in Run. Stop all in-flight servers - // to unblock them, then keep draining resC before retry cleanup can destroy - // their data directories. + var ( + primaryErr error + cleanupTimer *time.Timer + cleanupC <-chan time.Time + ) + remaining := len(servers) + for remaining > 0 { + var result runResult + select { + case result = <-resC: + remaining-- + case <-cleanupC: + cleanupC = nil + // The remaining Run calls did not finish naturally. Cancel them to + // bound cleanup time, then drain their results below. var wg sync.WaitGroup for _, s := range servers { if s.State() != Running { @@ -760,12 +796,26 @@ func runTestServers(servers []testServerRunner) error { }() } wg.Wait() + continue + } + if result.err != nil && primaryErr == nil { + primaryErr = result.err + cleanupTimer = time.NewTimer(runServersCleanupGracePeriod) + cleanupC = cleanupTimer.C } } - for _, err := range errs { - if err != nil { - return errors.WithStack(err) + if cleanupTimer != nil { + cleanupTimer.Stop() + } + if primaryErr != nil { + // Every Run call has returned, so closing these servers cannot race with + // partial startup. Do this before retry cleanup can remove their data dirs. + for _, s := range servers { + if s.State() == Running { + _ = s.Stop() + } } + return errors.WithStack(primaryErr) } return nil } diff --git a/tests/cluster_test.go b/tests/cluster_test.go index 229ade9c3ad..ac7c6f4f2c8 100644 --- a/tests/cluster_test.go +++ b/tests/cluster_test.go @@ -18,6 +18,7 @@ import ( "os" "sync/atomic" "testing" + "time" "github.com/stretchr/testify/require" "go.uber.org/goleak" @@ -63,6 +64,7 @@ type stubTestServer struct { state atomic.Int32 stopCh chan struct{} runErr error + stopErr error stopped atomic.Bool finished atomic.Bool } @@ -85,7 +87,7 @@ func (s *stubTestServer) runWithStartSignal(started chan<- struct{}) error { } <-s.stopCh s.finished.Store(true) - return nil + return s.stopErr } func (s *stubTestServer) Stop() error { @@ -102,16 +104,18 @@ func (s *stubTestServer) State() int32 { } func TestRunServersWaitsForInFlightRuns(t *testing.T) { - t.Parallel() - re := require.New(t) + oldCleanupGracePeriod := runServersCleanupGracePeriod + runServersCleanupGracePeriod = time.Millisecond + t.Cleanup(func() { runServersCleanupGracePeriod = oldCleanupGracePeriod }) + blockedServer := newStubTestServer() + blockedServer.stopErr = errors.New("start canceled") failedServer := newStubTestServer() failedServer.runErr = errors.New("start failed") - blockedServer := newStubTestServer() - // Start the failed server first to verify that RunServers waits until the - // following server is stoppable before handling the failure. - err := runTestServers([]testServerRunner{failedServer, blockedServer}) + // The cleanup error from the lower-index server must not replace the startup + // error that triggered cleanup. + err := runTestServers([]testServerRunner{blockedServer, failedServer}) re.EqualError(err, "start failed") re.True(blockedServer.stopped.Load()) re.True(blockedServer.finished.Load()) diff --git a/tests/integrations/tso/consistency_test.go b/tests/integrations/tso/consistency_test.go index 766106993a1..578c8cbc844 100644 --- a/tests/integrations/tso/consistency_test.go +++ b/tests/integrations/tso/consistency_test.go @@ -16,10 +16,13 @@ package tso import ( "context" + "net" + "net/url" "sync" "testing" "time" + "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" "google.golang.org/grpc" @@ -68,6 +71,55 @@ func TestMicroserviceTSOConsistencySuite(t *testing.T) { }) } +func TestRunInitialServersClosesStartingServersBeforeRetry(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + cluster, err := tests.NewTestCluster(ctx, serverCount) + re.NoError(err) + t.Cleanup(func() { + cancel() + cluster.Destroy() + }) + + oldServer := cluster.GetServer("pd1") + conflictURL, err := url.Parse(cluster.GetServer("pd3").GetConfig().ClientUrls) + re.NoError(err) + listener, err := net.Listen("tcp", conflictURL.Host) + re.NoError(err) + t.Cleanup(func() { re.NoError(listener.Close()) }) + + re.NoError(cluster.RunInitialServers()) + re.True(oldServer.GetServer().IsClosed()) +} + +func TestRunFailureAfterEtcdStartClosesServer(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + cluster, err := tests.NewTestCluster(ctx, 1) + re.NoError(err) + t.Cleanup(func() { + cancel() + cluster.Destroy() + }) + + const failpointName = "github.com/tikv/pd/server/failAfterStartEtcd" + re.NoError(failpoint.Enable(failpointName, "return(true)")) + t.Cleanup(func() { re.NoError(failpoint.Disable(failpointName)) }) + + testServer := cluster.GetServer("pd1") + clientURL, err := url.Parse(testServer.GetConfig().ClientUrls) + re.NoError(err) + err = cluster.RunInitialServers() + re.ErrorContains(err, "injected error after etcd startup") + re.NoError(testServer.Destroy()) + + conn, err := net.DialTimeout("tcp", clientURL.Host, time.Second) + re.Error(err) + if conn != nil { + re.NoError(conn.Close()) + } +} + func (suite *tsoConsistencyTestSuite) SetupSuite() { re := suite.Require() From 644088f497cc48a97a16d3ffd09ac653d327e4ee Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 13 Jul 2026 18:27:56 +0800 Subject: [PATCH 12/15] server: warn on failed-start cleanup errors Signed-off-by: Ryan Leung --- server/server.go | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/server/server.go b/server/server.go index 56acc5a9afd..2f0782d4829 100644 --- a/server/server.go +++ b/server/server.go @@ -723,12 +723,12 @@ func (s *Server) cleanupFailedStart() { } if s.client != nil { if err := s.client.Close(); err != nil { - log.Error("close etcd client meet error", errs.ZapError(errs.ErrCloseEtcdClient, err)) + log.Warn("close etcd client meet error", errs.ZapError(errs.ErrCloseEtcdClient, err)) } } if s.electionClient != nil { if err := s.electionClient.Close(); err != nil { - log.Error("close election client meet error", errs.ZapError(errs.ErrCloseEtcdClient, err)) + log.Warn("close election client meet error", errs.ZapError(errs.ErrCloseEtcdClient, err)) } } if s.httpClient != nil { @@ -742,12 +742,12 @@ func (s *Server) cleanupFailedStart() { } if s.storage != nil { if err := s.storage.Close(); err != nil { - log.Error("close storage meet error", errs.ZapError(err)) + log.Warn("close storage meet error", errs.ZapError(err)) } } if s.hotRegionStorage != nil { if err := s.hotRegionStorage.Close(); err != nil { - log.Error("close hot region storage meet error", errs.ZapError(err)) + log.Warn("close hot region storage meet error", errs.ZapError(err)) } } } From 26b38d8cd414896bfe00042f03b614481fbb5114 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Wed, 15 Jul 2026 18:47:20 +0800 Subject: [PATCH 13/15] *: bound startup cleanup and preserve watch semantics Cancel failed server starts without calling the blocking Stop path, and return a non-retryable error when a startup goroutine still does not exit. Keep keyspace-group revision updates on the failpoint-aware watcher callback and strengthen the startup, snapshot, and listener cleanup coverage. Signed-off-by: Ryan Leung --- pkg/tso/keyspace_group_manager.go | 9 -- pkg/tso/keyspace_group_manager_test.go | 36 +++++++ pkg/utils/etcdutil/etcdutil.go | 11 --- pkg/utils/etcdutil/etcdutil_test.go | 10 +- tests/cluster.go | 105 ++++++++++++++------- tests/cluster_test.go | 99 ++++++++++++++++++- tests/integrations/tso/consistency_test.go | 15 ++- 7 files changed, 224 insertions(+), 61 deletions(-) diff --git a/pkg/tso/keyspace_group_manager.go b/pkg/tso/keyspace_group_manager.go index e746994d109..c5e922667dc 100644 --- a/pkg/tso/keyspace_group_manager.go +++ b/pkg/tso/keyspace_group_manager.go @@ -621,15 +621,6 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { kgm.Close() return errs.ErrLoadKeyspaceGroupsTerminated.Wrap(err) } - if loadedRevision := kgm.groupWatcher.GetLoadedRevision(); loadedRevision > 0 { - // The watch loop can advance the manager after the initial load finishes - // but before WaitLoad returns. In that case, keeping the newer revision is - // correct and SetModRevision returns false. - if !kgm.SetModRevision(uint64(loadedRevision)) { - log.Debug("keyspace group revision is already at or ahead of the loaded snapshot", - zap.Int64("loaded-revision", loadedRevision)) - } - } if !defaultKGConfigured { log.Info("initializing default keyspace group") diff --git a/pkg/tso/keyspace_group_manager_test.go b/pkg/tso/keyspace_group_manager_test.go index 9d42be9e252..216dd167578 100644 --- a/pkg/tso/keyspace_group_manager_test.go +++ b/pkg/tso/keyspace_group_manager_test.go @@ -261,6 +261,42 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsSetsModRevisio re.GreaterOrEqual(loadedRevision, deletedRevision) } +func (suite *keyspaceGroupManagerTestSuite) TestInitialSkipKeyspaceWatchDoesNotAdvanceRevision() { + re := suite.Require() + + mgr := suite.newUniqueKeyspaceGroupManager(1) + re.NotNil(mgr) + defer mgr.Close() + + point := fmt.Sprintf("return(\"%s\")", mgr.electionNamePrefix) + re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/tso/SkipKeyspaceWatch", point)) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/tso/SkipKeyspaceWatch")) + }() + + const ( + groupID = uint32(1) + keyspaceID = uint32(101) + ) + re.NoError(addKeyspaceGroupAssignment( + suite.ctx, + suite.etcdClient, + groupID, + []string{mgr.tsoServiceID.ServiceAddr}, + []int{mcs.DefaultKeyspaceGroupReplicaPriority}, + []uint32{keyspaceID}, + )) + re.NoError(mgr.Initialize()) + + mgr.RLock() + loadedGroup := mgr.kgs[groupID] + loadedRevision := mgr.modRevision + mgr.RUnlock() + + re.Nil(loadedGroup) + re.Zero(loadedRevision, "a skipped initial watch must not advertise an unapplied revision") +} + // TestLoadWithDifferentBatchSize tests the loading of the keyspace group assignment with the different batch size. func (suite *keyspaceGroupManagerTestSuite) TestLoadWithDifferentBatchSize() { re := suite.Require() diff --git a/pkg/utils/etcdutil/etcdutil.go b/pkg/utils/etcdutil/etcdutil.go index a10792b122f..aa1e2c52f99 100644 --- a/pkg/utils/etcdutil/etcdutil.go +++ b/pkg/utils/etcdutil/etcdutil.go @@ -22,7 +22,6 @@ import ( "os" "strings" "sync" - "sync/atomic" "time" "github.com/gogo/protobuf/proto" @@ -375,8 +374,6 @@ type LoopWatcher struct { postEventsFn func([]*clientv3.Event) error // preEventsFn is used to call before handling all events. preEventsFn func([]*clientv3.Event) error - // loadedRevision is the etcd snapshot revision loaded during initialization. - loadedRevision atomic.Int64 // revisionUpdatedFn is called when the watcher has caught up to a newer etcd revision. revisionUpdatedFn func(int64) @@ -496,9 +493,6 @@ func (lw *LoopWatcher) initFromEtcd(ctx context.Context) int64 { if err != nil { log.Warn("meet error when loading in watch loop", zap.String("name", lw.name), zap.String("key", lw.key), zap.Error(err)) } else { - if watchStartRevision > 0 { - lw.loadedRevision.Store(watchStartRevision - 1) - } log.Info("load finished in watch loop", zap.String("name", lw.name), zap.String("key", lw.key)) } lw.isLoadedCh <- err @@ -781,11 +775,6 @@ func (lw *LoopWatcher) WaitLoad() error { return <-lw.isLoadedCh } -// GetLoadedRevision returns the etcd snapshot revision loaded during initialization. -func (lw *LoopWatcher) GetLoadedRevision() int64 { - return lw.loadedRevision.Load() -} - // SetRevisionUpdatedCallback sets a callback that is invoked after the watcher catches up to an etcd revision. func (lw *LoopWatcher) SetRevisionUpdatedCallback(fn func(int64)) { if fn == nil { diff --git a/pkg/utils/etcdutil/etcdutil_test.go b/pkg/utils/etcdutil/etcdutil_test.go index ba46ef8eb14..c76f645c230 100644 --- a/pkg/utils/etcdutil/etcdutil_test.go +++ b/pkg/utils/etcdutil/etcdutil_test.go @@ -467,12 +467,12 @@ func (suite *loopWatcherTestSuite) TestLoadNoExistedKey() { re.Empty(cache) } -func (suite *loopWatcherTestSuite) TestGetLoadedRevision() { +func (suite *loopWatcherTestSuite) TestLoadUsesSingleSnapshotRevision() { re := suite.Require() ctx, cancel := context.WithCancel(suite.ctx) defer cancel() - prefix := "TestGetLoadedRevision/" + prefix := "TestLoadUsesSingleSnapshotRevision/" _, err := suite.client.Txn(ctx).Then( clientv3.OpPut(prefix+"a", ""), clientv3.OpPut(prefix+"b", ""), @@ -503,10 +503,10 @@ func (suite *loopWatcherTestSuite) TestGetLoadedRevision() { true, /* withPrefix */ ) watcher.SetLoadBatchSize(1) - watcher.StartWatchLoop() - err = watcher.WaitLoad() + nextRevision, err := watcher.load(ctx) re.NoError(err) - re.Equal(expectedRevision, watcher.GetLoadedRevision()) + re.NoError(putErr) + re.Equal(expectedRevision+1, nextRevision) } func (suite *loopWatcherTestSuite) TestLoadWithLimitChange() { diff --git a/tests/cluster.go b/tests/cluster.go index e234eb34625..31abadc20fc 100644 --- a/tests/cluster.go +++ b/tests/cluster.go @@ -16,6 +16,7 @@ package tests import ( "context" + stdErrors "errors" "net/http" "os" "strings" @@ -81,6 +82,9 @@ var ( // naturally after a peer fails. Canceling etcd before it becomes ready can // leave its embedded listeners blocked in shutdown. runServersCleanupGracePeriod = 10 * time.Second + // errRunServersCleanupTimeout prevents retry cleanup from destroying a data + // directory while its startup goroutine may still be using it. + errRunServersCleanupTimeout = errors.New("timed out waiting for starting servers to stop") ) type startServersRetryAction int @@ -103,6 +107,9 @@ func classifyInitialServersError(err error) startServersRetryAction { if err == nil { return startServersNoRetry } + if stdErrors.Is(err, errRunServersCleanupTimeout) { + return startServersNoRetry + } errMsg := err.Error() switch { case strings.Contains(errMsg, "address already in use") || strings.Contains(errMsg, "Etcd cluster ID mismatch"): @@ -218,6 +225,15 @@ func (s *TestServer) runWithStartSignal(started chan<- struct{}) error { return nil } +func (s *TestServer) cancelRun() { + s.RLock() + runCancel := s.runCancel + s.RUnlock() + if runCancel != nil { + runCancel() + } +} + // Stop is used to stop a TestServer. func (s *TestServer) Stop() error { s.Lock() @@ -745,6 +761,7 @@ func RunServers(servers []*TestServer) error { type testServerRunner interface { runWithStartSignal(chan<- struct{}) error + cancelRun() Stop() error State() int32 } @@ -769,52 +786,74 @@ func runTestServers(servers []testServerRunner) error { } var ( - primaryErr error - cleanupTimer *time.Timer - cleanupC <-chan time.Time + primaryErr error + graceTimer *time.Timer + graceC <-chan time.Time + postCancelTimer *time.Timer + postCancelC <-chan time.Time ) remaining := len(servers) + completed := make([]bool, len(servers)) + handleResult := func(result runResult) { + remaining-- + completed[result.index] = true + if result.err != nil && primaryErr == nil { + primaryErr = result.err + graceTimer = time.NewTimer(runServersCleanupGracePeriod) + graceC = graceTimer.C + } + } + stopCompletedServers := func() { + for i, s := range servers { + if completed[i] && s.State() == Running { + _ = s.Stop() + } + } + } for remaining > 0 { - var result runResult select { - case result = <-resC: - remaining-- - case <-cleanupC: - cleanupC = nil - // The remaining Run calls did not finish naturally. Cancel them to - // bound cleanup time, then drain their results below. - var wg sync.WaitGroup + case result := <-resC: + handleResult(result) + case <-graceC: + graceC = nil + // Cancel without joining here: Stop waits for runDone and can block + // forever when startup does not observe its context. for _, s := range servers { - if s.State() != Running { - continue + if s.State() == Running { + s.cancelRun() } - server := s - wg.Add(1) - go func() { - defer wg.Done() - _ = server.Stop() - }() } - wg.Wait() - continue - } - if result.err != nil && primaryErr == nil { - primaryErr = result.err - cleanupTimer = time.NewTimer(runServersCleanupGracePeriod) - cleanupC = cleanupTimer.C + postCancelTimer = time.NewTimer(runServersCleanupGracePeriod) + postCancelC = postCancelTimer.C + case <-postCancelC: + postCancelC = nil + // Prefer results that became ready at the timeout boundary. + drainResults: + for remaining > 0 { + select { + case result := <-resC: + handleResult(result) + default: + break drainResults + } + } + if remaining > 0 { + stopCompletedServers() + return errors.Wrapf(errRunServersCleanupTimeout, + "server startup cleanup did not complete after cancellation; original error: %v", primaryErr) + } } } - if cleanupTimer != nil { - cleanupTimer.Stop() + if graceTimer != nil { + graceTimer.Stop() + } + if postCancelTimer != nil { + postCancelTimer.Stop() } if primaryErr != nil { // Every Run call has returned, so closing these servers cannot race with // partial startup. Do this before retry cleanup can remove their data dirs. - for _, s := range servers { - if s.State() == Running { - _ = s.Stop() - } - } + stopCompletedServers() return errors.WithStack(primaryErr) } return nil diff --git a/tests/cluster_test.go b/tests/cluster_test.go index ac7c6f4f2c8..cf3346fcbad 100644 --- a/tests/cluster_test.go +++ b/tests/cluster_test.go @@ -56,6 +56,8 @@ func TestClassifyInitialServersError(t *testing.T) { 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.Wrap( + errRunServersCleanupTimeout, "listen tcp 127.0.0.1:2379: bind: address already in use"))) re.Equal(startServersNoRetry, classifyInitialServersError(errors.New("some other error"))) re.Equal(startServersNoRetry, classifyInitialServersError(nil)) } @@ -66,6 +68,8 @@ type stubTestServer struct { runErr error stopErr error stopped atomic.Bool + canceled atomic.Bool + stopOnce atomic.Bool finished atomic.Bool } @@ -87,15 +91,27 @@ func (s *stubTestServer) runWithStartSignal(started chan<- struct{}) error { } <-s.stopCh s.finished.Store(true) + if s.stopErr != nil { + s.state.Store(Initial) + } return s.stopErr } +func (s *stubTestServer) cancelRun() { + s.canceled.Store(true) + if s.stopOnce.CompareAndSwap(false, true) { + close(s.stopCh) + } +} + func (s *stubTestServer) Stop() error { if !s.state.CompareAndSwap(Running, Stop) { return errors.New("server is not running") } s.stopped.Store(true) - close(s.stopCh) + if s.stopOnce.CompareAndSwap(false, true) { + close(s.stopCh) + } return nil } @@ -103,6 +119,85 @@ func (s *stubTestServer) State() int32 { return s.state.Load() } +type uninterruptibleTestServer struct { + state atomic.Int32 + releaseCh chan struct{} + cancelCalled chan struct{} + cancelOnce atomic.Bool + releaseOnce atomic.Bool + finished atomic.Bool +} + +func newUninterruptibleTestServer() *uninterruptibleTestServer { + return &uninterruptibleTestServer{ + releaseCh: make(chan struct{}), + cancelCalled: make(chan struct{}), + } +} + +func (s *uninterruptibleTestServer) runWithStartSignal(started chan<- struct{}) error { + s.state.Store(Running) + close(started) + <-s.releaseCh + s.state.Store(Initial) + s.finished.Store(true) + return errors.New("start canceled") +} + +func (s *uninterruptibleTestServer) cancelRun() { + if s.cancelOnce.CompareAndSwap(false, true) { + close(s.cancelCalled) + } +} + +func (s *uninterruptibleTestServer) Stop() error { + // Model TestServer.Stop waiting for runDone while startup ignores cancellation. + <-s.releaseCh + s.state.Store(Stop) + return nil +} + +func (s *uninterruptibleTestServer) State() int32 { + return s.state.Load() +} + +func (s *uninterruptibleTestServer) release() { + if s.releaseOnce.CompareAndSwap(false, true) { + close(s.releaseCh) + } +} + +func TestRunServersCleanupIsBounded(t *testing.T) { + oldCleanupGracePeriod := runServersCleanupGracePeriod + runServersCleanupGracePeriod = time.Millisecond + t.Cleanup(func() { runServersCleanupGracePeriod = oldCleanupGracePeriod }) + + blocked := newUninterruptibleTestServer() + defer blocked.release() + failed := newStubTestServer() + failed.runErr = errors.New("start failed") + + done := make(chan error, 1) + go func() { + done <- runTestServers([]testServerRunner{blocked, failed}) + }() + + select { + case <-blocked.cancelCalled: + case <-time.After(time.Second): + t.Fatal("cleanup did not cancel the blocked server") + } + + select { + case err := <-done: + require.ErrorIs(t, err, errRunServersCleanupTimeout) + case <-time.After(50 * time.Millisecond): + t.Fatal("cleanup remained blocked after the grace period") + } + blocked.release() + require.Eventually(t, blocked.finished.Load, time.Second, time.Millisecond) +} + func TestRunServersWaitsForInFlightRuns(t *testing.T) { re := require.New(t) oldCleanupGracePeriod := runServersCleanupGracePeriod @@ -117,7 +212,7 @@ func TestRunServersWaitsForInFlightRuns(t *testing.T) { // error that triggered cleanup. err := runTestServers([]testServerRunner{blockedServer, failedServer}) re.EqualError(err, "start failed") - re.True(blockedServer.stopped.Load()) + re.True(blockedServer.canceled.Load()) re.True(blockedServer.finished.Load()) } diff --git a/tests/integrations/tso/consistency_test.go b/tests/integrations/tso/consistency_test.go index 578c8cbc844..b2a27ae5e0d 100644 --- a/tests/integrations/tso/consistency_test.go +++ b/tests/integrations/tso/consistency_test.go @@ -82,6 +82,10 @@ func TestRunInitialServersClosesStartingServersBeforeRetry(t *testing.T) { }) oldServer := cluster.GetServer("pd1") + oldClientURL, err := url.Parse(oldServer.GetConfig().ClientUrls) + re.NoError(err) + oldPeerURL, err := url.Parse(oldServer.GetConfig().PeerUrls) + re.NoError(err) conflictURL, err := url.Parse(cluster.GetServer("pd3").GetConfig().ClientUrls) re.NoError(err) listener, err := net.Listen("tcp", conflictURL.Host) @@ -89,7 +93,16 @@ func TestRunInitialServersClosesStartingServersBeforeRetry(t *testing.T) { t.Cleanup(func() { re.NoError(listener.Close()) }) re.NoError(cluster.RunInitialServers()) - re.True(oldServer.GetServer().IsClosed()) + re.Equal(tests.Destroy, oldServer.State()) + for _, addr := range []string{oldClientURL.Host, oldPeerURL.Host} { + testutil.Eventually(re, func() bool { + conn, err := net.DialTimeout("tcp", addr, 100*time.Millisecond) + if conn != nil { + re.NoError(conn.Close()) + } + return err != nil + }, testutil.WithWaitFor(10*time.Second), testutil.WithTickInterval(100*time.Millisecond)) + } } func TestRunFailureAfterEtcdStartClosesServer(t *testing.T) { From eaf95e96f893fcfac483c54a00608f0e4094f850 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Thu, 16 Jul 2026 10:54:12 +0800 Subject: [PATCH 14/15] tso: avoid advancing group revision on watch progress Etcd progress revisions include unrelated writes and can differ across TSO nodes. Update the keyspace-group revision only after a snapshot load; group watch events continue to use their own KV revisions. Signed-off-by: Ryan Leung --- pkg/tso/keyspace_group_manager.go | 2 +- pkg/utils/etcdutil/etcdutil.go | 18 +++++++-------- pkg/utils/etcdutil/etcdutil_test.go | 35 +++-------------------------- 3 files changed, 12 insertions(+), 43 deletions(-) diff --git a/pkg/tso/keyspace_group_manager.go b/pkg/tso/keyspace_group_manager.go index c5e922667dc..9ed9b40fefd 100644 --- a/pkg/tso/keyspace_group_manager.go +++ b/pkg/tso/keyspace_group_manager.go @@ -597,7 +597,7 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { postEventsFn, true, /* withPrefix */ ) - kgm.groupWatcher.SetRevisionUpdatedCallback(func(revision int64) { + kgm.groupWatcher.SetLoadRevisionUpdatedCallback(func(revision int64) { failpoint.Inject("SkipKeyspaceWatch", func(val failpoint.Value) { addr, ok := val.(string) if ok && addr == kgm.electionNamePrefix { diff --git a/pkg/utils/etcdutil/etcdutil.go b/pkg/utils/etcdutil/etcdutil.go index aa1e2c52f99..d288a992b82 100644 --- a/pkg/utils/etcdutil/etcdutil.go +++ b/pkg/utils/etcdutil/etcdutil.go @@ -374,8 +374,8 @@ type LoopWatcher struct { postEventsFn func([]*clientv3.Event) error // preEventsFn is used to call before handling all events. preEventsFn func([]*clientv3.Event) error - // revisionUpdatedFn is called when the watcher has caught up to a newer etcd revision. - revisionUpdatedFn func(int64) + // loadRevisionUpdatedFn is called after a snapshot has been loaded. + loadRevisionUpdatedFn func(int64) // forceLoadMu is used to ensure two force loads have minimal interval. forceLoadMu syncutil.RWMutex @@ -418,7 +418,7 @@ func NewLoopWatcher( deleteFn: deleteFn, postEventsFn: postEventsFn, preEventsFn: preEventsFn, - revisionUpdatedFn: func(int64) {}, + loadRevisionUpdatedFn: func(int64) {}, isWithPrefix: isWithPrefix, lastTimeForceLoad: time.Now(), loadRetryTimes: defaultLoadFromEtcdRetryTimes, @@ -600,7 +600,6 @@ func (lw *LoopWatcher) watch(ctx context.Context, revision int64) (nextRevision } else if wresp.IsProgressNotify() { log.Debug("watcher receives progress notify in watch loop", zap.Int64("revision", revision), zap.String("name", lw.name), zap.String("key", lw.key)) - lw.revisionUpdatedFn(wresp.Header.Revision) goto watchChanLoop } if err := lw.preEventsFn(wresp.Events); err != nil { @@ -634,7 +633,6 @@ func (lw *LoopWatcher) watch(ctx context.Context, revision int64) (nextRevision log.Error("run post event failed in watch loop", zap.Error(err), zap.Int64("revision", revision), zap.String("name", lw.name), zap.String("key", lw.key)) } - lw.revisionUpdatedFn(wresp.Header.Revision) revision = wresp.Header.Revision + 1 } goto watchChanLoop // Use goto to avoid creating a new watchChan @@ -716,7 +714,7 @@ func (lw *LoopWatcher) load(ctx context.Context) (nextRevision int64, err error) } // Note: if there are no keys in etcd, the resp.More is false. It also means the load is finished. if !resp.More { - lw.revisionUpdatedFn(snapshotRevision) + lw.loadRevisionUpdatedFn(snapshotRevision) return snapshotRevision + 1, err } } @@ -775,13 +773,13 @@ func (lw *LoopWatcher) WaitLoad() error { return <-lw.isLoadedCh } -// SetRevisionUpdatedCallback sets a callback that is invoked after the watcher catches up to an etcd revision. -func (lw *LoopWatcher) SetRevisionUpdatedCallback(fn func(int64)) { +// SetLoadRevisionUpdatedCallback sets a callback that is invoked after a snapshot has been loaded. +func (lw *LoopWatcher) SetLoadRevisionUpdatedCallback(fn func(int64)) { if fn == nil { - lw.revisionUpdatedFn = func(int64) {} + lw.loadRevisionUpdatedFn = func(int64) {} return } - lw.revisionUpdatedFn = fn + lw.loadRevisionUpdatedFn = fn } // SetLoadRetryTimes sets the retry times when loading data from etcd. diff --git a/pkg/utils/etcdutil/etcdutil_test.go b/pkg/utils/etcdutil/etcdutil_test.go index c76f645c230..cfecee6ab52 100644 --- a/pkg/utils/etcdutil/etcdutil_test.go +++ b/pkg/utils/etcdutil/etcdutil_test.go @@ -502,11 +502,14 @@ func (suite *loopWatcherTestSuite) TestLoadUsesSingleSnapshotRevision() { func([]*clientv3.Event) error { return nil }, true, /* withPrefix */ ) + var loadedRevision atomic.Int64 + watcher.SetLoadRevisionUpdatedCallback(loadedRevision.Store) watcher.SetLoadBatchSize(1) nextRevision, err := watcher.load(ctx) re.NoError(err) re.NoError(putErr) re.Equal(expectedRevision+1, nextRevision) + re.Equal(expectedRevision, loadedRevision.Load()) } func (suite *loopWatcherTestSuite) TestLoadWithLimitChange() { @@ -816,38 +819,6 @@ func (suite *loopWatcherTestSuite) TestWatcherRequestProgress() { checkWatcherRequestProgress(true) } -func (suite *loopWatcherTestSuite) TestWatcherReportsProgressRevision() { - re := suite.Require() - ctx, cancel := context.WithCancel(suite.ctx) - defer cancel() - - watcher := NewLoopWatcher( - ctx, - &suite.wg, - suite.client, - "test", - "TestWatcherReportsProgressRevision", - func([]*clientv3.Event) error { return nil }, - func(*mvccpb.KeyValue) error { return nil }, - func(*mvccpb.KeyValue) error { return nil }, - func([]*clientv3.Event) error { return nil }, - false, /* withPrefix */ - ) - var updatedRevision atomic.Int64 - watcher.SetRevisionUpdatedCallback(updatedRevision.Store) - - suite.wg.Add(1) - go func() { - defer suite.wg.Done() - _, err := watcher.watch(ctx, 0) - re.NoError(err) - }() - - testutil.Eventually(re, func() bool { - return updatedRevision.Load() > 0 - }) -} - func (suite *loopWatcherTestSuite) startEtcd(re *require.Assertions) { etcd1, err := embed.StartEtcd(suite.config) re.NoError(err) From 70252379e26de5d357cb3441b29d436795671452 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Thu, 16 Jul 2026 17:18:24 +0800 Subject: [PATCH 15/15] tso: persist comparable keyspace group revisions Update a durable revision marker atomically with keyspace group storage transactions. TSO watchers load and watch the marker alongside group metadata, preserving delete revisions without including unrelated etcd writes. Signed-off-by: Ryan Leung --- pkg/keyspace/tso_keyspace_group.go | 62 ++++++++++++++++++- pkg/keyspace/tso_keyspace_group_test.go | 32 ++++++++++ pkg/tso/keyspace_group_manager.go | 38 +++++++----- pkg/tso/keyspace_group_manager_test.go | 51 ++++++++++++++- pkg/utils/etcdutil/etcdutil.go | 14 ----- pkg/utils/etcdutil/etcdutil_test.go | 3 - pkg/utils/keypath/absolute_key_path.go | 24 +++++-- .../integrations/realcluster/etcd_key_test.go | 2 + 8 files changed, 186 insertions(+), 40 deletions(-) diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go index 08bec239d85..90e97c36cd7 100644 --- a/pkg/keyspace/tso_keyspace_group.go +++ b/pkg/keyspace/tso_keyspace_group.go @@ -58,6 +58,7 @@ const ( defaultKeyspaceCountSplitThreshold = 40000 // autoSplitKeyspaceGroupPatrolInterval is the interval for patrolling keyspace group size for auto-split. autoSplitKeyspaceGroupPatrolInterval = 15 * time.Minute + keyspaceGroupRevisionValue = "1" ) const ( @@ -90,6 +91,58 @@ type GroupManager struct { tsoNodesWatcher *etcdutil.LoopWatcher } +type keyspaceGroupRevisionStorage struct { + endpoint.KeyspaceGroupStorage +} + +type keyspaceGroupRevisionTxn struct { + kv.Txn + changed bool +} + +// Save records keyspace group writes so the transaction can advance the revision marker. +func (txn *keyspaceGroupRevisionTxn) Save(key, value string) error { + if err := txn.Txn.Save(key, value); err != nil { + return err + } + if strings.HasPrefix(key, keypath.KeyspaceGroupIDPrefix()) { + txn.changed = true + } + return nil +} + +// Remove records keyspace group deletions so the transaction can advance the revision marker. +func (txn *keyspaceGroupRevisionTxn) Remove(key string) error { + if err := txn.Txn.Remove(key); err != nil { + return err + } + if strings.HasPrefix(key, keypath.KeyspaceGroupIDPrefix()) { + txn.changed = true + } + return nil +} + +// RunInTxn atomically advances the revision marker when the transaction changes a keyspace group. +func (s *keyspaceGroupRevisionStorage) RunInTxn(ctx context.Context, f func(txn kv.Txn) error) error { + return s.KeyspaceGroupStorage.RunInTxn(ctx, func(txn kv.Txn) error { + revisionTxn := &keyspaceGroupRevisionTxn{Txn: txn} + if err := f(revisionTxn); err != nil { + return err + } + if !revisionTxn.changed { + return nil + } + return txn.Save(keypath.KeyspaceGroupRevisionPath(), keyspaceGroupRevisionValue) + }) +} + +func withKeyspaceGroupRevision(store endpoint.KeyspaceGroupStorage) endpoint.KeyspaceGroupStorage { + if _, ok := store.(*keyspaceGroupRevisionStorage); ok { + return store + } + return &keyspaceGroupRevisionStorage{KeyspaceGroupStorage: store} +} + // NewKeyspaceGroupManager creates a Manager of keyspace group related data. func NewKeyspaceGroupManager( ctx context.Context, @@ -104,7 +157,7 @@ func NewKeyspaceGroupManager( m := &GroupManager{ ctx: ctx, cancel: cancel, - store: store, + store: withKeyspaceGroupRevision(store), groups: groups, client: client, nodesBalancer: balancer.GenByPolicy[string](defaultBalancerPolicy), @@ -141,6 +194,13 @@ func (m *GroupManager) Bootstrap(ctx context.Context) error { if err != nil && err != errs.ErrKeyspaceGroupExists { return err } + // Persist a marker on every bootstrap so upgraded clusters retain a + // revision that is at least as new as all existing keyspace group state. + if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error { + return txn.Save(keypath.KeyspaceGroupRevisionPath(), keyspaceGroupRevisionValue) + }); err != nil { + return err + } // Load all the keyspace groups from the storage and add to the respective userKind groups. groups, err := m.store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) diff --git a/pkg/keyspace/tso_keyspace_group_test.go b/pkg/keyspace/tso_keyspace_group_test.go index 489dbbf7efd..6a950001dc7 100644 --- a/pkg/keyspace/tso_keyspace_group_test.go +++ b/pkg/keyspace/tso_keyspace_group_test.go @@ -33,6 +33,7 @@ import ( "github.com/tikv/pd/pkg/storage/endpoint" "github.com/tikv/pd/pkg/storage/kv" "github.com/tikv/pd/pkg/utils/etcdutil" + "github.com/tikv/pd/pkg/utils/keypath" "github.com/tikv/pd/pkg/versioninfo/kerneltype" ) @@ -131,6 +132,37 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupOperations() { re.Error(err) } +func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupChangesUpdateRevisionMarker() { + re := suite.Require() + + removeMarker := func() { + re.NoError(suite.kgm.store.RunInTxn(suite.ctx, func(txn kv.Txn) error { + return txn.Remove(keypath.KeyspaceGroupRevisionPath()) + })) + } + checkMarker := func() { + var marker string + re.NoError(suite.kgm.store.RunInTxn(suite.ctx, func(txn kv.Txn) error { + var err error + marker, err = txn.Load(keypath.KeyspaceGroupRevisionPath()) + return err + })) + re.NotEmpty(marker) + } + + removeMarker() + re.NoError(suite.kgm.CreateKeyspaceGroups([]*endpoint.KeyspaceGroup{{ + ID: 1, + UserKind: endpoint.Standard.String(), + }})) + checkMarker() + + removeMarker() + _, err := suite.kgm.DeleteKeyspaceGroupByID(1) + re.NoError(err) + checkMarker() +} + func (suite *keyspaceGroupTestSuite) TestKeyspaceAssignment() { re := suite.Require() diff --git a/pkg/tso/keyspace_group_manager.go b/pkg/tso/keyspace_group_manager.go index 9ed9b40fefd..8f1afcf1b40 100644 --- a/pkg/tso/keyspace_group_manager.go +++ b/pkg/tso/keyspace_group_manager.go @@ -534,9 +534,25 @@ func (kgm *KeyspaceGroupManager) InitializeTSOServerWatchLoop() error { // membership/distribution metadata. // Key: /pd/{cluster_id}/tso/keyspace_groups/membership/{group} // Value: endpoint.KeyspaceGroup +// Revision marker: /pd/{cluster_id}/tso/keyspace_groups/revision func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { defaultKGConfigured := false + maxLoadedModRevision := uint64(0) + preEventsFn := func([]*clientv3.Event) error { + maxLoadedModRevision = 0 + return nil + } putFn := func(kv *mvccpb.KeyValue) error { + if string(kv.Key) == keypath.KeyspaceGroupRevisionPath() { + failpoint.Inject("SkipKeyspaceWatch", func(val failpoint.Value) { + addr, ok := val.(string) + if ok && addr == kgm.electionNamePrefix { + failpoint.Return(nil) + } + }) + maxLoadedModRevision = max(maxLoadedModRevision, uint64(kv.ModRevision)) + return nil + } group := &endpoint.KeyspaceGroup{} if err := json.Unmarshal(kv.Value, group); err != nil { return errs.ErrJSONUnmarshal.Wrap(err) @@ -547,6 +563,7 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { failpoint.Return(nil) } }) + maxLoadedModRevision = max(maxLoadedModRevision, uint64(kv.ModRevision)) kgm.updateKeyspaceGroup(group) if group.ID == constant.DefaultKeyspaceGroupID { defaultKGConfigured = true @@ -554,6 +571,9 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { return nil } deleteFn := func(kv *mvccpb.KeyValue) error { + if string(kv.Key) == keypath.KeyspaceGroupRevisionPath() { + return nil + } groupID, err := ExtractKeyspaceGroupIDFromPath(kgm.compiledKGMembershipIDRegexp, string(kv.Key)) if err != nil { return err @@ -581,6 +601,8 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { zap.Uint64("new-mod-revision", uint64(last.Kv.ModRevision)), ) } + } else if maxLoadedModRevision > 0 { + kgm.SetModRevision(maxLoadedModRevision) } return nil } @@ -589,25 +611,13 @@ func (kgm *KeyspaceGroupManager) InitializeGroupWatchLoop() error { &kgm.wg, kgm.etcdClient, "keyspace-watcher", - // To keep the consistency with the previous code, we should trim the suffix `/`. - strings.TrimSuffix(keypath.KeyspaceGroupIDPrefix(), "/"), - func([]*clientv3.Event) error { return nil }, + keypath.KeyspaceGroupPrefix(), + preEventsFn, putFn, deleteFn, postEventsFn, true, /* withPrefix */ ) - kgm.groupWatcher.SetLoadRevisionUpdatedCallback(func(revision int64) { - failpoint.Inject("SkipKeyspaceWatch", func(val failpoint.Value) { - addr, ok := val.(string) - if ok && addr == kgm.electionNamePrefix { - failpoint.Return() - } - }) - if revision > 0 { - kgm.SetModRevision(uint64(revision)) - } - }) if kgm.loadFromEtcdMaxRetryTimes > 0 { kgm.groupWatcher.SetLoadRetryTimes(kgm.loadFromEtcdMaxRetryTimes) } diff --git a/pkg/tso/keyspace_group_manager_test.go b/pkg/tso/keyspace_group_manager_test.go index 216dd167578..195efd558ba 100644 --- a/pkg/tso/keyspace_group_manager_test.go +++ b/pkg/tso/keyspace_group_manager_test.go @@ -246,7 +246,10 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsSetsModRevisio []uint32{keyspaceID + 1}, ) re.NoError(err) - deleteResp, err := suite.etcdClient.Delete(suite.ctx, keypath.KeyspaceGroupIDPath(deletedGroupID)) + deleteResp, err := suite.etcdClient.Txn(suite.ctx).Then( + clientv3.OpDelete(keypath.KeyspaceGroupIDPath(deletedGroupID)), + clientv3.OpPut(keypath.KeyspaceGroupRevisionPath(), "1"), + ).Commit() re.NoError(err) deletedRevision := uint64(deleteResp.Header.Revision) re.Greater(deletedRevision, groupRevision) @@ -258,7 +261,51 @@ func (suite *keyspaceGroupManagerTestSuite) TestLoadKeyspaceGroupsSetsModRevisio re.NoError(err) re.NotNil(kg) re.Equal(groupID, loadedGroupID) - re.GreaterOrEqual(loadedRevision, deletedRevision) + re.Equal(deletedRevision, loadedRevision) +} + +func (suite *keyspaceGroupManagerTestSuite) TestSnapshotRevisionRemainsComparableAcrossManagers() { + re := suite.Require() + keypath.SetClusterID(rand.Uint64()) + + cfg1 := suite.createConfig() + cfg2 := suite.createConfig() + mgr1 := suite.newKeyspaceGroupManager(1, cfg1) + mgr2 := suite.newKeyspaceGroupManager(1, cfg2) + defer mgr1.Close() + defer mgr2.Close() + + const ( + groupID = uint32(1) + keyspaceID = uint32(101) + ) + re.NoError(addKeyspaceGroupAssignment( + suite.ctx, + suite.etcdClient, + groupID, + []string{ + mgr1.tsoServiceID.ServiceAddr, + mgr2.tsoServiceID.ServiceAddr, + }, + []int{ + mcs.DefaultKeyspaceGroupReplicaPriority, + mcs.DefaultKeyspaceGroupReplicaPriority, + }, + []uint32{keyspaceID}, + )) + re.NoError(mgr1.Initialize()) + + _, _, _, oldRevision, err := mgr1.FindGroupByKeyspaceID(keyspaceID) + re.NoError(err) + re.NotZero(oldRevision) + + _, err = suite.etcdClient.Put(suite.ctx, "/unrelated/revision-gap", "1") + re.NoError(err) + + re.NoError(mgr2.Initialize()) + _, _, _, newRevision, err := mgr2.FindGroupByKeyspaceID(keyspaceID) + re.NoError(err) + re.Equal(oldRevision, newRevision) } func (suite *keyspaceGroupManagerTestSuite) TestInitialSkipKeyspaceWatchDoesNotAdvanceRevision() { diff --git a/pkg/utils/etcdutil/etcdutil.go b/pkg/utils/etcdutil/etcdutil.go index d288a992b82..7b5e145f4a8 100644 --- a/pkg/utils/etcdutil/etcdutil.go +++ b/pkg/utils/etcdutil/etcdutil.go @@ -374,9 +374,6 @@ type LoopWatcher struct { postEventsFn func([]*clientv3.Event) error // preEventsFn is used to call before handling all events. preEventsFn func([]*clientv3.Event) error - // loadRevisionUpdatedFn is called after a snapshot has been loaded. - loadRevisionUpdatedFn func(int64) - // forceLoadMu is used to ensure two force loads have minimal interval. forceLoadMu syncutil.RWMutex // lastTimeForceLoad is used to record the last time force loading data from etcd. @@ -418,7 +415,6 @@ func NewLoopWatcher( deleteFn: deleteFn, postEventsFn: postEventsFn, preEventsFn: preEventsFn, - loadRevisionUpdatedFn: func(int64) {}, isWithPrefix: isWithPrefix, lastTimeForceLoad: time.Now(), loadRetryTimes: defaultLoadFromEtcdRetryTimes, @@ -714,7 +710,6 @@ func (lw *LoopWatcher) load(ctx context.Context) (nextRevision int64, err error) } // Note: if there are no keys in etcd, the resp.More is false. It also means the load is finished. if !resp.More { - lw.loadRevisionUpdatedFn(snapshotRevision) return snapshotRevision + 1, err } } @@ -773,15 +768,6 @@ func (lw *LoopWatcher) WaitLoad() error { return <-lw.isLoadedCh } -// SetLoadRevisionUpdatedCallback sets a callback that is invoked after a snapshot has been loaded. -func (lw *LoopWatcher) SetLoadRevisionUpdatedCallback(fn func(int64)) { - if fn == nil { - lw.loadRevisionUpdatedFn = func(int64) {} - return - } - lw.loadRevisionUpdatedFn = fn -} - // SetLoadRetryTimes sets the retry times when loading data from etcd. func (lw *LoopWatcher) SetLoadRetryTimes(times int) { lw.loadRetryTimes = times diff --git a/pkg/utils/etcdutil/etcdutil_test.go b/pkg/utils/etcdutil/etcdutil_test.go index cfecee6ab52..631df07d59c 100644 --- a/pkg/utils/etcdutil/etcdutil_test.go +++ b/pkg/utils/etcdutil/etcdutil_test.go @@ -502,14 +502,11 @@ func (suite *loopWatcherTestSuite) TestLoadUsesSingleSnapshotRevision() { func([]*clientv3.Event) error { return nil }, true, /* withPrefix */ ) - var loadedRevision atomic.Int64 - watcher.SetLoadRevisionUpdatedCallback(loadedRevision.Store) watcher.SetLoadBatchSize(1) nextRevision, err := watcher.load(ctx) re.NoError(err) re.NoError(putErr) re.Equal(expectedRevision+1, nextRevision) - re.Equal(expectedRevision, loadedRevision.Load()) } func (suite *loopWatcherTestSuite) TestLoadWithLimitChange() { diff --git a/pkg/utils/keypath/absolute_key_path.go b/pkg/utils/keypath/absolute_key_path.go index 28ca4e8dfaf..9cf23304d52 100644 --- a/pkg/utils/keypath/absolute_key_path.go +++ b/pkg/utils/keypath/absolute_key_path.go @@ -92,12 +92,14 @@ const ( minResolvedTSPathFormat = "/pd/%d/raft/min_resolved_ts" // "/pd/{cluster_id}/raft/min_resolved_ts" externalTimestampPathFormat = "/pd/%d/raft/external_timestamp" // "/pd/{cluster_id}/raft/external_timestamp" - keyspaceMetaPrefixFormat = "/pd/%d/keyspaces/meta/" // "/pd/{cluster_id}/keyspaces/meta/" - keyspaceMetaPathFormat = "/pd/%d/keyspaces/meta/%08d" // "/pd/{cluster_id}/keyspaces/meta/{keyspace_id}" - keyspaceIDPathFormat = "/pd/%d/keyspaces/id/%s" // "/pd/{cluster_id}/keyspaces/id/{keyspace_name}" - keyspaceGroupIDPrefixFormat = "/pd/%d/tso/keyspace_groups/membership/" // "/pd/{cluster_id}/tso/keyspace_groups/membership/" - keyspaceGroupIDPathFormat = "/pd/%d/tso/keyspace_groups/membership/%05d" // "/pd/{cluster_id}/tso/keyspace_groups/membership/{group_id}" - keyspaceGroupIDPattern = `tso/keyspace_groups/membership/(\d{5})$` + keyspaceMetaPrefixFormat = "/pd/%d/keyspaces/meta/" // "/pd/{cluster_id}/keyspaces/meta/" + keyspaceMetaPathFormat = "/pd/%d/keyspaces/meta/%08d" // "/pd/{cluster_id}/keyspaces/meta/{keyspace_id}" + keyspaceIDPathFormat = "/pd/%d/keyspaces/id/%s" // "/pd/{cluster_id}/keyspaces/id/{keyspace_name}" + keyspaceGroupPrefixFormat = "/pd/%d/tso/keyspace_groups/" // "/pd/{cluster_id}/tso/keyspace_groups/" + keyspaceGroupIDPrefixFormat = "/pd/%d/tso/keyspace_groups/membership/" // "/pd/{cluster_id}/tso/keyspace_groups/membership/" + keyspaceGroupIDPathFormat = "/pd/%d/tso/keyspace_groups/membership/%05d" // "/pd/{cluster_id}/tso/keyspace_groups/membership/{group_id}" + keyspaceGroupRevisionPathFormat = "/pd/%d/tso/keyspace_groups/revision" // "/pd/{cluster_id}/tso/keyspace_groups/revision" + keyspaceGroupIDPattern = `tso/keyspace_groups/membership/(\d{5})$` servicePathFormat = "/ms/%d/%s/registry/" // "/ms/{cluster_id}/{service_name}/registry/" registryPathFormat = "/ms/%d/%s/registry/%s" // "/ms/{cluster_id}/{service_name}/registry/{service_addr}" @@ -231,6 +233,11 @@ func KeyspaceIDPath(name string) string { return fmt.Sprintf(keyspaceIDPathFormat, ClusterID(), name) } +// KeyspaceGroupPrefix returns the prefix of keyspace group metadata. +func KeyspaceGroupPrefix() string { + return fmt.Sprintf(keyspaceGroupPrefixFormat, ClusterID()) +} + // KeyspaceGroupIDPrefix returns the prefix of keyspace group id. func KeyspaceGroupIDPrefix() string { return fmt.Sprintf(keyspaceGroupIDPrefixFormat, ClusterID()) @@ -241,6 +248,11 @@ func KeyspaceGroupIDPath(id uint32) string { return fmt.Sprintf(keyspaceGroupIDPathFormat, ClusterID(), id) } +// KeyspaceGroupRevisionPath returns the path of the durable keyspace group revision marker. +func KeyspaceGroupRevisionPath() string { + return fmt.Sprintf(keyspaceGroupRevisionPathFormat, ClusterID()) +} + // GetCompiledKeyspaceGroupIDRegexp returns the compiled regular expression for matching keyspace group id. func GetCompiledKeyspaceGroupIDRegexp() *regexp.Regexp { return regexp.MustCompile(keyspaceGroupIDPattern) diff --git a/tests/integrations/realcluster/etcd_key_test.go b/tests/integrations/realcluster/etcd_key_test.go index 77007b439bc..fb894fbd11e 100644 --- a/tests/integrations/realcluster/etcd_key_test.go +++ b/tests/integrations/realcluster/etcd_key_test.go @@ -78,6 +78,7 @@ var ( "/pd//scheduler_config/evict-stopping-store-scheduler", "/pd//timestamp", "/pd//tso/keyspace_groups/membership/", // ms + "/pd//tso/keyspace_groups/revision", // ms "/pd/cluster_id", } // The keys that prefix is `/ms`. @@ -95,6 +96,7 @@ var ( // These keys with `/pd` are only in `ms` mode. pdMSKeys = []string{ "/pd//tso/keyspace_groups/membership/", + "/pd//tso/keyspace_groups/revision", } )