Skip to content

keyspace: cache meta-service group status #10976

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bufferflies wants to merge 13 commits into
base: master
Choose a base branch
from
+611 −189
Open

keyspace: cache meta-service group status #10976

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
373334d
keyspace: cache meta-service group status (#463)
AmoebaProtozoa Feb 2, 2026
c4a0aa9
keyspace: refine meta-service group status cache
bufferflies Jul 7, 2026
13a9e89
Merge remote-tracking branch 'origin/master' into pr10976
bufferflies Jul 7, 2026
f01ce3e
Merge remote-tracking branch 'origin/master' into pr10976
bufferflies Jul 7, 2026
f514e59
keyspace: persist group changes before updating in-memory view
bufferflies Jul 8, 2026
77936b6
keyspace: add HasGroups to avoid map copy on create path
bufferflies Jul 8, 2026
e5262c6
keyspace: serialize meta-service group flush with synchronous writes
bufferflies Jul 13, 2026
e036819
keyspace: reset persisted status when re-adding a meta-service group
bufferflies Jul 13, 2026
6580fdc
keyspace: revert TSO keyspace group move on config txn commit failure
bufferflies Jul 13, 2026
88ae189
keyspace: merge patched fields into cached meta-service group status
bufferflies Jul 14, 2026
28bbf4b
keyspace: roll back config txn side effects under the meta-group lock
bufferflies Jul 14, 2026
52c082b
keyspace: flush meta-service group status off the mgm lock, fenced to…
bufferflies Jul 15, 2026
91caa0b
keyspace: mark added groups dirty when status reconcile fails
bufferflies Jul 15, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
116 changes: 70 additions & 46 deletions pkg/keyspace/keyspace.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()),
Expand Down Expand Up @@ -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()),
Expand Down Expand Up @@ -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
Expand All @@ -597,15 +581,15 @@ 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.
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
Expand All @@ -621,10 +605,12 @@ 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
// store operations and does not take the mgm lock, so it is safe to call
// while still holding the read lock.
manager.rollbackMetaServiceGroupAssignment(groupID)
// Roll back the reservation made by pickGroupLocked. updateAssignmentTxn
// only mutates the cached count under the leaf statusMu (never the mgm
// lock), so it is safe to call while still holding the mgm lock above.
if rollbackErr := manager.mgm.updateAssignmentTxn(nil, groupID, ""); rollbackErr != nil {
log.Error("failed to revert meta-service group assignment", zap.Error(rollbackErr))
}
return err
}
return nil
Expand Down Expand Up @@ -877,21 +863,29 @@ 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)
Comment on lines 870 to 875

@coderabbitai coderabbitai Bot Jul 3, 2026 •
edited
Loading

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ 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.

✅ Addressed in commits b9b3eed to f01ce3e

@lhy1024 lhy1024 Jul 8, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@coderabbitai coderabbitai Bot Jul 8, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
rg -nP 'func \(.*GroupManager\) UpdateKeyspaceGroup\(' pkg/keyspace -A5
rg -nP 'func StringUserKind' pkg/storage/endpoint

Length 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 runTxnWithMetaGroupLock to accept a rollback func() invoked on any RunInTxn error, before the lock is released.
  • In updateKeyspaceConfigTxn, hoist oldUserKind/newUserKind/oldID/newID to function scope and track whether UpdateKeyspaceGroup succeeded.
  • 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.

@rleungx rleungx Jul 13, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@coderabbitai coderabbitai Bot Jul 13, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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!

@bufferflies bufferflies Jul 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@bufferflies bufferflies Jul 15, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@rleungx rleungx Jul 15, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

