From 8ece8fb44bfc24fee83bfc4539fa5d9aaa12e215 Mon Sep 17 00:00:00 2001
From: Nishant Das <nishdas93@gmail.com>
Date: Mon, 30 Sep 2019 23:24:47 +0800
Subject: [PATCH] Expose DB Metrics (#3663)

* add in bolt metrics

* unregister in db teardown

* unregister in Close()

* fix clear db case

* fix test error

* gaz

* remove unregister

* remove gaz
---
 WORKSPACE                                              |  7 +++++++
 .../blockchain/forkchoice/lmd_ghost_yaml_test.go       |  2 ++
 beacon-chain/db/kv/BUILD.bazel                         |  2 ++
 beacon-chain/db/kv/kv.go                               | 10 ++++++++++
 4 files changed, 21 insertions(+)

diff --git a/WORKSPACE b/WORKSPACE
index b1e4da9e75..33a51dbe96 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -1231,3 +1231,10 @@ go_repository(
     sum = "h1:2BvfKmzob6Bmd4YsL0zygOqfdFnK7GR4QL06Do4/p7Y=",
     version = "v0.0.4",
 )
+
+go_repository(
+    name = "com_github_mdlayher_prombolt",
+    importpath = "github.com/mdlayher/prombolt",
+    sum = "h1:N257g6TTx0LxYoskSDFxvkSJ3NOZpy9IF1xQ7Gu+K8I=",
+    version = "v0.0.0-20161005185022-dfcf01d20ee9",
+)
diff --git a/beacon-chain/blockchain/forkchoice/lmd_ghost_yaml_test.go b/beacon-chain/blockchain/forkchoice/lmd_ghost_yaml_test.go
index 3375eb1f03..4d818090b0 100644
--- a/beacon-chain/blockchain/forkchoice/lmd_ghost_yaml_test.go
+++ b/beacon-chain/blockchain/forkchoice/lmd_ghost_yaml_test.go
@@ -134,5 +134,7 @@ func TestGetHeadFromYaml(t *testing.T) {
 		}
 
 		helpers.ClearAllCaches()
+		testDB.TeardownDB(t, db)
+
 	}
 }
diff --git a/beacon-chain/db/kv/BUILD.bazel b/beacon-chain/db/kv/BUILD.bazel
index 9c8caff11c..3117777fb5 100644
--- a/beacon-chain/db/kv/BUILD.bazel
+++ b/beacon-chain/db/kv/BUILD.bazel
@@ -27,7 +27,9 @@ go_library(
         "@com_github_ethereum_go_ethereum//common:go_default_library",
         "@com_github_gogo_protobuf//proto:go_default_library",
         "@com_github_karlseguin_ccache//:go_default_library",
+        "@com_github_mdlayher_prombolt//:go_default_library",
         "@com_github_pkg_errors//:go_default_library",
+        "@com_github_prometheus_client_golang//prometheus:go_default_library",
         "@com_github_prysmaticlabs_go_ssz//:go_default_library",
         "@io_opencensus_go//trace:go_default_library",
     ],
diff --git a/beacon-chain/db/kv/kv.go b/beacon-chain/db/kv/kv.go
index 9525aa83a4..3006a49cc3 100644
--- a/beacon-chain/db/kv/kv.go
+++ b/beacon-chain/db/kv/kv.go
@@ -7,7 +7,9 @@ import (
 
 	"github.com/boltdb/bolt"
 	"github.com/karlseguin/ccache"
+	"github.com/mdlayher/prombolt"
 	"github.com/pkg/errors"
+	"github.com/prometheus/client_golang/prometheus"
 )
 
 // BlockCacheSize specifies 4 epochs worth of blocks cached.
@@ -76,6 +78,7 @@ func NewKVStore(dirPath string) (*Store, error) {
 	}); err != nil {
 		return nil, err
 	}
+	err = prometheus.Register(createBoltCollector(kv.db))
 
 	return kv, err
 }
@@ -85,11 +88,13 @@ func (k *Store) ClearDB() error {
 	if _, err := os.Stat(k.databasePath); os.IsNotExist(err) {
 		return nil
 	}
+	prometheus.Unregister(createBoltCollector(k.db))
 	return os.RemoveAll(k.databasePath)
 }
 
 // Close closes the underlying BoltDB database.
 func (k *Store) Close() error {
+	prometheus.Unregister(createBoltCollector(k.db))
 	return k.db.Close()
 }
 
@@ -106,3 +111,8 @@ func createBuckets(tx *bolt.Tx, buckets ...[]byte) error {
 	}
 	return nil
 }
+
+// createBoltCollector returns a prometheus collector specifically configured for boltdb.
+func createBoltCollector(db *bolt.DB) prometheus.Collector {
+	return prombolt.New("boltDB", db)
+}