lint

cmd/mimir/config-descriptor.json

@@ -15177,6 +15177,17 @@
           "fieldType": "boolean",
           "fieldCategory": "experimental"
         },
+        {
+          "kind": "field",
+          "name": "prioritize_query_components",
+          "required": false,
+          "desc": "Prioritize rotation through query components over rotation through tenants during dequeueing.",
+          "fieldValue": null,
+          "fieldDefaultValue": false,
+          "fieldFlag": "query-scheduler.prioritize-query-components",
+          "fieldType": "boolean",
+          "fieldCategory": "experimental"
+        },
         {
           "kind": "field",
           "name": "querier_forget_delay",

cmd/mimir/help-all.txt.tmpl

@@ -2131,6 +2131,8 @@ Usage of ./cmd/mimir/mimir:
     	Maximum number of outstanding requests per tenant per query-scheduler. In-flight requests above this limit will fail with HTTP response status code 429. (default 100)
   -query-scheduler.max-used-instances int
     	The maximum number of query-scheduler instances to use, regardless how many replicas are running. This option can be set only when -query-scheduler.service-discovery-mode is set to 'ring'. 0 to use all available query-scheduler instances.
+  -query-scheduler.prioritize-query-components
+    	[experimental] Prioritize rotation through query components over rotation through tenants during dequeueing.
   -query-scheduler.querier-forget-delay duration
     	[experimental] If a querier disconnects without sending notification about graceful shutdown, the query-scheduler will keep the querier in the tenant's shard until the forget delay has passed. This feature is useful to reduce the blast radius when shuffle-sharding is enabled.
   -query-scheduler.ring.consul.acl-token string

docs/sources/mimir/configure/configuration-parameters/index.md

@@ -1611,6 +1611,11 @@ The `query_scheduler` block configures the query-scheduler.
 # CLI flag: -query-scheduler.additional-query-queue-dimensions-enabled
 [additional_query_queue_dimensions_enabled: <boolean> | default = false]
 
+# (experimental) Prioritize rotation through query components over rotation
+# through tenants during dequeueing.
+# CLI flag: -query-scheduler.prioritize-query-components
+[prioritize_query_components: <boolean> | default = false]
+
 # (experimental) If a querier disconnects without sending notification about
 # graceful shutdown, the query-scheduler will keep the querier in the tenant's
 # shard until the forget delay has passed. This feature is useful to reduce the

pkg/scheduler/queue/queue_test.go

@@ -9,7 +9,6 @@ import (
 	"context"
 	"errors"
 	"fmt"
-	promtest "github.com/prometheus/client_golang/prometheus/testutil"
 	"math/rand"
 	"strconv"
 	"strings"
@@ -22,6 +21,7 @@ import (
 	"github.com/grafana/dskit/services"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+	promtest "github.com/prometheus/client_golang/prometheus/testutil"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 	"golang.org/x/sync/errgroup"

pkg/scheduler/queue/tenant_querier_assignment.go

@@ -96,14 +96,6 @@ type tenantQuerierAssignments struct {
 	currentQuerier   *QuerierID
 }
 
-func (tqa *tenantQuerierAssignments) getTenant(tenantID TenantID) (*queueTenant, error) {
-	if tenantID == emptyTenantID {
-		return nil, ErrInvalidTenantID
-	}
-	tenant := tqa.tenantsByID[tenantID]
-	return tenant, nil
-}
-
 // createOrUpdateTenant creates or updates a tenant into the tenant-querier assignment state.
 //
 // New tenants are added to the tenant order list and tenant-querier shards are shuffled if needed.

pkg/scheduler/queue/tenant_queues_test.go

@@ -8,13 +8,14 @@ package queue
 import (
 	"context"
 	"fmt"
-	"github.com/grafana/dskit/httpgrpc"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
 	"math"
 	"math/rand"
 	"testing"
 	"time"
+
+	"github.com/grafana/dskit/httpgrpc"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
 )
 
 func (qb *queueBroker) enqueueObjectsForTests(tenantID TenantID, numObjects int) error {
@@ -111,7 +112,7 @@ func TestQueues(t *testing.T) {
 		{buildExpectedObject(tenantTwo.tenantID, 2), tenantTwo},
 		{buildExpectedObject(tenantOne.tenantID, 5), tenantOne},
 	}
-	lastTenantIndex = assertExpectedValuesOnDequeue(t, qb, -1, "querier-2", expectedDequeueVals)
+	_ = assertExpectedValuesOnDequeue(t, qb, -1, "querier-2", expectedDequeueVals)
 
 	//[one two three]
 	// confirm fifo by adding a third tenant queue and iterating to it
@@ -310,7 +311,7 @@ func TestQueuesOnTerminatingQuerier(t *testing.T) {
 
 	// After disconnecting querier-2, it's expected to own no queue.
 	qb.tenantQuerierAssignments.removeQuerier("querier-2")
-	req, tenant, qTwolastTenantIndex, err = qb.dequeueRequestForQuerier(qTwolastTenantIndex, "querier-2")
+	req, tenant, _, err = qb.dequeueRequestForQuerier(qTwolastTenantIndex, "querier-2")
 	assert.Nil(t, req)
 	assert.Nil(t, tenant)
 	assert.Equal(t, ErrQuerierShuttingDown, err)

pkg/scheduler/queue/tree_queue.go

@@ -118,9 +118,7 @@ func newNode(name string, depth int, da DequeueAlgorithm) (*Node, error) {
 	if da == nil {
 		return nil, fmt.Errorf("cannot create a node without a defined DequeueAlgorithm")
 	}
-	switch da.(type) {
-	case *tenantQuerierAssignments:
-		tqa := da.(*tenantQuerierAssignments)
+	if tqa, ok := da.(*tenantQuerierAssignments); ok {
 		tqa.tenantOrderIndex = localQueueIndex - 1 // start from -2 so that we first check local queue
 		if tqa.tenantNodes == nil {
 			tqa.tenantNodes = map[string][]*Node{}

pkg/scheduler/queue/tree_queue_test.go

@@ -5,10 +5,9 @@ package queue
 import (
 	"fmt"
 	"strings"
+	"testing"
 
-	//"fmt"
 	"github.com/stretchr/testify/require"
-	"testing"
 )
 
 func newTenantQuerierAssignments() *tenantQuerierAssignments {
@@ -620,12 +619,22 @@ func Test_ChangeTenantQuerierAssignments(t *testing.T) {
 	tree, err := NewTree(state, &roundRobinState{}, &roundRobinState{})
 	require.NoError(t, err)
 
-	err = tree.EnqueueBackByPath(QueuePath{"tenant-1", "query-component-1"}, "query-1")
-	err = tree.EnqueueBackByPath(QueuePath{"tenant-2", "query-component-1"}, "query-2")
-	err = tree.EnqueueBackByPath(QueuePath{"tenant-2", "query-component-1"}, "query-3")
-	err = tree.EnqueueBackByPath(QueuePath{"tenant-2", "query-component-1"}, "query-4")
-	err = tree.EnqueueBackByPath(QueuePath{"tenant-3", "query-component-1"}, "query-5")
-	require.NoError(t, err)
+	type enqueueObj struct {
+		obj  any
+		path QueuePath
+	}
+	enqueueObjs := []enqueueObj{
+		{"query-1", QueuePath{"tenant-1", "query-component-1"}},
+		{"query-2", QueuePath{"tenant-2", "query-component-1"}},
+		{"query-3", QueuePath{"tenant-2", "query-component-1"}},
+		{"query-4", QueuePath{"tenant-2", "query-component-1"}},
+		{"query-5", QueuePath{"tenant-3", "query-component-1"}},
+	}
+
+	for _, eo := range enqueueObjs {
+		err = tree.EnqueueBackByPath(eo.path, eo.obj)
+		require.NoError(t, err)
+	}
 
 	querier1 := QuerierID("querier-1")
 	querier2 := QuerierID("querier-2")

pkg/scheduler/scheduler.go

@@ -95,7 +95,7 @@ type connectedFrontend struct {
 type Config struct {
 	MaxOutstandingPerTenant               int           `yaml:"max_outstanding_requests_per_tenant"`
 	AdditionalQueryQueueDimensionsEnabled bool          `yaml:"additional_query_queue_dimensions_enabled" category:"experimental"`
-	PrioritizeQueryComponentsEnabled      bool          `yaml:"prioritize_query_components_enabled" category:"experimental"`
+	PrioritizeQueryComponents             bool          `yaml:"prioritize_query_components" category:"experimental"`
 	QuerierForgetDelay                    time.Duration `yaml:"querier_forget_delay" category:"experimental"`
 
 	GRPCClientConfig grpcclient.Config         `yaml:"grpc_client_config" doc:"description=This configures the gRPC client used to report errors back to the query-frontend."`
@@ -105,6 +105,7 @@ type Config struct {
 func (cfg *Config) RegisterFlags(f *flag.FlagSet, logger log.Logger) {
 	f.IntVar(&cfg.MaxOutstandingPerTenant, "query-scheduler.max-outstanding-requests-per-tenant", 100, "Maximum number of outstanding requests per tenant per query-scheduler. In-flight requests above this limit will fail with HTTP response status code 429.")
 	f.BoolVar(&cfg.AdditionalQueryQueueDimensionsEnabled, "query-scheduler.additional-query-queue-dimensions-enabled", false, "Enqueue query requests with additional queue dimensions to split tenant request queues into subqueues. This enables separate requests to proceed from a tenant's subqueues even when other subqueues are blocked on slow query requests. Must be set on both query-frontend and scheduler to take effect. (default false)")
+	f.BoolVar(&cfg.PrioritizeQueryComponents, "query-scheduler.prioritize-query-components", false, "Prioritize rotation through query components over rotation through tenants during dequeueing.")
 	f.DurationVar(&cfg.QuerierForgetDelay, "query-scheduler.querier-forget-delay", 0, "If a querier disconnects without sending notification about graceful shutdown, the query-scheduler will keep the querier in the tenant's shard until the forget delay has passed. This feature is useful to reduce the blast radius when shuffle-sharding is enabled.")
 
 	cfg.GRPCClientConfig.RegisterFlagsWithPrefix("query-scheduler.grpc-client-config", f)
@@ -157,7 +158,7 @@ func NewScheduler(cfg Config, limits Limits, log log.Logger, registerer promethe
 		s.log,
 		cfg.MaxOutstandingPerTenant,
 		cfg.AdditionalQueryQueueDimensionsEnabled,
-		cfg.PrioritizeQueryComponentsEnabled,
+		cfg.PrioritizeQueryComponents,
 		cfg.QuerierForgetDelay,
 		s.queueLength,
 		s.discardedRequests,