diff --git a/pkg/querier/queryrange/roundtrip.go b/pkg/querier/queryrange/roundtrip.go index 3e658f18c86f4..21961d279ae61 100644 --- a/pkg/querier/queryrange/roundtrip.go +++ b/pkg/querier/queryrange/roundtrip.go @@ -28,6 +28,7 @@ import ( "github.com/grafana/loki/v3/pkg/storage/config" "github.com/grafana/loki/v3/pkg/util" "github.com/grafana/loki/v3/pkg/util/constants" + "github.com/grafana/loki/v3/pkg/util/httpreq" logutil "github.com/grafana/loki/v3/pkg/util/log" "github.com/grafana/loki/v3/pkg/util/validation" ) @@ -400,8 +401,9 @@ func (r roundTripper) Do(ctx context.Context, req base.Request) (base.Response, return nil, httpgrpc.Errorf(http.StatusBadRequest, err.Error()) } - // Only filter expressions are query sharded - if !e.HasFilter() { + // Some queries we don't want to parallelize as aggressively, like limited queries and `datasample` queries + tags := httpreq.ExtractQueryTagsFromContext(ctx) + if !e.HasFilter() || strings.Contains(tags, "datasample") { return r.limited.Do(ctx, req) } return r.log.Do(ctx, req) diff --git a/pkg/util/httpreq/tags.go b/pkg/util/httpreq/tags.go index a0ed6c0d05386..96e4091c7db42 100644 --- a/pkg/util/httpreq/tags.go +++ b/pkg/util/httpreq/tags.go @@ -40,6 +40,12 @@ func ExtractQueryTagsFromHTTP(req *http.Request) string { return safeQueryTags.ReplaceAllString(tags, "_") } +func ExtractQueryTagsFromContext(ctx context.Context) string { + // if the cast fails then v will be an empty string + v, _ := ctx.Value(QueryTagsHTTPHeader).(string) + return v +} + func InjectQueryTags(ctx context.Context, tags string) context.Context { tags = safeQueryTags.ReplaceAllString(tags, "_") return context.WithValue(ctx, QueryTagsHTTPHeader, tags)