keyspace: unassign tombstone keyspace from meta service group

pkg/keyspace/keyspace.go

@@ -1005,7 +1005,7 @@ func (manager *Manager) UpdateKeyspaceState(name string, newState keyspacepb.Key
 			return errs.ErrKeyspaceNotFound
 		}
 		// Update keyspace meta.
-		if err = manager.transformKeyspaceState(meta, newState, now); err != nil {
+		if err = manager.transformKeyspaceState(txn, meta, newState, now); err != nil {
 			return err
 		}
 		return manager.store.SaveKeyspaceMeta(txn, meta)
@@ -1052,18 +1052,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.
-	if manager.mgm != nil {
-		if groupID := meta.GetConfig()[MetaServiceGroupIDKey]; groupID != "" {
-			if err := manager.mgm.updateAssignmentTxn(txn, groupID, ""); err != nil {
-				return err
-			}
-		}
-	}
-	return nil
+	// 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.
+	return manager.unassignKeyspaceFromMetaServiceGroup(txn, meta)
 }
 
 // UpdateKeyspaceStateByID updates target keyspace to the given state if it's not already in that state.
@@ -1088,7 +1080,7 @@ func (manager *Manager) UpdateKeyspaceStateByID(id uint32, newState keyspacepb.K
 			return errs.ErrKeyspaceNotFound
 		}
 		// Update keyspace meta.
