From 79cb508d833f05a8db88741d61b5bdcf31c97122 Mon Sep 17 00:00:00 2001
From: i-norden <iansnorden@gmail.com>
Date: Mon, 25 Oct 2021 18:54:54 -0500
Subject: [PATCH] make DB stats collector registration optional

---
 cmd/env.go                |  2 ++
 cmd/root.go               |  7 ++++---
 cmd/util.go               | 14 ++++++++++----
 environments/example.toml |  1 +
 4 files changed, 17 insertions(+), 7 deletions(-)

diff --git a/cmd/env.go b/cmd/env.go
index 1525602..4f81f65 100644
--- a/cmd/env.go
+++ b/cmd/env.go
@@ -44,6 +44,7 @@ const (
 	PROM_HTTP      = "PROM_HTTP"
 	PROM_HTTP_ADDR = "PROM_HTTP_ADDR"
 	PROM_HTTP_PORT = "PROM_HTTP_PORT"
+	PROM_DB_STATS  = "PROM_DB_STATS"
 )
 
 // Bind env vars for eth node and DB configuration
@@ -76,6 +77,7 @@ func init() {
 	viper.BindEnv("prom.http", PROM_HTTP)
 	viper.BindEnv("prom.httpAddr", PROM_HTTP_ADDR)
 	viper.BindEnv("prom.httpPort", PROM_HTTP_PORT)
+	viper.BindEnv("prom.dbStats", PROM_DB_STATS)
 
 	viper.BindEnv("statediff.serviceWorkers", STATEDIFF_SERVICE_WORKERS)
 	viper.BindEnv("statediff.trieWorkers", STATEDIFF_TRIE_WORKERS)
diff --git a/cmd/root.go b/cmd/root.go
index 56739a8..6356e0b 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -136,8 +136,8 @@ func init() {
 	rootCmd.PersistentFlags().Bool("prom-http", false, "enable prometheus http service")
 	rootCmd.PersistentFlags().String("prom-http-addr", "127.0.0.1", "prometheus http host")
 	rootCmd.PersistentFlags().String("prom-http-port", "8080", "prometheus http port")
-
-	rootCmd.PersistentFlags().Bool("metrics", false, "enable metrics")
+	rootCmd.PersistentFlags().Bool("prom-db-stats", false, "enables prometheus db stats")
+	rootCmd.PersistentFlags().Bool("prom-metrics", false, "enable prometheus metrics")
 
 	viper.BindPFlag("server.httpPath", rootCmd.PersistentFlags().Lookup("http-path"))
 	viper.BindPFlag("server.ipcPath", rootCmd.PersistentFlags().Lookup("ipc-path"))
@@ -164,7 +164,8 @@ func init() {
 	viper.BindPFlag("prom.http", rootCmd.PersistentFlags().Lookup("prom-http"))
 	viper.BindPFlag("prom.httpAddr", rootCmd.PersistentFlags().Lookup("prom-http-addr"))
 	viper.BindPFlag("prom.httpPort", rootCmd.PersistentFlags().Lookup("prom-http-port"))
-	viper.BindPFlag("prom.metrics", rootCmd.PersistentFlags().Lookup("metrics"))
+	viper.BindPFlag("prom.dbStats", rootCmd.PersistentFlags().Lookup("prom-db-stats"))
+	viper.BindPFlag("prom.metrics", rootCmd.PersistentFlags().Lookup("prom-metrics"))
 }
 
 func initConfig() {
diff --git a/cmd/util.go b/cmd/util.go
index b55e5cd..7d38bd9 100644
--- a/cmd/util.go
+++ b/cmd/util.go
@@ -52,15 +52,17 @@ func createStateDiffService() (sd.StateDiffService, error) {
 	}
 
 	// create statediff service
-	logWithCommand.Info("Creating statediff service")
+	logWithCommand.Info("Setting up Postgres DB")
 	db, err := setupPostgres(nodeInfo)
 	if err != nil {
 		logWithCommand.Fatal(err)
 	}
+	logWithCommand.Info("Creating statediff indexer")
 	indexer, err := ind.NewStateDiffIndexer(chainConf, db)
 	if err != nil {
 		logWithCommand.Fatal(err)
 	}
+	logWithCommand.Info("Creating statediff service")
 	sdConf := sd.Config{
 		ServiceWorkers:  viper.GetUint("statediff.serviceWorkers"),
 		TrieWorkers:     viper.GetUint("statediff.trieWorkers"),
@@ -71,12 +73,16 @@ func createStateDiffService() (sd.StateDiffService, error) {
 }
 
 func setupPostgres(nodeInfo node.Info) (*postgres.DB, error) {
-	params := GetDBParams()
-	db, err := postgres.NewDB(postgres.DbConnectionString(params), GetDBConfig(), nodeInfo)
+	p := GetDBParams()
+	logWithCommand.Info("initializing DB connection pool")
+	db, err := postgres.NewDB(postgres.DbConnectionString(p), GetDBConfig(), nodeInfo)
 	if err != nil {
 		return nil, err
 	}
-	prom.RegisterDBCollector(params.Name, db.DB)
+	if viper.GetBool("prom.dbStats") {
+		logWithCommand.Info("registering DB collector")
+		prom.RegisterDBCollector(p.Name, db.DB)
+	}
 	return db, nil
 }
 
diff --git a/environments/example.toml b/environments/example.toml
index 0b1caa5..0ee5b6f 100644
--- a/environments/example.toml
+++ b/environments/example.toml
@@ -45,6 +45,7 @@
     trie = 1024
 
 [prom]
+    dbStats = false
     metrics = true
     http = true
     httpAddr = "localhost"