From 256f11d3b7f6d6fb1585337d28c25d4654121ce5 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 13 Jul 2026 11:33:43 +0800 Subject: [PATCH 1/5] server: clean up microservice metadata in PD mode Signed-off-by: Ryan Leung --- server/microservice_cleanup.go | 256 ++++++++++++++++++++++ server/microservice_cleanup_test.go | 220 +++++++++++++++++++ server/server.go | 1 + tests/integrations/mcs/tso/server_test.go | 133 +++++++++++ 4 files changed, 610 insertions(+) create mode 100644 server/microservice_cleanup.go create mode 100644 server/microservice_cleanup_test.go diff --git a/server/microservice_cleanup.go b/server/microservice_cleanup.go new file mode 100644 index 00000000000..3df4838ae2f --- /dev/null +++ b/server/microservice_cleanup.go @@ -0,0 +1,256 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "context" + "fmt" + "strconv" + "time" + + clientv3 "go.etcd.io/etcd/client/v3" + "go.uber.org/zap" + + "github.com/pingcap/errors" + "github.com/pingcap/failpoint" + "github.com/pingcap/log" + + "github.com/tikv/pd/pkg/errs" + "github.com/tikv/pd/pkg/keyspace" + "github.com/tikv/pd/pkg/keyspace/constant" + mcs "github.com/tikv/pd/pkg/mcs/utils/constant" + "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/utils/logutil" +) + +// microserviceMetadataCleanupDelay lets clients observe the service mode change +// before stale microservice metadata disappears from etcd. +const ( + microserviceMetadataCleanupDelay = 10 * time.Second + microserviceMetadataCleanupRetryInterval = 5 * time.Second +) + +func (s *Server) scheduleMicroserviceMetadataCleanup(ctx context.Context) { + if s.IsKeyspaceGroupEnabled() { + return + } + s.serverLoopWg.Add(1) + go func() { + defer logutil.LogPanic() + defer s.serverLoopWg.Done() + if !waitMicroserviceMetadataCleanupDelay(ctx) { + return + } + for { + if s.member != nil && !s.member.IsServing() { + return + } + if err := s.cleanupMicroserviceMetadataInPDMode(ctx); err == nil { + return + } else { + log.Warn("failed to clean up microservice metadata in PD mode, retry later", + errs.ZapError(err)) + } + timer := time.NewTimer(microserviceMetadataCleanupRetryInterval) + select { + case <-ctx.Done(): + timer.Stop() + return + case <-timer.C: + } + } + }() +} + +func waitMicroserviceMetadataCleanupDelay(ctx context.Context) bool { + delay := microserviceMetadataCleanupDelay + failpoint.Inject("skipMicroserviceMetadataCleanupDelay", func() { + delay = 0 + }) + if delay == 0 { + return true + } + timer := time.NewTimer(delay) + defer timer.Stop() + select { + case <-ctx.Done(): + return false + case <-timer.C: + return true + } +} + +func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error { + if s.IsKeyspaceGroupEnabled() { + return nil + } + start := time.Now() + + defaultGroupExists, err := s.checkDefaultOnlyTSOKeyspaceGroup() + if err != nil { + return err + } + cleanedKeyspaces, err := s.cleanupDefaultTSOKeyspaceGroupConfig(ctx) + if err != nil { + return err + } + deletedDefaultGroup, err := s.deleteDefaultTSOKeyspaceGroup(ctx) + if err != nil { + return err + } + deletedMicroserviceKeys, err := s.deleteMicroserviceEtcdKeys(ctx) + if err != nil { + return err + } + + if defaultGroupExists || deletedDefaultGroup || cleanedKeyspaces > 0 || deletedMicroserviceKeys > 0 { + log.Info("cleaned up microservice metadata in PD mode", + zap.Bool("default-keyspace-group-exists", defaultGroupExists), + zap.Bool("deleted-default-keyspace-group", deletedDefaultGroup), + zap.Int("cleaned-keyspace-configs", cleanedKeyspaces), + zap.Int64("deleted-microservice-keys", deletedMicroserviceKeys), + zap.Duration("cost", time.Since(start))) + } + return nil +} + +func (s *Server) checkDefaultOnlyTSOKeyspaceGroup() (bool, error) { + groups, err := s.storage.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) + if err != nil { + return false, err + } + defaultGroupExists := false + for _, group := range groups { + if group == nil { + continue + } + if group.ID != constant.DefaultKeyspaceGroupID { + return false, errors.Errorf("found non-default TSO keyspace group %d when cleaning up PD mode microservice metadata", group.ID) + } + if group.IsSplitting() { + return false, errors.Errorf("default TSO keyspace group is splitting when cleaning up PD mode microservice metadata") + } + if group.IsMerging() { + return false, errors.Errorf("default TSO keyspace group is merging when cleaning up PD mode microservice metadata") + } + defaultGroupExists = true + } + return defaultGroupExists, nil +} + +func (s *Server) cleanupDefaultTSOKeyspaceGroupConfig(ctx context.Context) (int, error) { + var ( + cleaned int + startID uint32 = constant.StartKeyspaceID + ) + for { + var ( + batchCleaned int + keyspaceIDs []uint32 + loaded int + ) + err := s.storage.RunInTxn(ctx, func(txn kv.Txn) error { + metas, err := s.storage.LoadRangeKeyspace(txn, startID, etcdutil.MaxEtcdTxnOps) + if err != nil { + return err + } + loaded = len(metas) + keyspaceIDs = make([]uint32, 0, len(metas)) + for _, meta := range metas { + if meta == nil { + continue + } + keyspaceIDs = append(keyspaceIDs, meta.GetId()) + if meta.Config == nil { + continue + } + groupIDText, ok := meta.Config[keyspace.TSOKeyspaceGroupIDKey] + if !ok { + continue + } + groupID, err := strconv.ParseUint(groupIDText, 10, 32) + if err != nil { + return errors.Errorf("keyspace %d has invalid TSO keyspace group ID %q", meta.GetId(), groupIDText) + } + if uint32(groupID) != constant.DefaultKeyspaceGroupID { + return errors.Errorf("keyspace %d is assigned to non-default TSO keyspace group %d", meta.GetId(), groupID) + } + delete(meta.Config, keyspace.TSOKeyspaceGroupIDKey) + if err := s.storage.SaveKeyspaceMeta(txn, meta); err != nil { + return err + } + batchCleaned++ + } + return nil + }) + if err != nil { + return cleaned, err + } + cleaned += batchCleaned + if loaded < etcdutil.MaxEtcdTxnOps { + break + } + if len(keyspaceIDs) == 0 { + break + } + lastID := keyspaceIDs[len(keyspaceIDs)-1] + if lastID == ^uint32(0) { + break + } + startID = lastID + 1 + } + return cleaned, nil +} + +func (s *Server) deleteDefaultTSOKeyspaceGroup(ctx context.Context) (bool, error) { + deleted := false + err := s.storage.RunInTxn(ctx, func(txn kv.Txn) error { + group, err := s.storage.LoadKeyspaceGroup(txn, constant.DefaultKeyspaceGroupID) + if err != nil || group == nil { + return err + } + if group.IsSplitting() { + return errors.Errorf("default TSO keyspace group is splitting when deleting PD mode microservice metadata") + } + if group.IsMerging() { + return errors.Errorf("default TSO keyspace group is merging when deleting PD mode microservice metadata") + } + if err := s.storage.DeleteKeyspaceGroup(txn, constant.DefaultKeyspaceGroupID); err != nil { + return err + } + deleted = true + return nil + }) + return deleted, err +} + +func (s *Server) deleteMicroserviceEtcdKeys(ctx context.Context) (int64, error) { + if s.client == nil { + return 0, nil + } + ctx, cancel := context.WithTimeout(ctx, etcdutil.DefaultRequestTimeout) + defer cancel() + resp, err := s.client.Delete(ctx, microserviceEtcdPrefix(), clientv3.WithPrefix()) + if err != nil { + return 0, errs.ErrEtcdKVDelete.Wrap(err).GenWithStackByCause() + } + return resp.Deleted, nil +} + +func microserviceEtcdPrefix() string { + return fmt.Sprintf("%s/%d/", mcs.MicroserviceRootPath, keypath.ClusterID()) +} diff --git a/server/microservice_cleanup_test.go b/server/microservice_cleanup_test.go new file mode 100644 index 00000000000..e22a548fcae --- /dev/null +++ b/server/microservice_cleanup_test.go @@ -0,0 +1,220 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package server + +import ( + "context" + "strconv" + "testing" + "time" + + "github.com/stretchr/testify/require" + clientv3 "go.etcd.io/etcd/client/v3" + + "github.com/pingcap/failpoint" + "github.com/pingcap/kvproto/pkg/keyspacepb" + + "github.com/tikv/pd/pkg/keyspace" + "github.com/tikv/pd/pkg/keyspace/constant" + mcs "github.com/tikv/pd/pkg/mcs/utils/constant" + "github.com/tikv/pd/pkg/storage" + "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" +) + +func TestCleanupMicroserviceMetadataInPDMode(t *testing.T) { + re := require.New(t) + ctx := context.Background() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + keypath.SetClusterID(12345) + defer keypath.ResetClusterID() + + store := storage.NewStorageWithEtcdBackend(client) + svr := &Server{storage: store, client: client} + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + if err := store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + Keyspaces: []uint32{1}, + }); err != nil { + return err + } + if err := store.SaveKeyspaceMeta(txn, &keyspacepb.KeyspaceMeta{ + Id: 1, + Name: "keyspace-1", + Config: map[string]string{ + keyspace.TSOKeyspaceGroupIDKey: strconv.FormatUint(uint64(constant.DefaultKeyspaceGroupID), 10), + "gc_life_time": "10m", + }, + }); err != nil { + return err + } + return store.SaveKeyspaceMeta(txn, &keyspacepb.KeyspaceMeta{ + Id: 2, + Name: "keyspace-2", + Config: map[string]string{ + "custom": "value", + }, + }) + })) + _, err := client.Put(ctx, keypath.RegistryPath(mcs.TSOServiceName, "127.0.0.1:3379"), "tso") + re.NoError(err) + _, err = client.Put(ctx, keypath.RegistryPath(mcs.SchedulingServiceName, "127.0.0.1:3379"), "scheduling") + re.NoError(err) + _, err = client.Put(ctx, keypath.ElectionPath(&keypath.MsParam{ + ServiceName: mcs.TSOServiceName, + GroupID: constant.DefaultKeyspaceGroupID, + }), "primary") + re.NoError(err) + _, err = client.Put(ctx, keypath.TimestampPath(constant.DefaultKeyspaceGroupID), "timestamp") + re.NoError(err) + + re.NoError(svr.cleanupMicroserviceMetadataInPDMode(ctx)) + re.NoError(svr.cleanupMicroserviceMetadataInPDMode(ctx)) + + groups, err := store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) + re.NoError(err) + re.Empty(groups) + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + meta, err := store.LoadKeyspaceMeta(txn, 1) + re.NoError(err) + re.NotNil(meta) + re.NotContains(meta.GetConfig(), keyspace.TSOKeyspaceGroupIDKey) + re.Equal("10m", meta.GetConfig()["gc_life_time"]) + + meta, err = store.LoadKeyspaceMeta(txn, 2) + re.NoError(err) + re.NotNil(meta) + re.Equal("value", meta.GetConfig()["custom"]) + return nil + })) + resp, err := etcdutil.EtcdKVGet(client, microserviceEtcdPrefix(), clientv3.WithPrefix()) + re.NoError(err) + re.Empty(resp.Kvs) + resp, err = etcdutil.EtcdKVGet(client, keypath.TimestampPath(constant.DefaultKeyspaceGroupID)) + re.NoError(err) + re.Len(resp.Kvs, 1) + re.Equal("timestamp", string(resp.Kvs[0].Value)) +} + +func TestScheduleMicroserviceMetadataCleanupReturnsImmediately(t *testing.T) { + re := require.New(t) + ctx := context.Background() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + keypath.SetClusterID(12348) + defer keypath.ResetClusterID() + + store := storage.NewStorageWithEtcdBackend(client) + svr := &Server{storage: store, client: client} + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + return store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + }) + })) + re.NoError(failpoint.Enable("github.com/tikv/pd/server/skipMicroserviceMetadataCleanupDelay", "return(true)")) + re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/utils/etcdutil/SlowEtcdKVGet", "return(1)")) + t.Cleanup(func() { + _ = failpoint.Disable("github.com/tikv/pd/server/skipMicroserviceMetadataCleanupDelay") + _ = failpoint.Disable("github.com/tikv/pd/pkg/utils/etcdutil/SlowEtcdKVGet") + }) + + start := time.Now() + svr.scheduleMicroserviceMetadataCleanup(ctx) + re.Less(time.Since(start), 200*time.Millisecond) + re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/utils/etcdutil/SlowEtcdKVGet")) + svr.serverLoopWg.Wait() +} + +func TestCleanupMicroserviceMetadataInPDModeRejectsNonDefaultGroup(t *testing.T) { + re := require.New(t) + ctx := context.Background() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + keypath.SetClusterID(12346) + defer keypath.ResetClusterID() + + store := storage.NewStorageWithEtcdBackend(client) + svr := &Server{storage: store, client: client} + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + if err := store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + }); err != nil { + return err + } + return store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: 1, + UserKind: endpoint.Standard.String(), + }) + })) + _, err := client.Put(ctx, keypath.RegistryPath(mcs.TSOServiceName, "127.0.0.1:3379"), "tso") + re.NoError(err) + + err = svr.cleanupMicroserviceMetadataInPDMode(ctx) + re.ErrorContains(err, "non-default TSO keyspace group 1") + groups, err := store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) + re.NoError(err) + re.Len(groups, 2) + resp, err := etcdutil.EtcdKVGet(client, microserviceEtcdPrefix(), clientv3.WithPrefix()) + re.NoError(err) + re.Len(resp.Kvs, 1) +} + +func TestCleanupMicroserviceMetadataInPDModeRejectsNonDefaultKeyspaceAssignment(t *testing.T) { + re := require.New(t) + ctx := context.Background() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + keypath.SetClusterID(12347) + defer keypath.ResetClusterID() + + store := storage.NewStorageWithEtcdBackend(client) + svr := &Server{storage: store, client: client} + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + if err := store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + }); err != nil { + return err + } + return store.SaveKeyspaceMeta(txn, &keyspacepb.KeyspaceMeta{ + Id: 1, + Name: "keyspace-1", + Config: map[string]string{ + keyspace.TSOKeyspaceGroupIDKey: "1", + }, + }) + })) + _, err := client.Put(ctx, keypath.RegistryPath(mcs.TSOServiceName, "127.0.0.1:3379"), "tso") + re.NoError(err) + + err = svr.cleanupMicroserviceMetadataInPDMode(ctx) + re.ErrorContains(err, "keyspace 1 is assigned to non-default TSO keyspace group 1") + var group *endpoint.KeyspaceGroup + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + var err error + group, err = store.LoadKeyspaceGroup(txn, constant.DefaultKeyspaceGroupID) + return err + })) + re.NotNil(group) + resp, err := etcdutil.EtcdKVGet(client, microserviceEtcdPrefix(), clientv3.WithPrefix()) + re.NoError(err) + re.Len(resp.Kvs, 1) +} diff --git a/server/server.go b/server/server.go index fe0a2a37d1c..a3d9dd71ad8 100644 --- a/server/server.go +++ b/server/server.go @@ -2049,6 +2049,7 @@ func (s *Server) campaignLeader() { zap.String("leader-name", s.Name()), zap.Duration("total-cost", totalDuration), zap.Duration("cost", enableLeaderDuration)) + s.scheduleMicroserviceMetadataCleanup(ctx) leaderTicker := time.NewTicker(mcs.LeaderTickInterval) defer leaderTicker.Stop() diff --git a/tests/integrations/mcs/tso/server_test.go b/tests/integrations/mcs/tso/server_test.go index d20fd7642af..d569ea41bfa 100644 --- a/tests/integrations/mcs/tso/server_test.go +++ b/tests/integrations/mcs/tso/server_test.go @@ -34,11 +34,13 @@ import ( "github.com/pingcap/failpoint" "github.com/pingcap/kvproto/pkg/metapb" + "github.com/pingcap/kvproto/pkg/pdpb" pd "github.com/tikv/pd/client" "github.com/tikv/pd/client/opt" "github.com/tikv/pd/client/pkg/caller" "github.com/tikv/pd/pkg/core" + "github.com/tikv/pd/pkg/keyspace" "github.com/tikv/pd/pkg/keyspace/constant" "github.com/tikv/pd/pkg/mcs/discovery" tso "github.com/tikv/pd/pkg/mcs/tso/server" @@ -46,6 +48,7 @@ import ( mcs "github.com/tikv/pd/pkg/mcs/utils/constant" "github.com/tikv/pd/pkg/storage/endpoint" "github.com/tikv/pd/pkg/utils/etcdutil" + "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" @@ -769,6 +772,136 @@ func TestTSOServiceSwitch(t *testing.T) { re.NoError(failpoint.Disable("github.com/tikv/pd/client/servicediscovery/fastUpdateServiceMode")) } +func TestPDModeSwitchBetweenMicroserviceAndMonolithMultipleTimes(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + re.NoError(failpoint.Enable("github.com/tikv/pd/server/skipMicroserviceMetadataCleanupDelay", "return(true)")) + defer func() { + re.NoError(failpoint.Disable("github.com/tikv/pd/server/skipMicroserviceMetadataCleanupDelay")) + }() + + tc, err := tests.NewTestClusterWithKeyspaceGroup(ctx, 1, func(conf *config.Config, _ string) { + conf.Microservice.EnableTSODynamicSwitching = false + }) + re.NoError(err) + defer tc.Destroy() + + re.NoError(tc.RunInitialServers()) + pdServer := tc.GetServer(tc.WaitLeader()) + re.NotNil(pdServer) + re.NoError(pdServer.BootstrapCluster()) + + const switchRounds = 2 + for range switchRounds { + waitDefaultKeyspaceGroupReady(re, pdServer) + tsoCluster, err := tests.NewTestTSOCluster(ctx, 1, pdServer.GetAddr()) + re.NoError(err) + waitTSOServiceReady(re, tsoCluster) + checkMicroserviceTSOAvailable(ctx, re, pdServer) + tsoCluster.Destroy() + + pdServer = restartPDForServiceMode(ctx, re, tc, pdServer, nil) + checkPDTSOAvailable(ctx, re, pdServer) + waitMicroserviceMetadataCleaned(re, pdServer) + + pdServer = restartPDForServiceMode(ctx, re, tc, pdServer, []string{mcs.PDServiceName}) + } + + waitDefaultKeyspaceGroupReady(re, pdServer) + tsoCluster, err := tests.NewTestTSOCluster(ctx, 1, pdServer.GetAddr()) + re.NoError(err) + defer tsoCluster.Destroy() + waitTSOServiceReady(re, tsoCluster) + checkMicroserviceTSOAvailable(ctx, re, pdServer) +} + +func restartPDForServiceMode( + ctx context.Context, + re *require.Assertions, + tc *tests.TestCluster, + oldServer *tests.TestServer, + services []string, +) *tests.TestServer { + cfg := oldServer.GetConfig() + serverName := cfg.Name + re.NoError(oldServer.Stop()) + + newServer, err := tests.NewTestServer(ctx, cfg, services) + re.NoError(err) + tc.GetServers()[serverName] = newServer + re.NoError(newServer.Run()) + re.True(newServer.WaitLeader()) + return newServer +} + +func waitDefaultKeyspaceGroupReady(re *require.Assertions, pdServer *tests.TestServer) { + testutil.Eventually(re, func() bool { + resp, err := etcdutil.EtcdKVGet(pdServer.GetEtcdClient(), keypath.KeyspaceGroupIDPath(constant.DefaultKeyspaceGroupID)) + return err == nil && len(resp.Kvs) == 1 + }, testutil.WithWaitFor(10*time.Second), testutil.WithTickInterval(100*time.Millisecond)) +} + +func waitTSOServiceReady(re *require.Assertions, tsoCluster *tests.TestTSOCluster) { + primary := tsoCluster.WaitForDefaultPrimaryServing(re) + testutil.Eventually(re, func() bool { + resp, err := tests.TestDialClient.Get(primary.GetAddr() + tsoapi.APIPathPrefix + "/health") + if err != nil { + return false + } + defer resp.Body.Close() + return resp.StatusCode == http.StatusOK + }, testutil.WithWaitFor(10*time.Second), testutil.WithTickInterval(100*time.Millisecond)) +} + +func checkPDTSOAvailable(ctx context.Context, re *require.Assertions, pdServer *tests.TestServer) { + cli := utils.SetupClientWithAPIContext(ctx, re, pd.NewAPIContextV1(), []string{pdServer.GetAddr()}) + defer cli.Close() + physical, logical, err := cli.GetTS(ctx) + re.NoError(err) + re.NotZero(tsoutil.ComposeTS(physical, logical)) +} + +func checkMicroserviceTSOAvailable(ctx context.Context, re *require.Assertions, pdServer *tests.TestServer) { + addr := strings.TrimPrefix(pdServer.GetAddr(), "http://") + conn, err := grpc.Dial(addr, grpc.WithTransportCredentials(insecure.NewCredentials())) //nolint:staticcheck + re.NoError(err) + defer conn.Close() + + resp, err := pdpb.NewPDClient(conn).GetClusterInfo(ctx, &pdpb.GetClusterInfoRequest{}) + re.NoError(err) + re.Equal([]pdpb.ServiceMode{pdpb.ServiceMode_API_SVC_MODE}, resp.GetServiceModes()) + re.NotEmpty(resp.GetTsoUrls()) + + checkPDTSOAvailable(ctx, re, pdServer) +} + +func waitMicroserviceMetadataCleaned(re *require.Assertions, pdServer *tests.TestServer) { + testutil.Eventually(re, func() bool { + keyspaceGroupResp, err := etcdutil.EtcdKVGet(pdServer.GetEtcdClient(), keypath.KeyspaceGroupIDPrefix(), clientv3.WithPrefix()) + if err != nil || len(keyspaceGroupResp.Kvs) != 0 { + return false + } + microserviceResp, err := etcdutil.EtcdKVGet( + pdServer.GetEtcdClient(), + fmt.Sprintf("%s/%d/", mcs.MicroserviceRootPath, keypath.ClusterID()), + clientv3.WithPrefix()) + if err != nil || len(microserviceResp.Kvs) != 0 { + return false + } + keyspaces, err := pdServer.GetKeyspaceManager().LoadRangeKeyspace(constant.StartKeyspaceID, 0) + if err != nil { + return false + } + for _, meta := range keyspaces { + if _, ok := meta.GetConfig()[keyspace.TSOKeyspaceGroupIDKey]; ok { + return false + } + } + return true + }, testutil.WithWaitFor(30*time.Second), testutil.WithTickInterval(100*time.Millisecond)) +} + func checkTSOMonotonic(ctx context.Context, pdClient pd.Client, globalLastTS *uint64, count int) error { for range count { physical, logical, err := pdClient.GetTS(ctx) From e3fe646c7a27b8ee8c395810fa6a00472fdb25c6 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 13 Jul 2026 12:12:24 +0800 Subject: [PATCH 2/5] server: make mode cleanup safe and non-blocking Signed-off-by: Ryan Leung --- server/microservice_cleanup.go | 110 ++++++++++++++-------- server/microservice_cleanup_test.go | 95 ++++++++++++++++--- tests/integrations/mcs/tso/server_test.go | 4 - 3 files changed, 152 insertions(+), 57 deletions(-) diff --git a/server/microservice_cleanup.go b/server/microservice_cleanup.go index 3df4838ae2f..a98274c9d5b 100644 --- a/server/microservice_cleanup.go +++ b/server/microservice_cleanup.go @@ -24,7 +24,6 @@ import ( "go.uber.org/zap" "github.com/pingcap/errors" - "github.com/pingcap/failpoint" "github.com/pingcap/log" "github.com/tikv/pd/pkg/errs" @@ -37,11 +36,12 @@ import ( "github.com/tikv/pd/pkg/utils/logutil" ) -// microserviceMetadataCleanupDelay lets clients observe the service mode change -// before stale microservice metadata disappears from etcd. const ( - microserviceMetadataCleanupDelay = 10 * time.Second microserviceMetadataCleanupRetryInterval = 5 * time.Second + // Each loaded key adds one comparison and each updated key adds one write to + // the etcd transaction. Use half of the operation limit so a full batch stays + // within the transaction limit. + microserviceMetadataCleanupBatchSize = etcdutil.MaxEtcdTxnOps / 2 ) func (s *Server) scheduleMicroserviceMetadataCleanup(ctx context.Context) { @@ -52,19 +52,16 @@ func (s *Server) scheduleMicroserviceMetadataCleanup(ctx context.Context) { go func() { defer logutil.LogPanic() defer s.serverLoopWg.Done() - if !waitMicroserviceMetadataCleanupDelay(ctx) { - return - } for { if s.member != nil && !s.member.IsServing() { return } - if err := s.cleanupMicroserviceMetadataInPDMode(ctx); err == nil { + err := s.cleanupMicroserviceMetadataInPDMode(ctx) + if err == nil { return - } else { - log.Warn("failed to clean up microservice metadata in PD mode, retry later", - errs.ZapError(err)) } + log.Warn("failed to clean up microservice metadata in PD mode, retry later", + errs.ZapError(err)) timer := time.NewTimer(microserviceMetadataCleanupRetryInterval) select { case <-ctx.Done(): @@ -76,24 +73,6 @@ func (s *Server) scheduleMicroserviceMetadataCleanup(ctx context.Context) { }() } -func waitMicroserviceMetadataCleanupDelay(ctx context.Context) bool { - delay := microserviceMetadataCleanupDelay - failpoint.Inject("skipMicroserviceMetadataCleanupDelay", func() { - delay = 0 - }) - if delay == 0 { - return true - } - timer := time.NewTimer(delay) - defer timer.Stop() - select { - case <-ctx.Done(): - return false - case <-timer.C: - return true - } -} - func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error { if s.IsKeyspaceGroupEnabled() { return nil @@ -104,6 +83,9 @@ func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error if err != nil { return err } + if err := s.validateDefaultTSOKeyspaceGroupConfig(ctx); err != nil { + return err + } cleanedKeyspaces, err := s.cleanupDefaultTSOKeyspaceGroupConfig(ctx) if err != nil { return err @@ -128,6 +110,60 @@ func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error return nil } +func (s *Server) validateDefaultTSOKeyspaceGroupConfig(ctx context.Context) error { + startID := constant.StartKeyspaceID + for { + var ( + keyspaceIDs []uint32 + loaded int + ) + err := s.storage.RunInTxn(ctx, func(txn kv.Txn) error { + metas, err := s.storage.LoadRangeKeyspace(txn, startID, etcdutil.MaxEtcdTxnOps) + if err != nil { + return err + } + loaded = len(metas) + keyspaceIDs = make([]uint32, 0, len(metas)) + for _, meta := range metas { + if meta == nil { + continue + } + keyspaceIDs = append(keyspaceIDs, meta.GetId()) + if err := validateDefaultTSOKeyspaceGroupAssignment(meta.GetId(), meta.GetConfig()); err != nil { + return err + } + } + return nil + }) + if err != nil { + return err + } + if loaded < etcdutil.MaxEtcdTxnOps || len(keyspaceIDs) == 0 { + return nil + } + lastID := keyspaceIDs[len(keyspaceIDs)-1] + if lastID == ^uint32(0) { + return nil + } + startID = lastID + 1 + } +} + +func validateDefaultTSOKeyspaceGroupAssignment(keyspaceID uint32, config map[string]string) error { + groupIDText, ok := config[keyspace.TSOKeyspaceGroupIDKey] + if !ok { + return nil + } + groupID, err := strconv.ParseUint(groupIDText, 10, 32) + if err != nil { + return errors.Errorf("keyspace %d has invalid TSO keyspace group ID %q", keyspaceID, groupIDText) + } + if groupID != uint64(constant.DefaultKeyspaceGroupID) { + return errors.Errorf("keyspace %d is assigned to non-default TSO keyspace group %d", keyspaceID, groupID) + } + return nil +} + func (s *Server) checkDefaultOnlyTSOKeyspaceGroup() (bool, error) { groups, err := s.storage.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) if err != nil { @@ -155,7 +191,7 @@ func (s *Server) checkDefaultOnlyTSOKeyspaceGroup() (bool, error) { func (s *Server) cleanupDefaultTSOKeyspaceGroupConfig(ctx context.Context) (int, error) { var ( cleaned int - startID uint32 = constant.StartKeyspaceID + startID = constant.StartKeyspaceID ) for { var ( @@ -164,7 +200,7 @@ func (s *Server) cleanupDefaultTSOKeyspaceGroupConfig(ctx context.Context) (int, loaded int ) err := s.storage.RunInTxn(ctx, func(txn kv.Txn) error { - metas, err := s.storage.LoadRangeKeyspace(txn, startID, etcdutil.MaxEtcdTxnOps) + metas, err := s.storage.LoadRangeKeyspace(txn, startID, microserviceMetadataCleanupBatchSize) if err != nil { return err } @@ -178,16 +214,12 @@ func (s *Server) cleanupDefaultTSOKeyspaceGroupConfig(ctx context.Context) (int, if meta.Config == nil { continue } - groupIDText, ok := meta.Config[keyspace.TSOKeyspaceGroupIDKey] + _, ok := meta.Config[keyspace.TSOKeyspaceGroupIDKey] if !ok { continue } - groupID, err := strconv.ParseUint(groupIDText, 10, 32) - if err != nil { - return errors.Errorf("keyspace %d has invalid TSO keyspace group ID %q", meta.GetId(), groupIDText) - } - if uint32(groupID) != constant.DefaultKeyspaceGroupID { - return errors.Errorf("keyspace %d is assigned to non-default TSO keyspace group %d", meta.GetId(), groupID) + if err := validateDefaultTSOKeyspaceGroupAssignment(meta.GetId(), meta.GetConfig()); err != nil { + return err } delete(meta.Config, keyspace.TSOKeyspaceGroupIDKey) if err := s.storage.SaveKeyspaceMeta(txn, meta); err != nil { @@ -201,7 +233,7 @@ func (s *Server) cleanupDefaultTSOKeyspaceGroupConfig(ctx context.Context) (int, return cleaned, err } cleaned += batchCleaned - if loaded < etcdutil.MaxEtcdTxnOps { + if loaded < microserviceMetadataCleanupBatchSize { break } if len(keyspaceIDs) == 0 { diff --git a/server/microservice_cleanup_test.go b/server/microservice_cleanup_test.go index e22a548fcae..1f01d752cbf 100644 --- a/server/microservice_cleanup_test.go +++ b/server/microservice_cleanup_test.go @@ -128,10 +128,8 @@ func TestScheduleMicroserviceMetadataCleanupReturnsImmediately(t *testing.T) { UserKind: endpoint.Basic.String(), }) })) - re.NoError(failpoint.Enable("github.com/tikv/pd/server/skipMicroserviceMetadataCleanupDelay", "return(true)")) re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/utils/etcdutil/SlowEtcdKVGet", "return(1)")) t.Cleanup(func() { - _ = failpoint.Disable("github.com/tikv/pd/server/skipMicroserviceMetadataCleanupDelay") _ = failpoint.Disable("github.com/tikv/pd/pkg/utils/etcdutil/SlowEtcdKVGet") }) @@ -188,33 +186,102 @@ func TestCleanupMicroserviceMetadataInPDModeRejectsNonDefaultKeyspaceAssignment( store := storage.NewStorageWithEtcdBackend(client) svr := &Server{storage: store, client: client} re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { - if err := store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + return store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ ID: constant.DefaultKeyspaceGroupID, UserKind: endpoint.Basic.String(), - }); err != nil { - return err - } - return store.SaveKeyspaceMeta(txn, &keyspacepb.KeyspaceMeta{ - Id: 1, - Name: "keyspace-1", - Config: map[string]string{ - keyspace.TSOKeyspaceGroupIDKey: "1", - }, }) })) + metas := make([]*keyspacepb.KeyspaceMeta, 0, etcdutil.MaxEtcdTxnOps+1) + for id := uint32(1); id <= etcdutil.MaxEtcdTxnOps; id++ { + metas = append(metas, newKeyspaceMetaWithTSOGroup(id, "0")) + } + invalidID := uint32(etcdutil.MaxEtcdTxnOps + 1) + metas = append(metas, newKeyspaceMetaWithTSOGroup(invalidID, "1")) + saveKeyspaceMetas(ctx, re, store, metas) _, err := client.Put(ctx, keypath.RegistryPath(mcs.TSOServiceName, "127.0.0.1:3379"), "tso") re.NoError(err) err = svr.cleanupMicroserviceMetadataInPDMode(ctx) - re.ErrorContains(err, "keyspace 1 is assigned to non-default TSO keyspace group 1") + re.ErrorContains(err, "keyspace "+strconv.FormatUint(uint64(invalidID), 10)+" is assigned to non-default TSO keyspace group 1") var group *endpoint.KeyspaceGroup re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { var err error group, err = store.LoadKeyspaceGroup(txn, constant.DefaultKeyspaceGroupID) - return err + if err != nil { + return err + } + meta, err := store.LoadKeyspaceMeta(txn, 1) + re.NoError(err) + re.Equal("0", meta.GetConfig()[keyspace.TSOKeyspaceGroupIDKey]) + return nil })) re.NotNil(group) resp, err := etcdutil.EtcdKVGet(client, microserviceEtcdPrefix(), clientv3.WithPrefix()) re.NoError(err) re.Len(resp.Kvs, 1) } + +func TestCleanupMicroserviceMetadataInPDModeWithLargeKeyspaceBatch(t *testing.T) { + re := require.New(t) + ctx := context.Background() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + keypath.SetClusterID(12349) + defer keypath.ResetClusterID() + + store := storage.NewStorageWithEtcdBackend(client) + svr := &Server{storage: store, client: client} + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + return store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + }) + })) + metas := make([]*keyspacepb.KeyspaceMeta, 0, etcdutil.MaxEtcdTxnOps) + for id := uint32(1); id <= etcdutil.MaxEtcdTxnOps; id++ { + metas = append(metas, newKeyspaceMetaWithTSOGroup(id, "0")) + } + saveKeyspaceMetas(ctx, re, store, metas) + + re.NoError(svr.cleanupMicroserviceMetadataInPDMode(ctx)) + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + loaded, err := store.LoadRangeKeyspace(txn, constant.StartKeyspaceID, 0) + if err != nil { + return err + } + re.Len(loaded, etcdutil.MaxEtcdTxnOps) + for _, meta := range loaded { + re.NotContains(meta.GetConfig(), keyspace.TSOKeyspaceGroupIDKey) + } + return nil + })) +} + +func newKeyspaceMetaWithTSOGroup(id uint32, groupID string) *keyspacepb.KeyspaceMeta { + return &keyspacepb.KeyspaceMeta{ + Id: id, + Name: "keyspace-" + strconv.FormatUint(uint64(id), 10), + Config: map[string]string{ + keyspace.TSOKeyspaceGroupIDKey: groupID, + }, + } +} + +func saveKeyspaceMetas( + ctx context.Context, + re *require.Assertions, + store storage.Storage, + metas []*keyspacepb.KeyspaceMeta, +) { + for start := 0; start < len(metas); start += microserviceMetadataCleanupBatchSize { + end := min(start+microserviceMetadataCleanupBatchSize, len(metas)) + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + for _, meta := range metas[start:end] { + if err := store.SaveKeyspaceMeta(txn, meta); err != nil { + return err + } + } + return nil + })) + } +} diff --git a/tests/integrations/mcs/tso/server_test.go b/tests/integrations/mcs/tso/server_test.go index d569ea41bfa..eeec9f8ff25 100644 --- a/tests/integrations/mcs/tso/server_test.go +++ b/tests/integrations/mcs/tso/server_test.go @@ -776,10 +776,6 @@ func TestPDModeSwitchBetweenMicroserviceAndMonolithMultipleTimes(t *testing.T) { re := require.New(t) ctx, cancel := context.WithCancel(context.Background()) defer cancel() - re.NoError(failpoint.Enable("github.com/tikv/pd/server/skipMicroserviceMetadataCleanupDelay", "return(true)")) - defer func() { - re.NoError(failpoint.Disable("github.com/tikv/pd/server/skipMicroserviceMetadataCleanupDelay")) - }() tc, err := tests.NewTestClusterWithKeyspaceGroup(ctx, 1, func(conf *config.Config, _ string) { conf.Microservice.EnableTSODynamicSwitching = false From 3509b605d38cb3753184a938d17e17d3cc3e62ca Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 13 Jul 2026 12:21:50 +0800 Subject: [PATCH 3/5] server: simplify mode cleanup flow Signed-off-by: Ryan Leung --- server/microservice_cleanup.go | 185 +++++++++++++-------------------- 1 file changed, 74 insertions(+), 111 deletions(-) diff --git a/server/microservice_cleanup.go b/server/microservice_cleanup.go index a98274c9d5b..4045cb1cc7a 100644 --- a/server/microservice_cleanup.go +++ b/server/microservice_cleanup.go @@ -24,6 +24,7 @@ import ( "go.uber.org/zap" "github.com/pingcap/errors" + "github.com/pingcap/kvproto/pkg/keyspacepb" "github.com/pingcap/log" "github.com/tikv/pd/pkg/errs" @@ -79,11 +80,7 @@ func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error } start := time.Now() - defaultGroupExists, err := s.checkDefaultOnlyTSOKeyspaceGroup() - if err != nil { - return err - } - if err := s.validateDefaultTSOKeyspaceGroupConfig(ctx); err != nil { + if err := s.validateMicroserviceMetadataCleanup(ctx); err != nil { return err } cleanedKeyspaces, err := s.cleanupDefaultTSOKeyspaceGroupConfig(ctx) @@ -99,9 +96,8 @@ func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error return err } - if defaultGroupExists || deletedDefaultGroup || cleanedKeyspaces > 0 || deletedMicroserviceKeys > 0 { + if deletedDefaultGroup || cleanedKeyspaces > 0 || deletedMicroserviceKeys > 0 { log.Info("cleaned up microservice metadata in PD mode", - zap.Bool("default-keyspace-group-exists", defaultGroupExists), zap.Bool("deleted-default-keyspace-group", deletedDefaultGroup), zap.Int("cleaned-keyspace-configs", cleanedKeyspaces), zap.Int64("deleted-microservice-keys", deletedMicroserviceKeys), @@ -110,38 +106,89 @@ func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error return nil } -func (s *Server) validateDefaultTSOKeyspaceGroupConfig(ctx context.Context) error { +func (s *Server) validateMicroserviceMetadataCleanup(ctx context.Context) error { + groups, err := s.storage.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) + if err != nil { + return err + } + for _, group := range groups { + if group == nil { + continue + } + if group.ID != constant.DefaultKeyspaceGroupID { + return errors.Errorf("found non-default TSO keyspace group %d when cleaning up PD mode microservice metadata", group.ID) + } + if group.IsSplitting() { + return errors.Errorf("default TSO keyspace group is splitting when cleaning up PD mode microservice metadata") + } + if group.IsMerging() { + return errors.Errorf("default TSO keyspace group is merging when cleaning up PD mode microservice metadata") + } + } + return s.scanKeyspaceMetadata(ctx, etcdutil.MaxEtcdTxnOps, func(_ kv.Txn, metas []*keyspacepb.KeyspaceMeta) error { + for _, meta := range metas { + if meta == nil { + continue + } + if err := validateDefaultTSOKeyspaceGroupAssignment(meta.GetId(), meta.GetConfig()); err != nil { + return err + } + } + return nil + }) +} + +func (s *Server) cleanupDefaultTSOKeyspaceGroupConfig(ctx context.Context) (int, error) { + cleaned := 0 + err := s.scanKeyspaceMetadata(ctx, microserviceMetadataCleanupBatchSize, func(txn kv.Txn, metas []*keyspacepb.KeyspaceMeta) error { + for _, meta := range metas { + if meta == nil || meta.Config == nil { + continue + } + if _, ok := meta.Config[keyspace.TSOKeyspaceGroupIDKey]; !ok { + continue + } + if err := validateDefaultTSOKeyspaceGroupAssignment(meta.GetId(), meta.GetConfig()); err != nil { + return err + } + delete(meta.Config, keyspace.TSOKeyspaceGroupIDKey) + if err := s.storage.SaveKeyspaceMeta(txn, meta); err != nil { + return err + } + cleaned++ + } + return nil + }) + return cleaned, err +} + +func (s *Server) scanKeyspaceMetadata( + ctx context.Context, + batchSize int, + process func(kv.Txn, []*keyspacepb.KeyspaceMeta) error, +) error { startID := constant.StartKeyspaceID for { - var ( - keyspaceIDs []uint32 - loaded int - ) + var metas []*keyspacepb.KeyspaceMeta err := s.storage.RunInTxn(ctx, func(txn kv.Txn) error { - metas, err := s.storage.LoadRangeKeyspace(txn, startID, etcdutil.MaxEtcdTxnOps) + var err error + metas, err = s.storage.LoadRangeKeyspace(txn, startID, batchSize) if err != nil { return err } - loaded = len(metas) - keyspaceIDs = make([]uint32, 0, len(metas)) - for _, meta := range metas { - if meta == nil { - continue - } - keyspaceIDs = append(keyspaceIDs, meta.GetId()) - if err := validateDefaultTSOKeyspaceGroupAssignment(meta.GetId(), meta.GetConfig()); err != nil { - return err - } - } - return nil + return process(txn, metas) }) if err != nil { return err } - if loaded < etcdutil.MaxEtcdTxnOps || len(keyspaceIDs) == 0 { + if len(metas) < batchSize { + return nil + } + lastMeta := metas[len(metas)-1] + if lastMeta == nil { return nil } - lastID := keyspaceIDs[len(keyspaceIDs)-1] + lastID := lastMeta.GetId() if lastID == ^uint32(0) { return nil } @@ -164,90 +211,6 @@ func validateDefaultTSOKeyspaceGroupAssignment(keyspaceID uint32, config map[str return nil } -func (s *Server) checkDefaultOnlyTSOKeyspaceGroup() (bool, error) { - groups, err := s.storage.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) - if err != nil { - return false, err - } - defaultGroupExists := false - for _, group := range groups { - if group == nil { - continue - } - if group.ID != constant.DefaultKeyspaceGroupID { - return false, errors.Errorf("found non-default TSO keyspace group %d when cleaning up PD mode microservice metadata", group.ID) - } - if group.IsSplitting() { - return false, errors.Errorf("default TSO keyspace group is splitting when cleaning up PD mode microservice metadata") - } - if group.IsMerging() { - return false, errors.Errorf("default TSO keyspace group is merging when cleaning up PD mode microservice metadata") - } - defaultGroupExists = true - } - return defaultGroupExists, nil -} - -func (s *Server) cleanupDefaultTSOKeyspaceGroupConfig(ctx context.Context) (int, error) { - var ( - cleaned int - startID = constant.StartKeyspaceID - ) - for { - var ( - batchCleaned int - keyspaceIDs []uint32 - loaded int - ) - err := s.storage.RunInTxn(ctx, func(txn kv.Txn) error { - metas, err := s.storage.LoadRangeKeyspace(txn, startID, microserviceMetadataCleanupBatchSize) - if err != nil { - return err - } - loaded = len(metas) - keyspaceIDs = make([]uint32, 0, len(metas)) - for _, meta := range metas { - if meta == nil { - continue - } - keyspaceIDs = append(keyspaceIDs, meta.GetId()) - if meta.Config == nil { - continue - } - _, ok := meta.Config[keyspace.TSOKeyspaceGroupIDKey] - if !ok { - continue - } - if err := validateDefaultTSOKeyspaceGroupAssignment(meta.GetId(), meta.GetConfig()); err != nil { - return err - } - delete(meta.Config, keyspace.TSOKeyspaceGroupIDKey) - if err := s.storage.SaveKeyspaceMeta(txn, meta); err != nil { - return err - } - batchCleaned++ - } - return nil - }) - if err != nil { - return cleaned, err - } - cleaned += batchCleaned - if loaded < microserviceMetadataCleanupBatchSize { - break - } - if len(keyspaceIDs) == 0 { - break - } - lastID := keyspaceIDs[len(keyspaceIDs)-1] - if lastID == ^uint32(0) { - break - } - startID = lastID + 1 - } - return cleaned, nil -} - func (s *Server) deleteDefaultTSOKeyspaceGroup(ctx context.Context) (bool, error) { deleted := false err := s.storage.RunInTxn(ctx, func(txn kv.Txn) error { From ac16bb41d102b7679b6620e7c47fab153d77c623 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Mon, 13 Jul 2026 12:32:11 +0800 Subject: [PATCH 4/5] server: skip cleanup outside mode transitions Signed-off-by: Ryan Leung --- server/microservice_cleanup.go | 33 +++++++++++++++++++++++------ server/microservice_cleanup_test.go | 21 ++++++++++++++++++ 2 files changed, 47 insertions(+), 7 deletions(-) diff --git a/server/microservice_cleanup.go b/server/microservice_cleanup.go index 4045cb1cc7a..073261dcb3d 100644 --- a/server/microservice_cleanup.go +++ b/server/microservice_cleanup.go @@ -80,9 +80,13 @@ func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error } start := time.Now() - if err := s.validateMicroserviceMetadataCleanup(ctx); err != nil { + needsCleanup, err := s.validateMicroserviceMetadataCleanup(ctx) + if err != nil { return err } + if !needsCleanup { + return nil + } cleanedKeyspaces, err := s.cleanupDefaultTSOKeyspaceGroupConfig(ctx) if err != nil { return err @@ -106,26 +110,40 @@ func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error return nil } -func (s *Server) validateMicroserviceMetadataCleanup(ctx context.Context) error { +func (s *Server) validateMicroserviceMetadataCleanup(ctx context.Context) (bool, error) { groups, err := s.storage.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) if err != nil { - return err + return false, err } + needsCleanup := false for _, group := range groups { if group == nil { continue } + needsCleanup = true if group.ID != constant.DefaultKeyspaceGroupID { - return errors.Errorf("found non-default TSO keyspace group %d when cleaning up PD mode microservice metadata", group.ID) + return false, errors.Errorf("found non-default TSO keyspace group %d when cleaning up PD mode microservice metadata", group.ID) } if group.IsSplitting() { - return errors.Errorf("default TSO keyspace group is splitting when cleaning up PD mode microservice metadata") + return false, errors.Errorf("default TSO keyspace group is splitting when cleaning up PD mode microservice metadata") } if group.IsMerging() { - return errors.Errorf("default TSO keyspace group is merging when cleaning up PD mode microservice metadata") + return false, errors.Errorf("default TSO keyspace group is merging when cleaning up PD mode microservice metadata") } } - return s.scanKeyspaceMetadata(ctx, etcdutil.MaxEtcdTxnOps, func(_ kv.Txn, metas []*keyspacepb.KeyspaceMeta) error { + if !needsCleanup && s.client != nil { + getCtx, cancel := context.WithTimeout(ctx, etcdutil.DefaultRequestTimeout) + resp, err := s.client.Get(getCtx, microserviceEtcdPrefix(), clientv3.WithPrefix(), clientv3.WithLimit(1)) + cancel() + if err != nil { + return false, errs.ErrEtcdKVGet.Wrap(err).GenWithStackByCause() + } + needsCleanup = len(resp.Kvs) > 0 + } + if !needsCleanup { + return false, nil + } + err = s.scanKeyspaceMetadata(ctx, etcdutil.MaxEtcdTxnOps, func(_ kv.Txn, metas []*keyspacepb.KeyspaceMeta) error { for _, meta := range metas { if meta == nil { continue @@ -136,6 +154,7 @@ func (s *Server) validateMicroserviceMetadataCleanup(ctx context.Context) error } return nil }) + return true, err } func (s *Server) cleanupDefaultTSOKeyspaceGroupConfig(ctx context.Context) (int, error) { diff --git a/server/microservice_cleanup_test.go b/server/microservice_cleanup_test.go index 1f01d752cbf..2e9599621fd 100644 --- a/server/microservice_cleanup_test.go +++ b/server/microservice_cleanup_test.go @@ -140,6 +140,27 @@ func TestScheduleMicroserviceMetadataCleanupReturnsImmediately(t *testing.T) { svr.serverLoopWg.Wait() } +func TestMicroserviceMetadataCleanupTransitionDetection(t *testing.T) { + re := require.New(t) + ctx := context.Background() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + keypath.SetClusterID(12350) + defer keypath.ResetClusterID() + + store := storage.NewStorageWithEtcdBackend(client) + svr := &Server{storage: store, client: client} + needsCleanup, err := svr.validateMicroserviceMetadataCleanup(ctx) + re.NoError(err) + re.False(needsCleanup) + + _, err = client.Put(ctx, keypath.RegistryPath(mcs.TSOServiceName, "127.0.0.1:3379"), "tso") + re.NoError(err) + needsCleanup, err = svr.validateMicroserviceMetadataCleanup(ctx) + re.NoError(err) + re.True(needsCleanup) +} + func TestCleanupMicroserviceMetadataInPDModeRejectsNonDefaultGroup(t *testing.T) { re := require.New(t) ctx := context.Background() From ee7064875ae49dab2d6d2442fecb8a847a5a0509 Mon Sep 17 00:00:00 2001 From: Ryan Leung Date: Wed, 15 Jul 2026 14:50:25 +0800 Subject: [PATCH 5/5] server: protect microservice metadata cleanup Signed-off-by: Ryan Leung --- server/microservice_cleanup.go | 64 ++++++++---- server/microservice_cleanup_test.go | 145 ++++++++++++++++++++++++++++ 2 files changed, 192 insertions(+), 17 deletions(-) diff --git a/server/microservice_cleanup.go b/server/microservice_cleanup.go index 073261dcb3d..5c675f47e10 100644 --- a/server/microservice_cleanup.go +++ b/server/microservice_cleanup.go @@ -16,6 +16,7 @@ package server import ( "context" + goerrors "errors" "fmt" "strconv" "time" @@ -24,6 +25,7 @@ import ( "go.uber.org/zap" "github.com/pingcap/errors" + "github.com/pingcap/failpoint" "github.com/pingcap/kvproto/pkg/keyspacepb" "github.com/pingcap/log" @@ -45,6 +47,12 @@ const ( microserviceMetadataCleanupBatchSize = etcdutil.MaxEtcdTxnOps / 2 ) +var errMicroserviceMetadataCleanupRejected = errors.New("microservice metadata cleanup rejected") + +func rejectMicroserviceMetadataCleanup(format string, args ...any) error { + return errors.Wrapf(errMicroserviceMetadataCleanupRejected, format, args...) +} + func (s *Server) scheduleMicroserviceMetadataCleanup(ctx context.Context) { if s.IsKeyspaceGroupEnabled() { return @@ -61,6 +69,11 @@ func (s *Server) scheduleMicroserviceMetadataCleanup(ctx context.Context) { if err == nil { return } + if goerrors.Is(err, errMicroserviceMetadataCleanupRejected) { + log.Warn("cannot safely clean up microservice metadata in PD mode", + errs.ZapError(err)) + return + } log.Warn("failed to clean up microservice metadata in PD mode, retry later", errs.ZapError(err)) timer := time.NewTimer(microserviceMetadataCleanupRetryInterval) @@ -95,6 +108,10 @@ func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error if err != nil { return err } + failpoint.InjectCall("beforeDeleteMicroserviceEtcdKeys") + if _, err := s.checkMicroserviceEtcdKeysNotLeased(ctx); err != nil { + return err + } deletedMicroserviceKeys, err := s.deleteMicroserviceEtcdKeys(ctx) if err != nil { return err @@ -110,35 +127,48 @@ func (s *Server) cleanupMicroserviceMetadataInPDMode(ctx context.Context) error return nil } +func (s *Server) checkMicroserviceEtcdKeysNotLeased(ctx context.Context) (bool, error) { + if s.client == nil { + return false, nil + } + getCtx, cancel := context.WithTimeout(ctx, etcdutil.DefaultRequestTimeout) + defer cancel() + resp, err := s.client.Get(getCtx, microserviceEtcdPrefix(), clientv3.WithPrefix(), clientv3.WithKeysOnly()) + if err != nil { + return false, errs.ErrEtcdKVGet.Wrap(err).GenWithStackByCause() + } + for _, item := range resp.Kvs { + if item.Lease != 0 { + return true, errors.Errorf("microservice key %q is still leased", string(item.Key)) + } + } + return len(resp.Kvs) > 0, nil +} + func (s *Server) validateMicroserviceMetadataCleanup(ctx context.Context) (bool, error) { + microserviceKeysExist, err := s.checkMicroserviceEtcdKeysNotLeased(ctx) + if err != nil { + return false, err + } groups, err := s.storage.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) if err != nil { return false, err } - needsCleanup := false + needsCleanup := microserviceKeysExist for _, group := range groups { if group == nil { continue } needsCleanup = true if group.ID != constant.DefaultKeyspaceGroupID { - return false, errors.Errorf("found non-default TSO keyspace group %d when cleaning up PD mode microservice metadata", group.ID) + return false, rejectMicroserviceMetadataCleanup("found non-default TSO keyspace group %d when cleaning up PD mode microservice metadata", group.ID) } if group.IsSplitting() { - return false, errors.Errorf("default TSO keyspace group is splitting when cleaning up PD mode microservice metadata") + return false, rejectMicroserviceMetadataCleanup("default TSO keyspace group is splitting when cleaning up PD mode microservice metadata") } if group.IsMerging() { - return false, errors.Errorf("default TSO keyspace group is merging when cleaning up PD mode microservice metadata") - } - } - if !needsCleanup && s.client != nil { - getCtx, cancel := context.WithTimeout(ctx, etcdutil.DefaultRequestTimeout) - resp, err := s.client.Get(getCtx, microserviceEtcdPrefix(), clientv3.WithPrefix(), clientv3.WithLimit(1)) - cancel() - if err != nil { - return false, errs.ErrEtcdKVGet.Wrap(err).GenWithStackByCause() + return false, rejectMicroserviceMetadataCleanup("default TSO keyspace group is merging when cleaning up PD mode microservice metadata") } - needsCleanup = len(resp.Kvs) > 0 } if !needsCleanup { return false, nil @@ -222,10 +252,10 @@ func validateDefaultTSOKeyspaceGroupAssignment(keyspaceID uint32, config map[str } groupID, err := strconv.ParseUint(groupIDText, 10, 32) if err != nil { - return errors.Errorf("keyspace %d has invalid TSO keyspace group ID %q", keyspaceID, groupIDText) + return rejectMicroserviceMetadataCleanup("keyspace %d has invalid TSO keyspace group ID %q", keyspaceID, groupIDText) } if groupID != uint64(constant.DefaultKeyspaceGroupID) { - return errors.Errorf("keyspace %d is assigned to non-default TSO keyspace group %d", keyspaceID, groupID) + return rejectMicroserviceMetadataCleanup("keyspace %d is assigned to non-default TSO keyspace group %d", keyspaceID, groupID) } return nil } @@ -238,10 +268,10 @@ func (s *Server) deleteDefaultTSOKeyspaceGroup(ctx context.Context) (bool, error return err } if group.IsSplitting() { - return errors.Errorf("default TSO keyspace group is splitting when deleting PD mode microservice metadata") + return rejectMicroserviceMetadataCleanup("default TSO keyspace group is splitting when deleting PD mode microservice metadata") } if group.IsMerging() { - return errors.Errorf("default TSO keyspace group is merging when deleting PD mode microservice metadata") + return rejectMicroserviceMetadataCleanup("default TSO keyspace group is merging when deleting PD mode microservice metadata") } if err := s.storage.DeleteKeyspaceGroup(txn, constant.DefaultKeyspaceGroupID); err != nil { return err diff --git a/server/microservice_cleanup_test.go b/server/microservice_cleanup_test.go index 2e9599621fd..559458dadab 100644 --- a/server/microservice_cleanup_test.go +++ b/server/microservice_cleanup_test.go @@ -112,6 +112,112 @@ func TestCleanupMicroserviceMetadataInPDMode(t *testing.T) { re.Equal("timestamp", string(resp.Kvs[0].Value)) } +func TestCleanupMicroserviceMetadataWaitsForLeasedKeys(t *testing.T) { + re := require.New(t) + ctx := context.Background() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + keypath.SetClusterID(12351) + defer keypath.ResetClusterID() + + store := storage.NewStorageWithEtcdBackend(client) + svr := &Server{storage: store, client: client} + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + if err := store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + Keyspaces: []uint32{1}, + }); err != nil { + return err + } + return store.SaveKeyspaceMeta(txn, newKeyspaceMetaWithTSOGroup(1, "0")) + })) + registryPath := keypath.RegistryPath(mcs.TSOServiceName, "127.0.0.1:3379") + leaseID, err := etcdutil.EtcdKVPutWithTTL(ctx, client, registryPath, "tso", 60) + re.NoError(err) + re.NotZero(leaseID) + + err = svr.cleanupMicroserviceMetadataInPDMode(ctx) + re.ErrorContains(err, "is still leased") + + groups, err := store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) + re.NoError(err) + re.Len(groups, 1) + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + meta, err := store.LoadKeyspaceMeta(txn, 1) + if err != nil { + return err + } + re.Contains(meta.GetConfig(), keyspace.TSOKeyspaceGroupIDKey) + return nil + })) + resp, err := etcdutil.EtcdKVGet(client, registryPath) + re.NoError(err) + re.Len(resp.Kvs, 1) + re.NotZero(resp.Kvs[0].Lease) + + _, err = client.Revoke(ctx, leaseID) + re.NoError(err) + re.NoError(svr.cleanupMicroserviceMetadataInPDMode(ctx)) + groups, err = store.LoadKeyspaceGroups(constant.DefaultKeyspaceGroupID, 0) + re.NoError(err) + re.Empty(groups) + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + meta, err := store.LoadKeyspaceMeta(txn, 1) + if err != nil { + return err + } + re.NotContains(meta.GetConfig(), keyspace.TSOKeyspaceGroupIDKey) + return nil + })) +} + +func TestCleanupMicroserviceMetadataRechecksLeasesBeforeDelete(t *testing.T) { + re := require.New(t) + ctx := context.Background() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + keypath.SetClusterID(12352) + defer keypath.ResetClusterID() + + store := storage.NewStorageWithEtcdBackend(client) + svr := &Server{storage: store, client: client} + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + if err := store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + Keyspaces: []uint32{1}, + }); err != nil { + return err + } + return store.SaveKeyspaceMeta(txn, newKeyspaceMetaWithTSOGroup(1, "0")) + })) + registryPath := keypath.RegistryPath(mcs.TSOServiceName, "127.0.0.1:3379") + var leaseID clientv3.LeaseID + re.NoError(failpoint.EnableCall("github.com/tikv/pd/server/beforeDeleteMicroserviceEtcdKeys", func() { + if leaseID != clientv3.NoLease { + return + } + var err error + leaseID, err = etcdutil.EtcdKVPutWithTTL(ctx, client, registryPath, "tso", 60) + re.NoError(err) + })) + t.Cleanup(func() { + _ = failpoint.Disable("github.com/tikv/pd/server/beforeDeleteMicroserviceEtcdKeys") + }) + + err := svr.cleanupMicroserviceMetadataInPDMode(ctx) + re.ErrorContains(err, "is still leased") + re.NotZero(leaseID) + resp, err := etcdutil.EtcdKVGet(client, registryPath) + re.NoError(err) + re.Len(resp.Kvs, 1) + re.Equal(leaseID, clientv3.LeaseID(resp.Kvs[0].Lease)) + + _, err = client.Revoke(ctx, leaseID) + re.NoError(err) +} + func TestScheduleMicroserviceMetadataCleanupReturnsImmediately(t *testing.T) { re := require.New(t) ctx := context.Background() @@ -140,6 +246,45 @@ func TestScheduleMicroserviceMetadataCleanupReturnsImmediately(t *testing.T) { svr.serverLoopWg.Wait() } +func TestScheduleMicroserviceMetadataCleanupStopsOnRejectedState(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + _, client, clean := etcdutil.NewTestEtcdCluster(t, 1, nil) + defer clean() + keypath.SetClusterID(12353) + defer keypath.ResetClusterID() + + store := storage.NewStorageWithEtcdBackend(client) + svr := &Server{storage: store, client: client} + re.NoError(store.RunInTxn(ctx, func(txn kv.Txn) error { + if err := store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: constant.DefaultKeyspaceGroupID, + UserKind: endpoint.Basic.String(), + }); err != nil { + return err + } + return store.SaveKeyspaceGroup(txn, &endpoint.KeyspaceGroup{ + ID: 1, + UserKind: endpoint.Standard.String(), + }) + })) + + svr.scheduleMicroserviceMetadataCleanup(ctx) + done := make(chan struct{}) + go func() { + svr.serverLoopWg.Wait() + close(done) + }() + select { + case <-done: + case <-time.After(time.Second): + cancel() + <-done + t.Fatal("cleanup task kept retrying a rejected metadata state") + } +} + func TestMicroserviceMetadataCleanupTransitionDetection(t *testing.T) { re := require.New(t) ctx := context.Background()