Export the current segment index as a metic. (#601)

* Export the current segment index as a metic.

Signed-off-by: Callum Styan <[email protected]>

wal/wal.go

@@ -171,6 +171,7 @@ type WAL struct {
 	pageCompletions prometheus.Counter
 	truncateFail    prometheus.Counter
 	truncateTotal   prometheus.Counter
+	currentSegment  prometheus.Gauge
 }
 
 // New returns a new WAL over the given directory.
@@ -218,8 +219,12 @@ func NewSize(logger log.Logger, reg prometheus.Registerer, dir string, segmentSi
 		Name: "prometheus_tsdb_wal_truncations_total",
 		Help: "Total number of WAL truncations attempted.",
 	})
+	w.currentSegment = prometheus.NewGauge(prometheus.GaugeOpts{
+		Name: "prometheus_tsdb_wal_segment_current",
+		Help: "WAL segment index that TSDB is currently writing to.",
+	})
 	if reg != nil {
-		reg.MustRegister(w.fsyncDuration, w.pageFlushes, w.pageCompletions, w.truncateFail, w.truncateTotal)
+		reg.MustRegister(w.fsyncDuration, w.pageFlushes, w.pageCompletions, w.truncateFail, w.truncateTotal, w.currentSegment)
 	}
 
 	_, j, err := w.Segments()
@@ -413,7 +418,7 @@ func (w *WAL) setSegment(segment *Segment) error {
 		return err
 	}
 	w.donePages = int(stat.Size() / pageSize)
-
+	w.currentSegment.Set(float64(segment.Index()))
 	return nil
 }
 

wal/wal_test.go

@@ -23,6 +23,7 @@ import (
 	"path/filepath"
 	"testing"
 
+	client_testutil "github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/prometheus/tsdb/testutil"
 )
 
@@ -318,6 +319,35 @@ func TestClose(t *testing.T) {
 	testutil.NotOk(t, w.Close())
 }
 
+func TestSegmentMetric(t *testing.T) {
+	var (
+		segmentSize = pageSize
+		recordSize  = (pageSize / 2) - recordHeaderSize
+	)
+
+	dir, err := ioutil.TempDir("", "segment_metric")
+	testutil.Ok(t, err)
+	defer func() {
+		testutil.Ok(t, os.RemoveAll(dir))
+	}()
+	w, err := NewSize(nil, nil, dir, segmentSize)
+	testutil.Ok(t, err)
+
+	initialSegment := client_testutil.ToFloat64(w.currentSegment)
+
+	// Write 3 records, each of which is half the segment size, meaning we should rotate to the next segment.
+	for i := 0; i < 3; i++ {
+		buf := make([]byte, recordSize)
+		_, err := rand.Read(buf)
+		testutil.Ok(t, err)
+
+		err = w.Log(buf)
+		testutil.Ok(t, err)
+	}
+	testutil.Assert(t, client_testutil.ToFloat64(w.currentSegment) == initialSegment+1, "segment metric did not increment after segment rotation")
+	testutil.Ok(t, w.Close())
+}
+
 func BenchmarkWAL_LogBatched(b *testing.B) {
 	dir, err := ioutil.TempDir("", "bench_logbatch")
 	testutil.Ok(b, err)