From 185a4a817e1244f91aadbe2ad52e155b41416939 Mon Sep 17 00:00:00 2001
From: Anthony Romano <anthony.romano@coreos.com>
Date: Wed, 21 Jun 2017 22:27:28 -0700
Subject: [PATCH] mvcc: use GaugeFunc metric to load db size when requested

Relying on mvcc to set the db size metric can cause it to
miss size changes when a txn commits after the last write
completes before a quiescent period. Instead, load the
db size on demand.

Fixes #8146
---
 mvcc/kvstore.go     |  6 ++++--
 mvcc/kvstore_txn.go |  1 -
 mvcc/metrics.go     | 15 +++++++++++++--
 3 files changed, 17 insertions(+), 5 deletions(-)

diff --git a/mvcc/kvstore.go b/mvcc/kvstore.go
index d7e42d164c9b..a243cade94fd 100644
--- a/mvcc/kvstore.go
+++ b/mvcc/kvstore.go
@@ -131,6 +131,10 @@ func NewStore(b backend.Backend, le lease.Lessor, ig ConsistentIndexGetter) *sto
 	tx.Unlock()
 	s.b.ForceCommit()
 
+	reportDbTotalSizeInBytesMu.Lock()
+	reportDbTotalSizeInBytes = func() float64 { return float64(s.b.Size()) }
+	reportDbTotalSizeInBytesMu.Unlock()
+
 	if err := s.restore(); err != nil {
 		// TODO: return the error instead of panic here?
 		panic("failed to recover store from backend")
@@ -261,8 +265,6 @@ func (s *store) restore() error {
 	tx := s.b.BatchTx()
 	tx.Lock()
 
-	dbTotalSize.Set(float64(s.b.Size()))
-
 	_, finishedCompactBytes := tx.UnsafeRange(metaBucketName, finishedCompactKeyName, nil, 0)
 	if len(finishedCompactBytes) != 0 {
 		s.compactMainRev = bytesToRev(finishedCompactBytes[0]).main
diff --git a/mvcc/kvstore_txn.go b/mvcc/kvstore_txn.go
index ae8d37238369..ae89d0f7124c 100644
--- a/mvcc/kvstore_txn.go
+++ b/mvcc/kvstore_txn.go
@@ -105,7 +105,6 @@ func (tw *storeTxnWrite) End() {
 	if len(tw.changes) != 0 {
 		tw.s.revMu.Unlock()
 	}
-	dbTotalSize.Set(float64(tw.s.b.Size()))
 	tw.s.mu.RUnlock()
 }
 
diff --git a/mvcc/metrics.go b/mvcc/metrics.go
index aa8af6aa5525..a65fe59b996d 100644
--- a/mvcc/metrics.go
+++ b/mvcc/metrics.go
@@ -15,6 +15,8 @@
 package mvcc
 
 import (
+	"sync"
+
 	"github.com/prometheus/client_golang/prometheus"
 )
 
@@ -129,12 +131,21 @@ var (
 			Buckets: prometheus.ExponentialBuckets(100, 2, 14),
 		})
 
-	dbTotalSize = prometheus.NewGauge(prometheus.GaugeOpts{
+	dbTotalSize = prometheus.NewGaugeFunc(prometheus.GaugeOpts{
 		Namespace: "etcd_debugging",
 		Subsystem: "mvcc",
 		Name:      "db_total_size_in_bytes",
 		Help:      "Total size of the underlying database in bytes.",
-	})
+	},
+		func() float64 {
+			reportDbTotalSizeInBytesMu.RLock()
+			defer reportDbTotalSizeInBytesMu.RUnlock()
+			return reportDbTotalSizeInBytes()
+		},
+	)
+	// overridden by mvcc initialization
+	reportDbTotalSizeInBytesMu sync.RWMutex
+	reportDbTotalSizeInBytes   func() float64 = func() float64 { return 0 }
 )
 
 func init() {