From 8be1a14f64f2e01d02720c434cb60f6e9e284469 Mon Sep 17 00:00:00 2001
From: Iceber Gu <caiwei95@hotmail.com>
Date: Thu, 19 Dec 2024 13:45:26 +0800
Subject: [PATCH] internalstorage: add metrics config

Signed-off-by: Iceber Gu <caiwei95@hotmail.com>
---
 pkg/storage/internalstorage/config.go   |  8 +++++-
 pkg/storage/internalstorage/metrics.go  | 33 +++++++++----------------
 pkg/storage/internalstorage/register.go |  6 +++--
 3 files changed, 23 insertions(+), 24 deletions(-)

diff --git a/pkg/storage/internalstorage/config.go b/pkg/storage/internalstorage/config.go
index 81823c32b..8304e0da7 100644
--- a/pkg/storage/internalstorage/config.go
+++ b/pkg/storage/internalstorage/config.go
@@ -48,7 +48,8 @@ type Config struct {
 
 	Params map[string]string `yaml:"params"`
 
-	Log *LogConfig `yaml:"log"`
+	Log     *LogConfig    `yaml:"log"`
+	Metrics MetricsConfig `yaml:"metrics"`
 }
 
 type LogConfig struct {
@@ -60,6 +61,11 @@ type LogConfig struct {
 	Logger                    *lumberjack.Logger `yaml:"logger"`
 }
 
+type MetricsConfig struct {
+	Disable                bool          `yaml:"disable"`
+	DBStatsRefreshInterval time.Duration `yaml:"dbStatsRefreshInterval" default:"15s"`
+}
+
 type MySQLConfig struct {
 	DialTimeout  *time.Duration `yaml:"dialTimeout"`
 	ReadTimeout  *time.Duration `yaml:"readTimeout"`
diff --git a/pkg/storage/internalstorage/metrics.go b/pkg/storage/internalstorage/metrics.go
index afecf023a..c8cd9b437 100644
--- a/pkg/storage/internalstorage/metrics.go
+++ b/pkg/storage/internalstorage/metrics.go
@@ -23,38 +23,29 @@ var (
 )
 
 const (
-	defaultRefreshInterval = 15 // the prometheus default pull metrics every 15 seconds
+	defaultRefreshInterval = 15 * time.Second // the prometheus default pull metrics every 15 seconds
 )
 
+// Prometheus implements the gorm.Plugin and provides db metrics.
+// The `Prometheus` structure name is preserved, and this file may be moved to a separate package
+// in the future to be made available to other gorm-based storage layers.
 type Prometheus struct {
 	*gorm.DB
 	*DBStats
 
-	refreshInterval uint32
+	refreshInterval time.Duration
 	refreshOnce     sync.Once
 	Labels          map[string]string
 }
 
-type MetricsConfig struct {
-	DBName          string            // use DBName as metrics label
-	RefreshInterval uint32            // refresh metrics interval.
-	Labels          map[string]string // metrics labels
-}
-
-func NewGormMetrics(config MetricsConfig) *Prometheus {
-	if config.RefreshInterval == 0 {
-		config.RefreshInterval = defaultRefreshInterval
-	}
-
-	labels := make(map[string]string)
-	if config.Labels != nil {
-		labels = config.Labels
-	}
-	if config.DBName != "" {
-		labels["db_name"] = config.DBName
+func NewGormMetrics(dbName string, refreshInterval time.Duration) *Prometheus {
+	// TODO(iceber): Use real-time mode to get DBStats when the refresh time is below a certain threshold
+	if refreshInterval == 0 {
+		refreshInterval = defaultRefreshInterval
 	}
 
-	return &Prometheus{refreshInterval: config.RefreshInterval, Labels: labels}
+	labels := map[string]string{"db_name": dbName}
+	return &Prometheus{refreshInterval: refreshInterval, Labels: labels}
 }
 
 func (p *Prometheus) Name() string {
@@ -67,7 +58,7 @@ func (p *Prometheus) Initialize(db *gorm.DB) error { // can be called repeatedly
 
 	p.refreshOnce.Do(func() {
 		go func() {
-			for range time.Tick(time.Duration(p.refreshInterval) * time.Second) {
+			for range time.Tick(p.refreshInterval) {
 				p.refresh()
 			}
 		}()
diff --git a/pkg/storage/internalstorage/register.go b/pkg/storage/internalstorage/register.go
index 2dbedffdb..2bfdb6542 100644
--- a/pkg/storage/internalstorage/register.go
+++ b/pkg/storage/internalstorage/register.go
@@ -82,8 +82,10 @@ func NewStorageFactory(configPath string) (storage.StorageFactory, error) {
 		return nil, err
 	}
 
-	if err := db.Use(NewGormMetrics(MetricsConfig{DBName: cfg.Database})); err != nil {
-		return nil, err
+	if !cfg.Metrics.Disable {
+		if err := db.Use(NewGormMetrics(cfg.Database, cfg.Metrics.DBStatsRefreshInterval)); err != nil {
+			return nil, err
+		}
 	}
 
 	sqlDB, err := db.DB()