diff --git a/pkg/keyspace/keyspace.go b/pkg/keyspace/keyspace.go index 71536d9eb0..323ecf4024 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()), @@ -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,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 @@ -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 @@ -877,14 +863,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 +878,14 @@ 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 + // 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) @@ -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]) @@ -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 @@ -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()), @@ -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) } @@ -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 == "" { diff --git a/pkg/keyspace/keyspace_test.go b/pkg/keyspace/keyspace_test.go index 3223d9b1b6..35c4fd85b7 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} @@ -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{ diff --git a/pkg/keyspace/meta_service_group.go b/pkg/keyspace/meta_service_group.go index 87d53ebe9c..6ff4c02cad 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,9 +30,34 @@ import ( "github.com/tikv/pd/server/config" ) +const ( + flushInterval = 5 * time.Minute + flushThreshold = 1000 +) + // MetaServiceGroupManager manages external meta-service groups. +// +// 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. @@ -41,6 +67,17 @@ 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) + 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 + cachedStatus map[string]*endpoint.MetaServiceGroupStatus + dirtyCount int + flushCh chan struct{} } // SetKeyspaceAssignmentCounter sets the authoritative keyspace assignment @@ -50,30 +87,145 @@ 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 +} + +// 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, 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 { + // 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 ( 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.statusMu.Lock() + m.cachedStatus = statusMap + m.dirtyCount = 0 + m.statusMu.Unlock() + 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 { + // 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 + // (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() + return nil + } + snapshot := copyStatusMap(m.cachedStatus) + flushed := m.dirtyCount + m.dirtyCount = 0 + m.statusMu.Unlock() + + 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 + } + } + return nil + }); err != nil { + m.statusMu.Lock() + m.dirtyCount += flushed + m.statusMu.Unlock() + return err + } + return nil +} + +// GetStatus returns the status of each meta-service group. +func (m *MetaServiceGroupManager) GetStatus(_ context.Context) (map[string]*endpoint.MetaServiceGroupStatus, error) { + m.statusMu.Lock() + defer m.statusMu.Unlock() + return copyStatusMap(m.cachedStatus), nil } // GetAssignmentCounts returns the count of each meta-service group. @@ -90,6 +242,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 +267,64 @@ func (m *MetaServiceGroupManager) PatchStatus(ctx context.Context, groupID strin if patch.AssignmentCount != nil && *patch.AssignmentCount < 0 { return ErrInvalidAssignmentCount } + // 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() - 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. - if _, ok := m.metaServiceGroups[groupID]; !ok { + _, known := m.metaServiceGroups[groupID] + m.RUnlock() + if !known { 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) - }) + m.statusMu.Lock() + newStatus := endpoint.MetaServiceGroupStatus{} + if currentStatus := m.cachedStatus[groupID]; currentStatus != nil { + newStatus = *currentStatus + } + m.statusMu.Unlock() + if patch.AssignmentCount != nil { + newStatus.AssignmentCount = *patch.AssignmentCount + } + if patch.Enabled != nil { + newStatus.Enabled = *patch.Enabled + } + // 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); 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 { + 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() + 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 } -func (m *MetaServiceGroupManager) findMinMetaGroup(txn kv.Txn) (string, error) { - statusMap, err := m.store.LoadMetaServiceGroupStatus(txn, m.metaServiceGroups) - if err != nil { - return "", err - } +// findMinMetaGroupStatusLocked returns the enabled group with the least assigned +// keyspaces. The caller must hold statusMu. +func (m *MetaServiceGroupManager) findMinMetaGroupStatusLocked() (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 +337,33 @@ 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. -func (m *MetaServiceGroupManager) PickGroup(ctx context.Context) (string, error) { - m.RLock() - defer m.RUnlock() - return m.pickGroupLocked(ctx) +// and updates the cached assignment count. +func (m *MetaServiceGroupManager) PickGroup(_ context.Context) (string, error) { + m.Lock() + defer m.Unlock() + return m.pickGroupLocked() } // 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 mgm 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. -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 { +// hold the mgm lock across both, e.g. via Manager.assignGroupAndSaveKeyspace. The +// selection and the reservation are done together under statusMu so a concurrent +// assignment cannot pick the same minimum twice. +func (m *MetaServiceGroupManager) pickGroupLocked() (string, error) { + m.statusMu.Lock() + defer m.statusMu.Unlock() + assignedGroup, err := m.findMinMetaGroupStatusLocked() + if err != nil { + return "", err + } + if err := m.applyAssignmentDeltaStatusLocked("", assignedGroup); err != nil { return "", err } return assignedGroup, nil @@ -180,98 +372,97 @@ 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) { +func (m *MetaServiceGroupManager) AssignToGroup(_ context.Context, count int) (string, error) { 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() + m.statusMu.Lock() + defer m.statusMu.Unlock() + assignedGroup, err := m.findMinMetaGroupStatusLocked() + if err != nil { return "", err } + m.cachedStatus[assignedGroup].AssignmentCount += count + m.markDirtyStatusLocked(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 -// concurrent UpdateGroupsSafely cannot delete a group between this validation -// and the persisted assignment count update. -func (m *MetaServiceGroupManager) reassignKeyspaceLocked(txn kv.Txn, oldGroupID, newGroupID string) error { +// moves a single keyspace assignment from oldGroupID to newGroupID. The caller +// must hold the mgm lock for the whole enclosing transaction so a concurrent +// UpdateGroupsSafely cannot delete a group between this validation and the cached +// assignment count update; metaServiceGroups is therefore read without taking +// the mgm lock again. statusMu guards the cache reads and the count update. +func (m *MetaServiceGroupManager) reassignKeyspaceLocked(_ kv.Txn, oldGroupID, newGroupID string) error { + m.statusMu.Lock() + defer m.statusMu.Unlock() if newGroupID != "" { if _, ok := m.metaServiceGroups[newGroupID]; !ok { return ErrUnknownMetaServiceGroup } // 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.applyAssignmentDeltaStatusLocked(oldGroupID, newGroupID) +} + +// updateAssignmentTxn applies an assignment count delta to the cache. It takes +// only statusMu (the leaf lock), deliberately NOT the mgm lock: callers such as +// RemoveKeyspace and the tombstone unassignment already hold the keyspace +// metaLock, and the create/config paths take the mgm lock before metaLock, so +// acquiring the mgm lock here would invert that order and deadlock. 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) updateAssignmentTxn(_ kv.Txn, oldGroupID, newGroupID string) error { + m.statusMu.Lock() + defer m.statusMu.Unlock() + return m.applyAssignmentDeltaStatusLocked(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. +// applyAssignmentDeltaStatusLocked moves one assignment from oldGroupID to +// newGroupID in the cached status and marks it dirty for the flush loop to +// persist later. The caller must hold statusMu. +func (m *MetaServiceGroupManager) applyAssignmentDeltaStatusLocked(oldGroupID, newGroupID string) error { + 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.markDirtyStatusLocked(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 +// markDirtyStatusLocked accumulates dirty count and signals the flush loop once +// the threshold is reached. The caller must hold statusMu. +func (m *MetaServiceGroupManager) markDirtyStatusLocked(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. @@ -298,6 +489,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( @@ -321,13 +521,17 @@ 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, 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 { @@ -344,33 +548,70 @@ func (m *MetaServiceGroupManager) persistGroupsLocked( if err := persist(); err != nil { return err } - m.metaServiceGroups = metaServiceGroups - // 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 { + // 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. + reconcileFailed := false + 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)) + reconcileFailed = true + } + } + // 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) + } + // 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 } // 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,23 +619,15 @@ func (m *MetaServiceGroupManager) assignedKeyspaceCounts(ctx context.Context, gr } return m.keyspaceAssignmentCounter(set) } - // Fallback path: derive counts from the persisted 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 { + // Fallback path: derive counts from the cached status under statusMu. + counts := make(map[string]int, len(groupIDs)) + m.statusMu.Lock() + for _, id := range groupIDs { + if status := m.cachedStatus[id]; status != nil { counts[id] = status.AssignmentCount } - return nil - }); err != nil { - return nil, err } + m.statusMu.Unlock() return counts, nil } @@ -403,4 +636,16 @@ func (m *MetaServiceGroupManager) updateGroups(metaServiceGroups map[string]stri m.Lock() defer m.Unlock() m.metaServiceGroups = metaServiceGroups + m.statusMu.Lock() + defer m.statusMu.Unlock() + 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..db93da723d 100644 --- a/pkg/keyspace/meta_service_group_test.go +++ b/pkg/keyspace/meta_service_group_test.go @@ -16,8 +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" @@ -46,7 +49,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() { @@ -263,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 @@ -284,6 +316,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 @@ -293,3 +361,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) +} diff --git a/server/server.go b/server/server.go index fe0a2a37d1..e5d26eee4c 100644 --- a/server/server.go +++ b/server/server.go @@ -534,7 +534,20 @@ 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(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( s.ctx, s.storage,