keyspace: cache meta-service group status

[bufferflies]

What problem does this PR solve?

Issue Number: Close #10892

This is a cherry-pick of 2baf2a5 from the PD microservices branch.

Original author: @AmoebaProtozoa

Problem Summary:

Meta service group status was loaded and persisted through storage on assignment paths. This adds an in-memory status cache and flush loop so assignment count updates can be handled locally and persisted in batches by the serving leader.

What is changed and how does it work?

• Cache meta-service group status in MetaServiceGroupManager and refresh it from storage on initialization/service readiness.
• Update assignment counts in the cached status for group selection, reassignment, rollback, and removal paths.
• Flush dirty cached status to storage periodically or when the dirty threshold is reached, guarded by a leader checker.
• Keep PatchStatus synchronous with storage and cache so API updates are immediately visible.
• Adapt the cherry-pick to current master after PR meta-service-group: add status #10904, preserving the current UpdateGroupsSafely delete guard and authoritative keyspace assignment scanner.
• Add regression coverage for cache refresh and threshold-triggered flush.

Check List

Tests

• Unit test

Code changes

• Has persistent data change

Side effects

• No compatibility impact

Related changes

• Need to cherry-pick to the release branch

Release note

Cache meta-service group status and flush assignment count updates to storage in batches.

Validation

• GOFLAGS=-buildvcs=false make gotest GOTEST_ARGS='./pkg/keyspace -run TestMetaServiceGroupTestSuite -count=1 -timeout=5m'
• GOFLAGS=-buildvcs=false make gotest GOTEST_ARGS='./pkg/keyspace -run "TestMetaServiceGroupTestSuite|TestAssignGroupAndSaveKeyspace|TestKeyspaceTestSuite/(TestUpdateKeyspaceState|TestTombstoneKeyspaceUnassignsMetaServiceGroup)" -count=1 -timeout=5m'
• GOFLAGS=-buildvcs=false make gotest GOTEST_ARGS='./pkg/keyspace -count=1 -timeout=5m'
• GOFLAGS=-buildvcs=false make gotest GOTEST_ARGS='./server -run TestNonExistent -count=0 -timeout=5m'
• git diff --check && git diff --cached --check
• GOFLAGS=-buildvcs=false make check

Summary by CodeRabbit

• New Features
  • Added cached in-memory meta-service group status with batched persistence and leader-aware flushing.
  • Added the ability to refresh group status from storage and expose whether groups are available.
• Bug Fixes
  • Tightened assignment/config concurrency handling to avoid races during group changes, including safer rollback when updates fail after reassignment.
• Tests
  • Updated manager initialization to use context/error handling.
  • Added coverage for cache refresh and batched persistence after reaching write thresholds.

