Skip to content

Commit

Permalink
internalstorage: add metrics config
Browse files Browse the repository at this point in the history
Signed-off-by: Iceber Gu <[email protected]>
  • Loading branch information
Iceber committed Dec 19, 2024
1 parent c2451ac commit 8be1a14
Show file tree
Hide file tree
Showing 3 changed files with 23 additions and 24 deletions.
8 changes: 7 additions & 1 deletion pkg/storage/internalstorage/config.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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"`
Expand Down
33 changes: 12 additions & 21 deletions pkg/storage/internalstorage/metrics.go
Original file line number Diff line number Diff line change
Expand Up @@ -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 {
Expand All @@ -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()
}
}()
Expand Down
6 changes: 4 additions & 2 deletions pkg/storage/internalstorage/register.go
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Expand Down

0 comments on commit 8be1a14

Please sign in to comment.