From 1334d15846199a509f61ffb9234ce8fa8d62ed4f Mon Sep 17 00:00:00 2001
From: guo-shaoge <shaoge1994@163.com>
Date: Tue, 19 Sep 2023 23:37:34 +0800
Subject: [PATCH] add metric for tiflash ru consumption

Signed-off-by: guo-shaoge <shaoge1994@163.com>
---
 go.mod                                        |  1 +
 go.sum                                        |  1 +
 .../resourcemanager/server/grpc_service.go    |  4 +++-
 pkg/mcs/resourcemanager/server/manager.go     | 21 ++++++++++++-------
 pkg/mcs/resourcemanager/server/metrics.go     |  1 +
 5 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/go.mod b/go.mod
index bb72c0d5dba..14d3d799df8 100644
--- a/go.mod
+++ b/go.mod
@@ -5,6 +5,7 @@ go 1.21
 require (
 	github.com/AlekSi/gocov-xml v1.0.0
 	github.com/BurntSushi/toml v0.3.1
+	github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751
 	github.com/aws/aws-sdk-go-v2/config v1.18.19
 	github.com/aws/aws-sdk-go-v2/credentials v1.13.18
 	github.com/aws/aws-sdk-go-v2/service/kms v1.20.8
diff --git a/go.sum b/go.sum
index a007acf869e..14ae66f0832 100644
--- a/go.sum
+++ b/go.sum
@@ -21,6 +21,7 @@ github.com/Xeoncross/go-aesctr-with-hmac v0.0.0-20200623134604-12b17a7ff502 h1:L
 github.com/Xeoncross/go-aesctr-with-hmac v0.0.0-20200623134604-12b17a7ff502/go.mod h1:pmnBM9bxWSiHvC/gSWunUIyDvGn33EkP2CUjxFKtTTM=
 github.com/agiledragon/gomonkey/v2 v2.3.1/go.mod h1:ap1AmDzcVOAz1YpeJ3TCzIgstoaWLA6jbbgxfB4w2iY=
 github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
+github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751 h1:JYp7IbQjafoB+tBA3gMyHYHrpOtNuDiK/uB5uXxq5wM=
 github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
diff --git a/pkg/mcs/resourcemanager/server/grpc_service.go b/pkg/mcs/resourcemanager/server/grpc_service.go
index 5c1b5f0e458..dfb25b2023f 100644
--- a/pkg/mcs/resourcemanager/server/grpc_service.go
+++ b/pkg/mcs/resourcemanager/server/grpc_service.go
@@ -191,11 +191,13 @@ func (s *Service) AcquireTokenBuckets(stream rmpb.ResourceManager_AcquireTokenBu
 			}
 			// Send the consumption to update the metrics.
 			isBackground := req.GetIsBackground()
+			isTiFlash := req.GetIsTiflash()
 			s.manager.consumptionDispatcher <- struct {
 				resourceGroupName string
 				*rmpb.Consumption
 				isBackground bool
-			}{resourceGroupName, req.GetConsumptionSinceLastRequest(), isBackground}
+				isTiFlash bool
+			}{resourceGroupName, req.GetConsumptionSinceLastRequest(), isBackground, isTiFlash}
 			if isBackground {
 				continue
 			}
diff --git a/pkg/mcs/resourcemanager/server/manager.go b/pkg/mcs/resourcemanager/server/manager.go
index 21866ee1156..109104ab56c 100644
--- a/pkg/mcs/resourcemanager/server/manager.go
+++ b/pkg/mcs/resourcemanager/server/manager.go
@@ -60,6 +60,7 @@ type Manager struct {
 		resourceGroupName string
 		*rmpb.Consumption
 		isBackground bool
+		isTiFlash bool
 	}
 	// record update time of each resource group
 	consumptionRecord map[string]time.Time
@@ -81,6 +82,7 @@ func NewManager[T ConfigProvider](srv bs.Server) *Manager {
 			resourceGroupName string
 			*rmpb.Consumption
 			isBackground bool
+			isTiFlash bool
 		}, defaultConsumptionChanSize),
 		consumptionRecord: make(map[string]time.Time),
 	}
@@ -361,20 +363,23 @@ func (m *Manager) backgroundMetricsFlush(ctx context.Context) {
 			if consumption == nil {
 				continue
 			}
-			backgroundType := ""
+			bgOrTiFlashType := ""
 			if consumptionInfo.isBackground {
-				backgroundType = backgroundTypeLabel
+				bgOrTiFlashType = backgroundTypeLabel
+			}
+			if consumptionInfo.isTiFlash {
+				bgOrTiFlashType = tiflashTypeLabel
 			}
 
 			var (
 				name                     = consumptionInfo.resourceGroupName
-				rruMetrics               = readRequestUnitCost.WithLabelValues(name, backgroundType)
-				wruMetrics               = writeRequestUnitCost.WithLabelValues(name, backgroundType)
+				rruMetrics               = readRequestUnitCost.WithLabelValues(name, bgOrTiFlashType)
+				wruMetrics               = writeRequestUnitCost.WithLabelValues(name, bgOrTiFlashType)
 				sqlLayerRuMetrics        = sqlLayerRequestUnitCost.WithLabelValues(name)
-				readByteMetrics          = readByteCost.WithLabelValues(name, backgroundType)
-				writeByteMetrics         = writeByteCost.WithLabelValues(name, backgroundType)
-				kvCPUMetrics             = kvCPUCost.WithLabelValues(name, backgroundType)
-				sqlCPUMetrics            = sqlCPUCost.WithLabelValues(name, backgroundType)
+				readByteMetrics          = readByteCost.WithLabelValues(name, bgOrTiFlashType)
+				writeByteMetrics         = writeByteCost.WithLabelValues(name, bgOrTiFlashType)
+				kvCPUMetrics             = kvCPUCost.WithLabelValues(name, bgOrTiFlashType)
+				sqlCPUMetrics            = sqlCPUCost.WithLabelValues(name, bgOrTiFlashType)
 				readRequestCountMetrics  = requestCount.WithLabelValues(name, readTypeLabel)
 				writeRequestCountMetrics = requestCount.WithLabelValues(name, writeTypeLabel)
 			)
diff --git a/pkg/mcs/resourcemanager/server/metrics.go b/pkg/mcs/resourcemanager/server/metrics.go
index 083c44894ef..fd5559fe352 100644
--- a/pkg/mcs/resourcemanager/server/metrics.go
+++ b/pkg/mcs/resourcemanager/server/metrics.go
@@ -26,6 +26,7 @@ const (
 	readTypeLabel          = "read"
 	writeTypeLabel         = "write"
 	backgroundTypeLabel    = "background"
+	tiflashTypeLabel       = "tiflash"
 )
 
 var (