Add initialization to wal stripeSeries to prune inactive series

Author: kgeckhart

internal/static/metrics/wal/series.go

@@ -10,6 +10,7 @@ import (
 	"github.com/prometheus/prometheus/model/exemplar"
 	"github.com/prometheus/prometheus/model/labels"
 	"github.com/prometheus/prometheus/tsdb/chunks"
+	"go.uber.org/atomic"
 )
 
 // memSeries is a chunkless version of tsdb.memSeries.
@@ -80,11 +81,12 @@ func (m seriesHashmap) Delete(hash uint64, ref chunks.HeadSeriesRef) {
 // Filling the padded space with the maps was profiled to be slower -
 // likely due to the additional pointer dereferences.
 type stripeSeries struct {
-	size      int
-	series    []map[chunks.HeadSeriesRef]*memSeries
-	hashes    []seriesHashmap
-	exemplars []map[chunks.HeadSeriesRef]*exemplar.Exemplar
-	locks     []stripeLock
+	size        int
+	series      []map[chunks.HeadSeriesRef]*memSeries
+	hashes      []seriesHashmap
+	exemplars   []map[chunks.HeadSeriesRef]*exemplar.Exemplar
+	locks       []stripeLock
+	initialized *atomic.Bool
 
 	gcMut sync.Mutex
 }
@@ -95,13 +97,18 @@ type stripeLock struct {
 	_ [40]byte
 }
 
+// newStripeSeries creates a new stripeSeries with the given stripe size in an uninitialized state.
+// When in an uninitialized state, reads and writes are not lock protected. After loading any
+// initial data, a call to MarkInitialized() must be made before using the stripeSeries for
+// ensuring proper function of stripeSeries.gc().
 func newStripeSeries(stripeSize int) *stripeSeries {
 	s := &stripeSeries{
-		size:      stripeSize,
-		series:    make([]map[chunks.HeadSeriesRef]*memSeries, stripeSize),
-		hashes:    make([]seriesHashmap, stripeSize),
-		exemplars: make([]map[chunks.HeadSeriesRef]*exemplar.Exemplar, stripeSize),
-		locks:     make([]stripeLock, stripeSize),
+		size:        stripeSize,
+		series:      make([]map[chunks.HeadSeriesRef]*memSeries, stripeSize),
+		hashes:      make([]seriesHashmap, stripeSize),
+		exemplars:   make([]map[chunks.HeadSeriesRef]*exemplar.Exemplar, stripeSize),
+		locks:       make([]stripeLock, stripeSize),
+		initialized: atomic.NewBool(false),
 	}
 	for i := range s.series {
 		s.series[i] = map[chunks.HeadSeriesRef]*memSeries{}
@@ -115,9 +122,63 @@ func newStripeSeries(stripeSize int) *stripeSeries {
 	return s
 }
 
+// MarkInitialized marks the stripeSeries initialized, allowing usage of stripeSeries.gc(). Returns
+// true if the stripeSeries was not initialized before, false otherwise.
+func (s *stripeSeries) MarkInitialized() bool {
+	return s.initialized.CompareAndSwap(false, true)
+}
+
+// RemoveInactiveSeries removes all series that have a lastTs of 0 while the stripeSeries is still in
+// an uninitialized state. If the stripeSeries is already initialized, it returns 0 and false. Otherwise,
+// it returns the number of series that were removed and true.
+//
+// The stripeSeries assumes that a chunks.HeadSeriesRef uniquely refers to a series in the stripeSeries.
+// But in practice, a chunks.HeadSeriesRef can remain on a WAL even after it has been removed from the
+// stripeSeries. If the series comes back before the original is removed from the WAL we are left with
+// multiple chunks.HeadSeriesRef for the same series. If the WAL is reloaded in this state, we end up with a
+// series leak. A call to stripeSeries.gc() is only capable of removing one instance of the chunks.HeadSeriesRef
+// as it assumes there can only be one chunks.HeadSeriesRef for a series. The remaining chunks.HeadSeriesRefs are
+// left in the stripeSeries and overtime can accumulate to consume a very large amount of memory.
+func (s *stripeSeries) RemoveInactiveSeries() (int, bool) {
+	if s.initialized.Load() {
+		return 0, false
+	}
+
+	inactiveSeries := 0
+	// Start with hashes first because it's easier to get to a series from the hash than a hash from a series.
+	for _, hashSeries := range s.hashes {
+		for hash, seriesForHash := range hashSeries {
+			for _, series := range seriesForHash {
+				if series.lastTs == 0 {
+					hashSeries.Delete(hash, series.ref)
+					inactiveSeries++
+
+					// Get the seriesRef lock to delete the series from s.series.
+					refLock := s.refLock(series.ref)
+					delete(s.series[refLock], series.ref)
+				}
+			}
+		}
+	}
+
+	for _, seriesRefs := range s.series {
+		for head, series := range seriesRefs {
+			if series.lastTs == 0 {
+				delete(s.series[inactiveSeries], head)
+				inactiveSeries++
+			}
+		}
+	}
+
+	return inactiveSeries, true
+}
+
 // gc garbage collects old chunks that are strictly before mint and removes
 // series entirely that have no chunks left.
 func (s *stripeSeries) gc(mint int64) map[chunks.HeadSeriesRef]struct{} {
+	if !s.initialized.Load() {
+		return nil
+	}
 	// NOTE(rfratto): GC will grab two locks, one for the hash and the other for
 	// series. It's not valid for any other function to grab both locks,
 	// otherwise a deadlock might occur when running GC in parallel with
@@ -141,7 +202,7 @@ func (s *stripeSeries) gc(mint int64) map[chunks.HeadSeriesRef]struct{} {
 
 				// The series is stale. We need to obtain a second lock for the
 				// ref if it's different than the hash lock.
-				refLock := int(series.ref) & (s.size - 1)
+				refLock := int(s.refLock(series.ref))
 				if hashLock != refLock {
 					s.locks[refLock].Lock()
 				}
@@ -168,14 +229,14 @@ func (s *stripeSeries) gc(mint int64) map[chunks.HeadSeriesRef]struct{} {
 }
 
 func (s *stripeSeries) GetByID(id chunks.HeadSeriesRef) *memSeries {
-	refLock := uint64(id) & uint64(s.size-1)
+	refLock := s.refLock(id)
 	s.locks[refLock].RLock()
 	defer s.locks[refLock].RUnlock()
 	return s.series[refLock][id]
 }
 
 func (s *stripeSeries) GetByHash(hash uint64, lset labels.Labels) *memSeries {
-	hashLock := hash & uint64(s.size-1)
+	hashLock := s.hashLock(hash)
 
 	s.locks[hashLock].RLock()
 	defer s.locks[hashLock].RUnlock()
@@ -184,8 +245,8 @@ func (s *stripeSeries) GetByHash(hash uint64, lset labels.Labels) *memSeries {
 
 func (s *stripeSeries) Set(hash uint64, series *memSeries) {
 	var (
-		hashLock = hash & uint64(s.size-1)
-		refLock  = uint64(series.ref) & uint64(s.size-1)
+		hashLock = s.hashLock(hash)
+		refLock  = s.refLock(series.ref)
 	)
 
 	// We can't hold both locks at once otherwise we might deadlock with a
@@ -203,7 +264,7 @@ func (s *stripeSeries) Set(hash uint64, series *memSeries) {
 }
 
 func (s *stripeSeries) GetLatestExemplar(ref chunks.HeadSeriesRef) *exemplar.Exemplar {
-	i := uint64(ref) & uint64(s.size-1)
+	i := s.refLock(ref)
 
 	s.locks[i].RLock()
 	exemplar := s.exemplars[i][ref]
@@ -213,7 +274,7 @@ func (s *stripeSeries) GetLatestExemplar(ref chunks.HeadSeriesRef) *exemplar.Exe
 }
 
 func (s *stripeSeries) SetLatestExemplar(ref chunks.HeadSeriesRef, exemplar *exemplar.Exemplar) {
-	i := uint64(ref) & uint64(s.size-1)
+	i := s.refLock(ref)
 
 	// Make sure that's a valid series id and record its latest exemplar
 	s.locks[i].Lock()
@@ -227,6 +288,14 @@ func (s *stripeSeries) iterator() *stripeSeriesIterator {
 	return &stripeSeriesIterator{s}
 }
 
+func (s *stripeSeries) hashLock(hash uint64) uint64 {
+	return hash & uint64(s.size-1)
+}
+
+func (s *stripeSeries) refLock(ref chunks.HeadSeriesRef) uint64 {
+	return uint64(ref) & uint64(s.size-1)
+}
+
 // stripeSeriesIterator allows to iterate over series through a channel.
 // The channel should always be completely consumed to not leak.
 type stripeSeriesIterator struct {
@@ -243,7 +312,7 @@ func (it *stripeSeriesIterator) Channel() <-chan *memSeries {
 			for _, series := range it.s.series[i] {
 				series.Lock()
 
-				j := int(series.lset.Hash()) & (it.s.size - 1)
+				j := int(it.s.hashLock(series.lset.Hash()))
 				if i != j {
 					it.s.locks[j].RLock()
 				}

internal/static/metrics/wal/wal.go

@@ -15,7 +15,6 @@ import (
 
 	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
-	"github.com/grafana/alloy/internal/util"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/prometheus/model/exemplar"
 	"github.com/prometheus/prometheus/model/histogram"
@@ -29,6 +28,8 @@ import (
 	"github.com/prometheus/prometheus/tsdb/record"
 	"github.com/prometheus/prometheus/tsdb/wlog"
 	"go.uber.org/atomic"
+
+	"github.com/grafana/alloy/internal/util"
 )
 
 // ErrWALClosed is an error returned when a WAL operation can't run because the
@@ -198,6 +199,10 @@ func NewStorage(logger log.Logger, registerer prometheus.Registerer, path string
 		}
 	}
 
+	inactiveSeriesRemoved, _ := storage.series.RemoveInactiveSeries()
+	storage.metrics.numActiveSeries.Sub(float64(inactiveSeriesRemoved))
+	storage.series.MarkInitialized()
+
 	return storage, nil
 }
 

internal/static/metrics/wal/wal_test.go

@@ -363,6 +363,12 @@ func TestStorage_WillNotLeaveUnGCableSeries(t *testing.T) {
 	}
 	require.NoError(t, app.Commit())
 
+	// Make sure we have 4 active series left on the WAL
+	var activeSeries io_prometheus_client.Metric
+	err = s.metrics.numActiveSeries.Write(&activeSeries)
+	require.NoError(t, err)
+	require.Equal(t, 4.0, activeSeries.Gauge.GetValue())
+
 	// Forcefully create a bunch of new segments so when we truncate
 	// there's enough segments to be considered for truncation.
 	for i := 0; i < 3; i++ {
@@ -375,15 +381,24 @@ func TestStorage_WillNotLeaveUnGCableSeries(t *testing.T) {
 	err = s.Truncate(keepTs)
 	require.NoError(t, err)
 
+	// Make sure we have no active series left on the WAL
+	err = s.metrics.numActiveSeries.Write(&activeSeries)
+	require.NoError(t, err)
+	require.Equal(t, 0.0, activeSeries.Gauge.GetValue())
+
 	// Publish new samples that will create new RefIDs
 	for _, metric := range payload {
-		// Write a new sample for the same metric and see what happens
 		metric.samples = metric.samples[1:]
 		metric.samples[0].ts = metric.samples[0].ts * 10
 		metric.Write(t, app)
 	}
 	require.NoError(t, app.Commit())
 
+	// Make sure we have 4 active series left on the WAL
+	err = s.metrics.numActiveSeries.Write(&activeSeries)
+	require.NoError(t, err)
+	require.Equal(t, 4.0, activeSeries.Gauge.GetValue())
+
 	// Close the WAL before we have a chance to remove the first RefIDs
 	err = s.Close()
 	require.NoError(t, err)
@@ -401,11 +416,11 @@ func TestStorage_WillNotLeaveUnGCableSeries(t *testing.T) {
 		err = s.Truncate(keepTs)
 		require.NoError(t, err)
 	}
-	// We should 0 active series but will have 4 instead because the first 4 RefIDs are lost in the WAL
-	var metric io_prometheus_client.Metric
-	err = s.metrics.numActiveSeries.Write(&metric)
+
+	// Make sure we have no active series left on the WAL
+	err = s.metrics.numActiveSeries.Write(&activeSeries)
 	require.NoError(t, err)
-	require.Equal(t, 0.0, metric.Gauge.GetValue())
+	require.Equal(t, 0.0, activeSeries.Gauge.GetValue())
 }
 
 func TestStorage_WriteStalenessMarkers(t *testing.T) {