Enable @ modifier with result cache handling (#3744)

* Enable @ modifier with result cache handling

Signed-off-by: Ganesh Vernekar <cs15btech11018@iith.ac.in>

* Fix review comments

Signed-off-by: Ganesh Vernekar <cs15btech11018@iith.ac.in>

* Add querier.at-modifier-enabled flag

Signed-off-by: Ganesh Vernekar <cs15btech11018@iith.ac.in>

* Consider frontend.max-cache-freshness for caching @ modifier

Signed-off-by: Ganesh Vernekar <cs15btech11018@iith.ac.in>

CHANGELOG.md

@@ -8,6 +8,7 @@
   - `-<prefix>.s3.sse.type`
   - `-<prefix>.s3.sse.kms-key-id`
   - `-<prefix>.s3.sse.kms-encryption-context`
+* [FEATURE] Querier: Enable `@ <timestamp>` modifier in PromQL using the new `-querier.at-modifier-enabled` flag. #3744
 * [ENHANCEMENT] Ingester: exposed metric `cortex_ingester_oldest_unshipped_block_timestamp_seconds`, tracking the unix timestamp of the oldest TSDB block not shipped to the storage yet. #3705
 * [ENHANCEMENT] Prometheus upgraded. #3739
   * Avoid unnecessary `runtime.GC()` during compactions.

docs/blocks-storage/querier.md

@@ -132,6 +132,10 @@ querier:
   # CLI flag: -querier.query-store-for-labels-enabled
   [query_store_for_labels_enabled: <boolean> | default = false]
 
+  # Enable the @ modifier in PromQL.
+  # CLI flag: -querier.at-modifier-enabled
+  [at_modifier_enabled: <boolean> | default = false]
+
   # The time after which a metric should be queried from storage and not just
   # ingesters. 0 means all queries are sent to store. When running the blocks
   # storage, if this option is enabled, the time range of the query sent to the

docs/configuration/config-file-reference.md

@@ -800,6 +800,10 @@ The `querier_config` configures the Cortex querier.
 # CLI flag: -querier.query-store-for-labels-enabled
 [query_store_for_labels_enabled: <boolean> | default = false]
 
+# Enable the @ modifier in PromQL.
+# CLI flag: -querier.at-modifier-enabled
+[at_modifier_enabled: <boolean> | default = false]
+
 # The time after which a metric should be queried from storage and not just
 # ingesters. 0 means all queries are sent to store. When running the blocks
 # storage, if this option is enabled, the time range of the query sent to the

pkg/cortex/modules.go

@@ -487,10 +487,11 @@ func (t *Cortex) initQueryFrontendTripperware() (serv services.Service, err erro
 		queryrange.PrometheusResponseExtractor{},
 		t.Cfg.Schema,
 		promql.EngineOpts{
-			Logger:     util_log.Logger,
-			Reg:        prometheus.DefaultRegisterer,
-			MaxSamples: t.Cfg.Querier.MaxSamples,
-			Timeout:    t.Cfg.Querier.Timeout,
+			Logger:           util_log.Logger,
+			Reg:              prometheus.DefaultRegisterer,
+			MaxSamples:       t.Cfg.Querier.MaxSamples,
+			Timeout:          t.Cfg.Querier.Timeout,
+			EnableAtModifier: t.Cfg.Querier.AtModifierEnabled,
 			NoStepSubqueryIntervalFn: func(int64) int64 {
 				return t.Cfg.Querier.DefaultEvaluationInterval.Milliseconds()
 			},

pkg/querier/querier.go

@@ -42,6 +42,7 @@ type Config struct {
 	MaxSamples           int           `yaml:"max_samples"`
 	QueryIngestersWithin time.Duration `yaml:"query_ingesters_within"`
 	QueryStoreForLabels  bool          `yaml:"query_store_for_labels_enabled"`
+	AtModifierEnabled    bool          `yaml:"at_modifier_enabled"`
 
 	// QueryStoreAfter the time after which queries should also be sent to the store and not just ingesters.
 	QueryStoreAfter    time.Duration `yaml:"query_store_after"`
@@ -88,6 +89,7 @@ func (cfg *Config) RegisterFlags(f *flag.FlagSet) {
 	f.IntVar(&cfg.MaxSamples, "querier.max-samples", 50e6, "Maximum number of samples a single query can load into memory.")
 	f.DurationVar(&cfg.QueryIngestersWithin, "querier.query-ingesters-within", 0, "Maximum lookback beyond which queries are not sent to ingester. 0 means all queries are sent to ingester.")
 	f.BoolVar(&cfg.QueryStoreForLabels, "querier.query-store-for-labels-enabled", false, "Query long-term store for series, label values and label names APIs. Works only with blocks engine.")
+	f.BoolVar(&cfg.AtModifierEnabled, "querier.at-modifier-enabled", false, "Enable the @ modifier in PromQL.")
 	f.DurationVar(&cfg.MaxQueryIntoFuture, "querier.max-query-into-future", 10*time.Minute, "Maximum duration into the future you can query. 0 to disable.")
 	f.DurationVar(&cfg.DefaultEvaluationInterval, "querier.default-evaluation-interval", time.Minute, "The default evaluation interval or step size for subqueries.")
 	f.DurationVar(&cfg.QueryStoreAfter, "querier.query-store-after", 0, "The time after which a metric should be queried from storage and not just ingesters. 0 means all queries are sent to store. When running the blocks storage, if this option is enabled, the time range of the query sent to the store will be manipulated to ensure the query end is not more recent than 'now - query-store-after'.")
@@ -170,6 +172,7 @@ func New(cfg Config, limits *validation.Overrides, distributor Distributor, stor
 		MaxSamples:         cfg.MaxSamples,
 		Timeout:            cfg.Timeout,
 		LookbackDelta:      cfg.LookbackDelta,
+		EnableAtModifier:   cfg.AtModifierEnabled,
 		NoStepSubqueryIntervalFn: func(int64) int64 {
 			return cfg.DefaultEvaluationInterval.Milliseconds()
 		},

pkg/querier/queryrange/results_cache.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"sort"
+	"strings"
 	"time"
 
 	"github.com/go-kit/kit/log"
@@ -17,6 +18,7 @@ import (
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/model"
+	"github.com/prometheus/prometheus/promql/parser"
 	"github.com/uber/jaeger-client-go"
 	"github.com/weaveworks/common/httpgrpc"
 
@@ -209,9 +211,9 @@ func (s resultsCache) Do(ctx context.Context, r Request) (Response, error) {
 
 	cached, ok := s.get(ctx, key)
 	if ok {
-		response, extents, err = s.handleHit(ctx, r, cached)
+		response, extents, err = s.handleHit(ctx, r, cached, maxCacheTime)
 	} else {
-		response, extents, err = s.handleMiss(ctx, r)
+		response, extents, err = s.handleMiss(ctx, r, maxCacheTime)
 	}
 
 	if err == nil && len(extents) > 0 {
@@ -226,7 +228,7 @@ func (s resultsCache) Do(ctx context.Context, r Request) (Response, error) {
 }
 
 // shouldCacheResponse says whether the response should be cached or not.
-func (s resultsCache) shouldCacheResponse(ctx context.Context, r Response) bool {
+func (s resultsCache) shouldCacheResponse(ctx context.Context, req Request, r Response, maxCacheTime int64) bool {
 	headerValues := getHeaderValuesWithName(r, cacheControlHeader)
 	for _, v := range headerValues {
 		if v == noStoreValue {
@@ -235,6 +237,10 @@ func (s resultsCache) shouldCacheResponse(ctx context.Context, r Response) bool
 		}
 	}
 
+	if !isAtModifierCachable(req, maxCacheTime) {
+		return false
+	}
+
 	if s.cacheGenNumberLoader == nil {
 		return true
 	}
@@ -257,6 +263,54 @@ func (s resultsCache) shouldCacheResponse(ctx context.Context, r Response) bool
 	return true
 }
 
+var errAtModifierAfterEnd = errors.New("at modifier after end")
+
+// isAtModifierCachable returns true if the @ modifier result
+// is safe to cache.
+func isAtModifierCachable(r Request, maxCacheTime int64) bool {
+	// There are 2 cases when @ modifier is not safe to cache:
+	//   1. When @ modifier points to time beyond the maxCacheTime.
+	//   2. If the @ modifier time is > the query range end while being
+	//      below maxCacheTime. In such cases if any tenant is intentionally
+	//      playing with old data, we could cache empty result if we look
+	//      beyond query end.
+	query := r.GetQuery()
+	if !strings.Contains(query, "@") {
+		return true
+	}
+	expr, err := parser.ParseExpr(query)
+	if err != nil {
+		// We are being pessimistic in such cases.
+		return false
+	}
+
+	end := r.GetEnd()
+	atModCachable := true
+	parser.Inspect(expr, func(n parser.Node, _ []parser.Node) error {
+		switch e := n.(type) {
+		case *parser.VectorSelector:
+			if e.Timestamp != nil && (*e.Timestamp > end || *e.Timestamp > maxCacheTime) {
+				atModCachable = false
+				return errAtModifierAfterEnd
+			}
+		case *parser.MatrixSelector:
+			ts := e.VectorSelector.(*parser.VectorSelector).Timestamp
+			if ts != nil && (*ts > end || *ts > maxCacheTime) {
+				atModCachable = false
+				return errAtModifierAfterEnd
+			}
+		case *parser.SubqueryExpr:
+			if e.Timestamp != nil && (*e.Timestamp > end || *e.Timestamp > maxCacheTime) {
+				atModCachable = false
+				return errAtModifierAfterEnd
+			}
+		}
+		return nil
+	})
+
+	return atModCachable
+}
+
 func getHeaderValuesWithName(r Response, headerName string) (headerValues []string) {
 	for _, hv := range r.GetHeaders() {
 		if hv.GetName() != headerName {
@@ -269,13 +323,13 @@ func getHeaderValuesWithName(r Response, headerName string) (headerValues []stri
 	return
 }
 
-func (s resultsCache) handleMiss(ctx context.Context, r Request) (Response, []Extent, error) {
+func (s resultsCache) handleMiss(ctx context.Context, r Request, maxCacheTime int64) (Response, []Extent, error) {
 	response, err := s.next.Do(ctx, r)
 	if err != nil {
 		return nil, nil, err
 	}
 
-	if !s.shouldCacheResponse(ctx, response) {
+	if !s.shouldCacheResponse(ctx, r, response, maxCacheTime) {
 		return response, []Extent{}, nil
 	}
 
@@ -290,7 +344,7 @@ func (s resultsCache) handleMiss(ctx context.Context, r Request) (Response, []Ex
 	return response, extents, nil
 }
 
-func (s resultsCache) handleHit(ctx context.Context, r Request, extents []Extent) (Response, []Extent, error) {
+func (s resultsCache) handleHit(ctx context.Context, r Request, extents []Extent, maxCacheTime int64) (Response, []Extent, error) {
 	var (
 		reqResps []RequestResponse
 		err      error
@@ -315,7 +369,7 @@ func (s resultsCache) handleHit(ctx context.Context, r Request, extents []Extent
 
 	for _, reqResp := range reqResps {
 		responses = append(responses, reqResp.Response)
-		if !s.shouldCacheResponse(ctx, reqResp.Response) {
+		if !s.shouldCacheResponse(ctx, r, reqResp.Response, maxCacheTime) {
 			continue
 		}
 		extent, err := toExtent(ctx, reqResp.Request, s.extractor.ResponseWithoutHeaders(reqResp.Response))