Streaming PromQL engine: extract linked list of series type

CHANGELOG.md

@@ -10,8 +10,8 @@
 * [FEATURE] Continuous-test: now runable as a module with `mimir -target=continuous-test`. #7747
 * [FEATURE] Store-gateway: Allow specific tenants to be enabled or disabled via `-store-gateway.enabled-tenants` or `-store-gateway.disabled-tenants` CLI flags or their corresponding YAML settings. #7653
 * [FEATURE] New `-<prefix>.s3.bucket-lookup-type` flag configures lookup style type, used to access bucket in s3 compatible providers. #7684
+* [FEATURE] Querier: add experimental streaming PromQL engine, enabled with `-querier.promql-engine=streaming`. #7693 #7899
 * [FEATURE] New `/ingester/unregister-on-shutdown` HTTP endpoint allows dynamic access to ingesters' `-ingester.ring.unregister-on-shutdown` configuration. #7739
-* [FEATURE] Querier: add experimental streaming PromQL engine, enabled with `-querier.promql-engine=streaming`. #7693
 * [FEATURE] Server: added experimental [PROXY protocol support](https://www.haproxy.org/download/2.3/doc/proxy-protocol.txt). The PROXY protocol support can be enabled via `-server.proxy-protocol-enabled=true`. When enabled, the support is added both to HTTP and gRPC listening ports. #7698
 * [ENHANCEMENT] Store-gateway: merge series from different blocks concurrently. #7456
 * [ENHANCEMENT] Store-gateway: Add `stage="wait_max_concurrent"` to `cortex_bucket_store_series_request_stage_duration_seconds` which records how long the query had to wait for its turn for `-blocks-storage.bucket-store.max-concurrent`. #7609

pkg/streamingpromql/operator/selector.go

@@ -27,23 +27,12 @@ type Selector struct {
 	// Set for range vector selectors, otherwise 0.
 	Range time.Duration
 
-	querier                   storage.Querier
-	currentSeriesBatch        *seriesBatch
-	seriesIndexInCurrentBatch int
+	querier storage.Querier
+	series  *seriesList
 }
 
-// There's not too much science behind this number: this is based on the batch size used for chunks streaming.
-const seriesBatchSize = 256
-
-var seriesBatchPool = sync.Pool{New: func() any {
-	return &seriesBatch{
-		series: make([]storage.Series, 0, seriesBatchSize),
-		next:   nil,
-	}
-}}
-
 func (s *Selector) SeriesMetadata(ctx context.Context) ([]SeriesMetadata, error) {
-	if s.currentSeriesBatch != nil {
+	if s.series != nil {
 		return nil, errors.New("should not call Selector.SeriesMetadata() multiple times")
 	}
 
@@ -85,23 +74,80 @@ func (s *Selector) SeriesMetadata(ctx context.Context) ([]SeriesMetadata, error)
 	}
 
 	ss := s.querier.Select(ctx, true, hints, s.Matchers...)
-	s.currentSeriesBatch = seriesBatchPool.Get().(*seriesBatch)
-	incompleteBatch := s.currentSeriesBatch
-	totalSeries := 0
+	s.series = newSeriesList()
 
 	for ss.Next() {
-		if len(incompleteBatch.series) == cap(incompleteBatch.series) {
-			nextBatch := seriesBatchPool.Get().(*seriesBatch)
-			incompleteBatch.next = nextBatch
-			incompleteBatch = nextBatch
-		}
+		s.series.Add(ss.At())
+	}
+
+	return s.series.ToSeriesMetadata(), ss.Err()
+}
 
-		incompleteBatch.series = append(incompleteBatch.series, ss.At())
-		totalSeries++
+func (s *Selector) Next(existing chunkenc.Iterator) (chunkenc.Iterator, error) {
+	if s.series.Len() == 0 {
+		return nil, EOS
+	}
+
+	return s.series.Pop().Iterator(existing), nil
+}
+
+func (s *Selector) Close() {
+	if s.series != nil {
+		s.series.Close()
+	}
+
+	if s.querier != nil {
+		_ = s.querier.Close()
+		s.querier = nil
 	}
+}
+
+// seriesList is a FIFO queue of storage.Series.
+//
+// It is implemented as a linked list of slices of storage.Series, to allow O(1) insertion
+// without the memory overhead of a linked list node per storage.Series.
+type seriesList struct {
+	currentSeriesBatch        *seriesBatch
+	seriesIndexInCurrentBatch int
+
+	lastSeriesBatch *seriesBatch
+
+	length int
+}
+
+func newSeriesList() *seriesList {
+	firstBatch := getSeriesBatch()
+
+	return &seriesList{
+		currentSeriesBatch: firstBatch,
+		lastSeriesBatch:    firstBatch,
+	}
+}
+
+// Add adds s to the end of this seriesList.
+func (l *seriesList) Add(s storage.Series) {
+	if len(l.lastSeriesBatch.series) == cap(l.lastSeriesBatch.series) {
+		nextBatch := getSeriesBatch()
+		l.lastSeriesBatch.next = nextBatch
+		l.lastSeriesBatch = nextBatch
+	}
+
+	l.lastSeriesBatch.series = append(l.lastSeriesBatch.series, s)
+	l.length++
+}
+
+// Len returns the total number of series ever added to this seriesList.
+func (l *seriesList) Len() int {
+	return l.length
+}
+
+// ToSeriesMetadata returns a SeriesMetadata value for each series added to this seriesList.
+//
+// Calling ToSeriesMetadata after calling Pop may return an incomplete list.
+func (l *seriesList) ToSeriesMetadata() []SeriesMetadata {
+	metadata := GetSeriesMetadataSlice(l.length)
+	batch := l.currentSeriesBatch
 
-	metadata := GetSeriesMetadataSlice(totalSeries)
-	batch := s.currentSeriesBatch
 	for batch != nil {
 		for _, s := range batch.series {
 			metadata = append(metadata, SeriesMetadata{Labels: s.Labels()})
@@ -110,45 +156,59 @@ func (s *Selector) SeriesMetadata(ctx context.Context) ([]SeriesMetadata, error)
 		batch = batch.next
 	}
 
-	return metadata, ss.Err()
+	return metadata
 }
 
-func (s *Selector) Next(existing chunkenc.Iterator) (chunkenc.Iterator, error) {
-	if s.currentSeriesBatch == nil || len(s.currentSeriesBatch.series) == 0 {
-		return nil, EOS
+// Pop returns the next series from the head of this seriesList, and advances
+// to the next item in this seriesList.
+func (l *seriesList) Pop() storage.Series {
+	if l.currentSeriesBatch == nil || len(l.currentSeriesBatch.series) == 0 {
+		panic("no more series to pop")
 	}
 
-	it := s.currentSeriesBatch.series[s.seriesIndexInCurrentBatch].Iterator(existing)
-	s.seriesIndexInCurrentBatch++
+	s := l.currentSeriesBatch.series[l.seriesIndexInCurrentBatch]
+	l.seriesIndexInCurrentBatch++
 
-	if s.seriesIndexInCurrentBatch == len(s.currentSeriesBatch.series) {
-		b := s.currentSeriesBatch
-		s.currentSeriesBatch = s.currentSeriesBatch.next
+	if l.seriesIndexInCurrentBatch == len(l.currentSeriesBatch.series) {
+		b := l.currentSeriesBatch
+		l.currentSeriesBatch = l.currentSeriesBatch.next
 		putSeriesBatch(b)
-		s.seriesIndexInCurrentBatch = 0
+		l.seriesIndexInCurrentBatch = 0
 	}
 
-	return it, nil
+	return s
 }
 
-func (s *Selector) Close() {
-	for s.currentSeriesBatch != nil {
-		b := s.currentSeriesBatch
-		s.currentSeriesBatch = s.currentSeriesBatch.next
+// Close releases resources associated with this seriesList.
+func (l *seriesList) Close() {
+	for l.currentSeriesBatch != nil {
+		b := l.currentSeriesBatch
+		l.currentSeriesBatch = l.currentSeriesBatch.next
 		putSeriesBatch(b)
 	}
 
-	if s.querier != nil {
-		_ = s.querier.Close()
-		s.querier = nil
-	}
+	l.lastSeriesBatch = nil // Should have been put back in the pool as part of the loop above.
 }
 
 type seriesBatch struct {
 	series []storage.Series
 	next   *seriesBatch
 }
 
+// There's not too much science behind this number: this is based on the batch size used for chunks streaming.
+const seriesBatchSize = 256
+
+var seriesBatchPool = sync.Pool{New: func() any {
+	return &seriesBatch{
+		series: make([]storage.Series, 0, seriesBatchSize),
+		next:   nil,
+	}
+}}
+
+func getSeriesBatch() *seriesBatch {
+	return seriesBatchPool.Get().(*seriesBatch)
+}
+
 func putSeriesBatch(b *seriesBatch) {
 	b.series = b.series[:0]
 	b.next = nil

pkg/streamingpromql/operator/selector_test.go

@@ -0,0 +1,97 @@
+// SPDX-License-Identifier: AGPL-3.0-only
+
+package operator
+
+import (
+	"fmt"
+	"strconv"
+	"testing"
+
+	"github.com/prometheus/prometheus/model/labels"
+	"github.com/prometheus/prometheus/storage"
+	"github.com/prometheus/prometheus/tsdb/chunkenc"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSeriesList_BasicListOperations(t *testing.T) {
+	list := newSeriesList()
+	require.Equal(t, 0, list.Len())
+
+	series1 := mockSeries{labels.FromStrings("series", "1")}
+	series2 := mockSeries{labels.FromStrings("series", "2")}
+	series3 := mockSeries{labels.FromStrings("series", "3")}
+	series4 := mockSeries{labels.FromStrings("series", "4")}
+
+	list.Add(series1)
+	requireSeriesListContents(t, list, series1)
+
+	list.Add(series2)
+	list.Add(series3)
+	list.Add(series4)
+	requireSeriesListContents(t, list, series1, series2, series3, series4)
+
+	require.Equal(t, series1, list.Pop())
+	require.Equal(t, series2, list.Pop())
+	require.Equal(t, series3, list.Pop())
+	require.Equal(t, series4, list.Pop())
+
+	list.Close()
+}
+
+func TestSeriesList_OperationsNearBatchBoundaries(t *testing.T) {
+	cases := []int{
+		seriesBatchSize - 1,
+		seriesBatchSize,
+		seriesBatchSize + 1,
+		(seriesBatchSize * 2) - 1,
+		seriesBatchSize * 2,
+		(seriesBatchSize * 2) + 1,
+	}
+
+	for _, seriesCount := range cases {
+		t.Run(fmt.Sprintf("N=%v", seriesCount), func(t *testing.T) {
+			list := newSeriesList()
+
+			seriesAdded := make([]storage.Series, 0, seriesCount)
+
+			for i := 0; i < seriesCount; i++ {
+				s := mockSeries{labels.FromStrings("series", strconv.Itoa(i))}
+				list.Add(s)
+				seriesAdded = append(seriesAdded, s)
+			}
+
+			requireSeriesListContents(t, list, seriesAdded...)
+
+			seriesPopped := make([]storage.Series, 0, seriesCount)
+
+			for i := 0; i < seriesCount; i++ {
+				seriesPopped = append(seriesPopped, list.Pop())
+			}
+
+			require.Equal(t, seriesAdded, seriesPopped)
+		})
+	}
+}
+
+func requireSeriesListContents(t *testing.T, list *seriesList, series ...storage.Series) {
+	require.Equal(t, len(series), list.Len())
+
+	expectedMetadata := make([]SeriesMetadata, 0, len(series))
+	for _, s := range series {
+		expectedMetadata = append(expectedMetadata, SeriesMetadata{Labels: s.Labels()})
+	}
+
+	require.Equal(t, expectedMetadata, list.ToSeriesMetadata())
+}
+
+type mockSeries struct {
+	labels labels.Labels
+}
+
+func (m mockSeries) Labels() labels.Labels {
+	return m.labels
+}
+
+func (m mockSeries) Iterator(chunkenc.Iterator) chunkenc.Iterator {
+	panic("mockSeries: Iterator() not supported")
+}