From c1c444504e4a58a7cc5861d5e9049649d4882a9f Mon Sep 17 00:00:00 2001
From: ansalamdaniel <98307038+ansalamdaniel@users.noreply.github.com>
Date: Thu, 2 Mar 2023 15:25:56 +0530
Subject: [PATCH] Feat: metrics for head_chunks & wal folders (#12013)

Signed-off-by: ansalamdaniel <ansalam.daniel@infracloud.io>
---
 tsdb/head.go      | 11 +++++++++++
 tsdb/wlog/wlog.go | 17 +++++++++++++++--
 2 files changed, 26 insertions(+), 2 deletions(-)

diff --git a/tsdb/head.go b/tsdb/head.go
index f539388f27..cf2c15d79b 100644
--- a/tsdb/head.go
+++ b/tsdb/head.go
@@ -502,6 +502,17 @@ func newHeadMetrics(h *Head, r prometheus.Registerer) *headMetrics {
 			}, func() float64 {
 				return float64(h.iso.lastAppendID())
 			}),
+			prometheus.NewGaugeFunc(prometheus.GaugeOpts{
+				Name: "prometheus_tsdb_head_chunks_storage_size_bytes",
+				Help: "Size of the chunks_head directory.",
+			}, func() float64 {
+				val, err := h.chunkDiskMapper.Size()
+				if err != nil {
+					level.Error(h.logger).Log("msg", "Failed to calculate size of \"chunks_head\" dir",
+						"err", err.Error())
+				}
+				return float64(val)
+			}),
 		)
 	}
 	return m
diff --git a/tsdb/wlog/wlog.go b/tsdb/wlog/wlog.go
index 5ae308d4ea..df8bab53ff 100644
--- a/tsdb/wlog/wlog.go
+++ b/tsdb/wlog/wlog.go
@@ -199,9 +199,10 @@ type wlMetrics struct {
 	truncateTotal   prometheus.Counter
 	currentSegment  prometheus.Gauge
 	writesFailed    prometheus.Counter
+	walFileSize     prometheus.GaugeFunc
 }
 
-func newWLMetrics(r prometheus.Registerer) *wlMetrics {
+func newWLMetrics(w *WL, r prometheus.Registerer) *wlMetrics {
 	m := &wlMetrics{}
 
 	m.fsyncDuration = prometheus.NewSummary(prometheus.SummaryOpts{
@@ -233,6 +234,17 @@ func newWLMetrics(r prometheus.Registerer) *wlMetrics {
 		Name: "writes_failed_total",
 		Help: "Total number of write log writes that failed.",
 	})
+	m.walFileSize = prometheus.NewGaugeFunc(prometheus.GaugeOpts{
+		Name: "storage_size_bytes",
+		Help: "Size of the write log directory.",
+	}, func() float64 {
+		val, err := w.Size()
+		if err != nil {
+			level.Error(w.logger).Log("msg", "Failed to calculate size of \"wal\" dir",
+				"err", err.Error())
+		}
+		return float64(val)
+	})
 
 	if r != nil {
 		r.MustRegister(
@@ -243,6 +255,7 @@ func newWLMetrics(r prometheus.Registerer) *wlMetrics {
 			m.truncateTotal,
 			m.currentSegment,
 			m.writesFailed,
+			m.walFileSize,
 		)
 	}
 
@@ -279,7 +292,7 @@ func NewSize(logger log.Logger, reg prometheus.Registerer, dir string, segmentSi
 	if filepath.Base(dir) == WblDirName {
 		prefix = "prometheus_tsdb_out_of_order_wbl_"
 	}
-	w.metrics = newWLMetrics(prometheus.WrapRegistererWithPrefix(prefix, reg))
+	w.metrics = newWLMetrics(w, prometheus.WrapRegistererWithPrefix(prefix, reg))
 
 	_, last, err := Segments(w.Dir())
 	if err != nil {