Skip to content

keyspace: add optional keyspace info metric #11176

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
nolouch wants to merge 15 commits into
base: master
Choose a base branch
from
Open

keyspace: add optional keyspace info metric #11176

Show file tree
Hide file tree
Changes from 8 commits
Commits
Show all changes
15 commits
Select commit Hold shift + click to select a range
77fa433
keyspace: add keyspace info metric
nolouch Aug 25, 2026
a5ebd3e
keyspace: generalize keyspace metrics config
nolouch Aug 25, 2026
8afd0b3
keyspace: populate info metric on demand
nolouch Aug 25, 2026
fd066d1
keyspace: fix info metric lifecycle
nolouch Aug 25, 2026
80a0a1e
keyspace: fix removal helper naming
nolouch Aug 25, 2026
c02f996
keyspace: fix info metric lifecycle paths
nolouch Aug 25, 2026
2c24a4c
keyspace: fix metrics helper lint
nolouch Aug 25, 2026
0db1607
keyspace: simplify info metric state
nolouch Aug 25, 2026
247ace8
keyspace: make level metrics startup-only
nolouch Aug 25, 2026
b0dc9cc
keyspace: maintain info metrics through watch
nolouch Aug 25, 2026
0ad0ff7
keyspace: fix metrics test license header
nolouch Aug 25, 2026
5f76688
keyspace: populate info metrics from client loads
nolouch Aug 25, 2026
c146ae5
keyspace: clarify info metric semantics
nolouch Aug 25, 2026
f23de01
keyspace: avoid redundant info metric deletion
nolouch Aug 26, 2026
873377b
keyspace: track info metric lifecycle
nolouch Aug 26, 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
82 changes: 72 additions & 10 deletions pkg/keyspace/keyspace.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,6 +76,7 @@ const (
// Config is the interface for keyspace config.
type Config interface {
GetPreAlloc() []string
IsKeyspaceLevelMetricsEnabled() bool
ToWaitRegionSplit() bool
GetWaitRegionSplitTimeout() time.Duration
GetCheckRegionSplitInterval() time.Duration
Expand Down Expand Up @@ -242,7 +243,13 @@ func (manager *Manager) initReserveKeyspace(id uint32, name string) error {

// UpdateConfig update keyspace manager's config.
func (manager *Manager) UpdateConfig(cfg Config) {
keyspaceInfoMetricsMu.Lock()
wasEnabled := manager.config.IsKeyspaceLevelMetricsEnabled()
manager.config = cfg
if wasEnabled && !cfg.IsKeyspaceLevelMetricsEnabled() {
resetKeyspaceInfoMetricsLocked()
}
keyspaceInfoMetricsMu.Unlock()
if manager.mgm != nil {
manager.mgm.updateGroups(cfg.GetMetaServiceGroups())
}
Expand Down Expand Up @@ -392,6 +399,7 @@ func (manager *Manager) CreateKeyspace(request *CreateKeyspaceRequest) (*keyspac
zap.Uint32("keyspace-id", keyspace.GetId()),
zap.String("keyspace-name", keyspace.GetName()),
)
manager.UpdateKeyspaceInfoMetrics(keyspace)

@lhy1024 lhy1024 Aug 25, 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.

Blocking: the metric is published later than the metadata it represents.

UpdateKeyspaceStateByID has already committed an ENABLED metadata record before UpdateKeyspaceForGroup runs. If that separate group write fails, this method returns here without rolling the metadata back and without reaching the metric update below. The durable keyspace cannot be retried under the same name, yet its mapping remains absent until an unrelated load or watch happens. CreateKeyspaceByID has the same ordering.

Please publish once the metadata becomes durable and remove the series only when a rollback deletion commits, or make every subsequent failure roll the metadata back. Whether a client has opened WatchKeyspaces should not change the result.

Unit-test reproducer 
var errReviewGroupSave = errors.New("review: save keyspace group")

type failingKeyspaceGroupStorage struct {
	endpoint.KeyspaceGroupStorage
	fail bool
}

func (s *failingKeyspaceGroupStorage) SaveKeyspaceGroup(
	txn kv.Txn, group *endpoint.KeyspaceGroup,
) error {
	if s.fail {
		return errReviewGroupSave
	}
	return s.KeyspaceGroupStorage.SaveKeyspaceGroup(txn, group)
}

func (suite *keyspaceTestSuite) TestEnabledMetadataKeepsInfoMetricWhenGroupUpdateFails() {
	re := suite.Require()
	resetKeyspaceInfoMetrics()

	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
	groupStore := &failingKeyspaceGroupStorage{KeyspaceGroupStorage: store}
	kgm := NewKeyspaceGroupManager(suite.ctx, groupStore, nil)
	re.NoError(kgm.Bootstrap(suite.ctx))
	manager := NewKeyspaceManager(
		suite.ctx, store, nil, mockid.NewIDAllocator(),
		&mockConfig{EnableKeyspaceLevelMetrics: true}, kgm, nil,
	)

	groupStore.fail = true
	const name = "review_group_failure"
	_, err := manager.CreateKeyspace(&CreateKeyspaceRequest{
		Name: name, CreateTime: time.Now().Unix(),
	})
	re.ErrorIs(err, errReviewGroupSave)

	var persisted *keyspacepb.KeyspaceMeta
	re.NoError(store.RunInTxn(suite.ctx, func(txn kv.Txn) error {
		loaded, id, err := store.LoadKeyspaceID(txn, name)
		re.NoError(err)
		re.True(loaded)
		persisted, err = store.LoadKeyspaceMeta(txn, id)
		return err
	}))
	re.Equal(keyspacepb.KeyspaceState_ENABLED, persisted.GetState())
	re.Equal(1, promtestutil.CollectAndCount(keyspaceInfo))
}

Run:

make gotest GOTEST_ARGS='./pkg/keyspace -run TestKeyspaceTestSuite/TestEnabledMetadataKeepsInfoMetricWhenGroupUpdateFails -count=1 -timeout=30s'

On 247ace82be, the persisted state is ENABLED, but the final assertion reports expected: 1, actual: 0.

@nolouch nolouch Aug 25, 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.

This no longer depends on the create flow. The metric is now populated from successful GetAllKeyspaces and LoadKeyspaceByID requests, so partial create failures won’t directly update it.

return keyspace, nil
}

Expand Down Expand Up @@ -529,6 +537,7 @@ func (manager *Manager) CreateKeyspaceByID(request *CreateKeyspaceByIDRequest) (
zap.String("keyspace-name", keyspace.GetName()),
zap.Any("keyspace", keyspace),
)
manager.UpdateKeyspaceInfoMetrics(keyspace)
return keyspace, nil
}

Expand Down Expand Up @@ -770,6 +779,9 @@ func (manager *Manager) LoadKeyspace(name string) (*keyspacepb.KeyspaceMeta, err
if manager.mgm != nil && meta != nil {
manager.mgm.AttachEndpoints(meta.GetConfig())
}
if err == nil {
manager.UpdateKeyspaceInfoMetrics(meta)
}
return meta, err
}

Expand All @@ -793,6 +805,9 @@ func (manager *Manager) LoadKeyspaceByID(spaceID uint32) (*keyspacepb.KeyspaceMe
if manager.mgm != nil && meta != nil {
manager.mgm.AttachEndpoints(meta.GetConfig())
}
if err == nil {
manager.UpdateKeyspaceInfoMetrics(meta)
Comment thread
nolouch marked this conversation as resolved.
Outdated

@lhy1024 lhy1024 Aug 25, 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.

Blocking: a load that read before deletion can permanently recreate the deleted series.

The read transaction finishes before this metric update. A concurrent RemoveKeyspace can therefore commit and run finishRemoveKeyspace after the load has read the old metadata but before this line executes. The load then recreates the series after the only cleanup event has finished, so it remains indefinitely. LoadKeyspace and LoadRangeKeyspace have the same post-transaction write.

The metric publication needs ordering or version validation against removal. A second unsynchronized existence check would leave the same check-then-write race.

Unit-test reproducer

Add sync/atomic to the imports.

type blockingKeyspaceStorage struct {
	endpoint.KeyspaceStorage
	blockNext atomic.Bool
	reached   chan struct{}
	release   chan struct{}
}

func (s *blockingKeyspaceStorage) RunInTxn(ctx context.Context, f func(kv.Txn) error) error {
	err := s.KeyspaceStorage.RunInTxn(ctx, f)
	if s.blockNext.CompareAndSwap(true, false) {
		close(s.reached)
		<-s.release
	}
	return err
}

func (suite *keyspaceTestSuite) TestStaleLoadCannotRestoreRemovedInfoMetric() {
	re := suite.Require()
	suite.manager.UpdateConfig(&mockConfig{EnableKeyspaceLevelMetrics: true})
	resetKeyspaceInfoMetrics()
	created, err := suite.manager.CreateKeyspace(&CreateKeyspaceRequest{
		Name: "review_stale_load", CreateTime: time.Now().Unix(),
	})
	re.NoError(err)
	for _, state := range []keyspacepb.KeyspaceState{
		keyspacepb.KeyspaceState_DISABLED,
		keyspacepb.KeyspaceState_ARCHIVED,
		keyspacepb.KeyspaceState_TOMBSTONE,
	} {
		_, err = suite.manager.UpdateKeyspaceState(created.GetName(), state, time.Now().Unix())
		re.NoError(err)
	}

	store := &blockingKeyspaceStorage{
		KeyspaceStorage: suite.manager.store,
		reached: make(chan struct{}),
		release: make(chan struct{}),
	}
	suite.manager.store = store
	store.blockNext.Store(true)
	loadDone := make(chan error, 1)
	go func() {
		_, loadErr := suite.manager.LoadKeyspaceByID(created.GetId())
		loadDone <- loadErr
	}()
	<-store.reached

	re.NoError(suite.manager.RemoveKeyspace(created.GetId()))
	re.Equal(0, promtestutil.CollectAndCount(keyspaceInfo))
	close(store.release)
	re.NoError(<-loadDone)
	re.Equal(0, promtestutil.CollectAndCount(keyspaceInfo))
}

Run:

make gotest GOTEST_ARGS='./pkg/keyspace -run TestKeyspaceTestSuite/TestStaleLoadCannotRestoreRemovedInfoMetric -count=1 -timeout=30s'

On 247ace82be, the final assertion reports expected: 0, actual: 1.

@nolouch nolouch Aug 25, 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.

This metric is now defined as a client-observed mapping, not a strongly consistent metadata snapshot. A concurrent deletion may leave a previously observed mapping temporarily, but deleted keyspaces won’t have active storage-size series to join with. Avoiding this window would require extra lifecycle synchronization, which we want to avoid here.

@lhy1024 lhy1024 Aug 26, 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 makes this broader than the concurrent window described above: no successful removal path deletes an observed series at all. RemoveKeyspace removes durable metadata, and RemoveKeyspacesFromGroup commits that transaction, but there is no metric cleanup caller anywhere.

func (suite *keyspaceTestSuite) TestRemovalClearsObservedInfoMetric() {
    re := suite.Require()
    resetKeyspaceInfoMetrics()
    suite.T().Cleanup(resetKeyspaceInfoMetrics)

    created, err := suite.manager.CreateKeyspace(&CreateKeyspaceRequest{
        Name: "review_rm_metric", CreateTime: time.Now().Unix(),
    })
    re.NoError(err)
    SetKeyspaceInfoMetrics(created.GetId(), created.GetName())
    re.Equal(1, testutil.CollectAndCount(keyspaceInfo))

    for _, state := range []keyspacepb.KeyspaceState{
        keyspacepb.KeyspaceState_DISABLED,
        keyspacepb.KeyspaceState_ARCHIVED,
    } {
        _, err = suite.manager.UpdateKeyspaceStateByID(
            created.GetId(), state, time.Now().Unix())
        re.NoError(err)
    }
    re.NoError(suite.manager.store.RunInTxn(suite.ctx, func(txn kv.Txn) error {
        return suite.manager.RemoveKeyspace(txn, created.GetId())
    }))
    re.Equal(0, testutil.CollectAndCount(keyspaceInfo))
}

Run:

make gotest GOTEST_ARGS="./pkg/keyspace -run TestKeyspaceTestSuite/TestRemovalClearsObservedInfoMetric -count=1 -timeout=30s"

On f23de01da, the removal commits but the final count is 1. I also reproduced the same result through the public GroupManager.RemoveKeyspacesFromGroup path after first populating the series through LoadKeyspaceByID.

This is explicitly required by Issue #11175. The fact that another storage-size metric may currently have no sample does not make stale high-cardinality labels safe: ID reuse makes the stale series ambiguous, and repeated create/remove cycles retain series indefinitely. Please restore post-commit cleanup.

}
return meta, err
}

Expand Down Expand Up @@ -1035,33 +1050,77 @@ func (manager *Manager) UpdateKeyspaceState(name string, newState keyspacepb.Key
return meta, nil
}

// RemoveKeyspace removes the keyspace specified by id if it's in proper state and not protected.
func (manager *Manager) RemoveKeyspace(txn kv.Txn, id uint32) error {
// RemoveKeyspace removes the keyspace specified by id and clears its caches after the transaction commits.
func (manager *Manager) RemoveKeyspace(id uint32) error {
var removed *keyspacepb.KeyspaceMeta
err := manager.store.RunInTxn(manager.ctx, func(txn kv.Txn) error {
var err error
removed, err = manager.stageRemoveKeyspace(txn, id)
return err
})
if err != nil {
return err
}
manager.finishRemoveKeyspace(removed)
return nil
}

// stageRemoveKeyspace stages keyspace removal and returns its metadata for post-commit cleanup.
func (manager *Manager) stageRemoveKeyspace(txn kv.Txn, id uint32) (*keyspacepb.KeyspaceMeta, error) {
manager.metaLock.Lock(id)
defer manager.metaLock.Unlock(id)
if isProtectedKeyspaceID(id) {
return newModifyProtectedKeyspaceError()
return nil, newModifyProtectedKeyspaceError()
}
meta, err := manager.store.LoadKeyspaceMeta(txn, id)
if err != nil {
return err
return nil, err
}
if meta == nil {
return errs.ErrKeyspaceNotFound
return nil, errs.ErrKeyspaceNotFound
}
if meta.GetState() == keyspacepb.KeyspaceState_ENABLED || meta.GetState() == keyspacepb.KeyspaceState_DISABLED {
return errors.Errorf("cannot remove keyspace in state %s", meta.GetState().String())
return nil, errors.Errorf("cannot remove keyspace in state %s", meta.GetState().String())
}
err = manager.store.RemoveKeyspace(txn, id, meta.GetName())
if err != nil {
return err
return nil, err
}
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.
return manager.unassignKeyspaceFromMetaServiceGroup(txn, meta)
if err := manager.unassignKeyspaceFromMetaServiceGroup(txn, meta); err != nil {
return nil, err
}
return meta, nil
}

// finishRemoveKeyspace clears the caches and metrics after a keyspace removal transaction commits.
func (manager *Manager) finishRemoveKeyspace(meta *keyspacepb.KeyspaceMeta) {
manager.keyspaceNameLookup.Delete(meta.GetId())
manager.keyspaceStateLookup.Delete(meta.GetId())
keyspaceInfoMetricsMu.Lock()
defer keyspaceInfoMetricsMu.Unlock()
if manager.config.IsKeyspaceLevelMetricsEnabled() {
deleteKeyspaceInfoMetricsLocked(meta.GetId(), meta.GetName())
}
}

// UpdateKeyspaceInfoMetrics updates the keyspace ID-to-name mapping metric when keyspace-level metrics are enabled.
func (manager *Manager) UpdateKeyspaceInfoMetrics(meta *keyspacepb.KeyspaceMeta) {
keyspaceInfoMetricsMu.Lock()
defer keyspaceInfoMetricsMu.Unlock()
if meta == nil || !manager.config.IsKeyspaceLevelMetricsEnabled() {
Comment thread
nolouch marked this conversation as resolved.
Outdated
return
}
setKeyspaceInfoMetricsLocked(meta.GetId(), meta.GetName())
}

// DeleteKeyspaceInfoMetrics removes the keyspace info series for id.
func (*Manager) DeleteKeyspaceInfoMetrics(id uint32) {
keyspaceInfoMetricsMu.Lock()
defer keyspaceInfoMetricsMu.Unlock()
deleteKeyspaceInfoMetricsByIDLocked(id)
}

// UpdateKeyspaceStateByID updates target keyspace to the given state if it's not already in that state.
Expand Down Expand Up @@ -1190,6 +1249,9 @@ func (manager *Manager) LoadRangeKeyspace(startID uint32, limit int) ([]*keyspac
}
}
}
for _, meta := range keyspaces {
manager.UpdateKeyspaceInfoMetrics(meta)
}
return keyspaces, nil
}

