From 373334d6b98bd3f86ede9b630b281e5e59ebff5a Mon Sep 17 00:00:00 2001 From: David <8039876+AmoebaProtozoa@users.noreply.github.com> Date: Mon, 2 Feb 2026 15:15:25 +0800 Subject: [PATCH 01/11] keyspace: cache meta-service group status (#463) * cache meta-service group status Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * fix lint Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * fix lint Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * fix deadlock Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * cleanup Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * assign to group no longer error Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * lint Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * fix integration test Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * leader callback Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * add more tests Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * lint Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * don't error keyspace udpate when assignment count failed Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * add some comments Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * don't flush on follower Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * Apply suggestions from code review Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> * add comments for leader check Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> * improve comments Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> --------- Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> (cherry picked from commit 2baf2a502b864633cd6198a90094033fd28caff4) Signed-off-by: bufferflies <1045931706@qq.com> # Conflicts: # pkg/keyspace/keyspace.go # pkg/keyspace/keyspace_test.go # pkg/keyspace/meta_service_group.go # pkg/keyspace/meta_service_group_test.go # server/server.go # tests/integrations/client/meta_service_group_test.go --- pkg/keyspace/keyspace.go | 60 ++--- pkg/keyspace/keyspace_test.go | 9 +- pkg/keyspace/meta_service_group.go | 318 +++++++++++++++--------- pkg/keyspace/meta_service_group_test.go | 41 ++- server/server.go | 12 +- 5 files changed, 283 insertions(+), 157 deletions(-) diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go index af992c397d..18c1e31b6a 100644 --- a/pkg/keyspace/keyspace.go +++ b/pkg/keyspace/keyspace.go @@ -569,26 +569,10 @@ func (manager *Manager) saveNewKeyspace(keyspace *keyspacepb.KeyspaceMeta) error }) } -// rollbackMetaServiceGroupAssignment decrements the assignment count that -// PickGroup incremented for a keyspace whose creation failed before its metadata -// was persisted, keeping the persisted counter in sync with actual keyspaces. -func (manager *Manager) rollbackMetaServiceGroupAssignment(groupID string) { - if manager.mgm == nil || groupID == "" { - return - } - if err := manager.mgm.store.RunInTxn(manager.ctx, func(txn kv.Txn) error { - return manager.mgm.updateAssignmentTxn(txn, groupID, "") - }); err != nil { - log.Warn("[keyspace] failed to roll back meta-service group assignment count", - zap.String("meta-service-group", groupID), - zap.Error(err)) - } -} - // assignGroupAndSaveKeyspace assigns a meta-service group to the keyspace (when // assign is true) and persists the keyspace metadata while holding the -// meta-service group manager's read lock across both steps. This keeps the -// selection and the persisted assignment atomic with respect to group deletion: +// meta-service group manager's lock across both steps. This keeps the +// selection and the cached assignment atomic with respect to group deletion: // UpdateGroupsSafely takes the write lock, so a group cannot be removed in the // window between assignment and the keyspace being saved, which would otherwise // leave the keyspace referencing a non-existent group. config must point to @@ -597,8 +581,8 @@ func (manager *Manager) assignGroupAndSaveKeyspace(assign bool, config *map[stri if !assign { return manager.saveNewKeyspace(keyspace) } - manager.mgm.RLock() - defer manager.mgm.RUnlock() + manager.mgm.Lock() + defer manager.mgm.Unlock() // Re-check under the lock: the pre-lock check may be stale if a concurrent // PATCH deleted the last group in the meantime. In that case create the // keyspace without a meta-service group instead of failing the creation. @@ -622,9 +606,9 @@ func (manager *Manager) assignGroupAndSaveKeyspace(assign bool, config *map[stri keyspace.Config = *config if err := manager.saveNewKeyspace(keyspace); err != nil { // Roll back the reservation made by pickGroupLocked. This only performs - // store operations and does not take the mgm lock, so it is safe to call - // while still holding the read lock. - manager.rollbackMetaServiceGroupAssignment(groupID) + if rollbackErr := manager.mgm.updateAssignmentLockedTxn(nil, groupID, ""); rollbackErr != nil { + log.Error("failed to revert meta-service group assignment", zap.Error(rollbackErr)) + } return err } return nil @@ -877,14 +861,14 @@ func applyKeyspaceConfigMutations(config map[string]string, mutations []*Mutatio } // runTxnWithMetaGroupLock runs f inside a storage transaction while holding the -// meta-service group manager's read lock for the whole transaction. This keeps -// keyspace assignment validation and the persisted assignment count update +// meta-service group manager's lock for the whole transaction. This keeps +// keyspace assignment validation and the cached assignment count update // atomic with respect to MetaServiceGroupManager.UpdateGroupsSafely, which takes // the write lock before deleting a group. func (manager *Manager) runTxnWithMetaGroupLock(f func(txn kv.Txn) error) error { if manager.mgm != nil { - manager.mgm.RLock() - defer manager.mgm.RUnlock() + manager.mgm.Lock() + defer manager.mgm.Unlock() } return manager.store.RunInTxn(manager.ctx, f) } @@ -892,6 +876,8 @@ func (manager *Manager) runTxnWithMetaGroupLock(f func(txn kv.Txn) error) error func (manager *Manager) updateKeyspaceConfigTxn(name string, update func(meta *keyspacepb.KeyspaceMeta) error) (*keyspacepb.KeyspaceMeta, error) { var meta *keyspacepb.KeyspaceMeta oldConfig := make(map[string]string) + var oldMetaServiceGroup, newMetaServiceGroup string + metaServiceGroupReassigned := false txnFunc := func(txn kv.Txn) error { // First get KeyspaceID from Name. loaded, id, err := manager.store.LoadKeyspaceID(txn, name) @@ -933,14 +919,15 @@ func (manager *Manager) updateKeyspaceConfigTxn(name string, update func(meta *k // txn doesn't commit), while UpdateKeyspaceGroup persists immediately. Doing // the fallible meta-service validation first avoids leaving the TSO group move // persisted but unreverted when the meta-service reassignment fails. - oldMetaServiceGroup := oldConfig[MetaServiceGroupIDKey] - newMetaServiceGroup := newConfig[MetaServiceGroupIDKey] + oldMetaServiceGroup = oldConfig[MetaServiceGroupIDKey] + newMetaServiceGroup = newConfig[MetaServiceGroupIDKey] if manager.mgm != nil && oldMetaServiceGroup != newMetaServiceGroup { - // The read lock held by runTxnWithMetaGroupLock keeps this validation and + // The lock held by runTxnWithMetaGroupLock keeps this validation and // the assignment update atomic with respect to UpdateGroupsSafely. if err := manager.mgm.reassignKeyspaceLocked(txn, oldMetaServiceGroup, newMetaServiceGroup); err != nil { return err } + metaServiceGroupReassigned = true } oldUserKind := endpoint.StringUserKind(oldConfig[UserKindKey]) newUserKind := endpoint.StringUserKind(newConfig[UserKindKey]) @@ -965,6 +952,11 @@ func (manager *Manager) updateKeyspaceConfigTxn(name string, update func(meta *k } err := manager.runTxnWithMetaGroupLock(txnFunc) if err != nil { + if metaServiceGroupReassigned { + if rollbackErr := manager.mgm.updateAssignmentTxn(nil, newMetaServiceGroup, oldMetaServiceGroup); rollbackErr != nil { + log.Error("failed to revert meta-service group assignment", zap.Error(rollbackErr)) + } + } log.Warn("[keyspace] failed to update keyspace config", zap.Uint32("keyspace-id", meta.GetId()), zap.String("name", meta.GetName()), @@ -1059,10 +1051,10 @@ func (manager *Manager) RemoveKeyspace(txn kv.Txn, id uint32) error { } manager.keyspaceNameLookup.Delete(id) manager.keyspaceStateLookup.Delete(id) - // Decrement the meta-service group assignment count in the same txn so the - // persisted counter stays in sync with the keyspaces actually referencing the - // group. Without this, removed keyspaces leak count and could permanently - // block deleting an otherwise-empty group. + // Decrement the meta-service group assignment count so the cached counter + // stays in sync with the keyspaces actually referencing the group. Without + // this, removed keyspaces leak count and could permanently block deleting an + // otherwise-empty group in tests without the authoritative keyspace scanner. if manager.mgm != nil { if groupID := meta.GetConfig()[MetaServiceGroupIDKey]; groupID != "" { if err := manager.mgm.updateAssignmentTxn(txn, groupID, ""); err != nil { diff --git a/pkg/keyspace/keyspace_test.go b/pkg/keyspace/keyspace_test.go index 8a7153ddff..c21c280b9d 100644 --- a/pkg/keyspace/keyspace_test.go +++ b/pkg/keyspace/keyspace_test.go @@ -1095,7 +1095,8 @@ func TestAssignGroupAndSaveKeyspace(t *testing.T) { kgm := NewKeyspaceGroupManager(ctx, store, nil) // No groups available: assign=true (stale pre-check) must not fail creation. - emptyMgm := NewMetaServiceGroupManager(store, map[string]string{}) + emptyMgm, err := NewMetaServiceGroupManager(ctx, store, map[string]string{}) + re.NoError(err) managerNoGroup := NewKeyspaceManager(ctx, store, nil, mockid.NewIDAllocator(), &mockConfig{}, kgm, emptyMgm) cfg := map[string]string{} ks := &keyspacepb.KeyspaceMeta{Id: 100, Name: "ks-stale-precheck", Config: cfg} @@ -1107,7 +1108,8 @@ func TestAssignGroupAndSaveKeyspace(t *testing.T) { // A present, enabled group is still assigned. Groups are disabled by // default, so it must be enabled before it is eligible for assignment. - mgm := NewMetaServiceGroupManager(store, map[string]string{"g1": "addr1"}) + mgm, err := NewMetaServiceGroupManager(ctx, store, map[string]string{"g1": "addr1"}) + re.NoError(err) enabled := true re.NoError(mgm.PatchStatus(ctx, "g1", &MetaServiceGroupStatusPatch{Enabled: &enabled})) managerWithGroup := NewKeyspaceManager(ctx, store, nil, mockid.NewIDAllocator(), &mockConfig{}, kgm, mgm) @@ -1118,7 +1120,8 @@ func TestAssignGroupAndSaveKeyspace(t *testing.T) { // A group that exists but is disabled must not fail creation: the keyspace is // created without a meta-service group assignment instead. - disabledMgm := NewMetaServiceGroupManager(store, map[string]string{"g2": "addr2"}) + disabledMgm, err := NewMetaServiceGroupManager(ctx, store, map[string]string{"g2": "addr2"}) + re.NoError(err) managerDisabled := NewKeyspaceManager(ctx, store, nil, mockid.NewIDAllocator(), &mockConfig{}, kgm, disabledMgm) cfg3 := map[string]string{} ks3 := &keyspacepb.KeyspaceMeta{Id: 102, Name: "ks-disabled-group", Config: cfg3} diff --git a/pkg/keyspace/meta_service_group.go b/pkg/keyspace/meta_service_group.go index 87d53ebe9c..298a197499 100644 --- a/pkg/keyspace/meta_service_group.go +++ b/pkg/keyspace/meta_service_group.go @@ -18,6 +18,7 @@ import ( "context" "fmt" "math" + "time" "go.uber.org/zap" @@ -29,8 +30,14 @@ import ( "github.com/tikv/pd/server/config" ) +const ( + flushInterval = 5 * time.Minute + flushThreshold = 1000 +) + // MetaServiceGroupManager manages external meta-service groups. type MetaServiceGroupManager struct { + ctx context.Context store endpoint.MetaServiceGroupStorage syncutil.RWMutex // metaServiceGroups is the available external meta-service groups. @@ -41,6 +48,10 @@ type MetaServiceGroupManager struct { // the authoritative source for the delete guard so a stale persisted counter // cannot permanently block removing an actually-empty group. keyspaceAssignmentCounter func(groupIDs map[string]struct{}) (map[string]int, error) + cachedStatus map[string]*endpoint.MetaServiceGroupStatus + dirtyCount int + flushCh chan struct{} + isLeader func() bool } // SetKeyspaceAssignmentCounter sets the authoritative keyspace assignment @@ -50,30 +61,101 @@ func (m *MetaServiceGroupManager) SetKeyspaceAssignmentCounter(counter func(grou m.keyspaceAssignmentCounter = counter } +// SetLeaderChecker sets a function to determine leader ownership. +func (m *MetaServiceGroupManager) SetLeaderChecker(isLeader func() bool) { + m.Lock() + defer m.Unlock() + m.isLeader = isLeader +} + // NewMetaServiceGroupManager creates a new MetaServiceGroupManager. func NewMetaServiceGroupManager( + ctx context.Context, store endpoint.MetaServiceGroupStorage, metaServiceGroups map[string]string, -) *MetaServiceGroupManager { - return &MetaServiceGroupManager{ +) (*MetaServiceGroupManager, error) { + m := &MetaServiceGroupManager{ + ctx: ctx, store: store, metaServiceGroups: metaServiceGroups, + flushCh: make(chan struct{}, 1), + } + if err := m.RefreshCache(); err != nil { + return nil, err } + go m.flushLoop() + return m, nil } -// GetStatus returns the status of each meta-service group. -func (m *MetaServiceGroupManager) GetStatus(ctx context.Context) (map[string]*endpoint.MetaServiceGroupStatus, error) { - m.RLock() - defer m.RUnlock() +// RefreshCache loads the current status from storage into memory. +func (m *MetaServiceGroupManager) RefreshCache() error { + m.Lock() + defer m.Unlock() var ( err error statusMap map[string]*endpoint.MetaServiceGroupStatus ) - err = m.store.RunInTxn(ctx, func(txn kv.Txn) error { + if err = m.store.RunInTxn(m.ctx, func(txn kv.Txn) error { statusMap, err = m.store.LoadMetaServiceGroupStatus(txn, m.metaServiceGroups) return err - }) - return statusMap, err + }); err != nil { + log.Error("[keyspace] failed to load meta-service group status from storage", zap.Error(err)) + return err + } + m.cachedStatus = statusMap + m.dirtyCount = 0 + log.Info("[keyspace] meta-service group status loaded from storage", zap.Any("meta-service-group-status", statusMap)) + return nil +} + +func (m *MetaServiceGroupManager) flushLoop() { + ticker := time.NewTicker(flushInterval) + defer ticker.Stop() + for { + select { + case <-m.ctx.Done(): + return + case <-ticker.C: + if err := m.flushToStorage(); err != nil { + log.Error("[keyspace] failed to flush meta-service group status to storage", zap.Error(err)) + } + case <-m.flushCh: + if err := m.flushToStorage(); err != nil { + log.Error("[keyspace] failed to flush meta-service group status to storage", zap.Error(err)) + } + } + } +} + +func (m *MetaServiceGroupManager) flushToStorage() error { + m.Lock() + defer m.Unlock() + if m.isLeader != nil && !m.isLeader() { + return nil + } + if m.dirtyCount == 0 { + return nil + } + if err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) error { + for id, status := range m.cachedStatus { + if err := m.store.SaveMetaServiceGroupStatus(txn, id, status); err != nil { + return err + } + } + return nil + }); err != nil { + return err + } + m.dirtyCount = 0 + return nil +} + +// GetStatus returns the status of each meta-service group. +func (m *MetaServiceGroupManager) GetStatus(ctx context.Context) (map[string]*endpoint.MetaServiceGroupStatus, error) { + _ = ctx + m.RLock() + defer m.RUnlock() + return copyStatusMap(m.cachedStatus), nil } // GetAssignmentCounts returns the count of each meta-service group. @@ -90,6 +172,19 @@ func (m *MetaServiceGroupManager) GetAssignmentCounts(ctx context.Context) (map[ return counts, nil } +func copyStatusMap(statusMap map[string]*endpoint.MetaServiceGroupStatus) map[string]*endpoint.MetaServiceGroupStatus { + statuses := make(map[string]*endpoint.MetaServiceGroupStatus, len(statusMap)) + for groupID, status := range statusMap { + if status == nil { + statuses[groupID] = &endpoint.MetaServiceGroupStatus{} + continue + } + copiedStatus := *status + statuses[groupID] = &copiedStatus + } + return statuses +} + // MetaServiceGroupStatusPatch represents a patch operation for a meta-service group. // NOTE: This type is exported by HTTP API. Please pay more attention when modifying it. type MetaServiceGroupStatusPatch struct { @@ -102,38 +197,35 @@ func (m *MetaServiceGroupManager) PatchStatus(ctx context.Context, groupID strin if patch.AssignmentCount != nil && *patch.AssignmentCount < 0 { return ErrInvalidAssignmentCount } - m.RLock() - defer m.RUnlock() - // Validate existence against the in-memory group set under the lock, then - // touch only the target group's status in the txn. Loading every group would - // widen the etcd compare set so an unrelated concurrent assignment could make - // this patch fail with a spurious txn conflict. + m.Lock() + defer m.Unlock() if _, ok := m.metaServiceGroups[groupID]; !ok { return ErrUnknownMetaServiceGroup } - return m.store.RunInTxn(ctx, func(txn kv.Txn) error { - status, err := m.loadGroupStatus(txn, groupID) - if err != nil { - return err - } - if patch.AssignmentCount != nil { - status.AssignmentCount = *patch.AssignmentCount - } - if patch.Enabled != nil { - status.Enabled = *patch.Enabled - } - return m.store.SaveMetaServiceGroupStatus(txn, groupID, status) - }) + currentStatus := m.cachedStatus[groupID] + if currentStatus == nil { + currentStatus = &endpoint.MetaServiceGroupStatus{} + } + newStatus := *currentStatus + if patch.AssignmentCount != nil { + newStatus.AssignmentCount = *patch.AssignmentCount + } + if patch.Enabled != nil { + newStatus.Enabled = *patch.Enabled + } + if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error { + return m.store.SaveMetaServiceGroupStatus(txn, groupID, &newStatus) + }); err != nil { + return err + } + m.cachedStatus[groupID] = &newStatus + return nil } -func (m *MetaServiceGroupManager) findMinMetaGroup(txn kv.Txn) (string, error) { - statusMap, err := m.store.LoadMetaServiceGroupStatus(txn, m.metaServiceGroups) - if err != nil { - return "", err - } +func (m *MetaServiceGroupManager) findMinMetaGroupLocked() (string, error) { minCount := math.MaxInt var assignedGroup string - for currentGroup, status := range statusMap { + for currentGroup, status := range m.cachedStatus { if status.Enabled && status.AssignmentCount < minCount { minCount = status.AssignmentCount assignedGroup = currentGroup @@ -146,32 +238,31 @@ func (m *MetaServiceGroupManager) findMinMetaGroup(txn kv.Txn) (string, error) { } // PickGroup returns the meta-service group with the least assigned keyspaces -// without updating the persisted assignment count. +// and updates the cached assignment count. func (m *MetaServiceGroupManager) PickGroup(ctx context.Context) (string, error) { - m.RLock() - defer m.RUnlock() + _ = ctx + m.Lock() + defer m.Unlock() return m.pickGroupLocked(ctx) } // hasGroupsLocked reports whether any meta-service group is currently available. -// The caller must hold the read lock. +// The caller must hold the lock. func (m *MetaServiceGroupManager) hasGroupsLocked() bool { return len(m.metaServiceGroups) > 0 } -// pickGroupLocked is PickGroup with the read lock already held by the caller. +// pickGroupLocked is PickGroup with the lock already held by the caller. // Callers that need group selection and the subsequent keyspace metadata save to // be atomic with respect to group deletion (which takes the write lock) must -// hold the read lock across both, e.g. via Manager.assignGroupAndSaveKeyspace. +// hold the lock across both, e.g. via Manager.assignGroupAndSaveKeyspace. func (m *MetaServiceGroupManager) pickGroupLocked(ctx context.Context) (string, error) { - var assignedGroup string - if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error { - var err error - if assignedGroup, err = m.findMinMetaGroup(txn); err != nil { - return err - } - return m.updateAssignmentTxn(txn, "", assignedGroup) - }); err != nil { + _ = ctx + assignedGroup, err := m.findMinMetaGroupLocked() + if err != nil { + return "", err + } + if err := m.updateAssignmentLockedTxn(nil, "", assignedGroup); err != nil { return "", err } return assignedGroup, nil @@ -181,36 +272,26 @@ func (m *MetaServiceGroupManager) pickGroupLocked(ctx context.Context) (string, // It returns the assigned meta-service group and an error if any. // only used for testing now, as it doesn't guarantee the atomicity of select and update. UpdateAssignment should be used in production code instead. func (m *MetaServiceGroupManager) AssignToGroup(ctx context.Context, count int) (string, error) { + _ = ctx if count < 0 { return "", ErrInvalidAssignmentCount } - m.RLock() - defer m.RUnlock() - var assignedGroup string - if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error { - var err error - assignedGroup, err = m.findMinMetaGroup(txn) - if err != nil { - return err - } - statusMap, err := m.store.LoadMetaServiceGroupStatus(txn, m.metaServiceGroups) - if err != nil { - return err - } - status := statusMap[assignedGroup] - status.AssignmentCount += count - return m.store.SaveMetaServiceGroupStatus(txn, assignedGroup, status) - }); err != nil { + m.Lock() + defer m.Unlock() + assignedGroup, err := m.findMinMetaGroupLocked() + if err != nil { return "", err } + m.cachedStatus[assignedGroup].AssignmentCount += count + m.markDirtyLocked(count) return assignedGroup, nil } // reassignKeyspaceLocked validates that newGroupID (if any) still exists and // moves a single keyspace assignment from oldGroupID to newGroupID within txn. -// The caller must hold the read lock for the whole enclosing transaction so a +// The caller must hold the lock for the whole enclosing transaction so a // concurrent UpdateGroupsSafely cannot delete a group between this validation -// and the persisted assignment count update. +// and the cached assignment count update. func (m *MetaServiceGroupManager) reassignKeyspaceLocked(txn kv.Txn, oldGroupID, newGroupID string) error { if newGroupID != "" { if _, ok := m.metaServiceGroups[newGroupID]; !ok { @@ -218,60 +299,51 @@ func (m *MetaServiceGroupManager) reassignKeyspaceLocked(txn kv.Txn, oldGroupID, } // Disabled groups are skipped by automatic assignment, so reject moving a // keyspace into one to keep manual reassignment consistent with it. - statusMap, err := m.store.LoadMetaServiceGroupStatus(txn, map[string]string{newGroupID: ""}) - if err != nil { - return err - } - if status := statusMap[newGroupID]; status == nil || !status.Enabled { + if status := m.cachedStatus[newGroupID]; status == nil || !status.Enabled { return ErrMetaServiceGroupDisabled } } - return m.updateAssignmentTxn(txn, oldGroupID, newGroupID) + return m.updateAssignmentLockedTxn(txn, oldGroupID, newGroupID) } func (m *MetaServiceGroupManager) updateAssignmentTxn(txn kv.Txn, oldGroupID, newGroupID string) error { - // Load only the affected groups instead of the whole m.metaServiceGroups map: - // some callers (e.g. RemoveKeyspace) reach this without holding the - // meta-service group lock, so reading the shared map here would race with - // UpdateGroupsSafely. + m.Lock() + defer m.Unlock() + return m.updateAssignmentLockedTxn(txn, oldGroupID, newGroupID) +} + +func (m *MetaServiceGroupManager) updateAssignmentLockedTxn(txn kv.Txn, oldGroupID, newGroupID string) error { + _ = txn + dirtyCount := 0 if oldGroupID != "" { - status, err := m.loadGroupStatus(txn, oldGroupID) - if err != nil { - return err - } - // Only persist a decrement when the group still has assignments. A deleted - // group's status key is already removed, so skipping avoids recreating a - // stale zero-value status; a count of 0 needs no change anyway. This also - // guards against underflow after a manual reset via PatchStatus, which - // would otherwise make findMinMetaGroup prefer the group. - if status.AssignmentCount > 0 { + if status := m.cachedStatus[oldGroupID]; status != nil && status.AssignmentCount > 0 { status.AssignmentCount-- - if err := m.store.SaveMetaServiceGroupStatus(txn, oldGroupID, status); err != nil { - return err - } + dirtyCount++ } } if newGroupID != "" { - status, err := m.loadGroupStatus(txn, newGroupID) - if err != nil { - return err + status := m.cachedStatus[newGroupID] + if status == nil { + return ErrUnknownMetaServiceGroup } status.AssignmentCount++ - if err := m.store.SaveMetaServiceGroupStatus(txn, newGroupID, status); err != nil { - return err - } + dirtyCount++ + } + if dirtyCount > 0 { + m.markDirtyLocked(dirtyCount) } return nil } -// loadGroupStatus loads the persisted status of a single meta-service group -// within txn, without touching the shared m.metaServiceGroups map. -func (m *MetaServiceGroupManager) loadGroupStatus(txn kv.Txn, groupID string) (*endpoint.MetaServiceGroupStatus, error) { - statusMap, err := m.store.LoadMetaServiceGroupStatus(txn, map[string]string{groupID: ""}) - if err != nil { - return nil, err +func (m *MetaServiceGroupManager) markDirtyLocked(count int) { + m.dirtyCount += count + if m.dirtyCount < flushThreshold { + return + } + select { + case m.flushCh <- struct{}{}: + default: } - return statusMap[groupID], nil } // AttachEndpoints append potential meta-service group endpoint to the given keyspace config map. @@ -321,7 +393,7 @@ func (m *MetaServiceGroupManager) UpdateGroupsSafely( // persistGroupsLocked performs the delete-guard check and persists the new // groups while holding the write lock, which blocks concurrent keyspace -// assignment (AssignToGroup/PickGroup/reassign all take the read lock). +// assignment (AssignToGroup/PickGroup/reassign all take the lock). func (m *MetaServiceGroupManager) persistGroupsLocked( ctx context.Context, metaServiceGroups map[string]string, @@ -345,12 +417,20 @@ func (m *MetaServiceGroupManager) persistGroupsLocked( return err } m.metaServiceGroups = metaServiceGroups + for groupID := range metaServiceGroups { + if m.cachedStatus[groupID] == nil { + m.cachedStatus[groupID] = &endpoint.MetaServiceGroupStatus{} + } + } // Clear the persisted status for deleted groups so re-adding a group with // the same ID does not inherit a stale assignment count or enabled state, // which would skew list output and PickGroup balancing. Best-effort: the // config deletion is already persisted and the delete guard relies on // actual keyspace scans, not this counter. if len(deletedGroups) > 0 { + for _, id := range deletedGroups { + delete(m.cachedStatus, id) + } if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error { for _, id := range deletedGroups { if err := m.store.RemoveMetaServiceGroupStatus(txn, id); err != nil { @@ -368,9 +448,9 @@ func (m *MetaServiceGroupManager) persistGroupsLocked( // assignedKeyspaceCounts returns the number of keyspaces assigned to each of the // given groups. It prefers the authoritative keyspace scan (immune to counter -// drift) and falls back to the persisted counter when no scanner is configured, +// drift) and falls back to the cached counter when no scanner is configured, // e.g. in unit tests without a keyspace manager. -func (m *MetaServiceGroupManager) assignedKeyspaceCounts(ctx context.Context, groupIDs []string) (map[string]int, error) { +func (m *MetaServiceGroupManager) assignedKeyspaceCounts(_ context.Context, groupIDs []string) (map[string]int, error) { if m.keyspaceAssignmentCounter != nil { set := make(map[string]struct{}, len(groupIDs)) for _, id := range groupIDs { @@ -378,22 +458,14 @@ func (m *MetaServiceGroupManager) assignedKeyspaceCounts(ctx context.Context, gr } return m.keyspaceAssignmentCounter(set) } - // Fallback path: derive counts from the persisted status. The caller holds + // Fallback path: derive counts from the cached status. The caller holds // the write lock, so m.metaServiceGroups is accessed without an extra read // lock (which would deadlock against the held write lock). - var counts map[string]int - if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error { - statusMap, err := m.store.LoadMetaServiceGroupStatus(txn, m.metaServiceGroups) - if err != nil { - return err - } - counts = make(map[string]int, len(statusMap)) - for id, status := range statusMap { + counts := make(map[string]int, len(groupIDs)) + for _, id := range groupIDs { + if status := m.cachedStatus[id]; status != nil { counts[id] = status.AssignmentCount } - return nil - }); err != nil { - return nil, err } return counts, nil } @@ -403,4 +475,14 @@ func (m *MetaServiceGroupManager) updateGroups(metaServiceGroups map[string]stri m.Lock() defer m.Unlock() m.metaServiceGroups = metaServiceGroups + for groupID := range metaServiceGroups { + if m.cachedStatus[groupID] == nil { + m.cachedStatus[groupID] = &endpoint.MetaServiceGroupStatus{} + } + } + for groupID := range m.cachedStatus { + if _, ok := metaServiceGroups[groupID]; !ok { + delete(m.cachedStatus, groupID) + } + } } diff --git a/pkg/keyspace/meta_service_group_test.go b/pkg/keyspace/meta_service_group_test.go index 80f956b4e3..448a156a4b 100644 --- a/pkg/keyspace/meta_service_group_test.go +++ b/pkg/keyspace/meta_service_group_test.go @@ -17,6 +17,7 @@ package keyspace import ( "context" "testing" + "time" "github.com/stretchr/testify/suite" @@ -46,7 +47,9 @@ func mockMetaServiceGroups() map[string]string { func (suite *metaServiceGroupTestSuite) SetupTest() { suite.ctx, suite.cancel = context.WithCancel(context.Background()) store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil) - suite.manager = NewMetaServiceGroupManager(store, mockMetaServiceGroups()) + var err error + suite.manager, err = NewMetaServiceGroupManager(suite.ctx, store, mockMetaServiceGroups()) + suite.Require().NoError(err) } func (suite *metaServiceGroupTestSuite) TearDownTest() { @@ -284,6 +287,42 @@ func (suite *metaServiceGroupTestSuite) TestReassignRejectsDisabledGroup() { re.NoError(err) } +func (suite *metaServiceGroupTestSuite) TestRefreshCacheLoadsFromStorage() { + re := suite.Require() + status := &endpoint.MetaServiceGroupStatus{ + AssignmentCount: 10, + Enabled: true, + } + err := suite.manager.store.RunInTxn(suite.ctx, func(txn kv.Txn) error { + return suite.manager.store.SaveMetaServiceGroupStatus(txn, "etcd-group-0", status) + }) + re.NoError(err) + + re.NoError(suite.manager.RefreshCache()) + statusMap, err := suite.manager.GetStatus(suite.ctx) + re.NoError(err) + re.Equal(10, statusMap["etcd-group-0"].AssignmentCount) + re.True(statusMap["etcd-group-0"].Enabled) +} + +func (suite *metaServiceGroupTestSuite) TestFlushAfterWriteThreshold() { + re := suite.Require() + suite.enableAllGroups() + assigned, err := suite.manager.AssignToGroup(suite.ctx, flushThreshold) + re.NoError(err) + re.NotEmpty(assigned) + + re.Eventually(func() bool { + var statusMap map[string]*endpoint.MetaServiceGroupStatus + err := suite.manager.store.RunInTxn(suite.ctx, func(txn kv.Txn) error { + var err error + statusMap, err = suite.manager.store.LoadMetaServiceGroupStatus(txn, mockMetaServiceGroups()) + return err + }) + return err == nil && statusMap[assigned] != nil && statusMap[assigned].AssignmentCount == flushThreshold + }, 5*time.Second, 100*time.Millisecond) +} + func (suite *metaServiceGroupTestSuite) enableAllGroups() { re := suite.Require() enabled := true diff --git a/server/server.go b/server/server.go index fe0a2a37d1..7ab8a9d30e 100644 --- a/server/server.go +++ b/server/server.go @@ -534,7 +534,17 @@ func (s *Server) startServer(ctx context.Context) error { if s.IsKeyspaceGroupEnabled() { s.keyspaceGroupManager = keyspace.NewKeyspaceGroupManager(s.ctx, s.storage, s.client) } - s.metaServiceGroupManager = keyspace.NewMetaServiceGroupManager(s.storage, s.cfg.Keyspace.GetMetaServiceGroups()) + s.metaServiceGroupManager, err = keyspace.NewMetaServiceGroupManager(s.ctx, s.storage, s.cfg.Keyspace.GetMetaServiceGroups()) + if err != nil { + return err + } + s.metaServiceGroupManager.SetLeaderChecker(s.IsServing) + s.AddServiceReadyCallback(func(_ context.Context) error { + if s.metaServiceGroupManager == nil { + return nil + } + return s.metaServiceGroupManager.RefreshCache() + }) s.keyspaceManager = keyspace.NewKeyspaceManager( s.ctx, s.storage, From c4a0aa95c000dbdd2ea85570f41b54717a3f59ee Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Tue, 7 Jul 2026 14:29:36 +0800 Subject: [PATCH 02/11] keyspace: refine meta-service group status cache Flush cached status outside the mgm lock via a snapshot so a slow etcd write no longer blocks concurrent assignment operations, restoring the dirty count on failure. Document that the assignment count update is applied to the cache and is intentionally not part of the caller's txn, fix a truncated rollback comment, and drop unused ctx parameters. Co-Authored-By: Claude Opus 4.8 Signed-off-by: tongjian <1045931706@qq.com> --- pkg/keyspace/keyspace.go | 6 ++-- pkg/keyspace/meta_service_group.go | 48 +++++++++++++++++++++--------- 2 files changed, 38 insertions(+), 16 deletions(-) diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go index 18c1e31b6a..ed5e25930f 100644 --- a/pkg/keyspace/keyspace.go +++ b/pkg/keyspace/keyspace.go @@ -589,7 +589,7 @@ func (manager *Manager) assignGroupAndSaveKeyspace(assign bool, config *map[stri if !manager.mgm.hasGroupsLocked() { return manager.saveNewKeyspace(keyspace) } - groupID, err := manager.mgm.pickGroupLocked(manager.ctx) + groupID, err := manager.mgm.pickGroupLocked() if err != nil { if goerrors.Is(err, errNoAvailableMetaServiceGroups) { // Groups exist but none are enabled: create the keyspace without a @@ -605,7 +605,9 @@ func (manager *Manager) assignGroupAndSaveKeyspace(assign bool, config *map[stri (*config)[MetaServiceGroupIDKey] = groupID keyspace.Config = *config if err := manager.saveNewKeyspace(keyspace); err != nil { - // Roll back the reservation made by pickGroupLocked. This only performs + // Roll back the reservation made by pickGroupLocked. This only mutates the + // cached count and does not re-take the mgm lock, so it is safe to call + // while still holding the lock acquired above. if rollbackErr := manager.mgm.updateAssignmentLockedTxn(nil, groupID, ""); rollbackErr != nil { log.Error("failed to revert meta-service group assignment", zap.Error(rollbackErr)) } diff --git a/pkg/keyspace/meta_service_group.go b/pkg/keyspace/meta_service_group.go index 298a197499..b8587090e3 100644 --- a/pkg/keyspace/meta_service_group.go +++ b/pkg/keyspace/meta_service_group.go @@ -128,31 +128,43 @@ func (m *MetaServiceGroupManager) flushLoop() { } func (m *MetaServiceGroupManager) flushToStorage() error { + // Snapshot the dirty state under the lock and reset the dirty counter + // optimistically, then release the lock before the storage transaction so a + // slow etcd write does not block concurrent assignment operations. The + // snapshot is a deep copy, so later cache mutations are not observed by this + // flush. On failure the dirty count is restored so the next tick retries. m.Lock() - defer m.Unlock() if m.isLeader != nil && !m.isLeader() { + m.Unlock() return nil } if m.dirtyCount == 0 { + m.Unlock() return nil } + snapshot := copyStatusMap(m.cachedStatus) + flushed := m.dirtyCount + m.dirtyCount = 0 + m.Unlock() + if err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) error { - for id, status := range m.cachedStatus { + for id, status := range snapshot { if err := m.store.SaveMetaServiceGroupStatus(txn, id, status); err != nil { return err } } return nil }); err != nil { + m.Lock() + m.dirtyCount += flushed + m.Unlock() return err } - m.dirtyCount = 0 return nil } // GetStatus returns the status of each meta-service group. -func (m *MetaServiceGroupManager) GetStatus(ctx context.Context) (map[string]*endpoint.MetaServiceGroupStatus, error) { - _ = ctx +func (m *MetaServiceGroupManager) GetStatus(_ context.Context) (map[string]*endpoint.MetaServiceGroupStatus, error) { m.RLock() defer m.RUnlock() return copyStatusMap(m.cachedStatus), nil @@ -239,11 +251,10 @@ func (m *MetaServiceGroupManager) findMinMetaGroupLocked() (string, error) { // PickGroup returns the meta-service group with the least assigned keyspaces // and updates the cached assignment count. -func (m *MetaServiceGroupManager) PickGroup(ctx context.Context) (string, error) { - _ = ctx +func (m *MetaServiceGroupManager) PickGroup(_ context.Context) (string, error) { m.Lock() defer m.Unlock() - return m.pickGroupLocked(ctx) + return m.pickGroupLocked() } // hasGroupsLocked reports whether any meta-service group is currently available. @@ -256,8 +267,7 @@ func (m *MetaServiceGroupManager) hasGroupsLocked() bool { // Callers that need group selection and the subsequent keyspace metadata save to // be atomic with respect to group deletion (which takes the write lock) must // hold the lock across both, e.g. via Manager.assignGroupAndSaveKeyspace. -func (m *MetaServiceGroupManager) pickGroupLocked(ctx context.Context) (string, error) { - _ = ctx +func (m *MetaServiceGroupManager) pickGroupLocked() (string, error) { assignedGroup, err := m.findMinMetaGroupLocked() if err != nil { return "", err @@ -271,8 +281,7 @@ func (m *MetaServiceGroupManager) pickGroupLocked(ctx context.Context) (string, // AssignToGroup increments count of the meta-service group with least assigned keyspaces. // It returns the assigned meta-service group and an error if any. // only used for testing now, as it doesn't guarantee the atomicity of select and update. UpdateAssignment should be used in production code instead. -func (m *MetaServiceGroupManager) AssignToGroup(ctx context.Context, count int) (string, error) { - _ = ctx +func (m *MetaServiceGroupManager) AssignToGroup(_ context.Context, count int) (string, error) { if count < 0 { return "", ErrInvalidAssignmentCount } @@ -306,14 +315,25 @@ func (m *MetaServiceGroupManager) reassignKeyspaceLocked(txn kv.Txn, oldGroupID, return m.updateAssignmentLockedTxn(txn, oldGroupID, newGroupID) } +// updateAssignmentTxn is updateAssignmentLockedTxn with the mgm lock acquired +// for the caller. The txn is accepted for call-site symmetry with the storage +// transaction the caller runs, but see updateAssignmentLockedTxn: the count +// change is applied to the in-memory cache and is NOT part of txn, so it is not +// rolled back if txn fails to commit. func (m *MetaServiceGroupManager) updateAssignmentTxn(txn kv.Txn, oldGroupID, newGroupID string) error { m.Lock() defer m.Unlock() return m.updateAssignmentLockedTxn(txn, oldGroupID, newGroupID) } -func (m *MetaServiceGroupManager) updateAssignmentLockedTxn(txn kv.Txn, oldGroupID, newGroupID string) error { - _ = txn +// updateAssignmentLockedTxn moves one assignment from oldGroupID to newGroupID in +// the cached status and marks it dirty for the flush loop to persist later. The +// txn parameter is intentionally unused: unlike the pre-cache implementation the +// count is no longer written inside the caller's storage transaction, so callers +// must not assume the count update is atomic with their txn. Counts are +// best-effort load-balancing hints; the delete guard uses the authoritative +// keyspace scan, and RefreshCache reconciles any drift on the next leader term. +func (m *MetaServiceGroupManager) updateAssignmentLockedTxn(_ kv.Txn, oldGroupID, newGroupID string) error { dirtyCount := 0 if oldGroupID != "" { if status := m.cachedStatus[oldGroupID]; status != nil && status.AssignmentCount > 0 { From f514e59e5e72abca91b4312cc7db52a3cbbd3cc0 Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Wed, 8 Jul 2026 16:03:01 +0800 Subject: [PATCH 03/11] keyspace: persist group changes before updating in-memory view In persistGroupsLocked the deleted-group status removal ran after the in-memory metaServiceGroups/cache update. Move it ahead of the memory mutation so all storage writes (persist config + clear deleted status) complete before the in-memory view changes, keeping durable state and memory consistent on the failure path. Co-Authored-By: Claude Opus 4.8 Signed-off-by: tongjian <1045931706@qq.com> --- pkg/keyspace/meta_service_group.go | 31 +++++++++++++++--------------- 1 file changed, 16 insertions(+), 15 deletions(-) diff --git a/pkg/keyspace/meta_service_group.go b/pkg/keyspace/meta_service_group.go index 6777fc1300..54d0225f8b 100644 --- a/pkg/keyspace/meta_service_group.go +++ b/pkg/keyspace/meta_service_group.go @@ -477,21 +477,10 @@ func (m *MetaServiceGroupManager) persistGroupsLocked( if err := persist(); err != nil { return err } - m.metaServiceGroups = metaServiceGroups - // Sync the cache to the new group set. Clearing deleted groups here (and the - // persisted status below) keeps re-adding a group with the same ID from - // inheriting a stale assignment count or enabled state, which would skew list - // output and PickGroup balancing. - m.statusMu.Lock() - for groupID := range metaServiceGroups { - if m.cachedStatus[groupID] == nil { - m.cachedStatus[groupID] = &endpoint.MetaServiceGroupStatus{} - } - } - for _, id := range deletedGroups { - delete(m.cachedStatus, id) - } - m.statusMu.Unlock() + // Clear the persisted status for deleted groups before touching memory, so all + // storage writes complete before the in-memory view changes. This keeps + // re-adding a group with the same ID from inheriting a stale assignment count + // or enabled state, which would skew list output and PickGroup balancing. // Best-effort: the config deletion is already persisted and the delete guard // relies on actual keyspace scans, not this counter. if len(deletedGroups) > 0 { @@ -507,6 +496,18 @@ func (m *MetaServiceGroupManager) persistGroupsLocked( zap.Strings("deleted-groups", deletedGroups), zap.Error(err)) } } + // Apply the change to the in-memory view only after the storage writes above. + m.metaServiceGroups = metaServiceGroups + m.statusMu.Lock() + for groupID := range metaServiceGroups { + if m.cachedStatus[groupID] == nil { + m.cachedStatus[groupID] = &endpoint.MetaServiceGroupStatus{} + } + } + for _, id := range deletedGroups { + delete(m.cachedStatus, id) + } + m.statusMu.Unlock() return nil } From 77936b6da1cfc27e79ffc9fc4d25aa59f6e3a2c4 Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Wed, 8 Jul 2026 16:07:49 +0800 Subject: [PATCH 04/11] keyspace: add HasGroups to avoid map copy on create path The keyspace creation path only needs to know whether any meta-service group exists, but called len(GetGroups()) > 0, which allocates and copies the whole group map. Add HasGroups, a lock-guarded existence check that reuses hasGroupsLocked without copying, and use it on both create paths. Co-Authored-By: Claude Opus 4.8 Signed-off-by: tongjian <1045931706@qq.com> --- pkg/keyspace/keyspace.go | 4 ++-- pkg/keyspace/meta_service_group.go | 9 +++++++++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go index 0e6b3581ca..8282f8162e 100644 --- a/pkg/keyspace/keyspace.go +++ b/pkg/keyspace/keyspace.go @@ -329,7 +329,7 @@ func (manager *Manager) CreateKeyspace(request *CreateKeyspaceRequest) (*keyspac } // Assign a meta-service group (if any exist) and save the keyspace atomically // with respect to group deletion. The assignment is reflected in request.Config. - assignToMetaServiceGroup := manager.mgm != nil && len(manager.mgm.GetGroups()) > 0 + assignToMetaServiceGroup := manager.mgm != nil && manager.mgm.HasGroups() if err = manager.assignGroupAndSaveKeyspace(assignToMetaServiceGroup, &request.Config, keyspace); err != nil { log.Warn("[create-keyspace] failed to save keyspace before split", zap.Uint32("keyspace-id", keyspace.GetId()), @@ -476,7 +476,7 @@ func (manager *Manager) CreateKeyspaceByID(request *CreateKeyspaceByIDRequest) ( } // Assign a meta-service group (if any exist) and save the keyspace atomically // with respect to group deletion. The assignment is reflected in request.Config. - assignToMetaServiceGroup := manager.mgm != nil && len(manager.mgm.GetGroups()) > 0 + assignToMetaServiceGroup := manager.mgm != nil && manager.mgm.HasGroups() if err = manager.assignGroupAndSaveKeyspace(assignToMetaServiceGroup, &request.Config, keyspace); err != nil { log.Warn("[keyspace] failed to save keyspace before split", zap.Uint32("keyspace-id", keyspace.GetId()), diff --git a/pkg/keyspace/meta_service_group.go b/pkg/keyspace/meta_service_group.go index 54d0225f8b..815c797ee7 100644 --- a/pkg/keyspace/meta_service_group.go +++ b/pkg/keyspace/meta_service_group.go @@ -431,6 +431,15 @@ func (m *MetaServiceGroupManager) GetGroups() map[string]string { return groups } +// HasGroups reports whether any meta-service group is currently available. It +// avoids the map copy of GetGroups for callers that only need the existence +// check, e.g. the keyspace creation path. +func (m *MetaServiceGroupManager) HasGroups() bool { + m.RLock() + defer m.RUnlock() + return m.hasGroupsLocked() +} + // UpdateGroupsSafely persists and applies meta-service group changes while // blocking concurrent keyspace assignments. func (m *MetaServiceGroupManager) UpdateGroupsSafely( From e5262c6be590fbcc9ee39f980864e0451bc313df Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Mon, 13 Jul 2026 14:51:24 +0800 Subject: [PATCH 05/11] keyspace: serialize meta-service group flush with synchronous writes flushToStorage snapshotted cachedStatus, released the lock, then wrote the snapshot without any lock. A concurrent PatchStatus could persist a new value that the stale snapshot then overwrote, losing the patch after the next RefreshCache; the same race could recreate a status key that UpdateGroupsSafely had deleted. Hold the mgm write lock across the whole flush, including the storage write, so it serializes with PatchStatus and persistGroupsLocked. This does not deadlock: the metaLock-holding paths only take statusMu, never the mgm lock. Add a regression test that pauses a flush mid-write and checks a concurrent PatchStatus wins. Co-Authored-By: Claude Opus 4.8 Signed-off-by: tongjian <1045931706@qq.com> --- pkg/keyspace/meta_service_group.go | 33 +++++++----- pkg/keyspace/meta_service_group_test.go | 69 +++++++++++++++++++++++++ 2 files changed, 89 insertions(+), 13 deletions(-) diff --git a/pkg/keyspace/meta_service_group.go b/pkg/keyspace/meta_service_group.go index 815c797ee7..2b1ed18e26 100644 --- a/pkg/keyspace/meta_service_group.go +++ b/pkg/keyspace/meta_service_group.go @@ -143,21 +143,28 @@ func (m *MetaServiceGroupManager) flushLoop() { } func (m *MetaServiceGroupManager) flushToStorage() error { - // Only the serving leader persists; a stale read here is benign because the - // next leader reloads the authoritative status via RefreshCache. - m.RLock() - skip := m.isLeader != nil && !m.isLeader() - m.RUnlock() - if skip { + // Hold the mgm write lock for the whole flush, including the storage write, so + // it serializes with PatchStatus and persistGroupsLocked (both take the mgm + // lock). Without that, the snapshot below could overwrite a status those paths + // persisted synchronously in the meantime, or recreate a status key that + // persistGroupsLocked just deleted. This cannot deadlock: the metaLock-holding + // paths (RemoveKeyspace, tombstone unassignment) only take statusMu, never the + // mgm lock, so nothing waits on the mgm lock while holding metaLock, and flush + // never takes metaLock. + m.Lock() + defer m.Unlock() + // Only the serving leader persists; a follower's next leader term reloads the + // authoritative status via RefreshCache. + if m.isLeader != nil && !m.isLeader() { return nil } - // Snapshot the dirty state under statusMu and reset the dirty counter - // optimistically, then release the lock before the storage transaction so a - // slow etcd write does not block concurrent assignment operations. The - // snapshot is a deep copy, so later cache mutations are not observed by this - // flush. On failure the dirty count is restored so the next tick retries. - // Only flushLoop calls this, so there is no concurrent flush to race the - // reset/restore. + // Snapshot the dirty state under statusMu (a deep copy) and reset the dirty + // counter, without holding statusMu across the storage write. The snapshot + // keeps the write from reading cachedStatus while a concurrent statusMu holder + // (an assignment-count update under metaLock) mutates it; such updates re-mark + // the counter dirty and are flushed on the next tick. On failure the dirty + // count is restored. Only flushLoop calls this, so no concurrent flush races + // the reset/restore. m.statusMu.Lock() if m.dirtyCount == 0 { m.statusMu.Unlock() diff --git a/pkg/keyspace/meta_service_group_test.go b/pkg/keyspace/meta_service_group_test.go index 448a156a4b..a5510b7792 100644 --- a/pkg/keyspace/meta_service_group_test.go +++ b/pkg/keyspace/meta_service_group_test.go @@ -16,9 +16,11 @@ package keyspace import ( "context" + "sync/atomic" "testing" "time" + "github.com/stretchr/testify/require" "github.com/stretchr/testify/suite" "github.com/tikv/pd/pkg/storage/endpoint" @@ -332,3 +334,70 @@ func (suite *metaServiceGroupTestSuite) enableAllGroups() { })) } } + +// blockingMetaServiceGroupStorage pauses the first SaveMetaServiceGroupStatus so +// a flush can be held mid-write to exercise its concurrency with PatchStatus. +type blockingMetaServiceGroupStorage struct { + *endpoint.StorageEndpoint + blockNextSave atomic.Bool + saveStarted chan struct{} + unblockSave chan struct{} +} + +func (s *blockingMetaServiceGroupStorage) SaveMetaServiceGroupStatus(txn kv.Txn, id string, status *endpoint.MetaServiceGroupStatus) error { + if s.blockNextSave.CompareAndSwap(true, false) { + close(s.saveStarted) + <-s.unblockSave + } + return s.StorageEndpoint.SaveMetaServiceGroupStatus(txn, id, status) +} + +// TestFlushDoesNotOverwriteConcurrentPatchStatus verifies that an in-flight flush +// cannot clobber a status persisted by a concurrent PatchStatus. flushToStorage +// holds the mgm lock across its storage write, so PatchStatus serializes after it +// and its value is the persisted winner. +func TestFlushDoesNotOverwriteConcurrentPatchStatus(t *testing.T) { + re := require.New(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + store := &blockingMetaServiceGroupStorage{ + StorageEndpoint: endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil), + saveStarted: make(chan struct{}), + unblockSave: make(chan struct{}), + } + manager, err := NewMetaServiceGroupManager(ctx, store, mockMetaServiceGroups()) + re.NoError(err) + + const groupID = "etcd-group-0" + manager.statusMu.Lock() + manager.cachedStatus[groupID] = &endpoint.MetaServiceGroupStatus{AssignmentCount: 1, Enabled: true} + manager.dirtyCount = 1 + manager.statusMu.Unlock() + + // Start a flush and pause it inside the storage write, holding the mgm lock. + store.blockNextSave.Store(true) + flushErr := make(chan error, 1) + go func() { flushErr <- manager.flushToStorage() }() + <-store.saveStarted + + // PatchStatus blocks on the mgm lock the paused flush holds, so run it in a + // goroutine; it can only complete once the flush releases the lock. + const newCount = 99 + patchErr := make(chan error, 1) + go func() { + c := newCount + patchErr <- manager.PatchStatus(ctx, groupID, &MetaServiceGroupStatusPatch{AssignmentCount: &c}) + }() + + // Let the flush finish; PatchStatus then serializes strictly after it. + close(store.unblockSave) + re.NoError(<-flushErr) + re.NoError(<-patchErr) + + // The patch applied after the flush must be the persisted winner. + re.NoError(manager.RefreshCache()) + statusMap, err := manager.GetStatus(ctx) + re.NoError(err) + re.Equal(newCount, statusMap[groupID].AssignmentCount) +} From e036819f345a53e083a5f44fbdafb2b7f9eb327c Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Mon, 13 Jul 2026 15:44:11 +0800 Subject: [PATCH 06/11] keyspace: reset persisted status when re-adding a meta-service group persistGroupsLocked only cleared status for deleted groups; a newly added (or re-added) group got a zero status in the cache but nothing in storage. If an earlier best-effort deletion had failed, the stale enabled/count was resurrected by the next RefreshCache. Reset the persisted status of added groups to zero synchronously, in the same best-effort transaction that clears deleted ones. Add a regression test. Co-Authored-By: Claude Opus 4.8 Signed-off-by: tongjian <1045931706@qq.com> --- pkg/keyspace/meta_service_group.go | 35 ++++++++++++++++++------- pkg/keyspace/meta_service_group_test.go | 27 +++++++++++++++++++ 2 files changed, 53 insertions(+), 9 deletions(-) diff --git a/pkg/keyspace/meta_service_group.go b/pkg/keyspace/meta_service_group.go index 2b1ed18e26..93b28029a8 100644 --- a/pkg/keyspace/meta_service_group.go +++ b/pkg/keyspace/meta_service_group.go @@ -493,23 +493,40 @@ func (m *MetaServiceGroupManager) persistGroupsLocked( if err := persist(); err != nil { return err } - // Clear the persisted status for deleted groups before touching memory, so all - // storage writes complete before the in-memory view changes. This keeps - // re-adding a group with the same ID from inheriting a stale assignment count - // or enabled state, which would skew list output and PickGroup balancing. - // Best-effort: the config deletion is already persisted and the delete guard - // relies on actual keyspace scans, not this counter. - if len(deletedGroups) > 0 { + // Collect the groups that have no cached status yet (newly added or re-added) + // before mutating the cache, so their persisted status can be reset below. + m.statusMu.Lock() + var addedGroups []string + for groupID := range metaServiceGroups { + if m.cachedStatus[groupID] == nil { + addedGroups = append(addedGroups, groupID) + } + } + m.statusMu.Unlock() + // Reconcile persisted status before touching memory, so all storage writes + // complete before the in-memory view changes. Clearing deleted groups and + // resetting added ones to a zero status keeps re-adding a group with the same + // ID from inheriting a stale assignment count or enabled state (e.g. left by + // an earlier best-effort deletion that failed), which would skew list output + // and PickGroup balancing. Best-effort: the config change is already persisted + // and the delete guard relies on actual keyspace scans, not this counter. + if len(deletedGroups) > 0 || len(addedGroups) > 0 { if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error { for _, id := range deletedGroups { if err := m.store.RemoveMetaServiceGroupStatus(txn, id); err != nil { return err } } + for _, id := range addedGroups { + if err := m.store.SaveMetaServiceGroupStatus(txn, id, &endpoint.MetaServiceGroupStatus{}); err != nil { + return err + } + } return nil }); err != nil { - log.Warn("[keyspace] failed to clear status for deleted meta-service groups", - zap.Strings("deleted-groups", deletedGroups), zap.Error(err)) + log.Warn("[keyspace] failed to reconcile status for meta-service group changes", + zap.Strings("deleted-groups", deletedGroups), + zap.Strings("added-groups", addedGroups), zap.Error(err)) } } // Apply the change to the in-memory view only after the storage writes above. diff --git a/pkg/keyspace/meta_service_group_test.go b/pkg/keyspace/meta_service_group_test.go index a5510b7792..db93da723d 100644 --- a/pkg/keyspace/meta_service_group_test.go +++ b/pkg/keyspace/meta_service_group_test.go @@ -268,6 +268,33 @@ func (suite *metaServiceGroupTestSuite) TestAssignToGroupRejectsNegativeCount() re.ErrorIs(err, ErrInvalidAssignmentCount) } +// TestUpdateGroupsSafelyResetsStatusForReaddedGroup verifies that adding a group +// whose status still lingers in storage (e.g. from an earlier best-effort delete +// that failed) resets that persisted status to zero, so RefreshCache does not +// resurrect the stale count/enabled state. +func (suite *metaServiceGroupTestSuite) TestUpdateGroupsSafelyResetsStatusForReaddedGroup() { + re := suite.Require() + const groupID = "etcd-group-readded" + // Simulate a stale persisted status for a group the manager does not know yet. + re.NoError(suite.manager.store.RunInTxn(suite.ctx, func(txn kv.Txn) error { + return suite.manager.store.SaveMetaServiceGroupStatus(txn, groupID, + &endpoint.MetaServiceGroupStatus{AssignmentCount: 5, Enabled: true}) + })) + + groups := mockMetaServiceGroups() + groups[groupID] = "etcd-group-readded.tidb-serverless.cluster.svc.local" + re.NoError(suite.manager.UpdateGroupsSafely(suite.ctx, groups, nil, + func() error { return nil }, nil)) + + // Reloading from storage must see a zero status, not the stale one. + re.NoError(suite.manager.RefreshCache()) + statusMap, err := suite.manager.GetStatus(suite.ctx) + re.NoError(err) + re.NotNil(statusMap[groupID]) + re.Equal(0, statusMap[groupID].AssignmentCount) + re.False(statusMap[groupID].Enabled) +} + func (suite *metaServiceGroupTestSuite) TestReassignRejectsDisabledGroup() { re := suite.Require() // Groups are disabled by default, so reassigning a keyspace into one must be From 6580fdc0a91c6186f3efae4f0966a6b90adffc5c Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Mon, 13 Jul 2026 15:44:11 +0800 Subject: [PATCH 07/11] keyspace: revert TSO keyspace group move on config txn commit failure UpdateKeyspaceGroup persists the TSO keyspace group move immediately, but the inline revert only covered a SaveKeyspaceMeta error inside txnFunc. A commit-time failure after txnFunc returned nil left the move applied while the keyspace meta was rolled back. Track the move and revert it from the outer error path too, guarded so the inner revert is not duplicated. Co-Authored-By: Claude Opus 4.8 Signed-off-by: tongjian <1045931706@qq.com> --- pkg/keyspace/keyspace.go | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go index 8282f8162e..f2bb651442 100644 --- a/pkg/keyspace/keyspace.go +++ b/pkg/keyspace/keyspace.go @@ -880,6 +880,12 @@ func (manager *Manager) updateKeyspaceConfigTxn(name string, update func(meta *k oldConfig := make(map[string]string) var oldMetaServiceGroup, newMetaServiceGroup string metaServiceGroupReassigned := false + // Captured for the outer rollback below: UpdateKeyspaceGroup persists the TSO + // keyspace group move immediately, so a commit failure after txnFunc returns + // nil leaves it applied while the keyspace meta was not saved. + var oldTSOGroupID, newTSOGroupID string + var oldTSOUserKind, newTSOUserKind endpoint.UserKind + keyspaceGroupMoved := false txnFunc := func(txn kv.Txn) error { // First get KeyspaceID from Name. loaded, id, err := manager.store.LoadKeyspaceID(txn, name) @@ -940,12 +946,17 @@ func (manager *Manager) updateKeyspaceConfigTxn(name string, update func(meta *k if err := manager.kgm.UpdateKeyspaceGroup(oldID, newID, oldUserKind, newUserKind, meta.GetId()); err != nil { return err } + oldTSOGroupID, newTSOGroupID = oldID, newID + oldTSOUserKind, newTSOUserKind = oldUserKind, newUserKind + keyspaceGroupMoved = true } // Save the updated keyspace meta. if err := manager.store.SaveKeyspaceMeta(txn, meta); err != nil { - if needUpdate { + if keyspaceGroupMoved { if err := manager.kgm.UpdateKeyspaceGroup(newID, oldID, newUserKind, oldUserKind, meta.GetId()); err != nil { log.Error("failed to revert keyspace group", zap.Error(err)) + } else { + keyspaceGroupMoved = false } } return err @@ -954,6 +965,15 @@ func (manager *Manager) updateKeyspaceConfigTxn(name string, update func(meta *k } err := manager.runTxnWithMetaGroupLock(txnFunc) if err != nil { + // The txn can fail at commit after txnFunc returned nil, leaving the + // immediately-persisted TSO keyspace group move applied while the keyspace + // meta was rolled back. keyspaceGroupMoved stays true only when the inner + // error branch did not already revert it, so revert it here. + if keyspaceGroupMoved { + if rollbackErr := manager.kgm.UpdateKeyspaceGroup(newTSOGroupID, oldTSOGroupID, newTSOUserKind, oldTSOUserKind, meta.GetId()); rollbackErr != nil { + log.Error("failed to revert keyspace group", zap.Error(rollbackErr)) + } + } if metaServiceGroupReassigned { if rollbackErr := manager.mgm.updateAssignmentTxn(nil, newMetaServiceGroup, oldMetaServiceGroup); rollbackErr != nil { log.Error("failed to revert meta-service group assignment", zap.Error(rollbackErr)) From 88ae1893d06c2915f3d160f33e275085050efcb2 Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Tue, 14 Jul 2026 16:22:53 +0800 Subject: [PATCH 08/11] keyspace: merge patched fields into cached meta-service group status PatchStatus replaced the whole cached status with a snapshot taken before its storage write, dropping any assignment-count change a concurrent updateAssignmentTxn (statusMu only) applied during the I/O window. Apply only the patched fields to the live cached status instead. Co-Authored-By: Claude Opus 4.8 Signed-off-by: tongjian <1045931706@qq.com> --- pkg/keyspace/meta_service_group.go | 16 +++++++++++++++- 1 file changed, 15 insertions(+), 1 deletion(-) diff --git a/pkg/keyspace/meta_service_group.go b/pkg/keyspace/meta_service_group.go index 93b28029a8..8e77ed7f59 100644 --- a/pkg/keyspace/meta_service_group.go +++ b/pkg/keyspace/meta_service_group.go @@ -262,8 +262,22 @@ func (m *MetaServiceGroupManager) PatchStatus(ctx context.Context, groupID strin }); err != nil { return err } + // Apply only the patched fields to the live cached status instead of replacing + // it wholesale: an assignment-count update (updateAssignmentTxn, statusMu only) + // may have mutated the count during the I/O above, and overwriting with the + // pre-I/O snapshot would drop it. m.statusMu.Lock() - m.cachedStatus[groupID] = &newStatus + status := m.cachedStatus[groupID] + if status == nil { + status = &endpoint.MetaServiceGroupStatus{} + m.cachedStatus[groupID] = status + } + if patch.AssignmentCount != nil { + status.AssignmentCount = *patch.AssignmentCount + } + if patch.Enabled != nil { + status.Enabled = *patch.Enabled + } m.statusMu.Unlock() return nil } From 28bbf4b6aabdbfa6c89cbd05f188379249d3fa63 Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Tue, 14 Jul 2026 16:22:53 +0800 Subject: [PATCH 09/11] keyspace: roll back config txn side effects under the meta-group lock The TSO keyspace group and meta-service assignment rollbacks in updateKeyspaceConfigTxn ran after runTxnWithMetaGroupLock released the mgm lock. Take the mgm lock across both undos so they are atomic with respect to UpdateGroupsSafely, matching the locking of the forward path. Co-Authored-By: Claude Opus 4.8 Signed-off-by: tongjian <1045931706@qq.com> --- pkg/keyspace/keyspace.go | 33 +++++++++++++++++++++------------ 1 file changed, 21 insertions(+), 12 deletions(-) diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go index f2bb651442..323ecf4024 100644 --- a/pkg/keyspace/keyspace.go +++ b/pkg/keyspace/keyspace.go @@ -965,20 +965,29 @@ func (manager *Manager) updateKeyspaceConfigTxn(name string, update func(meta *k } err := manager.runTxnWithMetaGroupLock(txnFunc) if err != nil { - // The txn can fail at commit after txnFunc returned nil, leaving the - // immediately-persisted TSO keyspace group move applied while the keyspace - // meta was rolled back. keyspaceGroupMoved stays true only when the inner - // error branch did not already revert it, so revert it here. - if keyspaceGroupMoved { - if rollbackErr := manager.kgm.UpdateKeyspaceGroup(newTSOGroupID, oldTSOGroupID, newTSOUserKind, oldTSOUserKind, meta.GetId()); rollbackErr != nil { - log.Error("failed to revert keyspace group", zap.Error(rollbackErr)) + // Roll back side effects that were persisted immediately (and so are not + // undone by the failed txn): the TSO keyspace group move and the cached + // meta-service assignment. A commit-time failure after txnFunc returned nil + // leaves the move applied while the keyspace meta was rolled back; + // keyspaceGroupMoved stays true only when the inner branch did not already + // revert it. Hold the mgm lock across both undos so they are atomic with + // respect to UpdateGroupsSafely, mirroring runTxnWithMetaGroupLock. + func() { + if manager.mgm != nil { + manager.mgm.Lock() + defer manager.mgm.Unlock() } - } - if metaServiceGroupReassigned { - if rollbackErr := manager.mgm.updateAssignmentTxn(nil, newMetaServiceGroup, oldMetaServiceGroup); rollbackErr != nil { - log.Error("failed to revert meta-service group assignment", zap.Error(rollbackErr)) + if keyspaceGroupMoved { + if rollbackErr := manager.kgm.UpdateKeyspaceGroup(newTSOGroupID, oldTSOGroupID, newTSOUserKind, oldTSOUserKind, meta.GetId()); rollbackErr != nil { + log.Error("failed to revert keyspace group", zap.Error(rollbackErr)) + } } - } + if metaServiceGroupReassigned { + if rollbackErr := manager.mgm.updateAssignmentTxn(nil, newMetaServiceGroup, oldMetaServiceGroup); rollbackErr != nil { + log.Error("failed to revert meta-service group assignment", zap.Error(rollbackErr)) + } + } + }() log.Warn("[keyspace] failed to update keyspace config", zap.Uint32("keyspace-id", meta.GetId()), zap.String("name", meta.GetName()), From 52c082b984fa7a355b0dcbd1caf30d72d5901dc6 Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Wed, 15 Jul 2026 11:09:47 +0800 Subject: [PATCH 10/11] keyspace: flush meta-service group status off the mgm lock, fenced to term Two issues raised on the flush path: - It held the mgm write lock across the etcd transaction, so a slow flush blocked read paths (AttachEndpoints/GetGroups) and assignment operations on the serving PD. - It checked isLeader only before the write and used the server-lifetime context, so a former leader could commit its snapshot after a new leader had refreshed or patched the same status. Introduce persistMu, an outermost mutex that serializes the storage writers (flush, PatchStatus, persistGroupsLocked, RefreshCache) instead of the mgm write lock, so the flush no longer blocks mgm readers or assignments while still preventing writers from clobbering each other. PatchStatus now takes only the mgm read lock for its existence check. Flush uses the leadership term context (set from the service-ready callback and canceled when the lease is lost) so a lost-leadership PD stops writing. The residual TOCTOU between the leadership check and an in-flight commit would need a leader-key-fenced txn, left as a follow-up. Co-Authored-By: Claude Opus 4.8 Signed-off-by: tongjian <1045931706@qq.com> --- pkg/keyspace/meta_service_group.go | 95 +++++++++++++++++++++--------- server/server.go | 5 +- 2 files changed, 72 insertions(+), 28 deletions(-) diff --git a/pkg/keyspace/meta_service_group.go b/pkg/keyspace/meta_service_group.go index 8e77ed7f59..cde9670c21 100644 --- a/pkg/keyspace/meta_service_group.go +++ b/pkg/keyspace/meta_service_group.go @@ -37,18 +37,27 @@ const ( // MetaServiceGroupManager manages external meta-service groups. // -// Locking: the embedded RWMutex guards metaServiceGroups and isLeader, while -// statusMu is a dedicated leaf lock guarding cachedStatus and dirtyCount only. -// statusMu must always be the innermost lock: never acquire the RWMutex or the -// keyspace metaLock, and never perform storage I/O, while holding it. This lets -// the assignment-count update paths that already hold the keyspace metaLock -// (RemoveKeyspace, tombstone unassignment) mutate the cache without taking the -// RWMutex, which would otherwise invert the RWMutex->metaLock order used by the -// create/config paths and deadlock. Allowed orders are RWMutex->statusMu and -// metaLock->statusMu; both are safe because statusMu wraps nothing. +// Locking hierarchy, outermost to innermost: persistMu -> RWMutex -> statusMu. +// - persistMu serializes the storage writers (flushToStorage, PatchStatus, +// persistGroupsLocked, RefreshCache) so their status writes cannot clobber +// each other, without taking the RWMutex write lock across storage I/O. That +// keeps a slow flush from blocking RWMutex read paths (AttachEndpoints, +// GetGroups) or assignment operations. +// - the embedded RWMutex guards metaServiceGroups, isLeader and leaderCtx. +// - statusMu is a leaf lock guarding cachedStatus and dirtyCount only: never +// acquire the RWMutex, the keyspace metaLock or persistMu, and never perform +// storage I/O, while holding it. This lets the assignment-count update paths +// that already hold the keyspace metaLock (RemoveKeyspace, tombstone +// unassignment) mutate the cache without taking the RWMutex, which would +// otherwise invert the RWMutex->metaLock order used by the create/config +// paths and deadlock. Allowed orders are RWMutex->statusMu and +// metaLock->statusMu; both are safe because statusMu wraps nothing. type MetaServiceGroupManager struct { ctx context.Context store endpoint.MetaServiceGroupStorage + // persistMu serializes storage writers; see the type comment. It is the + // outermost lock, taken before the RWMutex. + persistMu syncutil.Mutex syncutil.RWMutex // metaServiceGroups is the available external meta-service groups. // The key is the meta-service group name, and the value is the corresponding endpoint. @@ -59,6 +68,10 @@ type MetaServiceGroupManager struct { // cannot permanently block removing an actually-empty group. keyspaceAssignmentCounter func(groupIDs map[string]struct{}) (map[string]int, error) isLeader func() bool + // leaderCtx is the current leadership term's context, canceled when this PD + // loses its lease. flushToStorage uses it so a lost-leadership PD does not keep + // writing status with the server-lifetime context. Guarded by the RWMutex. + leaderCtx context.Context // statusMu guards cachedStatus and dirtyCount. See the type comment for the // leaf-lock discipline it must follow. statusMu syncutil.Mutex @@ -81,6 +94,15 @@ func (m *MetaServiceGroupManager) SetLeaderChecker(isLeader func() bool) { m.isLeader = isLeader } +// SetLeaderContext sets the current leadership term's context, so flushToStorage +// stops writing once the term ends. It is called from the service-ready callback +// on each leadership acquisition with that term's context. +func (m *MetaServiceGroupManager) SetLeaderContext(ctx context.Context) { + m.Lock() + defer m.Unlock() + m.leaderCtx = ctx +} + // NewMetaServiceGroupManager creates a new MetaServiceGroupManager. func NewMetaServiceGroupManager( ctx context.Context, @@ -102,6 +124,10 @@ func NewMetaServiceGroupManager( // RefreshCache loads the current status from storage into memory. func (m *MetaServiceGroupManager) RefreshCache() error { + // Serialize with flush/PatchStatus so the reload does not race an in-flight + // status write and load a half-applied view. + m.persistMu.Lock() + defer m.persistMu.Unlock() m.Lock() defer m.Unlock() var ( @@ -143,21 +169,25 @@ func (m *MetaServiceGroupManager) flushLoop() { } func (m *MetaServiceGroupManager) flushToStorage() error { - // Hold the mgm write lock for the whole flush, including the storage write, so - // it serializes with PatchStatus and persistGroupsLocked (both take the mgm - // lock). Without that, the snapshot below could overwrite a status those paths - // persisted synchronously in the meantime, or recreate a status key that - // persistGroupsLocked just deleted. This cannot deadlock: the metaLock-holding - // paths (RemoveKeyspace, tombstone unassignment) only take statusMu, never the - // mgm lock, so nothing waits on the mgm lock while holding metaLock, and flush - // never takes metaLock. - m.Lock() - defer m.Unlock() - // Only the serving leader persists; a follower's next leader term reloads the - // authoritative status via RefreshCache. - if m.isLeader != nil && !m.isLeader() { + // Serialize the storage write with PatchStatus/persistGroupsLocked/RefreshCache + // on persistMu (not the mgm write lock), so a slow flush does not block mgm read + // paths (AttachEndpoints/GetGroups) or assignment operations, while still keeping + // the snapshot from clobbering a synchronous write or recreating a deleted key. + m.persistMu.Lock() + defer m.persistMu.Unlock() + // Only the serving leader persists, fenced to the current term: skip when not + // leader, and use the leadership context (canceled when the lease is lost) so a + // former leader's write is canceled instead of landing after the new leader. + m.RLock() + notLeader := m.isLeader != nil && !m.isLeader() + ctx := m.leaderCtx + m.RUnlock() + if notLeader { return nil } + if ctx == nil { + ctx = m.ctx + } // Snapshot the dirty state under statusMu (a deep copy) and reset the dirty // counter, without holding statusMu across the storage write. The snapshot // keeps the write from reading cachedStatus while a concurrent statusMu holder @@ -175,7 +205,7 @@ func (m *MetaServiceGroupManager) flushToStorage() error { m.dirtyCount = 0 m.statusMu.Unlock() - if err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) error { + if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error { for id, status := range snapshot { if err := m.store.SaveMetaServiceGroupStatus(txn, id, status); err != nil { return err @@ -237,9 +267,16 @@ func (m *MetaServiceGroupManager) PatchStatus(ctx context.Context, groupID strin if patch.AssignmentCount != nil && *patch.AssignmentCount < 0 { return ErrInvalidAssignmentCount } - m.Lock() - defer m.Unlock() - if _, ok := m.metaServiceGroups[groupID]; !ok { + // Serialize with flush/persistGroupsLocked on persistMu so this synchronous + // write is not clobbered by a concurrent flush, without holding the mgm write + // lock across the storage I/O. persistMu also blocks persistGroupsLocked, so + // the group cannot be deleted between the existence check and the write. + m.persistMu.Lock() + defer m.persistMu.Unlock() + m.RLock() + _, known := m.metaServiceGroups[groupID] + m.RUnlock() + if !known { return ErrUnknownMetaServiceGroup } m.statusMu.Lock() @@ -256,7 +293,7 @@ func (m *MetaServiceGroupManager) PatchStatus(ctx context.Context, groupID strin } // Persist synchronously so API updates are immediately durable, then reflect // them in the cache. The store I/O runs without statusMu (it is never held - // across I/O); the enclosing mgm write lock serializes concurrent patches. + // across I/O); persistMu serializes concurrent patches and flushes. if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error { return m.store.SaveMetaServiceGroupStatus(txn, groupID, &newStatus) }); err != nil { @@ -491,6 +528,10 @@ func (m *MetaServiceGroupManager) persistGroupsLocked( deletedGroups []string, persist func() error, ) error { + // Take persistMu (outermost) so the status reconcile writes below serialize + // with flush/PatchStatus and are not clobbered or resurrected by a flush. + m.persistMu.Lock() + defer m.persistMu.Unlock() m.Lock() defer m.Unlock() if len(deletedGroups) > 0 { diff --git a/server/server.go b/server/server.go index 7ab8a9d30e..e5d26eee4c 100644 --- a/server/server.go +++ b/server/server.go @@ -539,10 +539,13 @@ func (s *Server) startServer(ctx context.Context) error { return err } s.metaServiceGroupManager.SetLeaderChecker(s.IsServing) - s.AddServiceReadyCallback(func(_ context.Context) error { + s.AddServiceReadyCallback(func(ctx context.Context) error { if s.metaServiceGroupManager == nil { return nil } + // ctx is this leadership term's context (canceled when the lease is lost), + // so status flushes stop writing once the term ends. + s.metaServiceGroupManager.SetLeaderContext(ctx) return s.metaServiceGroupManager.RefreshCache() }) s.keyspaceManager = keyspace.NewKeyspaceManager( From 91caa0b58202f4c6da8a4c4879d27e0e12a58278 Mon Sep 17 00:00:00 2001 From: tongjian <1045931706@qq.com> Date: Wed, 15 Jul 2026 15:03:04 +0800 Subject: [PATCH 11/11] keyspace: mark added groups dirty when status reconcile fails The best-effort status reset in persistGroupsLocked only logged on failure, so a re-added group with stale persisted status was left with a zero cache entry that was never flushed; the next RefreshCache resurrected the stale enabled/count. Mark the zeroed added groups dirty on reconcile failure so a later flush re-persists their zero status and the reset invariant holds. Co-Authored-By: Claude Opus 4.8 Signed-off-by: tongjian <1045931706@qq.com> --- pkg/keyspace/meta_service_group.go | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/pkg/keyspace/meta_service_group.go b/pkg/keyspace/meta_service_group.go index cde9670c21..6ff4c02cad 100644 --- a/pkg/keyspace/meta_service_group.go +++ b/pkg/keyspace/meta_service_group.go @@ -565,6 +565,7 @@ func (m *MetaServiceGroupManager) persistGroupsLocked( // an earlier best-effort deletion that failed), which would skew list output // and PickGroup balancing. Best-effort: the config change is already persisted // and the delete guard relies on actual keyspace scans, not this counter. + reconcileFailed := false if len(deletedGroups) > 0 || len(addedGroups) > 0 { if err := m.store.RunInTxn(ctx, func(txn kv.Txn) error { for _, id := range deletedGroups { @@ -582,6 +583,7 @@ func (m *MetaServiceGroupManager) persistGroupsLocked( log.Warn("[keyspace] failed to reconcile status for meta-service group changes", zap.Strings("deleted-groups", deletedGroups), zap.Strings("added-groups", addedGroups), zap.Error(err)) + reconcileFailed = true } } // Apply the change to the in-memory view only after the storage writes above. @@ -595,6 +597,12 @@ func (m *MetaServiceGroupManager) persistGroupsLocked( for _, id := range deletedGroups { delete(m.cachedStatus, id) } + // If the synchronous reset above failed, mark the zeroed added groups dirty so a + // later flush re-persists their zero status; otherwise the stale value left in + // storage would be resurrected by the next RefreshCache, breaking the reset. + if reconcileFailed && len(addedGroups) > 0 { + m.markDirtyStatusLocked(len(addedGroups)) + } m.statusMu.Unlock() return nil }