chore: treat datasample queries the same as limited queries (#13700)

Signed-off-by: Edward Welch <edward.welch@grafana.com>

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)

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)