feat(metering): add storage size metering and keyspace name caching

go.mod

@@ -35,7 +35,7 @@ require (
 	github.com/pingcap/errcode v0.3.0
 	github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c
 	github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86
-	github.com/pingcap/kvproto v0.0.0-20250922052949-d6cbb9eb8638
+	github.com/pingcap/kvproto v0.0.0-20250923064352-8eeada0a8a03
 	github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3
 	github.com/pingcap/metering_sdk v0.0.0-20250918015914-468cd6feb1dc
 	github.com/pingcap/sysutil v1.0.1-0.20230407040306-fb007c5aff21

go.sum

@@ -469,8 +469,8 @@ github.com/pingcap/errors v0.11.5-0.20211224045212-9687c2b0f87c/go.mod h1:X2r9ue
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86 h1:tdMsjOqUR7YXHoBitzdebTvOjs/swniBTOLy5XiMtuE=
 github.com/pingcap/failpoint v0.0.0-20240528011301-b51a646c7c86/go.mod h1:exzhVYca3WRtd6gclGNErRWb1qEgff3LYta0LvRmON4=
 github.com/pingcap/kvproto v0.0.0-20191211054548-3c6b38ea5107/go.mod h1:WWLmULLO7l8IOcQG+t+ItJ3fEcrL5FxF0Wu+HrMy26w=
-github.com/pingcap/kvproto v0.0.0-20250922052949-d6cbb9eb8638 h1:H2PYrfqNBl5K4Y0OY3+FpfC5LTSX9ejiA65q50I3zhE=
-github.com/pingcap/kvproto v0.0.0-20250922052949-d6cbb9eb8638/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
+github.com/pingcap/kvproto v0.0.0-20250923064352-8eeada0a8a03 h1:G6lEpMBW42aNvstzIjNR3NnH6mx+3nIH42Ic1Sb8h/U=
+github.com/pingcap/kvproto v0.0.0-20250923064352-8eeada0a8a03/go.mod h1:rXxWk2UnwfUhLXha1jxRWPADw9eMZGWEWCg92Tgmb/8=
 github.com/pingcap/log v0.0.0-20210625125904-98ed8e2eb1c7/go.mod h1:8AanEdAHATuRurdGxZXBz0At+9avep+ub7U1AGYLIMM=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3 h1:HR/ylkkLmGdSSDaD8IDP+SZrdhV1Kibl9KrHxJ9eciw=
 github.com/pingcap/log v1.1.1-0.20221110025148-ca232912c9f3/go.mod h1:DWQW5jICDR7UJh4HtxXSM20Churx4CQL0fwL/SoOSA4=

pkg/core/region.go

@@ -63,26 +63,27 @@ func errRegionIsStale(region *metapb.Region, origin *metapb.Region) error {
 // the properties are Read-Only once created except buckets.
 // the `buckets` could be modified by the request `report buckets` with greater version.
 type RegionInfo struct {
-	meta              *metapb.Region
-	learners          []*metapb.Peer
-	witnesses         []*metapb.Peer
-	voters            []*metapb.Peer
-	leader            *metapb.Peer
-	downPeers         []*pdpb.PeerStats
-	pendingPeers      []*metapb.Peer
-	term              uint64
-	cpuUsage          uint64
-	writtenBytes      uint64
-	writtenKeys       uint64
-	readBytes         uint64
-	readKeys          uint64
-	approximateSize   int64
-	approximateKvSize int64
-	approximateKeys   int64
-	interval          *pdpb.TimeInterval
-	replicationStatus *replication_modepb.RegionReplicationStatus
-	queryStats        *pdpb.QueryStats
-	flowRoundDivisor  uint64
+	meta                      *metapb.Region
+	learners                  []*metapb.Peer
+	witnesses                 []*metapb.Peer
+	voters                    []*metapb.Peer
+	leader                    *metapb.Peer
+	downPeers                 []*pdpb.PeerStats
+	pendingPeers              []*metapb.Peer
+	term                      uint64
+	cpuUsage                  uint64
+	writtenBytes              uint64
+	writtenKeys               uint64
+	readBytes                 uint64
+	readKeys                  uint64
+	approximateSize           int64
+	approximateKvSize         int64 // Unit: MiB
+	approximateColumnarKvSize int64 // Unit: MiB
+	approximateKeys           int64
+	interval                  *pdpb.TimeInterval
+	replicationStatus         *replication_modepb.RegionReplicationStatus
+	queryStats                *pdpb.QueryStats
+	flowRoundDivisor          uint64
 	// buckets is not thread unsafe, it should be accessed by the request `report buckets` with greater version.
 	buckets unsafe.Pointer
 	// source is used to indicate region's source, such as Storage/Sync/Heartbeat.
@@ -247,6 +248,7 @@ func RegionFromHeartbeat(heartbeat RegionHeartbeatRequest, flowRoundDivisor int)
 	// scheduling service doesn't need the following fields.
 	if h, ok := heartbeat.(*pdpb.RegionHeartbeatRequest); ok {
 		region.approximateKvSize = int64(h.GetApproximateKvSize() / units.MiB)
+		region.approximateColumnarKvSize = int64(h.GetApproximateColumnarKvSize() / units.MiB)
 		region.replicationStatus = h.GetReplicationStatus()
 		region.cpuUsage = h.GetCpuUsage()
 	}
@@ -296,23 +298,24 @@ func (r *RegionInfo) Clone(opts ...RegionCreateOption) *RegionInfo {
 	}
 
 	region := &RegionInfo{
-		term:              r.term,
-		meta:              typeutil.DeepClone(r.meta, RegionFactory),
-		leader:            typeutil.DeepClone(r.leader, RegionPeerFactory),
-		downPeers:         downPeers,
-		pendingPeers:      pendingPeers,
-		cpuUsage:          r.cpuUsage,
-		writtenBytes:      r.writtenBytes,
-		writtenKeys:       r.writtenKeys,
-		readBytes:         r.readBytes,
-		readKeys:          r.readKeys,
-		approximateSize:   r.approximateSize,
-		approximateKvSize: r.approximateKvSize,
-		approximateKeys:   r.approximateKeys,
-		interval:          typeutil.DeepClone(r.interval, TimeIntervalFactory),
-		replicationStatus: r.replicationStatus,
-		buckets:           r.buckets,
-		queryStats:        typeutil.DeepClone(r.queryStats, QueryStatsFactory),
+		term:                      r.term,
+		meta:                      typeutil.DeepClone(r.meta, RegionFactory),
+		leader:                    typeutil.DeepClone(r.leader, RegionPeerFactory),
+		downPeers:                 downPeers,
+		pendingPeers:              pendingPeers,
+		cpuUsage:                  r.cpuUsage,
+		writtenBytes:              r.writtenBytes,
+		writtenKeys:               r.writtenKeys,
+		readBytes:                 r.readBytes,
+		readKeys:                  r.readKeys,
+		approximateSize:           r.approximateSize,
+		approximateKvSize:         r.approximateKvSize,
+		approximateColumnarKvSize: r.approximateColumnarKvSize,
+		approximateKeys:           r.approximateKeys,
+		interval:                  typeutil.DeepClone(r.interval, TimeIntervalFactory),
+		replicationStatus:         r.replicationStatus,
+		buckets:                   r.buckets,
+		queryStats:                typeutil.DeepClone(r.queryStats, QueryStatsFactory),
 	}
 
 	for _, opt := range opts {
@@ -656,6 +659,11 @@ func (r *RegionInfo) GetApproximateKvSize() int64 {
 	return r.approximateKvSize
 }
 
+// GetApproximateColumnarKvSize returns the approximate columnar kv size of the region.
+func (r *RegionInfo) GetApproximateColumnarKvSize() int64 {
+	return r.approximateColumnarKvSize
+}
+
 // GetApproximateKeys returns the approximate keys of the region.
 func (r *RegionInfo) GetApproximateKeys() int64 {
 	return r.approximateKeys

pkg/keyspace/keyspace.go

@@ -18,6 +18,7 @@ import (
 	"bytes"
 	"context"
 	"strconv"
+	"sync"
 	"time"
 
 	"go.uber.org/zap"
@@ -96,6 +97,8 @@ type Manager struct {
 	kgm *GroupManager
 	// nextPatrolStartID is the next start id of keyspace assignment patrol.
 	nextPatrolStartID uint32
+	// cached keyspace name for each keyspace ID.
+	keyspaceNameLookup sync.Map // store as ID(uint32) -> name(string)
 }
 
 // CreateKeyspaceRequest represents necessary arguments to create a keyspace.
@@ -411,6 +414,8 @@ func (manager *Manager) saveNewKeyspace(keyspace *keyspacepb.KeyspaceMeta) error
 		if err != nil {
 			return err
 		}
+		// Update the keyspace name cache.
+		manager.keyspaceNameLookup.Store(keyspace.Id, keyspace.Name)
 		// Save keyspace meta.
 		// Check if keyspace with that id already exists.
 		loadedMeta, err := manager.store.LoadKeyspaceMeta(txn, keyspace.Id)
@@ -944,3 +949,29 @@ func (manager *Manager) PatrolKeyspaceAssignment(startKeyspaceID, endKeyspaceID
 	}
 	return nil
 }
+
+// GetKeyspaceNameByID gets the keyspace name by ID, which will try to get it from the cache first.
+// If not found, it will try to get it from the storage.
+func (manager *Manager) GetKeyspaceNameByID(id uint32) (string, error) {
+	if id == constant.NullKeyspaceID {
+		return "", nil
+	}
+	// Try to get the keyspace name from the cache first.
+	name, ok := manager.keyspaceNameLookup.Load(id)
+	if ok {
+		return name.(string), nil
+	}
+	var loadedName string
+	// If the keyspace name is not in the cache, try to get it from the storage.
+	meta, err := manager.LoadKeyspaceByID(id)
+	if err != nil {
+		return "", err
+	}
+	loadedName = meta.GetName()
+	if len(loadedName) == 0 {
+		return "", errors.Errorf("got an empty keyspace name by id %d", id)
+	}
+	// Load or store the keyspace name to the cache.
+	actual, _ := manager.keyspaceNameLookup.LoadOrStore(id, loadedName)
+	return actual.(string), nil
+}

pkg/keyspace/keyspace_test.go

@@ -135,6 +135,10 @@ func (suite *keyspaceTestSuite) TestCreateKeyspace() {
 		re.Equal(uint32(i+1), created.Id)
 		checkCreateRequest(re, request, created)
 
+		name, err := manager.GetKeyspaceNameByID(created.Id)
+		re.NoError(err)
+		re.Equal(created.Name, name)
+
 		loaded, err := manager.LoadKeyspace(request.Name)
 		re.NoError(err)
 		re.Equal(uint32(i+1), loaded.Id)

pkg/keyspace/util.go

@@ -20,6 +20,7 @@ import (
 	"encoding/hex"
 	"regexp"
 	"strconv"
+	"strings"
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/kvproto/pkg/keyspacepb"
@@ -161,6 +162,39 @@ func MakeLabelRule(id uint32) *labeler.LabelRule {
 	}
 }
 
+// ParseKeyspaceIDFromLabelRule parses the keyspace ID from the label rule.
+// It will return the keyspace ID and a boolean indicating whether the label
+// rule is a keyspace label rule.
+func ParseKeyspaceIDFromLabelRule(rule *labeler.LabelRule) (uint32, bool) {
+	// Validate the ID matches the expected format "keyspaces/<id>".
+	if rule == nil || !strings.HasPrefix(rule.ID, regionLabelIDPrefix) {
+		return 0, false
+	}
+	// Retrieve the keyspace ID.
+	keyspaceID, err := strconv.ParseUint(
+		strings.TrimPrefix(rule.ID, regionLabelIDPrefix),
+		endpoint.SpaceIDBase, 32,
+	)
+	if err != nil {
+		return 0, false
+	}
+	// Double check the keyspace ID from the label rule.
+	var idFromLabel uint64
+	for _, label := range rule.Labels {
+		if label.Key == regionLabelKey {
+			idFromLabel, err = strconv.ParseUint(label.Value, endpoint.SpaceIDBase, 32)
+			if err != nil {
+				return 0, false
+			}
+			break
+		}
+	}
+	if keyspaceID != idFromLabel {
+		return 0, false
+	}
+	return uint32(keyspaceID), true
+}
+
 // indexedHeap is a heap with index.
 type indexedHeap struct {
 	items []*endpoint.KeyspaceGroup

pkg/keyspace/util_test.go

@@ -133,15 +133,15 @@ func TestMakeLabelRule(t *testing.T) {
 		{
 			id: 0,
 			expectedLabelRule: &labeler.LabelRule{
-				ID:    "keyspaces/0",
+				ID:    getRegionLabelID(0),
 				Index: 0,
 				Labels: []labeler.RegionLabel{
 					{
-						Key:   "id",
+						Key:   regionLabelKey,
 						Value: "0",
 					},
 				},
-				RuleType: "key-range",
+				RuleType: labeler.KeyRange,
 				Data: []any{
 					map[string]any{
 						"start_key": hex.EncodeToString(codec.EncodeBytes([]byte{'r', 0, 0, 0})),
@@ -157,15 +157,15 @@ func TestMakeLabelRule(t *testing.T) {
 		{
 			id: 4242,
 			expectedLabelRule: &labeler.LabelRule{
-				ID:    "keyspaces/4242",
+				ID:    getRegionLabelID(4242),
 				Index: 0,
 				Labels: []labeler.RegionLabel{
 					{
-						Key:   "id",
+						Key:   regionLabelKey,
 						Value: "4242",
 					},
 				},
-				RuleType: "key-range",
+				RuleType: labeler.KeyRange,
 				Data: []any{
 					map[string]any{
 						"start_key": hex.EncodeToString(codec.EncodeBytes([]byte{'r', 0, 0x10, 0x92})),
@@ -183,3 +183,98 @@ func TestMakeLabelRule(t *testing.T) {
 		re.Equal(testCase.expectedLabelRule, MakeLabelRule(testCase.id))
 	}
 }
+
+func TestParseKeyspaceIDFromLabelRule(t *testing.T) {
+	re := require.New(t)
+	testCases := []struct {
+		labelRule  *labeler.LabelRule
+		expectedID uint32
+		expectedOK bool
+	}{
+		// Valid keyspace label rule.
+		{
+			labelRule:  MakeLabelRule(1),
+			expectedID: 1,
+			expectedOK: true,
+		},
+		// Invalid keyspace label ID - unmatched prefix.
+		{
+			labelRule: &labeler.LabelRule{
+				ID:    "not-keyspaces/1",
+				Index: 0,
+				Labels: []labeler.RegionLabel{
+					{
+						Key:   regionLabelKey,
+						Value: "1",
+					},
+				},
+			},
+			expectedID: 0,
+			expectedOK: false,
+		},
+		// Invalid keyspace label ID - invalid keyspace ID.
+		{
+			labelRule: &labeler.LabelRule{
+				ID:    "keyspaces/id1",
+				Index: 0,
+				Labels: []labeler.RegionLabel{
+					{
+						Key:   regionLabelKey,
+						Value: "1",
+					},
+				},
+			},
+			expectedID: 0,
+			expectedOK: false,
+		},
+		{
+			labelRule: &labeler.LabelRule{
+				ID:    "keyspaces/1id",
+				Index: 0,
+				Labels: []labeler.RegionLabel{
+					{
+						Key:   regionLabelKey,
+						Value: "1",
+					},
+				},
+			},
+			expectedID: 0,
+			expectedOK: false,
+		},
+		// Invalid keyspace label ID - invalid keyspace ID label rule.
+		{
+			labelRule: &labeler.LabelRule{
+				ID:    getRegionLabelID(1),
+				Index: 0,
+				Labels: []labeler.RegionLabel{
+					{
+						Key:   "not-id",
+						Value: "1",
+					},
+				},
+			},
+			expectedID: 0,
+			expectedOK: false,
+		},
+		// Invalid keyspace label ID - unmatched keyspace ID with label rule.
+		{
+			labelRule: &labeler.LabelRule{
+				ID:    getRegionLabelID(1),
+				Index: 0,
+				Labels: []labeler.RegionLabel{
+					{
+						Key:   "id",
+						Value: "2",
+					},
+				},
+			},
+			expectedID: 0,
+			expectedOK: false,
+		},
+	}
+	for _, testCase := range testCases {
+		id, ok := ParseKeyspaceIDFromLabelRule(testCase.labelRule)
+		re.Equal(testCase.expectedID, id)
+		re.Equal(testCase.expectedOK, ok)
+	}
+}