wiring for multiple sharding configs

Signed-off-by: Owen Diehl <ow.diehl@gmail.com>

Author: owen-d

pkg/querier/querysharding/middleware.go

@@ -2,42 +2,98 @@ package querysharding
 
 import (
 	"context"
+	"github.com/cortexproject/cortex/pkg/chunk"
 	"github.com/cortexproject/cortex/pkg/querier/astmapper"
 	"github.com/cortexproject/cortex/pkg/querier/queryrange"
+	"github.com/pkg/errors"
 	"github.com/prometheus/prometheus/promql"
 	"time"
 )
 
 const (
 	downStreamErrType = "downstream error"
-	parseErrType      = "parse error"
 )
 
-func QueryShardMiddleware(engine *promql.Engine) queryrange.Middleware {
+var (
+	invalidShardingRange = errors.New("Query does not fit in a single sharding configuration")
+)
+
+// Sharding configuration. Will eventually support ranges like chunk.PeriodConfig for meshing togethe multiple queryShard(s).
+// This will enable compatibility for deployments which change their shard factors over time.
+type ShardingConfig struct {
+	From   chunk.DayTime `yaml:"from"`
+	Shards int           `yaml:"shards"`
+}
+
+type ShardingConfigs []ShardingConfig
+
+func (confs ShardingConfigs) ValidRange(start, end int64) (ShardingConfig, error) {
+	for i, conf := range confs {
+		if start < int64(conf.From.Time) {
+			// the query starts before this config's range
+			return ShardingConfig{}, invalidShardingRange
+		} else if i == len(confs)-1 {
+			// the last configuration has no upper bound
+			return conf, nil
+		} else if end < int64(confs[i+1].From.Time) {
+			// The request is entirely scoped into this shard config
+			return conf, nil
+		} else {
+			continue
+		}
+	}
+
+	return ShardingConfig{}, invalidShardingRange
+
+}
+
+func mapQuery(mapper astmapper.ASTMapper, query string) (promql.Node, error) {
+	expr, err := promql.ParseExpr(query)
+	if err != nil {
+		return nil, err
+	}
+	return mapper.Map(expr)
+}
+
+func QueryShardMiddleware(engine *promql.Engine, confs ShardingConfigs) queryrange.Middleware {
 	return queryrange.MiddlewareFunc(func(next queryrange.Handler) queryrange.Handler {
 		return &queryShard{
+			confs:  confs,
 			next:   next,
 			engine: engine,
-			mapper: astmapper.NewMultiMapper(
-				astmapper.NewShardSummer(astmapper.DEFAULT_SHARDS, astmapper.VectorSquasher),
-				astmapper.MapperFunc(astmapper.ShallowEmbedSelectors),
-			),
 		}
 	})
 }
 
 type queryShard struct {
+	confs  ShardingConfigs
 	next   queryrange.Handler
 	engine *promql.Engine
-	mapper astmapper.ASTMapper
 }
 
 func (qs *queryShard) Do(ctx context.Context, r *queryrange.Request) (*queryrange.APIResponse, error) {
 	queryable := &DownstreamQueryable{r, qs.next}
 
+	conf, err := qs.confs.ValidRange(r.Start, r.End)
+	if err != nil {
+		return nil, err
+	}
+
+	mappedQuery, err := mapQuery(
+		astmapper.NewMultiMapper(
+			astmapper.NewShardSummer(conf.Shards, astmapper.VectorSquasher),
+			astmapper.MapperFunc(astmapper.ShallowEmbedSelectors),
+		),
+		r.Query,
+	)
+
+	if err != nil {
+		return nil, err
+	}
+
 	qry, err := qs.engine.NewRangeQuery(
 		queryable,
-		r.Query,
+		mappedQuery.String(),
 		time.Unix(r.Start, 0), time.Unix(r.End, 0),
 		time.Duration(r.Step)*time.Second,
 	)

pkg/querier/querysharding/middleware_test.go

@@ -2,10 +2,12 @@ package querysharding
 
 import (
 	"context"
+	"github.com/cortexproject/cortex/pkg/chunk"
 	"github.com/cortexproject/cortex/pkg/ingester/client"
 	"github.com/cortexproject/cortex/pkg/querier/queryrange"
 	"github.com/cortexproject/cortex/pkg/util"
 	"github.com/pkg/errors"
+	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/promql"
 	"github.com/stretchr/testify/require"
 	"testing"
@@ -83,7 +85,14 @@ func TestMiddleware(t *testing.T) {
 				Timeout:       time.Minute,
 			})
 
-			handler := QueryShardMiddleware(engine).Wrap(c.next)
+			handler := QueryShardMiddleware(
+				engine,
+				ShardingConfigs{
+					{
+						Shards: 3,
+					},
+				},
+			).Wrap(c.next)
 
 			// escape hatch for custom tests
 			if c.override != nil {
@@ -168,3 +177,103 @@ func defaultReq() *queryrange.Request {
 		Query: `__embedded_query__{__cortex_query__="687474705f72657175657374735f746f74616c7b636c75737465723d2270726f64227d"}`,
 	}
 }
+
+func TestShardingConfigs_ValidRange(t *testing.T) {
+	reqWith := func(start, end string) *queryrange.Request {
+		r := defaultReq()
+
+		if start != "" {
+			r.Start = int64(parseDate(start))
+		}
+
+		if end != "" {
+			r.End = int64(parseDate(end))
+		}
+
+		return r
+	}
+
+	var testExpr = []struct {
+		name     string
+		confs    ShardingConfigs
+		req      *queryrange.Request
+		expected ShardingConfig
+		err      error
+	}{
+		{
+			name:  "0 ln configs fail",
+			confs: ShardingConfigs{},
+			req:   defaultReq(),
+			err:   invalidShardingRange,
+		},
+		{
+			name: "request starts before beginning config",
+			confs: ShardingConfigs{
+				{
+					From:   chunk.DayTime{parseDate("2019-10-16")},
+					Shards: 1,
+				},
+			},
+			req: reqWith("2019-10-15", ""),
+			err: invalidShardingRange,
+		},
+		{
+			name: "request spans multiple configs",
+			confs: ShardingConfigs{
+				{
+					From:   chunk.DayTime{parseDate("2019-10-16")},
+					Shards: 1,
+				},
+				{
+					From:   chunk.DayTime{parseDate("2019-11-16")},
+					Shards: 2,
+				},
+			},
+			req: reqWith("2019-10-15", "2019-11-17"),
+			err: invalidShardingRange,
+		},
+		{
+			name: "selects correct config ",
+			confs: ShardingConfigs{
+				{
+					From:   chunk.DayTime{parseDate("2019-10-16")},
+					Shards: 1,
+				},
+				{
+					From:   chunk.DayTime{parseDate("2019-11-16")},
+					Shards: 2,
+				},
+				{
+					From:   chunk.DayTime{parseDate("2019-12-16")},
+					Shards: 3,
+				},
+			},
+			req: reqWith("2019-11-20", "2019-11-25"),
+			expected: ShardingConfig{
+				From:   chunk.DayTime{parseDate("2019-11-16")},
+				Shards: 2,
+			},
+		},
+	}
+
+	for _, c := range testExpr {
+		t.Run(c.name, func(t *testing.T) {
+			out, err := c.confs.ValidRange(c.req.Start, c.req.End)
+
+			if c.err != nil {
+				require.EqualError(t, err, c.err.Error())
+			} else {
+				require.Nil(t, err)
+				require.Equal(t, c.expected, out)
+			}
+		})
+	}
+}
+
+func parseDate(in string) model.Time {
+	t, err := time.Parse("2006-01-02", in)
+	if err != nil {
+		panic(err)
+	}
+	return model.Time(t.UnixNano())
+}