[ti-chi-bot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign lhy1024 for approval. For more information see the Code Review Process.
Please ensure that each of them provides their approval before proceeding.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[bufferflies]

@coderabbitai full review

[coderabbitai]

✅ Action performed

Full review finished.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

MetaServiceGroupManager now caches status in memory with batched flushing. Keyspace assignment and config updates use write-locked accounting with rollback handling, and server startup plus tests were updated for the new constructor and cache refresh flow.

Changes

Meta-service group cache and locking refactor

Layer / File(s)	Summary
Cached status and flush loop pkg/keyspace/meta_service_group.go	MetaServiceGroupManager adds cached status storage, dirty tracking, leader-aware flushing, cache refresh, and deep-copy status reads.
Cached assignment updates pkg/keyspace/meta_service_group.go	PatchStatus, PickGroup, AssignToGroup, persistGroupsLocked, assignedKeyspaceCounts, and updateGroups update cached counts and reconcile cached entries with group changes.
Write-locked keyspace assignment and rollback pkg/keyspace/keyspace.go	assignGroupAndSaveKeyspace, runTxnWithMetaGroupLock, and updateKeyspaceConfigTxn use the meta-service group manager write lock, re-check group availability, and conditionally roll back assignment counters after failures.
Server startup and test updates server/server.go, pkg/keyspace/keyspace_test.go, pkg/keyspace/meta_service_group_test.go	Server startup handles the context-aware constructor error, sets leader checking, refreshes cache on readiness, and tests use the new constructor while covering cache refresh and flush-threshold behavior.

Estimated code review effort: 4 (Complex) | ~60 minutes

Possibly related PRs

• tikv/pd#10597: Introduced the meta-service group assignment plumbing that this PR refines.
• tikv/pd#10881: Uses the meta-service group manager paths that this PR updates for cached status handling.
• tikv/pd#10904: Touches the same assignment-count and meta-service group accounting flow updated here.

Suggested labels: lgtm, approved

Suggested reviewers: ystaticy, rleungx, lhy1024

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: caching meta-service group status in keyspace handling.
Description check	✅ Passed	The description includes the issue link, problem statement, change summary, tests, side effects, and release note.
Linked Issues check	✅ Passed	The changes implement #10892 by adding cached status, flush behavior, refresh on readiness, and regression tests.
Out of Scope Changes check	✅ Passed	The server and keyspace updates support the cache feature and do not introduce unrelated scope beyond the issue.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@pkg/keyspace/meta_service_group.go`:
- Around line 315-336: updateAssignmentLockedTxn is mutating cachedStatus and
dirtyCount without using the storage txn, so aborted transactions can
desynchronize cache state from persisted state. Update
MetaServiceGroupManager.updateAssignmentLockedTxn to make cache changes
transactional by tying the cachedStatus assignment count and markDirtyLocked
updates to the txn outcome (or otherwise rolling them back on error), and keep
the oldGroupID decrement/newGroupID increment behavior symmetric even when the
source count is zero.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 85dd8198-22b4-4b3e-a3fa-d656ff749e92

📥 Commits

Reviewing files that changed from the base of the PR and between 1de3200 and 373334d.

📒 Files selected for processing (5)

• pkg/keyspace/keyspace.go
• pkg/keyspace/keyspace_test.go
• pkg/keyspace/meta_service_group.go
• pkg/keyspace/meta_service_group_test.go
• server/server.go

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

pkg/keyspace/meta_service_group.go (1)

40-40: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Pass the lifecycle context through RefreshCache, flushLoop, and flushToStorage instead of keeping m.ctx on MetaServiceGroupManager. NewMetaServiceGroupManager can hand the constructor context to the goroutine, and the service-ready callback already has a context.Context to pass into RefreshCache.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@pkg/keyspace/meta_service_group.go` at line 40, Remove the stored lifecycle
context from MetaServiceGroupManager and thread context explicitly through
RefreshCache, flushLoop, and flushToStorage instead. Update
NewMetaServiceGroupManager to pass the constructor context into the flush
goroutine, and use the context already available in the service-ready callback
when calling RefreshCache. Locate the affected flow in MetaServiceGroupManager,
RefreshCache, flushLoop, and flushToStorage and update their signatures/call
sites accordingly.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@pkg/keyspace/keyspace.go`:
- Around line 868-873: The rollback logic around runTxnWithMetaGroupLock and
UpdateKeyspaceGroup needs to stay under the meta-service group lock, and commit
failures must revert all side effects, not just the meta-service cache. Move the
non-transactional rollback work into the locked helper/error path so
UpdateGroupsSafely cannot see transient reassignment counts after the lock is
released, and track the successful UpdateKeyspaceGroup/Tso group move state
alongside metaServiceGroupReassigned so the rollback callback can restore both
before unlocking.

In `@pkg/keyspace/meta_service_group.go`:
- Around line 420-433: Persist the reset status for re-added groups instead of
only creating it in memory. In meta_service_group.go, update the cached-status
initialization around the metaServiceGroups loop so newly created
endpoint.MetaServiceGroupStatus entries are also added to newStatuses (or
otherwise marked dirty) before the cleanup transaction runs. Keep the
deletedGroups cleanup behavior, but ensure the status persistence path in
MetaServiceGroup metadata writes the zeroed state so a restart cannot reload
stale assignment/enabled values.

---

Nitpick comments:
In `@pkg/keyspace/meta_service_group.go`:
- Line 40: Remove the stored lifecycle context from MetaServiceGroupManager and
thread context explicitly through RefreshCache, flushLoop, and flushToStorage
instead. Update NewMetaServiceGroupManager to pass the constructor context into
the flush goroutine, and use the context already available in the service-ready
callback when calling RefreshCache. Locate the affected flow in
MetaServiceGroupManager, RefreshCache, flushLoop, and flushToStorage and update
their signatures/call sites accordingly.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 4fc3ed40-bd11-4452-a951-20fd54e535fc

📥 Commits

Reviewing files that changed from the base of the PR and between 1de3200 and 373334d.

📒 Files selected for processing (5)

• pkg/keyspace/keyspace.go
• pkg/keyspace/keyspace_test.go
• pkg/keyspace/meta_service_group.go
• pkg/keyspace/meta_service_group_test.go
• server/server.go

[codecov]

Codecov Report

❌ Patch coverage is 79.06977% with 63 lines in your changes missing coverage. Please review.
✅ Project coverage is 79.24%. Comparing base (aa5a988) to head (5ffb5bf).
⚠️ Report is 10 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #10976      +/-   ##
==========================================
+ Coverage   79.21%   79.24%   +0.02%     
==========================================
  Files         541      542       +1     
  Lines       75677    76247     +570     
==========================================
+ Hits        59949    60423     +474     
- Misses      11491    11561      +70     
- Partials     4237     4263      +26     

Flag	Coverage Δ
unittests	79.24% <79.06%> (+0.02%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[lhy1024]

flushToStorage can overwrite concurrent synchronous status updates. It snapshots the whole cachedStatus, releases statusMu, and then blindly writes every status from that snapshot. A concurrent PatchStatus can persist a new enabled/count value while this flush is in flight; if the older snapshot commits afterward, the patch is lost on the next RefreshCache or restart. The same race can also recreate a status key removed by UpdateGroupsSafely.

Please serialize flush persistence with PatchStatus/group deletion, or make the flush conditional/merge-safe and skip groups deleted after the snapshot.

[lhy1024]

A concrete regression test for this race could look like this:

import (
    "sync/atomic"

    "github.com/stretchr/testify/require"
)

type blockingMetaServiceGroupStorage struct {
    *endpoint.StorageEndpoint
    blockNextSave atomic.Bool
    saveStarted   chan struct{}
    unblockSave   chan struct{}
}

func (s *blockingMetaServiceGroupStorage) SaveMetaServiceGroupStatus(txn kv.Txn, id string, status *endpoint.MetaServiceGroupStatus) error {
    if s.blockNextSave.CompareAndSwap(true, false) {
        close(s.saveStarted)
        <-s.unblockSave
    }
    return s.StorageEndpoint.SaveMetaServiceGroupStatus(txn, id, status)
}

func TestFlushDoesNotOverwriteConcurrentPatchStatus(t *testing.T) {
    re := require.New(t)
    ctx, cancel := context.WithCancel(context.Background())
    defer cancel()

    store := &blockingMetaServiceGroupStorage{
        StorageEndpoint: endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil),
        saveStarted:     make(chan struct{}),
        unblockSave:     make(chan struct{}),
    }
    manager, err := NewMetaServiceGroupManager(ctx, store, mockMetaServiceGroups())
    re.NoError(err)

    groupID := "etcd-group-0"
    manager.statusMu.Lock()
    manager.cachedStatus[groupID] = &endpoint.MetaServiceGroupStatus{AssignmentCount: 1, Enabled: true}
    manager.dirtyCount = 1
    manager.statusMu.Unlock()

    store.blockNextSave.Store(true)
    flushErr := make(chan error, 1)
    go func() {
        flushErr <- manager.flushToStorage()
    }()
    <-store.saveStarted

    newCount := 99
    re.NoError(manager.PatchStatus(ctx, groupID, &MetaServiceGroupStatusPatch{
        AssignmentCount: &newCount,
    }))

    close(store.unblockSave)
    re.NoError(<-flushErr)
    re.NoError(manager.RefreshCache())

    statusMap, err := manager.GetStatus(ctx)
    re.NoError(err)
    re.Equal(newCount, statusMap[groupID].AssignmentCount)
}

I ran this focused test locally against the current implementation. It fails with expected: 99 and actual: 1, which means the old flush snapshot overwrites the synchronous PatchStatus update.

[bufferflies]

Fixed in e5262c6. flushToStorage now holds the mgm write lock across the storage write, serializing it with PatchStatus and persistGroupsLocked, so the snapshot can no longer overwrite a synchronous write or recreate a deleted status key. It stays deadlock-free because the metaLock-holding paths (RemoveKeyspace, tombstone unassignment) only take statusMu, never the mgm lock.

Also fixed an adjacent read-modify-write race in PatchStatus (88ae189): it now applies only the patched fields to the live cached status instead of overwriting it with a pre-I/O snapshot, so a concurrent updateAssignmentTxn count change is not dropped.

[bufferflies]

Thanks for the repro — added TestFlushDoesNotOverwriteConcurrentPatchStatus based on it. One tweak: with the fix, PatchStatus blocks on the mgm lock the paused flush holds, so calling it synchronously (as in the snippet) deadlocks against the unblockSave that comes after it. So I run PatchStatus in a goroutine and assert its value wins once the flush releases the lock.

[rleungx]

PatchStatus releases statusMu before persisting and then replaces the whole cached entry with the earlier snapshot. updateAssignmentTxn intentionally skips the manager lock, so a concurrent tombstone or removal can update AssignmentCount between those critical sections; an enabled-only patch then overwrites that successful delta while dirtyCount remains set. The next flush persists the stale count, even though both operations returned success.

[bufferflies]

Fixed in 88ae189. PatchStatus now applies only the patched fields to the live cached status instead of replacing the whole entry with the pre-I/O snapshot, so a concurrent updateAssignmentTxn count delta is preserved. The synchronous store write still lands under the same serialization as the flush (see below).

[rleungx]

isLeader is checked only before this transaction, so the write is not fenced to the leadership term. If this PD loses its lease after line 158, the flush continues with the server-lifetime context; these status saves are unconditional puts with no loaded-key comparison, so the former leader can commit its snapshot after the new leader has refreshed or patched the same status. A leader transfer can therefore overwrite a newer synchronous status update.

[bufferflies]

Mitigated in 52c082b: the flush now runs on the leadership-term context (set from the service-ready callback and canceled when the lease is lost) instead of the server-lifetime context, and re-checks isLeader right before the write. That cancels a former leader's in-flight flush and stops new ones after the term ends.

You are right that this does not fully close it: an already-issued commit can still land after cancellation. Fully fencing it needs a leader-key comparison inside the flush txn (so the put aborts when the lease key changed), which requires plumbing the leadership guard into the keyspace storage writes. I left that as a follow-up and noted it in the commit message — happy to do it here if you prefer.

[rleungx]

The term context narrows the window but does not fence a transaction whose commit has already been issued, so the former leader can still overwrite a newer synchronous PatchStatus after a transfer. Because the stale write changes API-visible enabled/count state, the original correctness issue remains.

[lhy1024]

This is still reproducible on the current head (52c082b98). The term context narrows the window, but it is canceled only after the old leader loop observes lease expiry and returns. The election key can already be gone, and a new leader can already be serving, while that context is still live.

I added a deterministic regression test that pauses the old flush after its leader check, marks the old instance non-serving while leaving the term context live (the interval before the leader ticker returns), lets the new leader persist count 99, and then resumes the old flush. The final persisted count is 1, not 99:

TestFormerLeaderFlushIsFencedBeforeTermContextCancellation
expected: 99
actual:   1

So context cancellation is not a leadership fence. Please include the election leader-key/value comparison in the same etcd transaction, equivalent to Leadership.LeaderTxn; otherwise an unconditional write that passed the preflight check immediately before lease expiry can still land in the next term.

[rleungx]

flushToStorage holds the manager write lock across the etcd transaction. AttachEndpoints on keyspace read paths and the group configuration/automatic-assignment paths need this lock, so a slow flush stalls otherwise read-only keyspace lookups and unrelated group operations on the serving PD. Routine threshold or periodic flushes can therefore turn one slow storage write into a broad request-latency spike.

[bufferflies]

Fixed in 52c082b. Introduced persistMu, an outermost mutex that serializes the storage writers (flush, PatchStatus, persistGroupsLocked, RefreshCache) instead of the mgm write lock. The flush now holds only persistMu across the etcd transaction and takes the mgm read lock just briefly to read isLeader/leaderCtx, so a slow flush no longer blocks AttachEndpoints/GetGroups read paths or assignment operations. PatchStatus likewise no longer holds the mgm write lock across its I/O.

[lhy1024]

RefreshCache replaces the cache with the persisted batching snapshot and resets dirtyCount, but it never reconciles AssignmentCount with authoritative keyspace metadata. A leader change before a successful flush therefore permanently drops all assignment deltas since the previous flush; subsequent PickGroup decisions continue from the stale baseline.

The manager already has keyspaceAssignmentCounter. A focused test configures that authoritative counter to return 1, while storage contains the last flushed value 0; after RefreshCache, the current implementation still returns 0:

TestRefreshCacheReconcilesAuthoritativeAssignmentCounts
expected: 1
actual:   0

Please preserve persisted administrative fields such as Enabled, while rebuilding assignment counts from keyspace metadata during leader readiness.

[lhy1024]

This is not atomic with the real callers. persist() is SetKeyspaceConfigWithoutKeyspaceManagerUpdate in both API paths, which CASes and writes the keyspace config before the independent status-cleanup transaction below. If RemoveMetaServiceGroupStatus fails, this method returns an error and skips afterPersist, but the config is already durable while metaServiceGroups and the keyspace manager still retain the old group. The serving PD can then assign new keyspaces to a group that has already been removed from persisted config.

Please make the config/status updates atomic, or compensate the config rollback under the same lock. Returning the cleanup error after committing the config leaves a split-brain state.

A focused regression test:

var errInjectedStatusCleanup = errors.New("injected status cleanup failure")

type cleanupFailingStorage struct {
    *endpoint.StorageEndpoint
}

func (*cleanupFailingStorage) RemoveMetaServiceGroupStatus(kv.Txn, string) error {
    return errInjectedStatusCleanup
}

func TestUpdateGroupsSafelyDoesNotCommitConfigBeforeStatusCleanup(t *testing.T) {
    re := require.New(t)
    ctx := context.Background()
    store := &cleanupFailingStorage{
        StorageEndpoint: endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil),
    }
    manager, err := NewMetaServiceGroupManager(ctx, store, mockMetaServiceGroups())
    re.NoError(err)

    newGroups := mockMetaServiceGroups()
    delete(newGroups, "etcd-group-0")
    configPersisted := false
    err = manager.UpdateGroupsSafely(ctx, newGroups, []string{"etcd-group-0"}, func() error {
        configPersisted = true
        return nil
    }, nil)

    re.ErrorIs(err, errInjectedStatusCleanup)
    re.False(configPersisted, "a failed group update must not leave config committed")
}

This fails on 1bc6846 because configPersisted is true.

[ti-chi-bot]

@bufferflies: The following tests failed, say /retest to rerun all failed tests or /retest-required to rerun all mandatory failed tests:

Test name	Commit	Details	Required	Rerun command
pull-error-log-review	5ffb5bf	link	true	/test pull-error-log-review
pull-unit-test-next-gen-2	5ffb5bf	link	true	/test pull-unit-test-next-gen-2
pull-unit-test-next-gen-3	5ffb5bf	link	true	/test pull-unit-test-next-gen-3

Full PR test history. Your PR dashboard.

Details

Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. I understand the commands that are listed here.

[rleungx]

An unassignment only advances the epoch when the cached count is greater than zero. A fresh leader starts at zero, so a tombstone/remove during the scan can commit without invalidating a result that still counted the old binding; that stale result is then marked ready. Please advance the epoch for every logical assignment mutation, even when the decrement is clamped at zero.

[rleungx]

An epoch change immediately starts another full scan. With 1M keyspaces, steady assignment or tombstone traffic can keep invalidating the scan and turn this into an unthrottled loop of about 8,334 range reads, while assignment_count_ready never becomes true. Please add backoff/rate limiting or reconcile the deltas without restarting the full scan.

[rleungx]

A canceled rebuild from the previous leader term can reach this branch after the new term rebuild has completed and unconditionally clear the shared ready/refreshing flags. That leaves the current leader reporting not ready with no worker running. Please guard these updates with a term/rebuild generation, or wait for the previous worker to exit.

[lhy1024]

This turns the cached assignment path into a global etcd-latency queue. assignGroupAndSaveKeyspace holds mgm.Lock() until saveNewKeyspace commits, so a slow transaction blocks every other keyspace creation (and PickGroup), even though their reservations are independent. statusMu already serializes selecting and reserving a group; the outer lock only needs to prevent UpdateGroupsSafely from deleting a selected group, which a shared RLock can do.

I reproduced this with a deterministic unit test: block the first RunInTxn after it has reserved a group, start a second assignGroupAndSaveKeyspace, and assert that it finishes before the first transaction is unblocked. The current head times out at 100ms:

store.blockNextTxn()
go assign(first)
<-store.txnStarted
go assign(second)
select {
case <-secondDone:
case <-time.After(100 * time.Millisecond):
    t.Fatal("a blocked storage save must not serialize another group assignment")
}

Please keep group deletion as the writer, but let normal assignment transactions take RLock while statusMu retains the atomic reservation.