From 142e64974bd16eb7046cdf2305c4d4aae2b6ec1e Mon Sep 17 00:00:00 2001 From: JmPotato Date: Wed, 26 Apr 2023 14:31:58 +0800 Subject: [PATCH 1/7] Patrol keyspace assignment before the first split Signed-off-by: JmPotato --- pkg/keyspace/tso_keyspace_group.go | 52 +++++++++++++++++++++++-- pkg/keyspace/tso_keyspace_group_test.go | 26 +++++++++++++ 2 files changed, 74 insertions(+), 4 deletions(-) diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go index b9e8eb311ff..5a5f6f5addb 100644 --- a/pkg/keyspace/tso_keyspace_group.go +++ b/pkg/keyspace/tso_keyspace_group.go @@ -22,6 +22,7 @@ import ( "time" "github.com/pingcap/errors" + "github.com/pingcap/kvproto/pkg/keyspacepb" "github.com/pingcap/log" "github.com/tikv/pd/pkg/balancer" "github.com/tikv/pd/pkg/mcs/discovery" @@ -56,6 +57,7 @@ type GroupManager struct { ctx context.Context cancel context.CancelFunc wg sync.WaitGroup + // the lock for the groups sync.RWMutex // groups is the cache of keyspace group related information. @@ -63,7 +65,10 @@ type GroupManager struct { groups map[endpoint.UserKind]*indexedHeap // store is the storage for keyspace group related information. - store endpoint.KeyspaceGroupStorage + store interface { + endpoint.KeyspaceGroupStorage + endpoint.KeyspaceStorage + } client *clientv3.Client @@ -77,10 +82,21 @@ type GroupManager struct { nodesBalancer balancer.Balancer[string] // serviceRegistryMap stores the mapping from the service registry key to the service address. serviceRegistryMap map[string]string + + // patrolKeyspaceAssignmentOnce is used to patrol all keyspaces and assign them to the keyspace groups. + patrolKeyspaceAssignmentOnce sync.Once } // NewKeyspaceGroupManager creates a Manager of keyspace group related data. -func NewKeyspaceGroupManager(ctx context.Context, store endpoint.KeyspaceGroupStorage, client *clientv3.Client, clusterID uint64) *GroupManager { +func NewKeyspaceGroupManager( + ctx context.Context, + store interface { + endpoint.KeyspaceGroupStorage + endpoint.KeyspaceStorage + }, + client *clientv3.Client, + clusterID uint64, +) *GroupManager { ctx, cancel := context.WithCancel(ctx) key := discovery.TSOPath(clusterID) groups := make(map[endpoint.UserKind]*indexedHeap) @@ -156,6 +172,31 @@ func (m *GroupManager) Close() { m.wg.Wait() } +// patrolKeyspaceAssignment is used to patrol all keyspaces and assign them to the keyspace groups. +func (m *GroupManager) patrolKeyspaceAssignment() (err error) { + m.patrolKeyspaceAssignmentOnce.Do(func() { + var keyspaces []*keyspacepb.KeyspaceMeta + keyspaces, err = m.store.LoadRangeKeyspace(utils.DefaultKeyspaceID, 0) + if err != nil { + return + } + config, err := m.GetKeyspaceConfigByKind(endpoint.Basic) + if err != nil { + return + } + for _, ks := range keyspaces { + if ks == nil { + continue + } + err = m.UpdateKeyspaceForGroup(endpoint.Basic, config[TSOKeyspaceGroupIDKey], ks.GetId(), opAdd) + if err != nil { + return + } + } + }) + return err +} + func (m *GroupManager) allocNodesToAllKeyspaceGroups() { defer logutil.LogPanic() defer m.wg.Done() @@ -535,11 +576,14 @@ func (m *GroupManager) UpdateKeyspaceGroup(oldGroupID, newGroupID string, oldUse // SplitKeyspaceGroupByID splits the keyspace group by ID into a new keyspace group with the given new ID. // And the keyspaces in the old keyspace group will be moved to the new keyspace group. func (m *GroupManager) SplitKeyspaceGroupByID(splitSourceID, splitTargetID uint32, keyspaces []uint32) error { + err := m.patrolKeyspaceAssignment() + if err != nil { + return err + } var splitSourceKg, splitTargetKg *endpoint.KeyspaceGroup m.Lock() defer m.Unlock() - // TODO: avoid to split when the keyspaces is empty. - if err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) (err error) { + if err = m.store.RunInTxn(m.ctx, func(txn kv.Txn) (err error) { // Load the old keyspace group first. splitSourceKg, err = m.store.LoadKeyspaceGroup(txn, splitSourceID) if err != nil { diff --git a/pkg/keyspace/tso_keyspace_group_test.go b/pkg/keyspace/tso_keyspace_group_test.go index 6286a71b3aa..0b08ce3927e 100644 --- a/pkg/keyspace/tso_keyspace_group_test.go +++ b/pkg/keyspace/tso_keyspace_group_test.go @@ -20,6 +20,7 @@ import ( "testing" "time" + "github.com/pingcap/kvproto/pkg/keyspacepb" "github.com/stretchr/testify/suite" "github.com/tikv/pd/pkg/mcs/utils" "github.com/tikv/pd/pkg/mock/mockcluster" @@ -317,3 +318,28 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupSplit() { err = suite.kgm.SplitKeyspaceGroupByID(2, 5, []uint32{111, 222, 444}) re.ErrorIs(err, ErrKeyspaceNotInKeyspaceGroup) } + +func (suite *keyspaceGroupTestSuite) TestPatrolKeyspaceAssignment() { + re := suite.Require() + // Create a keyspace group without any keyspace. + err := suite.kgm.CreateKeyspaceGroups([]*endpoint.KeyspaceGroup{ + { + ID: uint32(1), + UserKind: endpoint.Basic.String(), + }, + }) + re.NoError(err) + // Create a keyspace without any keyspace group. + now := time.Now().Unix() + err = suite.kg.saveNewKeyspace(&keyspacepb.KeyspaceMeta{ + Id: 111, + Name: "111", + State: keyspacepb.KeyspaceState_ENABLED, + CreatedAt: now, + StateChangedAt: now, + }) + re.NoError(err) + // Split to see if the keyspace is attached to the group. + err = suite.kgm.SplitKeyspaceGroupByID(1, 2, []uint32{111}) + re.NoError(err) +} From f3777c3380b52e33a54b7dc4f954929f645ac9e6 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Wed, 26 Apr 2023 14:43:31 +0800 Subject: [PATCH 2/7] Check keyspace num Signed-off-by: JmPotato --- pkg/keyspace/tso_keyspace_group.go | 10 ++++++---- pkg/keyspace/tso_keyspace_group_test.go | 3 +++ 2 files changed, 9 insertions(+), 4 deletions(-) diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go index 5a5f6f5addb..b508cfb1f9c 100644 --- a/pkg/keyspace/tso_keyspace_group.go +++ b/pkg/keyspace/tso_keyspace_group.go @@ -608,13 +608,15 @@ func (m *GroupManager) SplitKeyspaceGroupByID(splitSourceID, splitTargetID uint3 if splitTargetKg != nil { return ErrKeyspaceGroupExists } + keyspaceNum := len(keyspaces) + sourceKeyspaceNum := len(splitSourceKg.Keyspaces) // Check if the keyspaces are all in the old keyspace group. - if len(keyspaces) > len(splitSourceKg.Keyspaces) { + if keyspaceNum == 0 || keyspaceNum > sourceKeyspaceNum { return ErrKeyspaceNotInKeyspaceGroup } var ( - oldKeyspaceMap = make(map[uint32]struct{}, len(splitSourceKg.Keyspaces)) - newKeyspaceMap = make(map[uint32]struct{}, len(keyspaces)) + oldKeyspaceMap = make(map[uint32]struct{}, sourceKeyspaceNum) + newKeyspaceMap = make(map[uint32]struct{}, keyspaceNum) ) for _, keyspace := range splitSourceKg.Keyspaces { oldKeyspaceMap[keyspace] = struct{}{} @@ -626,7 +628,7 @@ func (m *GroupManager) SplitKeyspaceGroupByID(splitSourceID, splitTargetID uint3 newKeyspaceMap[keyspace] = struct{}{} } // Get the split keyspace group for the old keyspace group. - splitKeyspaces := make([]uint32, 0, len(splitSourceKg.Keyspaces)-len(keyspaces)) + splitKeyspaces := make([]uint32, 0, sourceKeyspaceNum-keyspaceNum) for _, keyspace := range splitSourceKg.Keyspaces { if _, ok := newKeyspaceMap[keyspace]; !ok { splitKeyspaces = append(splitKeyspaces, keyspace) diff --git a/pkg/keyspace/tso_keyspace_group_test.go b/pkg/keyspace/tso_keyspace_group_test.go index 0b08ce3927e..cc643613cfe 100644 --- a/pkg/keyspace/tso_keyspace_group_test.go +++ b/pkg/keyspace/tso_keyspace_group_test.go @@ -252,6 +252,9 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupSplit() { // split the keyspace group 1 to 4 err = suite.kgm.SplitKeyspaceGroupByID(1, 4, []uint32{333}) re.ErrorIs(err, ErrKeyspaceGroupNotEnoughReplicas) + // split the keyspace group 2 to 4 without giving any keyspace + err = suite.kgm.SplitKeyspaceGroupByID(2, 3, []uint32{}) + re.ErrorIs(err, ErrKeyspaceNotInKeyspaceGroup) // split the keyspace group 2 to 4 err = suite.kgm.SplitKeyspaceGroupByID(2, 4, []uint32{333}) re.NoError(err) From 0e634d9dd37f7c541d13c002ca54cb278e2c79ed Mon Sep 17 00:00:00 2001 From: JmPotato Date: Wed, 26 Apr 2023 16:10:26 +0800 Subject: [PATCH 3/7] Use lock to make sure the success of patrol Signed-off-by: JmPotato --- pkg/keyspace/tso_keyspace_group.go | 63 ++++++++++++++++++------------ 1 file changed, 38 insertions(+), 25 deletions(-) diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go index b508cfb1f9c..7f6a1bc86a8 100644 --- a/pkg/keyspace/tso_keyspace_group.go +++ b/pkg/keyspace/tso_keyspace_group.go @@ -22,7 +22,6 @@ import ( "time" "github.com/pingcap/errors" - "github.com/pingcap/kvproto/pkg/keyspacepb" "github.com/pingcap/log" "github.com/tikv/pd/pkg/balancer" "github.com/tikv/pd/pkg/mcs/discovery" @@ -58,11 +57,13 @@ type GroupManager struct { cancel context.CancelFunc wg sync.WaitGroup - // the lock for the groups sync.RWMutex // groups is the cache of keyspace group related information. // user kind -> keyspace group groups map[endpoint.UserKind]*indexedHeap + // patrolKeyspaceAssignmentOnce is used to indicate whether we have patrolled all keyspaces + // and assign them to the keyspace groups. + patrolKeyspaceAssignmentOnce bool // store is the storage for keyspace group related information. store interface { @@ -82,9 +83,6 @@ type GroupManager struct { nodesBalancer balancer.Balancer[string] // serviceRegistryMap stores the mapping from the service registry key to the service address. serviceRegistryMap map[string]string - - // patrolKeyspaceAssignmentOnce is used to patrol all keyspaces and assign them to the keyspace groups. - patrolKeyspaceAssignmentOnce sync.Once } // NewKeyspaceGroupManager creates a Manager of keyspace group related data. @@ -173,28 +171,35 @@ func (m *GroupManager) Close() { } // patrolKeyspaceAssignment is used to patrol all keyspaces and assign them to the keyspace groups. -func (m *GroupManager) patrolKeyspaceAssignment() (err error) { - m.patrolKeyspaceAssignmentOnce.Do(func() { - var keyspaces []*keyspacepb.KeyspaceMeta - keyspaces, err = m.store.LoadRangeKeyspace(utils.DefaultKeyspaceID, 0) - if err != nil { - return +func (m *GroupManager) patrolKeyspaceAssignment() error { + m.Lock() + defer m.Unlock() + if m.patrolKeyspaceAssignmentOnce { + return nil + } + keyspaces, err := m.store.LoadRangeKeyspace(utils.DefaultKeyspaceID, 0) + if err != nil { + return err + } + config, err := m.getKeyspaceConfigByKindLocked(endpoint.Basic) + if err != nil { + return err + } + for _, ks := range keyspaces { + if ks == nil { + continue } - config, err := m.GetKeyspaceConfigByKind(endpoint.Basic) + groupID, err := strconv.ParseUint(config[TSOKeyspaceGroupIDKey], 10, 64) if err != nil { - return + return err } - for _, ks := range keyspaces { - if ks == nil { - continue - } - err = m.UpdateKeyspaceForGroup(endpoint.Basic, config[TSOKeyspaceGroupIDKey], ks.GetId(), opAdd) - if err != nil { - return - } + err = m.updateKeyspaceForGroupLocked(endpoint.Basic, groupID, ks.GetId(), opAdd) + if err != nil { + return err } - }) - return err + } + m.patrolKeyspaceAssignmentOnce = true + return nil } func (m *GroupManager) allocNodesToAllKeyspaceGroups() { @@ -467,6 +472,10 @@ func (m *GroupManager) GetKeyspaceConfigByKind(userKind endpoint.UserKind) (map[ } m.RLock() defer m.RUnlock() + return m.getKeyspaceConfigByKindLocked(userKind) +} + +func (m *GroupManager) getKeyspaceConfigByKindLocked(userKind endpoint.UserKind) (map[string]string, error) { groups, ok := m.groups[userKind] if !ok { return map[string]string{}, errors.Errorf("user kind %s not found", userKind) @@ -493,9 +502,13 @@ func (m *GroupManager) UpdateKeyspaceForGroup(userKind endpoint.UserKind, groupI m.Lock() defer m.Unlock() - kg := m.groups[userKind].Get(uint32(id)) + return m.updateKeyspaceForGroupLocked(userKind, id, keyspaceID, mutation) +} + +func (m *GroupManager) updateKeyspaceForGroupLocked(userKind endpoint.UserKind, groupID uint64, keyspaceID uint32, mutation int) error { + kg := m.groups[userKind].Get(uint32(groupID)) if kg == nil { - return errors.Errorf("keyspace group %d not found", id) + return errors.Errorf("keyspace group %d not found", groupID) } if kg.IsSplitting() { return ErrKeyspaceGroupInSplit From 552d26ff20ae5b35282f99a71cd75f218e7a32ca Mon Sep 17 00:00:00 2001 From: JmPotato Date: Wed, 26 Apr 2023 16:20:53 +0800 Subject: [PATCH 4/7] Call patrolKeyspaceAssignment during the bootstrap Signed-off-by: JmPotato --- pkg/keyspace/tso_keyspace_group.go | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go index 7f6a1bc86a8..b10db4766d4 100644 --- a/pkg/keyspace/tso_keyspace_group.go +++ b/pkg/keyspace/tso_keyspace_group.go @@ -161,7 +161,8 @@ func (m *GroupManager) Bootstrap() error { m.groups[userKind].Put(group) } - return nil + // Load all the keyspaces from the storage and assign them to the respective keyspace groups. + return m.patrolKeyspaceAssignmentLocked() } // Close closes the manager. @@ -170,10 +171,14 @@ func (m *GroupManager) Close() { m.wg.Wait() } -// patrolKeyspaceAssignment is used to patrol all keyspaces and assign them to the keyspace groups. func (m *GroupManager) patrolKeyspaceAssignment() error { m.Lock() defer m.Unlock() + return m.patrolKeyspaceAssignmentLocked() +} + +// patrolKeyspaceAssignment is used to patrol all keyspaces and assign them to the keyspace groups. +func (m *GroupManager) patrolKeyspaceAssignmentLocked() error { if m.patrolKeyspaceAssignmentOnce { return nil } From 29d48fddc74278e68ec5f5331ae20131e31fe38e Mon Sep 17 00:00:00 2001 From: JmPotato Date: Wed, 26 Apr 2023 16:23:53 +0800 Subject: [PATCH 5/7] Fix a typo Signed-off-by: JmPotato --- pkg/keyspace/tso_keyspace_group_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/keyspace/tso_keyspace_group_test.go b/pkg/keyspace/tso_keyspace_group_test.go index cc643613cfe..b81afb98b5b 100644 --- a/pkg/keyspace/tso_keyspace_group_test.go +++ b/pkg/keyspace/tso_keyspace_group_test.go @@ -253,7 +253,7 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupSplit() { err = suite.kgm.SplitKeyspaceGroupByID(1, 4, []uint32{333}) re.ErrorIs(err, ErrKeyspaceGroupNotEnoughReplicas) // split the keyspace group 2 to 4 without giving any keyspace - err = suite.kgm.SplitKeyspaceGroupByID(2, 3, []uint32{}) + err = suite.kgm.SplitKeyspaceGroupByID(2, 4, []uint32{}) re.ErrorIs(err, ErrKeyspaceNotInKeyspaceGroup) // split the keyspace group 2 to 4 err = suite.kgm.SplitKeyspaceGroupByID(2, 4, []uint32{333}) From e9b9ec1e8f3ae73ef67f83997a20343b5dfdad46 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Wed, 26 Apr 2023 16:31:14 +0800 Subject: [PATCH 6/7] Fix the test Signed-off-by: JmPotato --- pkg/keyspace/tso_keyspace_group_test.go | 3 +++ 1 file changed, 3 insertions(+) diff --git a/pkg/keyspace/tso_keyspace_group_test.go b/pkg/keyspace/tso_keyspace_group_test.go index b81afb98b5b..cfe035578e2 100644 --- a/pkg/keyspace/tso_keyspace_group_test.go +++ b/pkg/keyspace/tso_keyspace_group_test.go @@ -324,11 +324,14 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupSplit() { func (suite *keyspaceGroupTestSuite) TestPatrolKeyspaceAssignment() { re := suite.Require() + // Force the patrol to run once. + suite.kgm.patrolKeyspaceAssignmentOnce = false // Create a keyspace group without any keyspace. err := suite.kgm.CreateKeyspaceGroups([]*endpoint.KeyspaceGroup{ { ID: uint32(1), UserKind: endpoint.Basic.String(), + Members: make([]endpoint.KeyspaceGroupMember, 2), }, }) re.NoError(err) From fdbe9d03421a7f2d8f1d641b2cb41584f26828a1 Mon Sep 17 00:00:00 2001 From: JmPotato Date: Wed, 26 Apr 2023 17:41:03 +0800 Subject: [PATCH 7/7] Fix tsoKeyspaceGroupManagerTestSuite Signed-off-by: JmPotato --- pkg/keyspace/tso_keyspace_group.go | 12 +++++------- .../mcs/tso/keyspace_group_manager_test.go | 4 ++++ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/pkg/keyspace/tso_keyspace_group.go b/pkg/keyspace/tso_keyspace_group.go index b10db4766d4..36b45f885d6 100644 --- a/pkg/keyspace/tso_keyspace_group.go +++ b/pkg/keyspace/tso_keyspace_group.go @@ -161,8 +161,7 @@ func (m *GroupManager) Bootstrap() error { m.groups[userKind].Put(group) } - // Load all the keyspaces from the storage and assign them to the respective keyspace groups. - return m.patrolKeyspaceAssignmentLocked() + return nil } // Close closes the manager. @@ -171,14 +170,10 @@ func (m *GroupManager) Close() { m.wg.Wait() } +// patrolKeyspaceAssignment is used to patrol all keyspaces and assign them to the keyspace groups. func (m *GroupManager) patrolKeyspaceAssignment() error { m.Lock() defer m.Unlock() - return m.patrolKeyspaceAssignmentLocked() -} - -// patrolKeyspaceAssignment is used to patrol all keyspaces and assign them to the keyspace groups. -func (m *GroupManager) patrolKeyspaceAssignmentLocked() error { if m.patrolKeyspaceAssignmentOnce { return nil } @@ -486,6 +481,9 @@ func (m *GroupManager) getKeyspaceConfigByKindLocked(userKind endpoint.UserKind) return map[string]string{}, errors.Errorf("user kind %s not found", userKind) } kg := groups.Top() + if kg == nil { + return map[string]string{}, errors.Errorf("no keyspace group for user kind %s", userKind) + } id := strconv.FormatUint(uint64(kg.ID), 10) config := map[string]string{ UserKindKey: userKind.String(), diff --git a/tests/integrations/mcs/tso/keyspace_group_manager_test.go b/tests/integrations/mcs/tso/keyspace_group_manager_test.go index 9cd46ad2d41..38000c257ce 100644 --- a/tests/integrations/mcs/tso/keyspace_group_manager_test.go +++ b/tests/integrations/mcs/tso/keyspace_group_manager_test.go @@ -83,6 +83,10 @@ func (suite *tsoKeyspaceGroupManagerTestSuite) TearDownTest() { func cleanupKeyspaceGroups(re *require.Assertions, server *tests.TestServer) { for _, group := range handlersutil.MustLoadKeyspaceGroups(re, server, "0", "0") { + // Do not delete default keyspace group. + if group.ID == mcsutils.DefaultKeyspaceGroupID { + continue + } handlersutil.MustDeleteKeyspaceGroup(re, server, group.ID) } }