Expand Down
113 changes: 108 additions & 5 deletions pkg/keyspace/keyspace_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,17 @@ package keyspace

import (
"context"
"errors"
"fmt"
"math"
"strconv"
"strings"
"sync"
"testing"
"time"

"github.com/prometheus/client_golang/prometheus"
promtestutil "github.com/prometheus/client_golang/prometheus/testutil"
"github.com/stretchr/testify/require"
"github.com/stretchr/testify/suite"
"go.uber.org/goleak"
Expand Down Expand Up @@ -64,14 +67,19 @@ func TestKeyspaceTestSuite(t *testing.T) {
}

type mockConfig struct {
PreAlloc []string
WaitRegionSplit bool
WaitRegionSplitTimeout typeutil.Duration
CheckRegionSplitInterval typeutil.Duration
PreAlloc []string
EnableKeyspaceLevelMetrics bool
WaitRegionSplit bool
WaitRegionSplitTimeout typeutil.Duration
CheckRegionSplitInterval typeutil.Duration
// MetaServiceGroups is used to mock the meta-service groups for keyspace assignment.
MetaServiceGroups map[string]string
}

func (m *mockConfig) IsKeyspaceLevelMetricsEnabled() bool {
return m.EnableKeyspaceLevelMetrics
}

func (m *mockConfig) GetPreAlloc() []string {
return m.PreAlloc
}
Expand Down Expand Up @@ -108,9 +116,100 @@ func (suite *keyspaceTestSuite) SetupTest() {
}

func (suite *keyspaceTestSuite) TearDownTest() {
resetKeyspaceInfoMetrics()
suite.cancel()
}

func (suite *keyspaceTestSuite) TestKeyspaceInfoMetricsLifecycle() {
re := suite.Require()
existing, err := suite.manager.CreateKeyspace(&CreateKeyspaceRequest{
Name: "metrics_existing",
CreateTime: time.Now().Unix(),
})
re.NoError(err)

cfg := &mockConfig{EnableKeyspaceLevelMetrics: true}
suite.manager.UpdateConfig(cfg)
re.Equal(0, promtestutil.CollectAndCount(keyspaceInfo))

existing, err = suite.manager.LoadKeyspace(existing.GetName())
re.NoError(err)
re.Equal(float64(1), promtestutil.ToFloat64(keyspaceInfo.WithLabelValues(
strconv.FormatUint(uint64(existing.GetId()), 10), existing.GetName())))

created, err := suite.manager.CreateKeyspace(&CreateKeyspaceRequest{
Name: "metrics_new",
CreateTime: time.Now().Unix(),
})
re.NoError(err)
re.Equal(float64(1), promtestutil.ToFloat64(keyspaceInfo.WithLabelValues(
strconv.FormatUint(uint64(created.GetId()), 10), created.GetName())))

_, err = suite.manager.UpdateKeyspaceState(created.GetName(), keyspacepb.KeyspaceState_DISABLED, time.Now().Unix())
re.NoError(err)
_, err = suite.manager.UpdateKeyspaceState(created.GetName(), keyspacepb.KeyspaceState_ARCHIVED, time.Now().Unix())
re.NoError(err)
_, err = suite.manager.UpdateKeyspaceState(created.GetName(), keyspacepb.KeyspaceState_TOMBSTONE, time.Now().Unix())
re.NoError(err)

errRollback := errors.New("rollback keyspace removal")
err = suite.manager.store.RunInTxn(suite.ctx, func(txn kv.Txn) error {
_, err := suite.manager.stageRemoveKeyspace(txn, created.GetId())
re.NoError(err)
return errRollback
})
re.ErrorIs(err, errRollback)
re.Equal(float64(1), promtestutil.ToFloat64(keyspaceInfo.WithLabelValues(
strconv.FormatUint(uint64(created.GetId()), 10), created.GetName())))
_, err = suite.manager.LoadKeyspace(created.GetName())
re.NoError(err)

re.NoError(suite.manager.RemoveKeyspace(created.GetId()))
expected := fmt.Sprintf(`# HELP pd_keyspace_info Keyspace metadata. The value is always 1.
# TYPE pd_keyspace_info gauge
pd_keyspace_info{keyspace_id="%d",keyspace_name="%s"} 1
`, existing.GetId(), existing.GetName())
re.NoError(promtestutil.CollectAndCompare(keyspaceInfo, strings.NewReader(expected), "pd_keyspace_info"))

suite.manager.UpdateConfig(&mockConfig{})
re.Equal(0, promtestutil.CollectAndCount(keyspaceInfo))
}

func (suite *keyspaceTestSuite) TestKeyspaceInfoMetricsCreateRollback() {
re := suite.Require()
const skipSplitRegion = "github.com/tikv/pd/pkg/keyspace/skipSplitRegion"
re.NoError(failpoint.Disable(skipSplitRegion))
defer func() { re.NoError(failpoint.Enable(skipSplitRegion, "return(true)")) }()

suite.manager.UpdateConfig(&mockConfig{EnableKeyspaceLevelMetrics: true})
resetKeyspaceInfoMetrics()
_, err := suite.manager.CreateKeyspace(&CreateKeyspaceRequest{
Name: "metrics_create_rollback",
CreateTime: time.Now().Unix(),
})
re.Error(err)
re.Equal(0, promtestutil.CollectAndCount(keyspaceInfo))
}

func (suite *keyspaceTestSuite) TestLoadRangeKeyspaceUpdatesInfoMetrics() {
re := suite.Require()
suite.manager.UpdateConfig(&mockConfig{EnableKeyspaceLevelMetrics: true})
created, err := suite.manager.CreateKeyspace(&CreateKeyspaceRequest{
Name: "metrics_range_load",
CreateTime: time.Now().Unix(),
})
re.NoError(err)
resetKeyspaceInfoMetrics()

keyspaces, err := suite.manager.LoadRangeKeyspace(created.GetId(), 1)
re.NoError(err)
re.Len(keyspaces, 1)
re.Equal(float64(1), promtestutil.ToFloat64(keyspaceInfo.WithLabelValues(
strconv.FormatUint(uint64(created.GetId()), 10), created.GetName())))
suite.manager.DeleteKeyspaceInfoMetrics(created.GetId())
re.Equal(0, promtestutil.CollectAndCount(keyspaceInfo))
}

func (suite *keyspaceTestSuite) SetupSuite() {
re := suite.Require()
re.NoError(failpoint.Enable("github.com/tikv/pd/pkg/keyspace/skipSplitRegion", "return(true)"))
Expand Down Expand Up @@ -1203,9 +1302,13 @@ func (suite *keyspaceTestSuite) TestTombstoneKeyspaceUnassignsMetaServiceGroup()
// 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.
var removed *keyspacepb.KeyspaceMeta
re.NoError(manager.store.RunInTxn(suite.ctx, func(txn kv.Txn) error {
return manager.RemoveKeyspace(txn, updated.GetId())
var err error
removed, err = manager.stageRemoveKeyspace(txn, updated.GetId())
return err
}))
manager.finishRemoveKeyspace(removed)
counts, err = manager.mgm.GetAssignmentCounts(suite.ctx)
re.NoError(err)
re.Equal(0, counts[groupID])
Expand Down
41 changes: 41 additions & 0 deletions pkg/keyspace/metrics.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@
package keyspace

import (
"strconv"
"sync"
"time"

"github.com/prometheus/client_golang/prometheus"
Expand All @@ -40,6 +42,15 @@ const (
)

var (
keyspaceInfo = prometheus.NewGaugeVec(

@lhy1024 lhy1024 Aug 25, 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.

Blocking: bootstrap never publishes the reserved keyspace mapping.

When enable-keyspace-level-metrics is true from startup, Bootstrap persists DEFAULT (or SYSTEM in NextGen) through initReserveKeyspace, but only the two public create methods call UpdateKeyspaceInfoMetrics. PD does not guarantee that an external client will subsequently open WatchKeyspaces or load this reserved keyspace, so a freshly bootstrapped cluster can expose no mapping for it. Bootstrap is already a creation path; publishing after the keyspace-group update succeeds completes that path without adding an etcd scan.

Unit-test reproducer 
func (suite *keyspaceTestSuite) TestBootstrapPublishesKeyspaceInfoMetric() {
	re := suite.Require()
	resetKeyspaceInfoMetrics()

	store := endpoint.NewStorageEndpoint(kv.NewMemoryKV(), nil)
	kgm := NewKeyspaceGroupManager(suite.ctx, store, nil)
	re.NoError(kgm.Bootstrap(suite.ctx))
	manager := NewKeyspaceManager(
		suite.ctx,
		store,
		nil,
		mockid.NewIDAllocator(),
		&mockConfig{EnableKeyspaceLevelMetrics: true},
		kgm,
		nil,
	)
	re.NoError(manager.Bootstrap())

	id, name := GetBootstrapKeyspaceID(), GetBootstrapKeyspaceName()
	re.Equal(1, promtestutil.CollectAndCount(keyspaceInfo))
	re.Equal(float64(1), promtestutil.ToFloat64(keyspaceInfo.WithLabelValues(
		strconv.FormatUint(uint64(id), 10), name)))
}

Run:

make gotest GOTEST_ARGS='./pkg/keyspace -run TestKeyspaceTestSuite/TestBootstrapPublishesKeyspaceInfoMetric -count=1 -timeout=30s'

On 247ace82be, the first assertion fails because the actual series count is 0.

@nolouch nolouch Aug 25, 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.

This metric is populated on demand. CSE calls GetAllKeyspaces at startup when name-based storage-size metrics are enabled, and that response includes the reserved keyspace. We intentionally don’t populate it during PD bootstrap to avoid another lifecycle path.

prometheus.GaugeOpts{
Namespace: namespace,
Subsystem: subsystem,
Name: "info",
Help: "Keyspace metadata. The value is always 1.",
}, []string{"keyspace_id", "keyspace_name"})
keyspaceInfoMetricsMu sync.Mutex

createKeyspaceStepDuration = prometheus.NewHistogramVec(
prometheus.HistogramOpts{
Namespace: namespace,
Expand All @@ -60,6 +71,7 @@ var (
)

func init() {
prometheus.MustRegister(keyspaceInfo)
prometheus.MustRegister(createKeyspaceStepDuration)
createKeyspaceStepDurationTotal = createKeyspaceStepDuration.WithLabelValues(StepTotal)
createKeyspaceStepDurationAllocateID = createKeyspaceStepDuration.WithLabelValues(StepAllocateID)
Expand All @@ -70,6 +82,35 @@ func init() {
createKeyspaceStepDurationUpdateKG = createKeyspaceStepDuration.WithLabelValues(StepUpdateKeyspaceGroup)
}

// setKeyspaceInfoMetricsLocked updates a keyspace info series while keyspaceInfoMetricsMu is held.
func setKeyspaceInfoMetricsLocked(id uint32, name string) {
keyspaceInfo.WithLabelValues(strconv.FormatUint(uint64(id), 10), name).Set(1)

@lhy1024 lhy1024 Aug 25, 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.

Blocking: a former leader can export both the old and new names for one keyspace ID.

These metrics are process-local, while every load/create/remove handler is leader-only. A PD can populate an ID, lose leadership, and miss the removal and CreateKeyspaceByID recreation performed by another PD. When it later loads the new metadata, this unconditional WithLabelValues call adds the new child without deleting the old child. Prometheus then sees two names for one ID, so the metric can no longer be used as a mapping.

With the current leader-only population model, reset these series when a PD stops serving as leader, matching the existing scheduling-metric lifecycle, and keep per-ID state so observing a different name replaces the previous exact child. Alternatively, a leader-independent watcher would have to maintain every PD's local registry. The per-ID state should retain the cached gauge and name: each client WatchKeyspaces stream invokes these callbacks, so the current code repeats WithLabelValues for every put and performs a DeletePartialMatch full-vector scan per stream for every delete.

Unit-test reproducer

The direct storage transactions below model changes committed to shared etcd by another PD process. That process can clean its own Prometheus registry, but not this process's registry.

func (suite *keyspaceTestSuite) TestExternalRecreationReplacesLocalInfoMetric() {
	re := suite.Require()
	suite.manager.UpdateConfig(&mockConfig{EnableKeyspaceLevelMetrics: true})
	created, err := suite.manager.CreateKeyspace(&CreateKeyspaceRequest{
		Name: "review_external", CreateTime: time.Now().Unix(),
	})
	re.NoError(err)

	for _, state := range []keyspacepb.KeyspaceState{
		keyspacepb.KeyspaceState_DISABLED,
		keyspacepb.KeyspaceState_ARCHIVED,
		keyspacepb.KeyspaceState_TOMBSTONE,
	} {
		_, err = suite.manager.UpdateKeyspaceState(created.GetName(), state, time.Now().Unix())
		re.NoError(err)
	}
	re.NoError(suite.manager.store.RunInTxn(suite.ctx, func(txn kv.Txn) error {
		return suite.manager.store.RemoveKeyspace(txn, created.GetId(), created.GetName())
	}))

	recreated := &keyspacepb.KeyspaceMeta{
		Keyspace: &keyspacepb.KeyspaceMeta_Id{Id: created.GetId()},
		Name: "review_recreated",
		State: keyspacepb.KeyspaceState_ENABLED,
	}
	re.NoError(suite.manager.store.RunInTxn(suite.ctx, func(txn kv.Txn) error {
		if err := suite.manager.store.SaveKeyspaceID(
			txn, recreated.GetId(), recreated.GetName(),
		); err != nil {
			return err
		}
		return suite.manager.store.SaveKeyspaceMeta(txn, recreated)
	}))
	_, err = suite.manager.LoadKeyspaceByID(recreated.GetId())
	re.NoError(err)

	expected := fmt.Sprintf(
		"# HELP pd_keyspace_info Keyspace metadata. The value is always 1.\n"+
			"# TYPE pd_keyspace_info gauge\n"+
			"pd_keyspace_info{keyspace_id=\"%d\",keyspace_name=\"%s\"} 1\n",
		recreated.GetId(), recreated.GetName(),
	)
	re.NoError(promtestutil.CollectAndCompare(
		keyspaceInfo, strings.NewReader(expected), "pd_keyspace_info",
	))
}

Run:

make gotest GOTEST_ARGS='./pkg/keyspace -run TestKeyspaceTestSuite/TestExternalRecreationReplacesLocalInfoMetric -count=1 -timeout=30s'

On 247ace82be, the comparison reports an extra series with keyspace_name="review_external" alongside review_recreated. A separate three-PD resign test also kept the old-leader series for the full 5-second assertion window.

@nolouch nolouch Aug 25, 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.

A new observation now removes any existing series for the same keyspace ID before setting the current name, so one ID won’t have two names. We intentionally allow a PD to retain its last observed mapping because this metric is client-observed rather than an authoritative snapshot.

@lhy1024 lhy1024 Aug 26, 2026 •
edited
Loading

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.

Blocking: the current implementation can export two names for one ID.

The latest commit removed the per-ID deletion and changed the test from old-name/new-name to the same name twice. The new comment's immutability assumption does not hold for a supported production lifecycle: after RemoveKeyspacesFromGroup commits, CreateKeyspaceByID accepts the removed ID with a different name.

Unit-test reproducer 
func (suite *keyspaceTestSuite) TestReusedIDReplacesObservedName() {
    re := suite.Require()
    resetKeyspaceInfoMetrics()
    suite.T().Cleanup(resetKeyspaceInfoMetrics)

    id := uint32(42)
    oldMeta, err := suite.manager.CreateKeyspaceByID(&CreateKeyspaceByIDRequest{
        ID: &id, Name: "review_old_name", CreateTime: time.Now().Unix(),
    })
    re.NoError(err)
    SetKeyspaceInfoMetrics(oldMeta.GetId(), oldMeta.GetName())

    for _, state := range []keyspacepb.KeyspaceState{
        keyspacepb.KeyspaceState_DISABLED,
        keyspacepb.KeyspaceState_ARCHIVED,
    } {
        _, err = suite.manager.UpdateKeyspaceStateByID(id, state, time.Now().Unix())
        re.NoError(err)
    }
    _, err = suite.manager.kgm.RemoveKeyspacesFromGroup(
        constant.DefaultKeyspaceGroupID, suite.manager, []uint32{id})
    re.NoError(err)

    newMeta, err := suite.manager.CreateKeyspaceByID(&CreateKeyspaceByIDRequest{
        ID: &id, Name: "review_new_name", CreateTime: time.Now().Unix(),
    })
    re.NoError(err)
    SetKeyspaceInfoMetrics(newMeta.GetId(), newMeta.GetName())

    expected := strings.NewReader(
        "# HELP pd_keyspace_info Keyspace ID-to-name mappings observed through client requests. The value is always 1.\n" +
            "# TYPE pd_keyspace_info gauge\n" +
            "pd_keyspace_info{keyspace_id=\"42\",keyspace_name=\"review_new_name\"} 1\n",
    )
    re.NoError(testutil.CollectAndCompare(
        keyspaceInfo, expected, "pd_keyspace_info"))
}

Run:

make gotest GOTEST_ARGS="./pkg/keyspace -run TestKeyspaceTestSuite/TestReusedIDReplacesObservedName -count=1 -timeout=30s"

On f23de01da, both CreateKeyspaceByID calls and the intervening removal succeed, but collection fails with the extra series:

pd_keyspace_info{keyspace_id="42",keyspace_name="review_old_name"} 1

Please retain per-ID {name, gauge} state. Repeated observations can reuse the cached gauge, while a changed name can delete the exact old child before creating the replacement. That avoids both stale names and a full-vector scan on every GetAllKeyspaces request.

@bufferflies bufferflies Aug 26, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

This regressed in the latest commit (f23de01, "avoid redundant info metric deletion"), which dropped the deleteKeyspaceInfoMetrics(id) call that used to run before this Set. The thread above states the fix keeps one name per ID, but that call is gone from the current head: calling SetKeyspaceInfoMetrics(id, "old") then SetKeyspaceInfoMetrics(id, "new") now leaves both series in the vector (I confirmed this locally: CollectAndCount returns 2, not 1), so the leader-loss/recreation scenario from the original report reproduces again.

}

// deleteKeyspaceInfoMetricsLocked deletes a keyspace info series while keyspaceInfoMetricsMu is held.
func deleteKeyspaceInfoMetricsLocked(id uint32, name string) {
keyspaceInfo.DeleteLabelValues(strconv.FormatUint(uint64(id), 10), name)
}

// deleteKeyspaceInfoMetricsByIDLocked deletes keyspace info series by ID while keyspaceInfoMetricsMu is held.
func deleteKeyspaceInfoMetricsByIDLocked(id uint32) {
keyspaceInfo.DeletePartialMatch(prometheus.Labels{
"keyspace_id": strconv.FormatUint(uint64(id), 10),
})
}

// resetKeyspaceInfoMetrics deletes all keyspace info series and cached gauges.
func resetKeyspaceInfoMetrics() {
keyspaceInfoMetricsMu.Lock()
defer keyspaceInfoMetricsMu.Unlock()
resetKeyspaceInfoMetricsLocked()
}

// resetKeyspaceInfoMetricsLocked resets the metric while keyspaceInfoMetricsMu is held.
func resetKeyspaceInfoMetricsLocked() {
keyspaceInfo.Reset()
}

// createKeyspaceTracer traces create-keyspace steps: one callback per step (same pattern as RegionHeartbeatProcessTracer), records metrics and logs per step.
type createKeyspaceTracer struct {
beginTime time.Time
Expand Down
Loading
Loading