@lhy1024 lhy1024 Jul 15, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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:

  1. request A applies old -> failed and releases the lock when its txn fails;
  2. request B acquires the lock and commits old -> committed;
  3. 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.

}

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
// 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)
Expand Down Expand Up @@ -933,14 +927,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])
Expand All @@ -951,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
Expand All @@ -965,6 +965,29 @@ func (manager *Manager) updateKeyspaceConfigTxn(name string, update func(meta *k
}
err := manager.runTxnWithMetaGroupLock(txnFunc)
if err != nil {
// 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 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()),
Expand Down Expand Up @@ -1059,9 +1082,9 @@ func (manager *Manager) RemoveKeyspace(txn kv.Txn, id uint32) error {
}
manager.keyspaceNameLookup.Delete(id)
manager.keyspaceStateLookup.Delete(id)
// Keep the meta-service group assignment accounting in sync within the same
// txn. Without this, removed keyspaces leak count and could permanently block
// deleting an otherwise-empty group.
// Keep the meta-service group assignment accounting in sync. Without this,
// removed keyspaces leak count and could permanently block deleting an
// otherwise-empty group in setups without the authoritative keyspace scanner.
return manager.unassignKeyspaceFromMetaServiceGroup(txn, meta)
}

Expand Down Expand Up @@ -1112,19 +1135,20 @@ func (manager *Manager) UpdateKeyspaceStateByID(id uint32, newState keyspacepb.K
}

// unassignKeyspaceFromMetaServiceGroup removes the keyspace's meta-service group
// binding within txn: it drops MetaServiceGroupIDKey from the config and, when a
// meta-service group manager is configured, decrements the persisted assignment
// binding: it drops MetaServiceGroupIDKey from the config within txn and, when a
// meta-service group manager is configured, decrements its cached assignment
// count. Once the config key is removed and persisted, a subsequent call is a
// no-op, so the removal and tombstone paths can both invoke it without
// double-counting.
//
// Callers already hold the keyspace metaLock (so meta is not mutated
// concurrently). This deliberately does NOT take mgm.RLock: the config-update
// path acquires mgm.RLock before metaLock (via runTxnWithMetaGroupLock), so
// grabbing mgm.RLock here while holding metaLock would invert the lock order and
// deadlock once UpdateGroupsSafely is waiting on mgm.Lock. The lock is
// unnecessary anyway — updateAssignmentTxn only touches the store, and the
// group delete guard relies on the authoritative keyspace scan, not this count.
// concurrently). This deliberately does NOT take the mgm lock: the create/config
// paths acquire the mgm lock before metaLock (via runTxnWithMetaGroupLock and
// assignGroupAndSaveKeyspace), so grabbing it here while holding metaLock would
// invert that order and deadlock. updateAssignmentTxn honors this by mutating the
// cache under its own leaf lock (statusMu) instead of the mgm lock. The count is
// best-effort anyway — the group delete guard relies on the authoritative
// keyspace scan, not this counter.
func (manager *Manager) unassignKeyspaceFromMetaServiceGroup(txn kv.Txn, meta *keyspacepb.KeyspaceMeta) error {
groupID := meta.GetConfig()[MetaServiceGroupIDKey]
if groupID == "" {
Expand Down
13 changes: 9 additions & 4 deletions pkg/keyspace/keyspace_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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}
Expand All @@ -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)
Expand All @@ -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}
Expand All @@ -1135,7 +1138,9 @@ func (suite *keyspaceTestSuite) TestTombstoneKeyspaceUnassignsMetaServiceGroup()
re.True(ok)
// Start without any group so creation never auto-assigns: meta-service groups
// are disabled by default, and this keeps the test independent of that.
manager.mgm = NewMetaServiceGroupManager(metaServiceGroupStore, map[string]string{})
mgm, err := NewMetaServiceGroupManager(suite.ctx, metaServiceGroupStore, map[string]string{})
re.NoError(err)
manager.mgm = mgm
manager.mgm.SetKeyspaceAssignmentCounter(manager.CountKeyspacesByMetaServiceGroup)

created, err := manager.CreateKeyspace(&CreateKeyspaceRequest{
Expand Down
Loading
Loading