chore: treat datasample queries the same as limited queries (grafana#…

…13700) Signed-off-by: Edward Welch <[email protected]>
cyriltovena · Jul 31, 2024 · fb6c63b · fb6c63b
1 parent 00b1f96
commit fb6c63b
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 2 deletions.
diff --git 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
@@ -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)