feat(resourcemanager): add keyspace name label to metrics

Signed-off-by: JmPotato <github@ipotato.me>

Author: JmPotato

pkg/mcs/resourcemanager/server/manager.go

@@ -17,7 +17,7 @@ package server
 import (
 	"context"
 	"encoding/json"
-	"math"
+	"fmt"
 	"strings"
 	"time"
 
@@ -54,20 +54,44 @@ type Manager struct {
 	srv              bs.Server
 	controllerConfig *ControllerConfig
 	krgms            map[uint32]*keyspaceResourceGroupManager
-	storage          endpoint.ResourceGroupStorage
+	storage          interface {
+		// Used to store the resource group settings and states.
+		endpoint.ResourceGroupStorage
+		// Used to get the keyspace meta info.
+		endpoint.KeyspaceStorage
+	}
 	// consumptionChan is used to send the consumption
 	// info to the background metrics flusher.
 	consumptionDispatcher chan *consumptionItem
 	// record update time of each resource group
 	consumptionRecord map[consumptionRecordKey]time.Time
+	// max per sec trackers for each keyspace and resource group.
+	maxPerSecTrackers map[trackerKey]*maxPerSecCostTracker
+	// cached counter metrics for each keyspace, resource group and RU type.
+	counterMetrics map[metricsKey]*counterMetrics
+	// cached gauge metrics for each keyspace, resource group and RU type.
+	gaugeMetrics map[metricsKey]*gaugeMetrics
+	// cached keyspace name for each keyspace ID.
+	keyspaceNameLookup map[uint32]string
 }
 
 type consumptionRecordKey struct {
 	keyspaceID uint32
-	name       string
+	groupName  string
+	ruType     string
+}
+
+type metricsKey struct {
+	keyspaceID uint32
+	groupName  string
 	ruType     string
 }
 
+type trackerKey struct {
+	keyspaceID uint32
+	groupName  string
+}
+
 // ConfigProvider is used to get resource manager config from the given
 // `bs.server` without modifying its interface.
 type ConfigProvider interface {
@@ -82,6 +106,10 @@ func NewManager[T ConfigProvider](srv bs.Server) *Manager {
 		krgms:                 make(map[uint32]*keyspaceResourceGroupManager),
 		consumptionDispatcher: make(chan *consumptionItem, defaultConsumptionChanSize),
 		consumptionRecord:     make(map[consumptionRecordKey]time.Time),
+		maxPerSecTrackers:     make(map[trackerKey]*maxPerSecCostTracker),
+		counterMetrics:        make(map[metricsKey]*counterMetrics),
+		gaugeMetrics:          make(map[metricsKey]*gaugeMetrics),
+		keyspaceNameLookup:    make(map[uint32]string),
 	}
 	// The first initialization after the server is started.
 	srv.AddStartCallback(func() {
@@ -346,6 +374,80 @@ func (m *Manager) dispatchConsumption(req *rmpb.TokenBucketRequest) error {
 	return nil
 }
 
+func (m *Manager) insertConsumptionRecord(keyspaceID uint32, groupName string, ruType string) {
+	m.consumptionRecord[consumptionRecordKey{keyspaceID: keyspaceID, groupName: groupName, ruType: ruType}] = time.Now()
+}
+
+func (m *Manager) deleteConsumptionRecord(record consumptionRecordKey) {
+	delete(m.consumptionRecord, record)
+}
+
+func (m *Manager) getMaxPerSecTracker(keyspaceID uint32, keyspaceName, groupName string) *maxPerSecCostTracker {
+	tracker := m.maxPerSecTrackers[trackerKey{keyspaceID, groupName}]
+	if tracker == nil {
+		tracker = newMaxPerSecCostTracker(keyspaceName, groupName, defaultCollectIntervalSec)
+		m.maxPerSecTrackers[trackerKey{keyspaceID, groupName}] = tracker
+	}
+	return tracker
+}
+
+func (m *Manager) deleteMaxPerSecTracker(keyspaceID uint32, groupName string) {
+	delete(m.maxPerSecTrackers, trackerKey{keyspaceID, groupName})
+}
+
+func (m *Manager) getCounterMetrics(keyspaceID uint32, keyspaceName, groupName, ruType string) *counterMetrics {
+	key := metricsKey{keyspaceID, groupName, ruType}
+	if m.counterMetrics[key] == nil {
+		m.counterMetrics[key] = newCounterMetrics(keyspaceName, groupName, ruType)
+	}
+	return m.counterMetrics[key]
+}
+
+func (m *Manager) getGaugeMetrics(keyspaceID uint32, keyspaceName, groupName string) *gaugeMetrics {
+	key := metricsKey{keyspaceID, groupName, ""}
+	if m.gaugeMetrics[key] == nil {
+		m.gaugeMetrics[key] = newGaugeMetrics(keyspaceName, groupName)
+	}
+	return m.gaugeMetrics[key]
+}
+
+func (m *Manager) deleteMetrics(keyspaceID uint32, keyspaceName, groupName, ruType string) {
+	delete(m.counterMetrics, metricsKey{keyspaceID, groupName, ruType})
+	delete(m.gaugeMetrics, metricsKey{keyspaceID, groupName, ""})
+	deleteLabelValues(keyspaceName, groupName, ruType)
+}
+
+func (m *Manager) getKeyspaceNameByID(ctx context.Context, id uint32) (string, error) {
+	if id == constant.NullKeyspaceID {
+		return "", nil
+	}
+	// Try to get the keyspace name from the cache first.
+	name, ok := m.keyspaceNameLookup[id]
+	if ok {
+		return name, nil
+	}
+	var loadedName string
+	// If the keyspace name is not in the cache, try to get it from the storage.
+	err := m.storage.RunInTxn(ctx, func(txn kv.Txn) error {
+		meta, err := m.storage.LoadKeyspaceMeta(txn, id)
+		if err != nil {
+			return err
+		}
+		loadedName = meta.GetName()
+		return nil
+	})
+	if err != nil {
+		log.Error("failed to get the keyspace name", zap.Uint32("keyspace-id", id), zap.Error(err))
+		return "", err
+	}
+	if len(loadedName) == 0 {
+		return "", fmt.Errorf("got an empty keyspace name by id %d", id)
+	}
+	// Update the cache.
+	m.keyspaceNameLookup[id] = loadedName
+	return loadedName, nil
+}
+
 func (m *Manager) backgroundMetricsFlush(ctx context.Context) {
 	defer logutil.LogPanic()
 	cleanUpTicker := time.NewTicker(metricsCleanupInterval)
@@ -355,31 +457,6 @@ func (m *Manager) backgroundMetricsFlush(ctx context.Context) {
 	recordMaxTicker := time.NewTicker(tickPerSecond)
 	defer recordMaxTicker.Stop()
 
-	maxPerSecTrackers := make(map[uint32]map[string]*maxPerSecCostTracker)
-	// getMaxPerSecTracker returns the max per sec tracker for the given keyspace ID and name.
-	// If the tracker doesn't exist, it will be created.
-	getMaxPerSecTracker := func(keyspaceID uint32, name string) *maxPerSecCostTracker {
-		trackers := maxPerSecTrackers[keyspaceID]
-		if trackers == nil {
-			trackers = make(map[string]*maxPerSecCostTracker)
-			maxPerSecTrackers[keyspaceID] = trackers
-		}
-		tracker, ok := trackers[name]
-		if !ok {
-			tracker = newMaxPerSecCostTracker(name, defaultCollectIntervalSec)
-			trackers[name] = tracker
-		}
-		return tracker
-	}
-	// deleteMaxPerSecTracker deletes the max per sec tracker for the given keyspace ID and name.
-	// If the tracker doesn't exist, it will do nothing.
-	deleteMaxPerSecTracker := func(keyspaceID uint32, name string) {
-		trackers := maxPerSecTrackers[keyspaceID]
-		if trackers == nil {
-			return
-		}
-		delete(trackers, name)
-	}
 	for {
 		select {
 		case <-ctx.Done():
@@ -389,6 +466,12 @@ func (m *Manager) backgroundMetricsFlush(ctx context.Context) {
 			if consumption == nil {
 				continue
 			}
+			keyspaceID := consumptionInfo.keyspaceID
+			keyspaceName, err := m.getKeyspaceNameByID(ctx, keyspaceID)
+			if err != nil {
+				continue
+			}
+
 			ruLabelType := defaultTypeLabel
 			if consumptionInfo.isBackground {
 				ruLabelType = backgroundTypeLabel
@@ -397,53 +480,10 @@ func (m *Manager) backgroundMetricsFlush(ctx context.Context) {
 				ruLabelType = tiflashTypeLabel
 			}
 
-			var (
-				// TODO: add keyspace name lable to the metrics.
-				keyspaceID               = consumptionInfo.keyspaceID
-				name                     = consumptionInfo.resourceGroupName
-				rruMetrics               = readRequestUnitCost.WithLabelValues(name, name, ruLabelType)
-				wruMetrics               = writeRequestUnitCost.WithLabelValues(name, name, ruLabelType)
-				sqlLayerRuMetrics        = sqlLayerRequestUnitCost.WithLabelValues(name, name)
-				readByteMetrics          = readByteCost.WithLabelValues(name, name, ruLabelType)
-				writeByteMetrics         = writeByteCost.WithLabelValues(name, name, ruLabelType)
-				kvCPUMetrics             = kvCPUCost.WithLabelValues(name, name, ruLabelType)
-				sqlCPUMetrics            = sqlCPUCost.WithLabelValues(name, name, ruLabelType)
-				readRequestCountMetrics  = requestCount.WithLabelValues(name, name, readTypeLabel)
-				writeRequestCountMetrics = requestCount.WithLabelValues(name, name, writeTypeLabel)
-			)
-			getMaxPerSecTracker(keyspaceID, name).CollectConsumption(consumption)
-
-			// RU info.
-			if consumption.RRU > 0 {
-				rruMetrics.Add(consumption.RRU)
-			}
-			if consumption.WRU > 0 {
-				wruMetrics.Add(consumption.WRU)
-			}
-			// Byte info.
-			if consumption.ReadBytes > 0 {
-				readByteMetrics.Add(consumption.ReadBytes)
-			}
-			if consumption.WriteBytes > 0 {
-				writeByteMetrics.Add(consumption.WriteBytes)
-			}
-			// CPU time info.
-			if consumption.TotalCpuTimeMs > 0 {
-				if consumption.SqlLayerCpuTimeMs > 0 {
-					sqlLayerRuMetrics.Add(consumption.SqlLayerCpuTimeMs * m.controllerConfig.RequestUnit.CPUMsCost)
-					sqlCPUMetrics.Add(consumption.SqlLayerCpuTimeMs)
-				}
-				kvCPUMetrics.Add(consumption.TotalCpuTimeMs - consumption.SqlLayerCpuTimeMs)
-			}
-			// RPC count info.
-			if consumption.KvReadRpcCount > 0 {
-				readRequestCountMetrics.Add(consumption.KvReadRpcCount)
-			}
-			if consumption.KvWriteRpcCount > 0 {
-				writeRequestCountMetrics.Add(consumption.KvWriteRpcCount)
-			}
-
-			m.consumptionRecord[consumptionRecordKey{keyspaceID: keyspaceID, name: name, ruType: ruLabelType}] = time.Now()
+			name := consumptionInfo.resourceGroupName
+			m.getMaxPerSecTracker(keyspaceID, keyspaceName, name).collect(consumption)
+			m.getCounterMetrics(keyspaceID, keyspaceName, name, ruLabelType).add(consumption, m.controllerConfig)
+			m.insertConsumptionRecord(keyspaceID, name, ruLabelType)
 
 			// TODO: maybe we need to distinguish background ru.
 			if rg := m.GetMutableResourceGroup(keyspaceID, name); rg != nil {
@@ -455,47 +495,44 @@ func (m *Manager) backgroundMetricsFlush(ctx context.Context) {
 				if time.Since(lastTime) <= metricsCleanupTimeout {
 					continue
 				}
-				readRequestUnitCost.DeleteLabelValues(r.name, r.name, r.ruType)
-				writeRequestUnitCost.DeleteLabelValues(r.name, r.name, r.ruType)
-				sqlLayerRequestUnitCost.DeleteLabelValues(r.name, r.name, r.ruType)
-				readByteCost.DeleteLabelValues(r.name, r.name, r.ruType)
-				writeByteCost.DeleteLabelValues(r.name, r.name, r.ruType)
-				kvCPUCost.DeleteLabelValues(r.name, r.name, r.ruType)
-				sqlCPUCost.DeleteLabelValues(r.name, r.name, r.ruType)
-				requestCount.DeleteLabelValues(r.name, r.name, readTypeLabel)
-				requestCount.DeleteLabelValues(r.name, r.name, writeTypeLabel)
-				availableRUCounter.DeleteLabelValues(r.name, r.name)
-				delete(m.consumptionRecord, r)
-				deleteMaxPerSecTracker(r.keyspaceID, r.name)
-				readRequestUnitMaxPerSecCost.DeleteLabelValues(r.name)
-				writeRequestUnitMaxPerSecCost.DeleteLabelValues(r.name)
-				resourceGroupConfigGauge.DeletePartialMatch(prometheus.Labels{newResourceGroupNameLabel: r.name})
+				keyspaceName, err := m.getKeyspaceNameByID(ctx, r.keyspaceID)
+				if err != nil {
+					continue
+				}
+				m.deleteConsumptionRecord(r)
+				m.deleteMetrics(r.keyspaceID, keyspaceName, r.groupName, r.ruType)
+				m.deleteMaxPerSecTracker(r.keyspaceID, r.groupName)
 			}
 		case <-availableRUTicker.C:
 			// Prevent from holding the lock too long when there're many keyspaces and resource groups.
 			for _, krgm := range m.getKeyspaceResourceGroupManagers() {
+				keyspaceName, err := m.getKeyspaceNameByID(ctx, krgm.keyspaceID)
+				if err != nil {
+					continue
+				}
 				for _, group := range krgm.getResourceGroupList(true, false) {
-					ru := math.Max(group.getRUToken(), 0)
-					availableRUCounter.WithLabelValues(group.Name, group.Name).Set(ru)
-					resourceGroupConfigGauge.WithLabelValues(group.Name, priorityLabel).Set(group.getPriority())
-					resourceGroupConfigGauge.WithLabelValues(group.Name, ruPerSecLabel).Set(group.getFillRate())
-					resourceGroupConfigGauge.WithLabelValues(group.Name, ruCapacityLabel).Set(group.getBurstLimit())
+					m.getGaugeMetrics(krgm.keyspaceID, keyspaceName, group.Name).set(group)
 				}
 			}
 		case <-recordMaxTicker.C:
 			// Record the sum of RRU and WRU every second.
 			for _, krgm := range m.getKeyspaceResourceGroupManagers() {
+				keyspaceName, err := m.getKeyspaceNameByID(ctx, krgm.keyspaceID)
+				if err != nil {
+					continue
+				}
 				names := krgm.getResourceGroupNames(true)
 				for _, name := range names {
-					getMaxPerSecTracker(krgm.keyspaceID, name).FlushMetrics()
+					m.getMaxPerSecTracker(krgm.keyspaceID, keyspaceName, name).flushMetrics()
 				}
 			}
 		}
 	}
 }
 
 type maxPerSecCostTracker struct {
-	name          string
+	keyspaceName  string
+	groupName     string
 	maxPerSecRRU  float64
 	maxPerSecWRU  float64
 	rruSum        float64
@@ -508,23 +545,22 @@ type maxPerSecCostTracker struct {
 	wruMaxMetrics prometheus.Gauge
 }
 
-func newMaxPerSecCostTracker(name string, flushPeriod int) *maxPerSecCostTracker {
+func newMaxPerSecCostTracker(keyspaceName, groupName string, flushPeriod int) *maxPerSecCostTracker {
 	return &maxPerSecCostTracker{
-		name:          name,
+		keyspaceName:  keyspaceName,
+		groupName:     groupName,
 		flushPeriod:   flushPeriod,
-		rruMaxMetrics: readRequestUnitMaxPerSecCost.WithLabelValues(name),
-		wruMaxMetrics: writeRequestUnitMaxPerSecCost.WithLabelValues(name),
+		rruMaxMetrics: readRequestUnitMaxPerSecCost.WithLabelValues(groupName, keyspaceName),
+		wruMaxMetrics: writeRequestUnitMaxPerSecCost.WithLabelValues(groupName, keyspaceName),
 	}
 }
 
-// CollectConsumption collects the consumption info.
-func (t *maxPerSecCostTracker) CollectConsumption(consume *rmpb.Consumption) {
+func (t *maxPerSecCostTracker) collect(consume *rmpb.Consumption) {
 	t.rruSum += consume.RRU
 	t.wruSum += consume.WRU
 }
 
-// FlushMetrics and set the maxPerSecRRU and maxPerSecWRU to the metrics.
-func (t *maxPerSecCostTracker) FlushMetrics() {
+func (t *maxPerSecCostTracker) flushMetrics() {
 	if t.lastRRUSum == 0 && t.lastWRUSum == 0 {
 		t.lastRRUSum = t.rruSum
 		t.lastWRUSum = t.wruSum

pkg/mcs/resourcemanager/server/manager_test.go

@@ -77,6 +77,7 @@ func TestBackgroundMetricsFlush(t *testing.T) {
 
 	// Verify consumption was added to the resource group.
 	testutil.Eventually(re, func() bool {
+		// TODO: test with the keyspace name.
 		updatedGroup := m.GetResourceGroup(constant.NullKeyspaceID, req.GetResourceGroupName(), true)
 		re.NotNil(updatedGroup)
 		return updatedGroup.RUConsumption.RRU == req.ConsumptionSinceLastRequest.RRU &&