Max bytes read limit

CHANGELOG.md

@@ -39,6 +39,7 @@
 * [6675](https://github.com/grafana/loki/pull/6675) **btaani**: Add logfmt expression parser for selective extraction of labels from logfmt formatted logs
 * [8474](https://github.com/grafana/loki/pull/8474) **farodin91**: Add support for short-lived S3 session tokens
 * [8774](https://github.com/grafana/loki/pull/8774) **slim-bean**: Add new logql template functions `bytes`, `duration`, `unixEpochMillis`, `unixEpochNanos`, `toDateInZone`, `b64Enc`, and `b64Dec`
+* [8670](https://github.com/grafana/loki/pull/8670) **salvacorts** Introduce two new limits to refuse log and metric queries that would read too much data.
 
 ##### Fixes
 

docs/sources/configuration/_index.md

@@ -2326,6 +2326,17 @@ The `limits_config` block configures global and per-tenant limits in Loki.
 # CLI flag: -frontend.min-sharding-lookback
 [min_sharding_lookback: <duration> | default = 0s]
 
+# Max number of bytes a query can fetch. Enforced in log and metric queries only
+# when TSDB is used. The default value of 0 disables this limit.
+# CLI flag: -frontend.max-query-bytes-read
+[max_query_bytes_read: <int> | default = 0B]
+
+# Max number of bytes a query can fetch after splitting and sharding. Enforced
+# in log and metric queries only when TSDB is used. The default value of 0
+# disables this limit.
+# CLI flag: -frontend.max-querier-bytes-read
+[max_querier_bytes_read: <int> | default = 0B]
+
 # Duration to delay the evaluation of rules to ensure the underlying metrics
 # have been pushed to Cortex.
 # CLI flag: -ruler.evaluation-delay-duration

pkg/querier/queryrange/codec.go

@@ -26,6 +26,7 @@ import (
 	"github.com/grafana/loki/pkg/logqlmodel"
 	"github.com/grafana/loki/pkg/logqlmodel/stats"
 	"github.com/grafana/loki/pkg/querier/queryrange/queryrangebase"
+	indexStats "github.com/grafana/loki/pkg/storage/stores/index/stats"
 	"github.com/grafana/loki/pkg/util"
 	"github.com/grafana/loki/pkg/util/httpreq"
 	"github.com/grafana/loki/pkg/util/marshal"
@@ -685,6 +686,20 @@ func (Codec) MergeResponse(responses ...queryrangebase.Response) (queryrangebase
 			Data:       names,
 			Statistics: mergedStats,
 		}, nil
+	case *IndexStatsResponse:
+		headers := responses[0].(*IndexStatsResponse).Headers
+		stats := make([]*indexStats.Stats, len(responses))
+		for i, res := range responses {
+			stats[i] = res.(*IndexStatsResponse).Response
+		}
+
+		mergedIndexStats := indexStats.MergeStats(stats...)
+
+		return &IndexStatsResponse{
+			Response: &mergedIndexStats,
+			Headers:  headers,
+		}, nil
+
 	default:
 		return nil, errors.New("unknown response in merging responses")
 	}

pkg/querier/queryrange/limits.go

@@ -5,29 +5,37 @@ import (
 	"fmt"
 	"net/http"
 	"sort"
+	"strings"
 	"sync"
 	"time"
 
+	"github.com/dustin/go-humanize"
+	"github.com/go-kit/log"
 	"github.com/go-kit/log/level"
 	"github.com/grafana/dskit/tenant"
 	"github.com/opentracing/opentracing-go"
+	"github.com/pkg/errors"
 	"github.com/prometheus/common/model"
 	"github.com/prometheus/prometheus/model/timestamp"
 	"github.com/weaveworks/common/httpgrpc"
 	"github.com/weaveworks/common/user"
 
 	"github.com/grafana/loki/pkg/logproto"
 	"github.com/grafana/loki/pkg/logql"
+	"github.com/grafana/loki/pkg/logql/syntax"
 	"github.com/grafana/loki/pkg/querier/queryrange/queryrangebase"
 	"github.com/grafana/loki/pkg/storage/config"
+	"github.com/grafana/loki/pkg/storage/stores/index/stats"
 	"github.com/grafana/loki/pkg/util"
 	util_log "github.com/grafana/loki/pkg/util/log"
 	"github.com/grafana/loki/pkg/util/spanlogger"
 	"github.com/grafana/loki/pkg/util/validation"
 )
 
 const (
-	limitErrTmpl = "maximum of series (%d) reached for a single query"
+	limitErrTmpl                  = "maximum of series (%d) reached for a single query"
+	limErrQueryTooManyBytesTmpl   = "the query would read too many bytes (query: %s, limit: %s)"
+	limErrQuerierTooManyBytesTmpl = "query too large to execute on a single querier, either because parallelization is not enabled, the query is unshardable, or a shard query is too big to execute: (query: %s, limit: %s)"
 )
 
 var (
@@ -45,13 +53,16 @@ type Limits interface {
 	// TSDBMaxQueryParallelism returns the limit to the number of split queries the
 	// frontend will process in parallel for TSDB queries.
 	TSDBMaxQueryParallelism(context.Context, string) int
+	MaxQueryBytesRead(context.Context, string) int
+	MaxQuerierBytesRead(context.Context, string) int
 }
 
 type limits struct {
 	Limits
 	// Use pointers so nil value can indicate if the value was set.
 	splitDuration       *time.Duration
 	maxQueryParallelism *int
+	maxQueryBytesRead   *int
 }
 
 func (l limits) QuerySplitDuration(user string) time.Duration {
@@ -179,6 +190,176 @@ func (l limitsMiddleware) Do(ctx context.Context, r queryrangebase.Request) (que
 	return l.next.Do(ctx, r)
 }
 
+type querySizeLimiter struct {
+	logger            log.Logger
+	next              queryrangebase.Handler
+	statsHandler      queryrangebase.Handler
+	cfg               []config.PeriodConfig
+	maxLookBackPeriod time.Duration
+	limitFunc         func(context.Context, string) int
+	limitErrorTmpl    string
+}
+
+func newQuerySizeLimiter(
+	next queryrangebase.Handler,
+	cfg []config.PeriodConfig,
+	logger log.Logger,
+	limits Limits,
+	codec queryrangebase.Codec,
+	limitFunc func(context.Context, string) int,
+	limitErrorTmpl string,
+	statsHandler ...queryrangebase.Handler,
+) *querySizeLimiter {
+	q := &querySizeLimiter{
+		logger:         logger,
+		next:           next,
+		cfg:            cfg,
+		limitFunc:      limitFunc,
+		limitErrorTmpl: limitErrorTmpl,
+	}
+
+	q.statsHandler = next
+	if len(statsHandler) > 0 {
+		q.statsHandler = statsHandler[0]
+	}
+
+	// Parallelize the index stats requests, so it doesn't send a huge request to a single index-gw (i.e. {app=~".+"} for 30d).
+	// Indices are sharded by 24 hours, so we split the stats request in 24h intervals.
+	statsSplitTimeMiddleware := SplitByIntervalMiddleware(cfg, WithSplitByLimits(limits, 24*time.Hour), codec, splitByTime, nil)
+	q.statsHandler = statsSplitTimeMiddleware.Wrap(q.statsHandler)
+
+	// Get MaxLookBackPeriod from downstream engine. This is needed for instant limited queries at getStatsForMatchers
+	ng := logql.NewDownstreamEngine(logql.EngineOpts{LogExecutingQuery: false}, DownstreamHandler{next: next, limits: limits}, limits, logger)
+	q.maxLookBackPeriod = ng.Opts().MaxLookBackPeriod
+
+	return q
+}
+
+// NewQuerierSizeLimiterMiddleware creates a new Middleware that enforces query size limits after sharding and splitting.
+// The errorTemplate should format two strings: the bytes that would be read and the bytes limit.
+func NewQuerierSizeLimiterMiddleware(
+	cfg []config.PeriodConfig,
+	logger log.Logger,
+	limits Limits,
+	codec queryrangebase.Codec,
+	statsHandler ...queryrangebase.Handler,
+) queryrangebase.Middleware {
+	return queryrangebase.MiddlewareFunc(func(next queryrangebase.Handler) queryrangebase.Handler {
+		return newQuerySizeLimiter(next, cfg, logger, limits, codec, limits.MaxQuerierBytesRead, limErrQuerierTooManyBytesTmpl, statsHandler...)
+	})
+}
+
+// NewQuerySizeLimiterMiddleware creates a new Middleware that enforces query size limits.
+// The errorTemplate should format two strings: the bytes that would be read and the bytes limit.
+func NewQuerySizeLimiterMiddleware(
+	cfg []config.PeriodConfig,
+	logger log.Logger,
+	limits Limits,
+	codec queryrangebase.Codec,
+	statsHandler ...queryrangebase.Handler,
+) queryrangebase.Middleware {
+	return queryrangebase.MiddlewareFunc(func(next queryrangebase.Handler) queryrangebase.Handler {
+		return newQuerySizeLimiter(next, cfg, logger, limits, codec, limits.MaxQueryBytesRead, limErrQueryTooManyBytesTmpl, statsHandler...)
+	})
+}
+
+// getBytesReadForRequest returns the number of bytes that would be read for the query in r.
+// Since the query expression may contain multiple stream matchers, this function sums up the
+// bytes that will be read for each stream.
+// E.g. for the following query:
+//
+//	count_over_time({job="foo"}[5m]) / count_over_time({job="bar"}[5m] offset 10m)
+//
+// this function will sum the bytes read for each of the following streams, taking into account
+// individual intervals and offsets
+//   - {job="foo"}
+//   - {job="bar"}
+func (q *querySizeLimiter) getBytesReadForRequest(ctx context.Context, r queryrangebase.Request) (uint64, error) {
+	sp, ctx := spanlogger.NewWithLogger(ctx, q.logger, "querySizeLimiter.getBytesReadForRequest")
+	defer sp.Finish()
+
+	expr, err := syntax.ParseExpr(r.GetQuery())
+	if err != nil {
+		return 0, err
+	}
+
+	matcherGroups, err := syntax.MatcherGroups(expr)
+	if err != nil {
+		return 0, err
+	}
+
+	// TODO: Set concurrency dynamically as in shardResolverForConf?
+	start := time.Now()
+	const maxConcurrentIndexReq = 10
+	matcherStats, err := getStatsForMatchers(ctx, q.logger, q.statsHandler, model.Time(r.GetStart()), model.Time(r.GetEnd()), matcherGroups, maxConcurrentIndexReq, q.maxLookBackPeriod)
+	if err != nil {
+		return 0, err
+	}
+
+	combinedStats := stats.MergeStats(matcherStats...)
+
+	level.Debug(sp).Log(
+		append(
+			combinedStats.LoggingKeyValues(),
+			"msg", "queried index",
+			"type", "combined",
+			"len", len(matcherStats),
+			"max_parallelism", maxConcurrentIndexReq,
+			"duration", time.Since(start),
+			"total_bytes", strings.Replace(humanize.Bytes(combinedStats.Bytes), " ", "", 1),
+		)...,
+	)
+
+	return combinedStats.Bytes, nil
+}
+
+func (q *querySizeLimiter) getSchemaCfg(r queryrangebase.Request) (config.PeriodConfig, error) {
+	maxRVDuration, maxOffset, err := maxRangeVectorAndOffsetDuration(r.GetQuery())
+	if err != nil {
+		return config.PeriodConfig{}, errors.New("failed to get range-vector and offset duration: " + err.Error())
+	}
+
+	adjustedStart := int64(model.Time(r.GetStart()).Add(-maxRVDuration).Add(-maxOffset))
+	adjustedEnd := int64(model.Time(r.GetEnd()).Add(-maxOffset))
+
+	return ShardingConfigs(q.cfg).ValidRange(adjustedStart, adjustedEnd)
+}
+
+func (q *querySizeLimiter) Do(ctx context.Context, r queryrangebase.Request) (queryrangebase.Response, error) {
+	log, ctx := spanlogger.New(ctx, "query_size_limits")
+	defer log.Finish()
+
+	// Only support TSDB
+	schemaCfg, err := q.getSchemaCfg(r)
+	if err != nil {
+		return nil, httpgrpc.Errorf(http.StatusInternalServerError, "Failed to get schema config: %s", err.Error())
+	}
+	if schemaCfg.IndexType != config.TSDBType {
+		return q.next.Do(ctx, r)
+	}
+
+	tenantIDs, err := tenant.TenantIDs(ctx)
+	if err != nil {
+		return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error())
+	}
+
+	limitFuncCapture := func(id string) int { return q.limitFunc(ctx, id) }
+	if maxBytesRead := validation.SmallestPositiveNonZeroIntPerTenant(tenantIDs, limitFuncCapture); maxBytesRead > 0 {
+		bytesRead, err := q.getBytesReadForRequest(ctx, r)
+		if err != nil {
+			return nil, httpgrpc.Errorf(http.StatusInternalServerError, "Failed to get bytes read stats for query: %s", err.Error())
+		}
+
+		if bytesRead > uint64(maxBytesRead) {
+			statsBytesStr := humanize.Bytes(bytesRead)
+			maxBytesReadStr := humanize.Bytes(uint64(maxBytesRead))
+			return nil, httpgrpc.Errorf(http.StatusBadRequest, q.limitErrorTmpl, statsBytesStr, maxBytesReadStr)
+		}
+	}
+
+	return q.next.Do(ctx, r)
+}
+
 type seriesLimiter struct {
 	hashes map[uint64]struct{}
 	rw     sync.RWMutex