Refactor Cortex querier to use Queryable interface directly (#2060)

* New function now requires Queryable for store, instead of ChunkStore.
This makes it easier to use non-chunk-based queryables. On the other
hand, chunk-based queryables are still optimized in the way they were
before (in mergeSeriesSets).

Merged ingester_streaming_queryable.go into distributor_queryable.go.
Merged unified_querier.go into querier.go.

Signed-off-by: Peter Štibraný <peter.stibrany@grafana.com>

* Don't wait for results to arrive if context is done.

Signed-off-by: Peter Štibraný <peter.stibrany@grafana.com>

* Rename parameters to NewQueryable.

Signed-off-by: Peter Štibraný <peter.stibrany@grafana.com>

* Include empty set with error. Unify checks in seriesSetWithFirstSeries methods.

Signed-off-by: Peter Štibraný <peter.stibrany@grafana.com>

Author: pstibrany

pkg/compactor/compactor_test.go

@@ -339,6 +339,7 @@ func TestCompactor_ShouldCompactAllUsersOnShardingEnabledButOnlyOneInstanceRunni
 	cfg := prepareConfig()
 	cfg.ShardingEnabled = true
 	cfg.ShardingRing.InstanceID = "compactor-1"
+	cfg.ShardingRing.InstanceAddr = "1.2.3.4"
 	cfg.ShardingRing.KVStore.Mock = consul.NewInMemoryClient(ring.GetCodec())
 
 	c, tsdbCompactor, logs, _, cleanup := prepare(t, cfg, bucketClient)

pkg/cortex/modules.go

@@ -238,7 +238,7 @@ func (t *Cortex) stopDistributor() (err error) {
 }
 
 func (t *Cortex) initQuerier(cfg *Config) (err error) {
-	queryable, engine := querier.New(cfg.Querier, t.distributor, t.querierChunkStore)
+	queryable, engine := querier.New(cfg.Querier, t.distributor, querier.NewChunkStoreQueryable(cfg.Querier, t.querierChunkStore))
 	api := v1.NewAPI(
 		engine,
 		queryable,
@@ -439,7 +439,7 @@ func (t *Cortex) stopTableManager() error {
 func (t *Cortex) initRuler(cfg *Config) (err error) {
 	cfg.Ruler.Ring.ListenPort = cfg.Server.GRPCListenPort
 	cfg.Ruler.Ring.KVStore.MemberlistKV = t.memberlistKVState.getMemberlistKV
-	queryable, engine := querier.New(cfg.Querier, t.distributor, t.querierChunkStore)
+	queryable, engine := querier.New(cfg.Querier, t.distributor, querier.NewChunkStoreQueryable(cfg.Querier, t.querierChunkStore))
 
 	t.ruler, err = ruler.NewRuler(cfg.Ruler, engine, queryable, t.distributor, prometheus.DefaultRegisterer, util.Logger)
 	if err != nil {

pkg/querier/chunk_store_queryable.go

@@ -44,10 +44,11 @@ func (q *chunkStoreQuerier) Select(sp *storage.SelectParams, matchers ...*labels
 		return nil, nil, promql.ErrStorage{Err: err}
 	}
 
-	return q.partitionChunks(chunks), nil, nil
+	return partitionChunks(chunks, q.mint, q.maxt, q.chunkIteratorFunc), nil, nil
 }
 
-func (q *chunkStoreQuerier) partitionChunks(chunks []chunk.Chunk) storage.SeriesSet {
+// Series in the returned set are sorted alphabetically by labels.
+func partitionChunks(chunks []chunk.Chunk, mint, maxt int64, iteratorFunc chunkIteratorFunc) storage.SeriesSet {
 	chunksBySeries := map[model.Fingerprint][]chunk.Chunk{}
 	for _, c := range chunks {
 		fp := client.Fingerprint(c.Metric)
@@ -59,9 +60,9 @@ func (q *chunkStoreQuerier) partitionChunks(chunks []chunk.Chunk) storage.Series
 		series = append(series, &chunkSeries{
 			labels:            chunksBySeries[i][0].Metric,
 			chunks:            chunksBySeries[i],
-			chunkIteratorFunc: q.chunkIteratorFunc,
-			mint:              q.mint,
-			maxt:              q.maxt,
+			chunkIteratorFunc: iteratorFunc,
+			mint:              mint,
+			maxt:              maxt,
 		})
 	}
 
@@ -95,3 +96,7 @@ func (s *chunkSeries) Labels() labels.Labels {
 func (s *chunkSeries) Iterator() storage.SeriesIterator {
 	return s.chunkIteratorFunc(s.chunks, model.Time(s.mint), model.Time(s.maxt))
 }
+
+func (s *chunkSeries) Chunks() []chunk.Chunk {
+	return s.chunks
+}

pkg/querier/chunk_store_queryable_test.go

@@ -3,6 +3,7 @@ package querier
 import (
 	"context"
 	"fmt"
+	"sort"
 	"testing"
 	"time"
 
@@ -65,3 +66,34 @@ func mkChunk(t require.TestingT, mint, maxt model.Time, step time.Duration, enco
 	}
 	return chunk.NewChunk(userID, fp, metric, pc, mint, maxt)
 }
+
+func TestPartitionChunksOutputIsSortedByLabels(t *testing.T) {
+	var allChunks []chunk.Chunk
+
+	const count = 10
+	// go down, to add series in reversed order
+	for i := count; i > 0; i-- {
+		ch := mkChunk(t, model.Time(0), model.Time(1000), time.Millisecond, promchunk.Bigchunk)
+		// mkChunk uses `foo` as metric name, so we rename metric to be unique
+		ch.Metric[0].Value = fmt.Sprintf("%02d", i)
+
+		allChunks = append(allChunks, ch)
+	}
+
+	res := partitionChunks(allChunks, 0, 1000, mergeChunks)
+
+	// collect labels from each series
+	var seriesLabels []labels.Labels
+	for res.Next() {
+		seriesLabels = append(seriesLabels, res.At().Labels())
+	}
+
+	require.Len(t, seriesLabels, count)
+	require.True(t, sort.IsSorted(sortedByLabels(seriesLabels)))
+}
+
+type sortedByLabels []labels.Labels
+
+func (b sortedByLabels) Len() int           { return len(b) }
+func (b sortedByLabels) Swap(i, j int)      { b[i], b[j] = b[j], b[i] }
+func (b sortedByLabels) Less(i, j int) bool { return labels.Compare(b[i], b[j]) < 0 }

pkg/querier/distributor_queryable.go

@@ -2,14 +2,17 @@ package querier
 
 import (
 	"context"
+	"sort"
 
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/pkg/labels"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/prometheus/prometheus/storage"
+	"github.com/weaveworks/common/user"
 
 	"github.com/cortexproject/cortex/pkg/ingester/client"
 	"github.com/cortexproject/cortex/pkg/prom1/storage/metric"
+	"github.com/cortexproject/cortex/pkg/util/chunkcompat"
 )
 
 // Distributor is the read interface to the distributor, made an interface here
@@ -22,13 +25,15 @@ type Distributor interface {
 	MetricsForLabelMatchers(ctx context.Context, from, through model.Time, matchers ...*labels.Matcher) ([]metric.Metric, error)
 }
 
-func newDistributorQueryable(distributor Distributor) storage.Queryable {
+func newDistributorQueryable(distributor Distributor, streaming bool, iteratorFn chunkIteratorFunc) storage.Queryable {
 	return storage.QueryableFunc(func(ctx context.Context, mint, maxt int64) (storage.Querier, error) {
 		return &distributorQuerier{
 			distributor: distributor,
 			ctx:         ctx,
 			mint:        mint,
 			maxt:        maxt,
+			streaming:   streaming,
+			chunkIterFn: iteratorFn,
 		}, nil
 	})
 }
@@ -37,13 +42,25 @@ type distributorQuerier struct {
 	distributor Distributor
 	ctx         context.Context
 	mint, maxt  int64
+	streaming   bool
+	chunkIterFn chunkIteratorFunc
 }
 
 func (q *distributorQuerier) Select(sp *storage.SelectParams, matchers ...*labels.Matcher) (storage.SeriesSet, storage.Warnings, error) {
-	mint, maxt := q.mint, q.maxt
-	if sp != nil {
-		mint = sp.Start
-		maxt = sp.End
+	// Kludge: Prometheus passes nil SelectParams if it is doing a 'series' operation,
+	// which needs only metadata.
+	if sp == nil {
+		ms, err := q.distributor.MetricsForLabelMatchers(q.ctx, model.Time(q.mint), model.Time(q.maxt), matchers...)
+		if err != nil {
+			return nil, nil, err
+		}
+		return metricsToSeriesSet(ms), nil, nil
+	}
+
+	mint, maxt := sp.Start, sp.End
+
+	if q.streaming {
+		return q.streamingSelect(*sp, matchers)
 	}
 
 	matrix, err := q.distributor.Query(q.ctx, model.Time(mint), model.Time(maxt), matchers...)
@@ -54,6 +71,45 @@ func (q *distributorQuerier) Select(sp *storage.SelectParams, matchers ...*label
 	return matrixToSeriesSet(matrix), nil, nil
 }
 
+func (q *distributorQuerier) streamingSelect(sp storage.SelectParams, matchers []*labels.Matcher) (storage.SeriesSet, storage.Warnings, error) {
+	userID, err := user.ExtractOrgID(q.ctx)
+	if err != nil {
+		return nil, nil, promql.ErrStorage{Err: err}
+	}
+
+	mint, maxt := sp.Start, sp.End
+
+	results, err := q.distributor.QueryStream(q.ctx, model.Time(mint), model.Time(maxt), matchers...)
+	if err != nil {
+		return nil, nil, promql.ErrStorage{Err: err}
+	}
+
+	serieses := make([]storage.Series, 0, len(results))
+	for _, result := range results {
+		// Sometimes the ingester can send series that have no data.
+		if len(result.Chunks) == 0 {
+			continue
+		}
+
+		ls := client.FromLabelAdaptersToLabels(result.Labels)
+		sort.Sort(ls)
+
+		chunks, err := chunkcompat.FromChunks(userID, ls, result.Chunks)
+		if err != nil {
+			return nil, nil, promql.ErrStorage{Err: err}
+		}
+
+		series := &chunkSeries{
+			labels:            ls,
+			chunks:            chunks,
+			chunkIteratorFunc: q.chunkIterFn,
+		}
+		serieses = append(serieses, series)
+	}
+
+	return newConcreteSeriesSet(serieses), nil, nil
+}
+
 func (q *distributorQuerier) LabelValues(name string) ([]string, storage.Warnings, error) {
 	lv, err := q.distributor.LabelValuesForLabelName(q.ctx, model.LabelName(name))
 	return lv, nil, err

pkg/querier/distributor_queryable_test.go

@@ -6,10 +6,15 @@ import (
 
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/pkg/labels"
+	"github.com/prometheus/prometheus/storage"
 	"github.com/stretchr/testify/require"
+	"github.com/weaveworks/common/user"
 
+	"github.com/cortexproject/cortex/pkg/chunk"
+	"github.com/cortexproject/cortex/pkg/chunk/encoding"
 	"github.com/cortexproject/cortex/pkg/ingester/client"
 	"github.com/cortexproject/cortex/pkg/prom1/storage/metric"
+	"github.com/cortexproject/cortex/pkg/util/chunkcompat"
 )
 
 const (
@@ -32,11 +37,58 @@ func TestDistributorQuerier(t *testing.T) {
 			},
 		},
 	}
-	queryable := newDistributorQueryable(d)
+	queryable := newDistributorQueryable(d, false, nil)
 	querier, err := queryable.Querier(context.Background(), mint, maxt)
 	require.NoError(t, err)
 
-	seriesSet, _, err := querier.Select(nil)
+	seriesSet, _, err := querier.Select(&storage.SelectParams{Start: mint, End: maxt})
+	require.NoError(t, err)
+
+	require.True(t, seriesSet.Next())
+	series := seriesSet.At()
+	require.Equal(t, labels.Labels{{Name: "bar", Value: "baz"}}, series.Labels())
+
+	require.True(t, seriesSet.Next())
+	series = seriesSet.At()
+	require.Equal(t, labels.Labels{{Name: "foo", Value: "bar"}}, series.Labels())
+
+	require.False(t, seriesSet.Next())
+	require.NoError(t, seriesSet.Err())
+}
+
+func TestIngesterStreaming(t *testing.T) {
+	// We need to make sure that there is atleast one chunk present,
+	// else no series will be selected.
+	promChunk, err := encoding.NewForEncoding(encoding.Bigchunk)
+	require.NoError(t, err)
+
+	clientChunks, err := chunkcompat.ToChunks([]chunk.Chunk{
+		chunk.NewChunk("", 0, nil, promChunk, model.Earliest, model.Earliest),
+	})
+	require.NoError(t, err)
+
+	d := &mockDistributor{
+		r: []client.TimeSeriesChunk{
+			{
+				Labels: []client.LabelAdapter{
+					{Name: "bar", Value: "baz"},
+				},
+				Chunks: clientChunks,
+			},
+			{
+				Labels: []client.LabelAdapter{
+					{Name: "foo", Value: "bar"},
+				},
+				Chunks: clientChunks,
+			},
+		},
+	}
+	ctx := user.InjectOrgID(context.Background(), "0")
+	queryable := newDistributorQueryable(d, true, mergeChunks)
+	querier, err := queryable.Querier(ctx, mint, maxt)
+	require.NoError(t, err)
+
+	seriesSet, _, err := querier.Select(&storage.SelectParams{Start: mint, End: maxt})
 	require.NoError(t, err)
 
 	require.True(t, seriesSet.Next())