-		if err = manager.transformKeyspaceState(meta, newState, now); err != nil {
+		if err = manager.transformKeyspaceState(txn, meta, newState, now); err != nil {
 			return err
 		}
 		return manager.store.SaveKeyspaceMeta(txn, meta)
@@ -1112,16 +1104,53 @@ func (manager *Manager) UpdateKeyspaceStateByID(id uint32, newState keyspacepb.K
 	return meta, nil
 }
 
+// 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
+// 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.
+func (manager *Manager) unassignKeyspaceFromMetaServiceGroup(txn kv.Txn, meta *keyspacepb.KeyspaceMeta) error {
+	groupID := meta.GetConfig()[MetaServiceGroupIDKey]
+	if groupID == "" {
+		return nil
+	}
+	delete(meta.Config, MetaServiceGroupIDKey)
+	if manager.mgm == nil {
+		return nil
+	}
+	return manager.mgm.updateAssignmentTxn(txn, groupID, "")
+}
+
 // transformKeyspaceState transforms the keyspace state to the target state and record the update time.
-func (manager *Manager) transformKeyspaceState(meta *keyspacepb.KeyspaceMeta, newState keyspacepb.KeyspaceState, now int64) error {
-	// If already in the target state, do nothing and return.
+func (manager *Manager) transformKeyspaceState(txn kv.Txn, meta *keyspacepb.KeyspaceMeta, newState keyspacepb.KeyspaceState, now int64) error {
+	// If already in the target state, do nothing and return. A TOMBSTONE keyspace
+	// still carrying a meta-service group binding (e.g. one tombstoned before this
+	// cleanup existed) is repaired here by re-applying the TOMBSTONE update; the
+	// unassignment is idempotent, so it is a no-op once the binding is cleared.
 	if meta.GetState() == newState {
+		if newState == keyspacepb.KeyspaceState_TOMBSTONE {
+			return manager.unassignKeyspaceFromMetaServiceGroup(txn, meta)
+		}
 		return nil
 	}
 	// Consult state transition table to check if the operation is legal.
 	if !slice.Contains(stateTransitionTable[meta.GetState()], newState) {
 		return errors.Errorf("cannot change keyspace state from %s to %s", meta.GetState().String(), newState.String())
 	}
+	if newState == keyspacepb.KeyspaceState_TOMBSTONE {
+		if err := manager.unassignKeyspaceFromMetaServiceGroup(txn, meta); err != nil {
+			return err
+		}
+	}
 	// If the operation is legal, update keyspace state and change time.
 	meta.State = newState
 	meta.StateChangedAt = now

pkg/keyspace/keyspace_test.go

@@ -1113,3 +1113,88 @@ func TestAssignGroupAndSaveKeyspace(t *testing.T) {
 	re.NoError(managerWithGroup.assignGroupAndSaveKeyspace(true, &cfg2, ks2))
 	re.Equal("g1", ks2.GetConfig()[MetaServiceGroupIDKey])
 }
+
+func (suite *keyspaceTestSuite) TestTombstoneKeyspaceUnassignsMetaServiceGroup() {
+	re := suite.Require()
+	manager := suite.manager
+	groupID := "etcd-group-0"
+	groupEndpoint := "etcd-group-0.tidb-serverless.cluster.svc.local"
+	metaServiceGroupStore, ok := manager.store.(endpoint.MetaServiceGroupStorage)
+	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{})
+	manager.mgm.SetKeyspaceAssignmentCounter(manager.CountKeyspacesByMetaServiceGroup)
+
+	created, err := manager.CreateKeyspace(&CreateKeyspaceRequest{
+		Name:       "test_ks_msg",
+		Config:     map[string]string{},
+		CreateTime: time.Now().Unix(),
+	})
+	re.NoError(err)
+	re.NotContains(created.GetConfig(), MetaServiceGroupIDKey)
+
+	// Make the group available and assign the keyspace to it explicitly, mirroring
+	// what a meta-service-group-enabled cluster persists.
+	assignKeyspaceToGroup := func(id uint32) {
+		re.NoError(manager.store.RunInTxn(suite.ctx, func(txn kv.Txn) error {
+			meta, err := manager.store.LoadKeyspaceMeta(txn, id)
+			if err != nil {
+				return err
+			}
+			if meta.Config == nil {
+				meta.Config = map[string]string{}
+			}
+			meta.Config[MetaServiceGroupIDKey] = groupID
+			if err := manager.store.SaveKeyspaceMeta(txn, meta); err != nil {
+				return err
+			}
+			return manager.mgm.updateAssignmentTxn(txn, "", groupID)
+		}))
+	}
+	manager.mgm.updateGroups(map[string]string{groupID: groupEndpoint})
+	assignKeyspaceToGroup(created.GetId())
+
+	counts, err := manager.mgm.GetAssignmentCounts(suite.ctx)
+	re.NoError(err)
+	re.Equal(1, counts[groupID])
+
+	_, err = manager.UpdateKeyspaceState(created.GetName(), keyspacepb.KeyspaceState_DISABLED, time.Now().Unix())
+	re.NoError(err)
+	_, err = manager.UpdateKeyspaceState(created.GetName(), keyspacepb.KeyspaceState_ARCHIVED, time.Now().Unix())
+	re.NoError(err)
+	updated, err := manager.UpdateKeyspaceState(created.GetName(), keyspacepb.KeyspaceState_TOMBSTONE, time.Now().Unix())
+	re.NoError(err)
+	re.NotContains(updated.GetConfig(), MetaServiceGroupIDKey)
+
+	loaded, err := manager.LoadKeyspace(created.GetName())
+	re.NoError(err)
+	re.NotContains(loaded.GetConfig(), MetaServiceGroupIDKey)
+	counts, err = manager.mgm.GetAssignmentCounts(suite.ctx)
+	re.NoError(err)
+	re.Equal(0, counts[groupID])
+
+	// A keyspace tombstoned before this cleanup existed still carries the group
+	// binding and an inflated counter. Re-applying the TOMBSTONE state (a
+	// same-state update) must repair it rather than being skipped.
+	assignKeyspaceToGroup(updated.GetId())
+	counts, err = manager.mgm.GetAssignmentCounts(suite.ctx)
+	re.NoError(err)
+	re.Equal(1, counts[groupID])
+	repaired, err := manager.UpdateKeyspaceState(created.GetName(), keyspacepb.KeyspaceState_TOMBSTONE, time.Now().Unix())
+	re.NoError(err)
+	re.NotContains(repaired.GetConfig(), MetaServiceGroupIDKey)
+	counts, err = manager.mgm.GetAssignmentCounts(suite.ctx)
+	re.NoError(err)
+	re.Equal(0, counts[groupID])
+
+	// Removing the already-tombstoned keyspace must not decrement the counter
+	// again: the group binding was cleared and persisted during the tombstone
+	// transition, so unassignment is a no-op and the count stays at zero.
+	re.NoError(manager.store.RunInTxn(suite.ctx, func(txn kv.Txn) error {
+		return manager.RemoveKeyspace(txn, updated.GetId())
+	}))
+	counts, err = manager.mgm.GetAssignmentCounts(suite.ctx)
+	re.NoError(err)
+	re.Equal(0, counts[groupID])
+}