keyspace: cache meta-service group status#10976
Conversation
* 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 2baf2a5) 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
ti-chi-bot
Bot
added
release-note
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
|
@coderabbitai full review |
ti-chi-bot
Bot
added
the
size/XL
✅ Action performedFull review finished. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughMetaServiceGroupManager now caches status in memory with batched flushing. Keyspace assignment and config updates use write-locked accounting with rollback handling, and server startup plus tests were updated for the new constructor and cache refresh flow. ChangesMeta-service group cache and locking refactor
Estimated code review effort: 4 (Complex) | ~60 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pkg/keyspace/meta_service_group.go`:
- Around line 315-336: updateAssignmentLockedTxn is mutating cachedStatus and
dirtyCount without using the storage txn, so aborted transactions can
desynchronize cache state from persisted state. Update
MetaServiceGroupManager.updateAssignmentLockedTxn to make cache changes
transactional by tying the cachedStatus assignment count and markDirtyLocked
updates to the txn outcome (or otherwise rolling them back on error), and keep
the oldGroupID decrement/newGroupID increment behavior symmetric even when the
source count is zero.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 85dd8198-22b4-4b3e-a3fa-d656ff749e92
📒 Files selected for processing (5)
pkg/keyspace/keyspace.gopkg/keyspace/keyspace_test.gopkg/keyspace/meta_service_group.gopkg/keyspace/meta_service_group_test.goserver/server.go
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (1)
pkg/keyspace/meta_service_group.go (1)
40-40: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winPass the lifecycle context through
RefreshCache,flushLoop, andflushToStorageinstead of keepingm.ctxonMetaServiceGroupManager.NewMetaServiceGroupManagercan hand the constructor context to the goroutine, and the service-ready callback already has acontext.Contextto pass intoRefreshCache.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pkg/keyspace/meta_service_group.go` at line 40, Remove the stored lifecycle context from MetaServiceGroupManager and thread context explicitly through RefreshCache, flushLoop, and flushToStorage instead. Update NewMetaServiceGroupManager to pass the constructor context into the flush goroutine, and use the context already available in the service-ready callback when calling RefreshCache. Locate the affected flow in MetaServiceGroupManager, RefreshCache, flushLoop, and flushToStorage and update their signatures/call sites accordingly.Source: Coding guidelines
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pkg/keyspace/keyspace.go`:
- Around line 868-873: The rollback logic around runTxnWithMetaGroupLock and
UpdateKeyspaceGroup needs to stay under the meta-service group lock, and commit
failures must revert all side effects, not just the meta-service cache. Move the
non-transactional rollback work into the locked helper/error path so
UpdateGroupsSafely cannot see transient reassignment counts after the lock is
released, and track the successful UpdateKeyspaceGroup/Tso group move state
alongside metaServiceGroupReassigned so the rollback callback can restore both
before unlocking.
In `@pkg/keyspace/meta_service_group.go`:
- Around line 420-433: Persist the reset status for re-added groups instead of
only creating it in memory. In meta_service_group.go, update the cached-status
initialization around the metaServiceGroups loop so newly created
endpoint.MetaServiceGroupStatus entries are also added to newStatuses (or
otherwise marked dirty) before the cleanup transaction runs. Keep the
deletedGroups cleanup behavior, but ensure the status persistence path in
MetaServiceGroup metadata writes the zeroed state so a restart cannot reload
stale assignment/enabled values.
---
Nitpick comments:
In `@pkg/keyspace/meta_service_group.go`:
- Line 40: Remove the stored lifecycle context from MetaServiceGroupManager and
thread context explicitly through RefreshCache, flushLoop, and flushToStorage
instead. Update NewMetaServiceGroupManager to pass the constructor context into
the flush goroutine, and use the context already available in the service-ready
callback when calling RefreshCache. Locate the affected flow in
MetaServiceGroupManager, RefreshCache, flushLoop, and flushToStorage and update
their signatures/call sites accordingly.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 4fc3ed40-bd11-4452-a951-20fd54e535fc
📒 Files selected for processing (5)
pkg/keyspace/keyspace.gopkg/keyspace/keyspace_test.gopkg/keyspace/meta_service_group.gopkg/keyspace/meta_service_group_test.goserver/server.go
| 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) |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | 🏗️ Heavy lift
Keep rollback under the meta-service group lock and cover commit failures.
Line 955 runs after runTxnWithMetaGroupLock has already deferred-unlocked, so UpdateGroupsSafely can observe transient reassignment counts before rollback. Also, RunInTxn can fail at commit after UpdateKeyspaceGroup has persisted, but this outer error path only reverts the meta-service cache. Move all non-transactional rollback into the locked helper/error path and include the TSO group move.
Suggested shape
-func (manager *Manager) runTxnWithMetaGroupLock(f func(txn kv.Txn) error) error {
+func (manager *Manager) runTxnWithMetaGroupLock(f func(txn kv.Txn) error, rollback func()) error {
if manager.mgm != nil {
manager.mgm.Lock()
defer manager.mgm.Unlock()
}
- return manager.store.RunInTxn(manager.ctx, f)
+ err := manager.store.RunInTxn(manager.ctx, f)
+ if err != nil && rollback != nil {
+ rollback()
+ }
+ return err
}Then track successful UpdateKeyspaceGroup side effects alongside metaServiceGroupReassigned, and revert both from the rollback callback before the lock is released.
Also applies to: 953-959
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@pkg/keyspace/keyspace.go` around lines 868 - 873, The rollback logic around
runTxnWithMetaGroupLock and UpdateKeyspaceGroup needs to stay under the
meta-service group lock, and commit failures must revert all side effects, not
just the meta-service cache. Move the non-transactional rollback work into the
locked helper/error path so UpdateGroupsSafely cannot see transient reassignment
counts after the lock is released, and track the successful
UpdateKeyspaceGroup/Tso group move state alongside metaServiceGroupReassigned so
the rollback callback can restore both before unlocking.
There was a problem hiding this comment.
This still looks unresolved on the current head (77936b6d). UpdateKeyspaceGroup is persisted before the keyspace-meta RunInTxn commits. The rollback inside the SaveKeyspaceMeta error branch only covers errors returned before RunInTxn commits; if RunInTxn returns a commit conflict/error after txnFunc returns nil, the outer error path only rolls back the meta-service cache and never moves the TSO keyspace group from newID back to oldID.
Please roll back successful UpdateKeyspaceGroup side effects from the outer RunInTxn error path as well, preferably while still holding the meta-service group manager lock.
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
rg -nP 'func \(.*GroupManager\) UpdateKeyspaceGroup\(' pkg/keyspace -A5
rg -nP 'func StringUserKind' pkg/storage/endpointLength of output: 8907
@lhy1024, you're right — this is still open. store.RunInTxn (in pkg/storage/kv/etcd_kv.go) runs txnFunc and then calls txn.commit() separately, so a commit-time failure after txnFunc returns nil bypasses the inline revert inside the SaveKeyspaceMeta error branch entirely, and the outer error path in updateKeyspaceConfigTxn (Line 955 area) only reverts the meta-service assignment count, never the UpdateKeyspaceGroup move.
I'll centralize both rollbacks into a single callback invoked by runTxnWithMetaGroupLock on any error while the mgm lock is still held, and drop the now-redundant inline revert to avoid double-reverting.
- Extend
runTxnWithMetaGroupLockto accept arollback func()invoked on anyRunInTxnerror, before the lock is released. - In
updateKeyspaceConfigTxn, hoistoldUserKind/newUserKind/oldID/newIDto function scope and track whetherUpdateKeyspaceGroupsucceeded. - Pass a rollback closure that reverts both the meta-service assignment count and the TSO keyspace group move, replacing the two separate/partial rollback sites.
This makes rollback fire exactly once per failure, whether the failure originates inside txnFunc (e.g. SaveKeyspaceMeta error) or from RunInTxn's commit step after txnFunc returns nil, and both reverts happen before the mgm lock is released.
🧬 Code Graph Analysis Results
Relevant code snippets
pkg/keyspace/keyspace.go
L318-L362 — runTxnWithMetaGroupLock (holds mgm lock across the storage txn)
func (manager *Manager) runTxnWithMetaGroupLock(f func(txn kv.Txn) error) error {
if manager.mgm != nil {
manager.mgm.Lock()
defer manager.mgm.Unlock()
}
return manager.store.RunInTxn(manager.ctx, f)
}L364-L463 — updateKeyspaceConfigTxn (shows UpdateKeyspaceGroup persistence timing and the outer RunInTxn error rollback behavior)
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)
if err != nil {
return err
}
if !loaded {
return errs.ErrKeyspaceNotFound
}
manager.metaLock.Lock(id)
defer manager.metaLock.Unlock(id)
// Load keyspace by id.
meta, err = manager.store.LoadKeyspaceMeta(txn, id)
if err != nil {
return err
}
if meta == nil {
return errs.ErrKeyspaceNotFound
}
// Only keyspace with state listed in allowChangeConfig are allowed to change their config.
if !slice.Contains(allowChangeConfig, meta.GetState()) {
return errors.Errorf("cannot change config for keyspace with state %s", meta.GetState().String())
}
// Initialize meta's config map if it's nil.
if meta.GetConfig() == nil {
meta.Config = map[string]string{}
}
for k, v := range meta.GetConfig() {
oldConfig[k] = v
}
// Update keyspace config.
if err := update(meta); err != nil {
return err
}
delete(meta.Config, MetaServiceGroupAddressesKey)
newConfig := meta.GetConfig()
// Reassign the meta-service group before moving the TSO keyspace group.
// reassignKeyspaceLocked only stages its changes in txn (discarded if the
// 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]
if manager.mgm != nil && oldMetaServiceGroup != newMetaServiceGroup {
// The lock held by runTxnWithMetaGroupLock keeps this validation and
// the assignment update atomic with respect to UpdateGroupsSafely, which takes
// the write lock before deleting a group.
if err := manager.mgm.reassignKeyspaceLocked(txn, oldMetaServiceGroup, newMetaServiceGroup); err != nil {
return err
}
metaServiceGroupReassigned = true
}
oldUserKind := endpoint.StringUserKind(oldConfig[UserKindKey])
newUserKind := endpoint.StringUserKind(newConfig[UserKindKey])
oldID := oldConfig[TSOKeyspaceGroupIDKey]
newID := newConfig[TSOKeyspaceGroupIDKey]
needUpdate := oldUserKind != newUserKind || oldID != newID
if needUpdate {
if err := manager.kgm.UpdateKeyspaceGroup(oldID, newID, oldUserKind, newUserKind, meta.GetId()); err != nil {
return err
}
}
// Save the updated keyspace meta.
if err := manager.store.SaveKeyspaceMeta(txn, meta); err != nil {
if needUpdate {
if err := manager.kgm.UpdateKeyspaceGroup(newID, oldID, newUserKind, oldUserKind, meta.GetId()); err != nil {
log.Error("failed to revert keyspace group", zap.Error(err))
}
}
return err
}
return nil
}
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()),
zap.Error(err),
)
return nil, err
}
if manager.mgm != nil {
manager.mgm.AttachEndpoints(meta.GetConfig())
}
log.Info("[keyspace] keyspace config updated",
zap.Uint32("keyspace-id", meta.GetId()),
zap.String("name", meta.GetName()),
zap.Any("new-config", meta.GetConfig()),
)
return meta, nil
}pkg/keyspace/tso_keyspace_group.go
L696-L762 — GroupManager.UpdateKeyspaceGroup (persists the keyspace group move immediately; used by updateKeyspaceConfigTxn)
func (m *GroupManager) UpdateKeyspaceGroup(oldGroupID, newGroupID string, oldUserKind, newUserKind endpoint.UserKind, keyspaceID uint32) error {
if m == nil {
return nil
}
oldID, err := strconv.ParseUint(oldGroupID, 10, 64)
if err != nil {
return err
}
newID, err := strconv.ParseUint(newGroupID, 10, 64)
if err != nil {
return err
}
m.Lock()
defer m.Unlock()
oldKG := m.groups[oldUserKind].Get(uint32(oldID))
if oldKG == nil {
return errors.Errorf("keyspace group %s not found in %s group", oldGroupID, oldUserKind)
}
newKG := m.groups[newUserKind].Get(uint32(newID))
if newKG == nil {
return errors.Errorf("keyspace group %s not found in %s group", newGroupID, newUserKind)
}
if oldKG.IsSplitting() {
return errs.ErrKeyspaceGroupInSplit.FastGenByArgs(uint32(oldID))
} else if newKG.IsSplitting() {
return errs.ErrKeyspaceGroupInSplit.FastGenByArgs(uint32(newID))
} else if oldKG.IsMerging() {
return errs.ErrKeyspaceGroupInMerging.FastGenByArgs(uint32(oldID))
} else if newKG.IsMerging() {
return errs.ErrKeyspaceGroupInMerging.FastGenByArgs(uint32(newID))
}
var updateOld, updateNew bool
if !slice.Contains(newKG.Keyspaces, keyspaceID) {
newKG.Keyspaces = append(newKG.Keyspaces, keyspaceID)
slices.Sort(newKG.Keyspaces)
updateNew = true
}
lenOfOldKeyspaces := len(oldKG.Keyspaces)
oldKG.Keyspaces = slice.Remove(oldKG.Keyspaces, keyspaceID)
if lenOfOldKeyspaces != len(oldKG.Keyspaces) {
updateOld = true
}
if err := m.saveKeyspaceGroups([]*endpoint.KeyspaceGroup{oldKG, newKG}, true); err != nil {
if updateOld {
oldKG.Keyspaces = append(oldKG.Keyspaces, keyspaceID)
slices.Sort(oldKG.Keyspaces)
}
if updateNew {
newKG.Keyspaces = slice.Remove(newKG.Keyspaces, keyspaceID)
}
return err
}
if updateOld {
m.groups[oldUserKind].Put(oldKG)
}
if updateNew {
m.groups[newUserKind].Put(newKG)
}
return nil
}pkg/storage/kv/etcd_kv.go
L212-L222 — RunInTxn (commit happens after txnFunc returns nil; commit conflict can surface via RunInTxn error)
func (kv *etcdKVBase) RunInTxn(ctx context.Context, f func(txn Txn) error) error {
txn := &etcdTxn{
kv: kv,
ctx: ctx,
}
err := f(txn)
if err != nil {
return err
}
return txn.commit()
}L56-L67 — StorageEndpoint.SaveKeyspaceMeta (the meta write happens inside the txn passed to RunInTxn)
func (*StorageEndpoint) SaveKeyspaceMeta(txn kv.Txn, meta *keyspacepb.KeyspaceMeta) error {
metaPath := keypath.KeyspaceMetaPath(meta.GetId())
metaVal, err := proto.Marshal(meta)
if err != nil {
return errs.ErrProtoMarshal.Wrap(err).GenWithStackByCause()
}
return txn.Save(metaPath, string(metaVal))
}pkg/keyspace/meta_service_group.go
L411-L421 — MetaServiceGroupManager.AttachEndpoints (used after successful config update)
func (m *MetaServiceGroupManager) AttachEndpoints(keyspaceConfig map[string]string) {
groupID := keyspaceConfig[MetaServiceGroupIDKey]
if groupID == "" {
return
}
m.RLock()
defer m.RUnlock()
if endpoints := m.metaServiceGroups[groupID]; endpoints != "" {
keyspaceConfig[MetaServiceGroupAddressesKey] = endpoints
}
}L437-L441 — MetaServiceGroupManager.HasGroups (guards assignment paths in other keyspace creation/config flows)
func (m *MetaServiceGroupManager) HasGroups() bool {
m.RLock()
defer m.RUnlock()
return m.hasGroupsLocked()
}Failed to handle agent chat message. Please try again.
There was a problem hiding this comment.
6580fdc0 now covers commit-time TSO rollback, but the rollback still runs after runTxnWithMetaGroupLock releases the manager lock. A second config update can acquire that lock and persist another TSO group move before lines 972–980 run; the first request then applies its stale rollback over the successful update, leaving group membership inconsistent with committed keyspace metadata. The meta-service count rollback has the same interleaving window.
There was a problem hiding this comment.
Seems like the humans are having a chat. I'll hop back into my burrow for now. If you need me again, just tag @coderabbitai in a new comment, and I'll come hopping out!
There was a problem hiding this comment.
Fixed in 6580fdc. The TSO keyspace group move is now tracked (keyspaceGroupMoved) and reverted from the outer error path as well, covering a commit-time failure after txnFunc returns nil; the flag prevents the inner revert from being duplicated.
Follow-up 28bbf4b takes the mgm lock across the outer rollback (both the TSO move and the meta-service assignment) so the undo is atomic with respect to UpdateGroupsSafely, per your note about keeping it under the lock.
There was a problem hiding this comment.
Fixed in 28bbf4b. The outer rollback (both the TSO keyspace group move and the meta-service assignment) now runs while holding the mgm lock, taken in a deferred-unlock closure, so it is atomic with respect to UpdateGroupsSafely and a second config update cannot interleave a TSO move between the failed commit and the rollback.
There was a problem hiding this comment.
The current head still releases the manager lock when runTxnWithMetaGroupLock returns at line 966 and only reacquires it at line 977. A second config update can therefore still run in that gap before the stale rollback, so the interleaving described above remains.
There was a problem hiding this comment.
This is still open after 28bbf4b6. The error path now reacquires mgm.Lock(), but runTxnWithMetaGroupLock has already released the original lock before returning. A later config update can acquire the lock and commit in that gap, after which the failed request reacquires it and applies a stale rollback over the successful update.
A deterministic state-transition test reproduces the same interleaving:
- request A applies
old -> failedand releases the lock when its txn fails; - request B acquires the lock and commits
old -> committed; - request A reacquires the lock and rolls back
failed -> old.
The expected source count is 0, but the stale rollback restores it to 1:
TestRollbackAfterReacquiringLockOverwritesLaterReassignment
expected: 0
actual: 1
Please run the rollback from inside runTxnWithMetaGroupLock, before its deferred unlock, for example via an error rollback callback. Releasing and reacquiring the same mutex does not preserve atomicity.
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 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
Resolve the conflict with tikv#10959 (tombstone keyspace unassignment) and fix a lock-order deadlock it surfaced. tikv#10959 added unassignKeyspaceFromMetaServiceGroup, which calls updateAssignmentTxn while holding the keyspace metaLock, on the invariant that updateAssignmentTxn takes no mgm lock. This PR's status cache had made updateAssignmentTxn acquire the mgm write lock, so the unassign paths ran metaLock -> mgm.Lock while the create/config paths run mgm.Lock -> metaLock: an AB-BA deadlock on the same keyspace id. Introduce a dedicated leaf mutex (statusMu) guarding cachedStatus and dirtyCount only. updateAssignmentTxn now mutates the cache under statusMu instead of the mgm lock, so metaLock -> statusMu and mgm.Lock -> statusMu coexist without a cycle. statusMu is always the innermost lock and never wraps storage I/O. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
ti-chi-bot
Bot
added
size/XXL
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #10976 +/- ##
==========================================
+ Coverage 79.21% 79.29% +0.07%
==========================================
Files 541 541
Lines 75677 76119 +442
==========================================
+ Hits 59949 60355 +406
- Misses 11491 11509 +18
- Partials 4237 4255 +18
Flags with carried forward coverage won't be shown. Click here to find out more. 🚀 New features to boost your workflow:
|
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 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
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 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
| m.statusMu.Unlock() | ||
| return nil | ||
| } | ||
| snapshot := copyStatusMap(m.cachedStatus) |
There was a problem hiding this comment.
flushToStorage can overwrite concurrent synchronous status updates. It snapshots the whole cachedStatus, releases statusMu, and then blindly writes every status from that snapshot. A concurrent PatchStatus can persist a new enabled/count value while this flush is in flight; if the older snapshot commits afterward, the patch is lost on the next RefreshCache or restart. The same race can also recreate a status key removed by UpdateGroupsSafely.
Please serialize flush persistence with PatchStatus/group deletion, or make the flush conditional/merge-safe and skip groups deleted after the snapshot.
There was a problem hiding this comment.
A concrete regression test for this race could look like this:
import (
"sync/atomic"
"github.com/stretchr/testify/require"
)
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)
}
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)
groupID := "etcd-group-0"
manager.statusMu.Lock()
manager.cachedStatus[groupID] = &endpoint.MetaServiceGroupStatus{AssignmentCount: 1, Enabled: true}
manager.dirtyCount = 1
manager.statusMu.Unlock()
store.blockNextSave.Store(true)
flushErr := make(chan error, 1)
go func() {
flushErr <- manager.flushToStorage()
}()
<-store.saveStarted
newCount := 99
re.NoError(manager.PatchStatus(ctx, groupID, &MetaServiceGroupStatusPatch{
AssignmentCount: &newCount,
}))
close(store.unblockSave)
re.NoError(<-flushErr)
re.NoError(manager.RefreshCache())
statusMap, err := manager.GetStatus(ctx)
re.NoError(err)
re.Equal(newCount, statusMap[groupID].AssignmentCount)
}I ran this focused test locally against the current implementation. It fails with expected: 99 and actual: 1, which means the old flush snapshot overwrites the synchronous PatchStatus update.
There was a problem hiding this comment.
Fixed in e5262c6. flushToStorage now holds the mgm write lock across the storage write, serializing it with PatchStatus and persistGroupsLocked, so the snapshot can no longer overwrite a synchronous write or recreate a deleted status key. It stays deadlock-free because the metaLock-holding paths (RemoveKeyspace, tombstone unassignment) only take statusMu, never the mgm lock.
Also fixed an adjacent read-modify-write race in PatchStatus (88ae189): it now applies only the patched fields to the live cached status instead of overwriting it with a pre-I/O snapshot, so a concurrent updateAssignmentTxn count change is not dropped.
There was a problem hiding this comment.
Thanks for the repro — added TestFlushDoesNotOverwriteConcurrentPatchStatus based on it. One tweak: with the fix, PatchStatus blocks on the mgm lock the paused flush holds, so calling it synchronously (as in the snippet) deadlocks against the unblockSave that comes after it. So I run PatchStatus in a goroutine and assert its value wins once the flush releases the lock.
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 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
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 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
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 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
| return err | ||
| } | ||
| m.statusMu.Lock() | ||
| m.cachedStatus[groupID] = &newStatus |
There was a problem hiding this comment.
PatchStatus releases statusMu before persisting and then replaces the whole cached entry with the earlier snapshot. updateAssignmentTxn intentionally skips the manager lock, so a concurrent tombstone or removal can update AssignmentCount between those critical sections; an enabled-only patch then overwrites that successful delta while dirtyCount remains set. The next flush persists the stale count, even though both operations returned success.
There was a problem hiding this comment.
Fixed in 88ae189. PatchStatus now applies only the patched fields to the live cached status instead of replacing the whole entry with the pre-I/O snapshot, so a concurrent updateAssignmentTxn count delta is preserved. The synchronous store write still lands under the same serialization as the flush (see below).
| m.dirtyCount = 0 | ||
| m.statusMu.Unlock() | ||
|
|
||
| if err := m.store.RunInTxn(m.ctx, func(txn kv.Txn) error { |
There was a problem hiding this comment.
isLeader is checked only before this transaction, so the write is not fenced to the leadership term. If this PD loses its lease after line 158, the flush continues with the server-lifetime context; these status saves are unconditional puts with no loaded-key comparison, so the former leader can commit its snapshot after the new leader has refreshed or patched the same status. A leader transfer can therefore overwrite a newer synchronous status update.
There was a problem hiding this comment.
Mitigated in 52c082b: the flush now runs on the leadership-term context (set from the service-ready callback and canceled when the lease is lost) instead of the server-lifetime context, and re-checks isLeader right before the write. That cancels a former leader's in-flight flush and stops new ones after the term ends.
You are right that this does not fully close it: an already-issued commit can still land after cancellation. Fully fencing it needs a leader-key comparison inside the flush txn (so the put aborts when the lease key changed), which requires plumbing the leadership guard into the keyspace storage writes. I left that as a follow-up and noted it in the commit message — happy to do it here if you prefer.
There was a problem hiding this comment.
The term context narrows the window but does not fence a transaction whose commit has already been issued, so the former leader can still overwrite a newer synchronous PatchStatus after a transfer. Because the stale write changes API-visible enabled/count state, the original correctness issue remains.
There was a problem hiding this comment.
This is still reproducible on the current head (52c082b98). The term context narrows the window, but it is canceled only after the old leader loop observes lease expiry and returns. The election key can already be gone, and a new leader can already be serving, while that context is still live.
I added a deterministic regression test that pauses the old flush after its leader check, marks the old instance non-serving while leaving the term context live (the interval before the leader ticker returns), lets the new leader persist count 99, and then resumes the old flush. The final persisted count is 1, not 99:
TestFormerLeaderFlushIsFencedBeforeTermContextCancellation
expected: 99
actual: 1
So context cancellation is not a leadership fence. Please include the election leader-key/value comparison in the same etcd transaction, equivalent to Leadership.LeaderTxn; otherwise an unconditional write that passed the preflight check immediately before lease expiry can still land in the next term.
| // 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() |
There was a problem hiding this comment.
flushToStorage holds the manager write lock across the etcd transaction. AttachEndpoints on keyspace read paths and the group configuration/automatic-assignment paths need this lock, so a slow flush stalls otherwise read-only keyspace lookups and unrelated group operations on the serving PD. Routine threshold or periodic flushes can therefore turn one slow storage write into a broad request-latency spike.
There was a problem hiding this comment.
Fixed in 52c082b. Introduced persistMu, an outermost mutex that serializes the storage writers (flush, PatchStatus, persistGroupsLocked, RefreshCache) instead of the mgm write lock. The flush now holds only persistMu across the etcd transaction and takes the mgm read lock just briefly to read isLeader/leaderCtx, so a slow flush no longer blocks AttachEndpoints/GetGroups read paths or assignment operations. PatchStatus likewise no longer holds the mgm write lock across its I/O.
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 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
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 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
… 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 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
| } | ||
| } | ||
| return nil | ||
| }); err != nil { |
There was a problem hiding this comment.
The reconciliation failure is logged and then the group change still returns success without marking the zeroed status dirty. If a re-added group has stale persisted status and this transaction fails, the in-memory group starts at zero but the next RefreshCache restores the stale enabled/count values, so the failure path still violates the reset invariant this block is intended to provide.
There was a problem hiding this comment.
Good catch. Fixed in 91caa0b: on reconcile failure the zeroed added groups are now marked dirty, so a later flush re-persists their zero status and RefreshCache no longer resurrects the stale value. The synchronous reset remains the primary path; the dirty mark is the self-healing fallback for the failure case.
There was a problem hiding this comment.
Marking it dirty does not schedule a flush unless dirtyCount reaches 1000; with a typical single re-added group, the write waits for the five-minute ticker. A RefreshCache caused by leader transfer or restart before then still reloads the stale status and clears dirtyCount. I reproduced this by injecting the reconcile save failure and calling RefreshCache immediately: count/enabled returned to 5/true.
There was a problem hiding this comment.
Confirmed with a deterministic unit test on 52c082b98. I preloaded a stale status with count 7, injected a failure into the zero-status save for the re-added group, let UpdateGroupsSafely return success, and then called RefreshCache. The stale value is restored:
TestReaddedGroupReconcileFailureDoesNotRestoreStaleStatus
expected: 0
actual: 7
The normal re-add path test passes, so the remaining issue is specifically this reconciliation failure path. Returning the storage error is the simplest fix; alternatively the zero status must remain dirty and be guaranteed to overwrite the stale value before a refresh can discard it.
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 <noreply@anthropic.com> Signed-off-by: tongjian <1045931706@qq.com>
|
@bufferflies: The following tests failed, say
Full PR test history. Your PR dashboard. DetailsInstructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. I understand the commands that are listed here. |
| return err | ||
| } | ||
| m.statusMu.Lock() | ||
| m.cachedStatus = statusMap |
There was a problem hiding this comment.
RefreshCache replaces the cache with the persisted batching snapshot and resets dirtyCount, but it never reconciles AssignmentCount with authoritative keyspace metadata. A leader change before a successful flush therefore permanently drops all assignment deltas since the previous flush; subsequent PickGroup decisions continue from the stale baseline.
The manager already has keyspaceAssignmentCounter. A focused test configures that authoritative counter to return 1, while storage contains the last flushed value 0; after RefreshCache, the current implementation still returns 0:
TestRefreshCacheReconcilesAuthoritativeAssignmentCounts
expected: 1
actual: 0
Please preserve persisted administrative fields such as Enabled, while rebuilding assignment counts from keyspace metadata during leader readiness.
4 participants
What problem does this PR solve?
Issue Number: Close #10892
This is a cherry-pick of 2baf2a5 from the PD microservices branch.
Original author: @AmoebaProtozoa
Problem Summary:
Meta service group status was loaded and persisted through storage on assignment paths. This adds an in-memory status cache and flush loop so assignment count updates can be handled locally and persisted in batches by the serving leader.
What is changed and how does it work?
MetaServiceGroupManagerand refresh it from storage on initialization/service readiness.PatchStatussynchronous with storage and cache so API updates are immediately visible.UpdateGroupsSafelydelete guard and authoritative keyspace assignment scanner.Check List
Tests
Code changes
Side effects
Related changes
Release note
Validation
GOFLAGS=-buildvcs=false make gotest GOTEST_ARGS='./pkg/keyspace -run TestMetaServiceGroupTestSuite -count=1 -timeout=5m'GOFLAGS=-buildvcs=false make gotest GOTEST_ARGS='./pkg/keyspace -run "TestMetaServiceGroupTestSuite|TestAssignGroupAndSaveKeyspace|TestKeyspaceTestSuite/(TestUpdateKeyspaceState|TestTombstoneKeyspaceUnassignsMetaServiceGroup)" -count=1 -timeout=5m'GOFLAGS=-buildvcs=false make gotest GOTEST_ARGS='./pkg/keyspace -count=1 -timeout=5m'GOFLAGS=-buildvcs=false make gotest GOTEST_ARGS='./server -run TestNonExistent -count=0 -timeout=5m'git diff --check && git diff --cached --checkGOFLAGS=-buildvcs=false make checkSummary by CodeRabbit