Restore stale series count from chunk snapshots

Signed-off-by: Ganesh Vernekar <[email protected]>

Author: codesome

tsdb/head_test.go

@@ -6988,6 +6988,10 @@ func TestHead_NumStaleSeries(t *testing.T) {
 	appendFloatHistogram(series5, 400, staleFH)
 	restartHeadAndVerifySeriesCounts(3, 5)
 
+	// This will test restarting with snapshot.
+	head.opts.EnableMemorySnapshotOnShutdown = true
+	restartHeadAndVerifySeriesCounts(3, 5)
+
 	// Test garbage collection behavior - stale series should be decremented when GC'd.
 	// Force a garbage collection by truncating old data.
 	require.NoError(t, head.Truncate(300))

tsdb/head_wal.go

@@ -16,7 +16,6 @@ package tsdb
 import (
 	"errors"
 	"fmt"
-	"github.com/prometheus/prometheus/model/value"
 	"maps"
 	"math"
 	"os"
@@ -33,6 +32,7 @@ import (
 	"github.com/prometheus/prometheus/model/histogram"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/model/metadata"
+	"github.com/prometheus/prometheus/model/value"
 	"github.com/prometheus/prometheus/storage"
 	"github.com/prometheus/prometheus/tsdb/chunkenc"
 	"github.com/prometheus/prometheus/tsdb/chunks"
@@ -1607,6 +1607,12 @@ func (h *Head) loadChunkSnapshot() (int, int, map[chunks.HeadSeriesRef]*memSerie
 				series.lastHistogramValue = csr.lastHistogramValue
 				series.lastFloatHistogramValue = csr.lastFloatHistogramValue
 
+				if value.IsStaleNaN(series.lastValue) ||
+					(series.lastHistogramValue != nil && value.IsStaleNaN(series.lastHistogramValue.Sum)) ||
+					(series.lastFloatHistogramValue != nil && value.IsStaleNaN(series.lastFloatHistogramValue.Sum)) {
+					h.numStaleSeries.Inc()
+				}
+
 				app, err := series.headChunks.chunk.Appender()
 				if err != nil {
 					errChan <- err