feat(unified-storage): round robin queue and scheduler package (#105544)

pkg/util/scheduler/queue.go

@@ -0,0 +1,363 @@
+package scheduler
+
+import (
+	"container/list"
+	"context"
+	"errors"
+	"time"
+
+	"github.com/grafana/dskit/services"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+)
+
+const (
+	// DefaultMaxSizePerTenant is the default maximum number of items per tenant in the queue.
+	DefaultMaxSizePerTenant = 100
+)
+
+var ErrQueueClosed = errors.New("queue closed")
+var ErrTenantQueueFull = errors.New("tenant queue full")
+var ErrNilRunnable = errors.New("cannot enqueue nil runnable")
+var ErrMissingTenantID = errors.New("item requires TenantID")
+
+type tenantQueue struct {
+	id       string
+	items    []func(ctx context.Context)
+	isActive bool
+}
+
+func (tq *tenantQueue) len() int {
+	return len(tq.items)
+}
+func (tq *tenantQueue) clear() {
+	tq.items = nil
+	tq.isActive = false
+}
+func (tq *tenantQueue) isEmpty() bool {
+	return len(tq.items) == 0
+}
+func (tq *tenantQueue) isFull(maxSize int) bool {
+	return maxSize > 0 && len(tq.items) >= maxSize
+}
+func (tq *tenantQueue) addItem(runnable func(ctx context.Context)) {
+	tq.items = append(tq.items, runnable)
+}
+
+type enqueueRequest struct {
+	tenantID string
+	runnable func(ctx context.Context)
+	respChan chan error
+}
+
+type dequeueRequest struct {
+	respChan chan dequeueResponse
+}
+
+type dequeueResponse struct {
+	runnable func(ctx context.Context)
+	err      error
+}
+
+type lenRequest struct {
+	respChan chan int
+}
+
+type activeTenantsLenRequest struct {
+	respChan chan int
+}
+
+type NoopQueue struct{}
+
+func (*NoopQueue) Enqueue(ctx context.Context, _ string, runnable func(ctx context.Context)) error {
+	runnable(ctx)
+	return nil
+}
+
+func NewNoopQueue() *NoopQueue {
+	return &NoopQueue{}
+}
+
+// Queue implements a multi-tenant qos with round-robin fairness using a dispatcher goroutine.
+type Queue struct {
+	services.Service
+
+	enqueueChan           chan enqueueRequest
+	dequeueChan           chan dequeueRequest
+	lenChan               chan lenRequest
+	activeTenantsLenChan  chan activeTenantsLenRequest
+	dispatcherStoppedChan chan struct{}
+
+	// tenantQueues stores the queues for each tenant
+	tenantQueues map[string]*tenantQueue
+	// activeTenants is a list of tenants with items in their queues
+	// used for round-robin dequeueing
+	activeTenants *list.List
+	// pendingDequeueRequests is a list of dequeue requests waiting for items
+	// used for notifying when items are available
+	pendingDequeueRequests *list.List
+	// maxSizePerTenant is the maximum number of items per tenant
+	maxSizePerTenant int
+
+	// Metrics
+	queueLength       *prometheus.GaugeVec
+	discardedRequests *prometheus.CounterVec
+	enqueueDuration   prometheus.Histogram
+}
+
+type QueueOptions struct {
+	MaxSizePerTenant int
+	Registerer       prometheus.Registerer
+}
+
+// NewQueue creates a new Queue and starts its dispatcher goroutine.
+func NewQueue(opts *QueueOptions) *Queue {
+	if opts.MaxSizePerTenant <= 0 {
+		opts.MaxSizePerTenant = DefaultMaxSizePerTenant
+	}
+
+	q := &Queue{
+		enqueueChan:           make(chan enqueueRequest),
+		dequeueChan:           make(chan dequeueRequest),
+		lenChan:               make(chan lenRequest),
+		activeTenantsLenChan:  make(chan activeTenantsLenRequest),
+		dispatcherStoppedChan: make(chan struct{}),
+
+		tenantQueues:           make(map[string]*tenantQueue),
+		activeTenants:          list.New(),
+		pendingDequeueRequests: list.New(),
+		maxSizePerTenant:       opts.MaxSizePerTenant,
+	}
+
+	q.queueLength = promauto.With(opts.Registerer).NewGaugeVec(prometheus.GaugeOpts{
+		Name: "queue_length",
+		Help: "Number of items in the queue",
+	}, []string{"namespace"})
+	q.discardedRequests = promauto.With(opts.Registerer).NewCounterVec(prometheus.CounterOpts{
+		Name: "discarded_requests_total",
+		Help: "Total number of discarded requests",
+	}, []string{"namespace", "reason"})
+	q.enqueueDuration = promauto.With(opts.Registerer).NewHistogram(prometheus.HistogramOpts{
+		Name:    "enqueue_duration_seconds",
+		Help:    "Duration of enqueue operation in seconds",
+		Buckets: prometheus.DefBuckets,
+	})
+
+	q.Service = services.NewBasicService(nil, q.dispatcherLoop, q.stopping)
+
+	return q
+}
+
+func (q *Queue) scheduleRoundRobin() {
+	// Process as long as we have both pending requests and active tenants
+	for {
+		// Get the front elements of both lists
+		reqElem := q.pendingDequeueRequests.Front()
+		tenantElem := q.activeTenants.Front()
+
+		// Exit when either list is empty
+		if reqElem == nil || tenantElem == nil {
+			break
+		}
+
+		req := reqElem.Value.(*dequeueRequest)
+		tq := tenantElem.Value.(*tenantQueue)
+
+		// Get and deliver the runnable item
+		item := tq.items[0]
+		req.respChan <- dequeueResponse{runnable: item, err: nil}
+
+		// Update bookkeeping
+		q.pendingDequeueRequests.Remove(reqElem)
+		tq.items = tq.items[1:]
+
+		// Update metrics
+		q.queueLength.WithLabelValues(tq.id).Set(float64(tq.len()))
+
+		// Round-robin: move to back if tenant still has items, otherwise remove
+		if tq.isEmpty() {
+			tq.clear()
+			q.activeTenants.Remove(tenantElem)
+		} else {
+			q.activeTenants.MoveToBack(tenantElem)
+		}
+	}
+}
+
+func (q *Queue) handleEnqueueRequest(req enqueueRequest) {
+	tq, exists := q.tenantQueues[req.tenantID]
+	if !exists {
+		tq = &tenantQueue{
+			id:    req.tenantID,
+			items: make([]func(ctx context.Context), 0, 8),
+		}
+		q.tenantQueues[req.tenantID] = tq
+	}
+
+	if tq.isFull(q.maxSizePerTenant) {
+		q.discardedRequests.WithLabelValues(req.tenantID, "queue_full").Inc()
+		req.respChan <- ErrTenantQueueFull
+		return
+	}
+
+	tq.addItem(req.runnable)
+	q.queueLength.WithLabelValues(req.tenantID).Set(float64(len(tq.items)))
+
+	if !tq.isActive {
+		q.activeTenants.PushBack(tq)
+		tq.isActive = true
+	}
+
+	req.respChan <- nil
+}
+
+func (q *Queue) handleDequeueRequest(req dequeueRequest) {
+	q.pendingDequeueRequests.PushBack(&req)
+}
+
+func (q *Queue) handleLenRequest(req lenRequest) {
+	total := 0
+	for _, tq := range q.tenantQueues {
+		total += tq.len()
+	}
+	req.respChan <- total
+}
+
+func (q *Queue) dispatcherLoop(ctx context.Context) error {
+	defer close(q.dispatcherStoppedChan)
+
+	for {
+		q.scheduleRoundRobin()
+
+		select {
+		case <-ctx.Done():
+			return nil
+
+		case req := <-q.enqueueChan:
+			q.handleEnqueueRequest(req)
+
+		case req := <-q.dequeueChan:
+			q.handleDequeueRequest(req)
+
+		case req := <-q.lenChan:
+			q.handleLenRequest(req)
+
+		case req := <-q.activeTenantsLenChan:
+			req.respChan <- q.activeTenants.Len()
+		}
+	}
+}
+
+// Enqueue adds a work item to the appropriate tenant's qos.
+// It blocks only if the dispatcher is busy or the tenant queue is full.
+func (q *Queue) Enqueue(ctx context.Context, tenantID string, runnable func(ctx context.Context)) error {
+	if runnable == nil {
+		return ErrNilRunnable
+	}
+	if tenantID == "" {
+		return ErrMissingTenantID
+	}
+
+	if q.State() != services.Running {
+		return ErrQueueClosed
+	}
+
+	start := time.Now()
+
+	respChan := make(chan error, 1)
+	req := enqueueRequest{
+		tenantID: tenantID,
+		runnable: runnable,
+		respChan: respChan,
+	}
+
+	var err error
+	select {
+	case q.enqueueChan <- req:
+		err = <-respChan
+		q.enqueueDuration.Observe(time.Since(start).Seconds())
+	case <-q.dispatcherStoppedChan:
+		q.discardedRequests.WithLabelValues(tenantID, "dispatcher_stopped").Inc()
+		err = ErrQueueClosed
+	case <-ctx.Done():
+		q.discardedRequests.WithLabelValues(tenantID, "context_canceled").Inc()
+		err = ctx.Err()
+	}
+
+	return err
+}
+
+// Dequeue removes and returns a work item from the qos using linked-list round-robin.
+// It blocks until an item is available for any tenant, the queue is closed,
+// or the context is cancelled.
+func (q *Queue) Dequeue(ctx context.Context) (func(ctx context.Context), error) {
+	if q.State() != services.Running {
+		return nil, ErrQueueClosed
+	}
+
+	respChan := make(chan dequeueResponse, 1)
+	req := dequeueRequest{
+		respChan: respChan,
+	}
+
+	select {
+	case q.dequeueChan <- req:
+		select {
+		case resp := <-respChan:
+			return resp.runnable, resp.err
+		case <-ctx.Done():
+			return nil, ctx.Err()
+		}
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	case <-q.dispatcherStoppedChan:
+		return nil, ErrQueueClosed
+	}
+}
+
+// Len returns the total number of items across all tenants in the queue.
+func (q *Queue) Len() int {
+	respChan := make(chan int, 1)
+	req := lenRequest{respChan: respChan}
+
+	select {
+	case q.lenChan <- req:
+		select {
+		case count := <-respChan:
+			return count
+		case <-q.dispatcherStoppedChan:
+			return 0
+		}
+	case <-q.dispatcherStoppedChan:
+		return 0
+	}
+}
+
+// ActiveTenantsLen returns the number of tenants with items currently in the queue.
+func (q *Queue) ActiveTenantsLen() int {
+	respChan := make(chan int, 1)
+	req := activeTenantsLenRequest{respChan: respChan}
+
+	select {
+	case q.activeTenantsLenChan <- req:
+		select {
+		case count := <-respChan:
+			return count
+		case <-q.dispatcherStoppedChan:
+			return 0
+		}
+	case <-q.dispatcherStoppedChan:
+		return 0
+	}
+}
+
+func (q *Queue) stopping(_ error) error {
+	q.queueLength.Reset()
+	q.discardedRequests.Reset()
+	for _, tq := range q.tenantQueues {
+		tq.clear()
+	}
+	q.activeTenants.Init()
+	q.pendingDequeueRequests.Init()
+	return nil
+}

pkg/util/scheduler/queue_test.go

@@ -0,0 +1,601 @@
+package scheduler
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/grafana/dskit/services"
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/stretchr/testify/require"
+)
+
+func QueueOptionsWithDefaults(opts *QueueOptions) *QueueOptions {
+	if opts == nil {
+		opts = &QueueOptions{}
+	}
+	if opts.MaxSizePerTenant <= 0 {
+		opts.MaxSizePerTenant = 10
+	}
+	if opts.Registerer == nil {
+		opts.Registerer = prometheus.NewRegistry()
+	}
+	return opts
+}
+
+//nolint:gocyclo
+func TestQueue(t *testing.T) {
+	t.Parallel()
+
+	t.Run("SimpleEnqueueAndDequeue", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, services.StartAndAwaitRunning(ctx, q))
+
+		defer func() {
+			require.NoError(t, services.StopAndAwaitTerminated(context.Background(), q))
+		}()
+
+		var wg sync.WaitGroup
+		const numItems = 5
+		const tenantID = "tenant-a"
+		var processed atomic.Int32
+
+		// Enqueue items
+		for i := 0; i < numItems; i++ {
+			err := q.Enqueue(ctx, tenantID, func(ctx context.Context) {
+				processed.Add(1)
+			})
+			require.NoError(t, err, "Enqueue should succeed")
+		}
+		require.Equal(t, numItems, q.Len(), "Queue length after enqueue")
+		require.Equal(t, 1, q.ActiveTenantsLen(), "Active tenants after enqueue")
+
+		// Dequeue items
+		for i := 0; i < numItems; i++ {
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				dequeueCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
+				defer cancel()
+				runnable, err := q.Dequeue(dequeueCtx)
+				require.NoError(t, err, "Dequeue should succeed")
+				require.NotNil(t, runnable, "Dequeued runnable should not be nil")
+				runnable(ctx)
+			}()
+		}
+
+		wg.Wait()
+
+		// Check that all items were processed
+		require.Equal(t, numItems, int(processed.Load()), "All items should have been processed")
+
+		// Let's simplify the dequeue check: Dequeue sequentially after enqueueing.
+		qSimple := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, services.StartAndAwaitRunning(ctx, qSimple))
+
+		for i := 0; i < numItems; i++ {
+			err := qSimple.Enqueue(ctx, tenantID, func(ctx context.Context) {})
+			require.NoError(t, err)
+		}
+		require.Equal(t, numItems, qSimple.Len(), "Queue length after enqueue (simple)")
+		require.Equal(t, 1, qSimple.ActiveTenantsLen(), "Active tenants after enqueue (simple)")
+
+		for i := 0; i < numItems; i++ {
+			dequeueCtx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
+			runnable, err := qSimple.Dequeue(dequeueCtx)
+			cancel() // Cancel context after use
+			require.NoError(t, err, "Dequeue %d should succeed (simple)", i)
+			require.NotNil(t, runnable, "Dequeued runnable %d should not be nil (simple)", i)
+		}
+
+		require.Equal(t, 0, qSimple.Len(), "Queue length after dequeue (simple)")
+		require.Equal(t, 0, qSimple.ActiveTenantsLen(), "Active tenants after dequeue (simple)")
+
+		// Check dequeue on empty queue times out
+		dequeueCtx, cancel := context.WithTimeout(ctx, 50*time.Millisecond)
+		_, err := qSimple.Dequeue(dequeueCtx)
+		cancel()
+		require.ErrorIs(t, err, context.DeadlineExceeded, "Dequeue on empty queue should time out")
+
+		// Stop the queue
+		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), qSimple))
+	})
+
+	t.Run("RoundRobinBetweenTenants", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, services.StartAndAwaitRunning(ctx, q))
+
+		tenantA := "tenant-a"
+		tenantB := "tenant-b"
+
+		// We'll use a very small test with just 2 items per tenant
+		// to reduce the chance of timeouts or other issues
+
+		// Enqueue items in tenant order: A, B, A, B
+		// Each item will record its tenant ID when executed
+		var results []string
+		var resultsMu sync.Mutex
+
+		makeRunnable := func(id string) func(ctx context.Context) {
+			return func(ctx context.Context) {
+				resultsMu.Lock()
+				results = append(results, id)
+				resultsMu.Unlock()
+			}
+		}
+
+		require.NoError(t, q.Enqueue(ctx, tenantA, makeRunnable(tenantA)))
+		require.NoError(t, q.Enqueue(ctx, tenantA, makeRunnable(tenantA)))
+		require.NoError(t, q.Enqueue(ctx, tenantB, makeRunnable(tenantB)))
+		require.NoError(t, q.Enqueue(ctx, tenantB, makeRunnable(tenantB)))
+
+		// Verify queue state
+		require.Equal(t, 4, q.Len(), "Queue should have 4 items")
+		require.Equal(t, 2, q.ActiveTenantsLen(), "Should have 2 active tenants")
+
+		// Use a longer timeout to handle CI environment variability
+		dequeueTimeout := 3 * time.Second
+
+		// Dequeue and execute the four items
+		for i := 0; i < 4; i++ {
+			// Use a more reliable context with longer timeout for CI environments
+			dequeueCtx, cancel := context.WithTimeout(ctx, dequeueTimeout)
+			runnable, err := q.Dequeue(dequeueCtx)
+			if err != nil {
+				t.Logf("Queue state: Len=%d, ActiveTenantsLen=%d", q.Len(), q.ActiveTenantsLen())
+			}
+
+			cancel()
+			require.NoError(t, err, "Dequeue %d should succeed", i)
+			require.NotNil(t, runnable, "Dequeued runnable %d should not be nil", i)
+			runnable(ctx) // Execute to record the tenant ID
+		}
+
+		// Check execution order - should alternate between tenants
+		resultsMu.Lock()
+		require.Equal(t, []string{tenantA, tenantB, tenantA, tenantB}, results)
+		resultsMu.Unlock()
+
+		// Verify queue is now empty
+		require.Equal(t, 0, q.Len())
+		require.Equal(t, 0, q.ActiveTenantsLen())
+	})
+
+	t.Run("TenantQueueFullError", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		q := NewQueue(QueueOptionsWithDefaults(&QueueOptions{MaxSizePerTenant: 2}))
+		require.NoError(t, q.StartAsync(context.Background()), "Queue should start")
+		require.NoError(t, q.AwaitRunning(context.Background()), "Queue should be running")
+
+		tenantID := "tenant-limited"
+
+		// Enqueue up to the limit
+		err := q.Enqueue(ctx, tenantID, func(ctx context.Context) {})
+		require.NoError(t, err)
+		err = q.Enqueue(ctx, tenantID, func(ctx context.Context) {})
+		require.NoError(t, err)
+
+		require.Equal(t, 2, q.Len())
+		require.Equal(t, 1, q.ActiveTenantsLen())
+
+		// Enqueue one more, expect error
+		err = q.Enqueue(ctx, tenantID, func(ctx context.Context) {})
+		require.ErrorIs(t, err, ErrTenantQueueFull, "Expected ErrTenantQueueFull")
+
+		// Len should still be 2
+		require.Equal(t, 2, q.Len())
+
+		// Dequeue one item
+		dequeueCtx, cancel := context.WithTimeout(ctx, 50*time.Millisecond)
+		_, err = q.Dequeue(dequeueCtx)
+		cancel()
+		require.NoError(t, err)
+		require.Equal(t, 1, q.Len())
+
+		// Now enqueue should succeed again
+		err = q.Enqueue(ctx, tenantID, func(ctx context.Context) {})
+		require.NoError(t, err, "Enqueue should succeed after dequeueing one item")
+		require.Equal(t, 2, q.Len(), "Length should be back to 2")
+	})
+
+	t.Run("DequeueContextCancellation", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, services.StartAndAwaitRunning(ctx, q))
+		require.NoError(t, q.AwaitRunning(ctx), "Queue should be running")
+
+		// Create an already canceled context instead of using timeout
+		cancelCtx, cancel := context.WithCancel(context.Background())
+		cancel() // Cancel immediately
+
+		// This should return with context.Canceled error immediately
+		runnable, err := q.Dequeue(cancelCtx)
+
+		// Verify the result - should be context.Canceled, not DeadlineExceeded
+		require.Nil(t, runnable, "Runnable should be nil on context cancellation")
+		require.ErrorIs(t, err, context.Canceled, "Expected context canceled error")
+	})
+
+	t.Run("DequeueWithTimeout", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, services.StartAndAwaitRunning(ctx, q))
+		require.NoError(t, q.AwaitRunning(ctx), "Queue should be running")
+
+		ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+		defer cancel()
+
+		// This should timeout since nothing is in the queue
+		runnable, err := q.Dequeue(ctx)
+		require.Nil(t, runnable, "Runnable should be nil on timeout")
+		require.ErrorIs(t, err, context.DeadlineExceeded, "Expected context deadline exceeded error")
+	})
+
+	t.Run("EnqueueAfterStopped", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, services.StartAndAwaitRunning(ctx, q))
+		require.NoError(t, q.AwaitRunning(ctx), "Queue should be running")
+
+		// Stop the queue first
+		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), q))
+
+		// Now try to enqueue - should return ErrQueueClosed
+		err := q.Enqueue(context.Background(), "tenant-id", func(ctx context.Context) {})
+		require.ErrorIs(t, err, ErrQueueClosed, "Enqueue after Stop should return ErrQueueClosed")
+	})
+
+	t.Run("EnqueueBeforeStarted", func(t *testing.T) {
+		t.Parallel()
+		ctx := context.Background()
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		err := q.Enqueue(ctx, "tenant-id", func(ctx context.Context) {})
+		require.ErrorIs(t, err, ErrQueueClosed, "Enqueue before Start should return ErrQueueClosed")
+	})
+
+	t.Run("DequeueAfterStopped", func(t *testing.T) {
+		t.Parallel()
+
+		ctx := context.Background()
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, services.StartAndAwaitRunning(ctx, q))
+		require.NoError(t, q.AwaitRunning(ctx), "Queue should be running")
+
+		// Stop the queue first
+		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), q))
+
+		// Now try to dequeue - should return ErrQueueClosed
+		runnable, err := q.Dequeue(context.Background())
+		require.Nil(t, runnable, "Runnable should be nil after Stop")
+		require.ErrorIs(t, err, ErrQueueClosed, "Dequeue after Stop should return ErrQueueClosed")
+	})
+
+	t.Run("DequeueBeforeStarted", func(t *testing.T) {
+		t.Parallel()
+		ctx := context.Background()
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		_, err := q.Dequeue(ctx)
+		require.ErrorIs(t, err, ErrQueueClosed, "Dequeue before Start should return ErrQueueClosed")
+	})
+
+	t.Run("ConcurrentEnqueuersAndDequeuers", func(t *testing.T) {
+		t.Parallel()
+
+		q := NewQueue(QueueOptionsWithDefaults(&QueueOptions{MaxSizePerTenant: 100}))
+		require.NoError(t, q.StartAsync(context.Background()), "Queue should start")
+		require.NoError(t, q.AwaitRunning(context.Background()), "Queue should be running")
+
+		const numProducers = 5
+		const numConsumers = 3
+		const itemsPerProducer = 20
+		const totalItems = numProducers * itemsPerProducer
+
+		// Track items to verify all are processed
+		processedItems := make(map[string]int)
+		var mu sync.Mutex
+		var wg sync.WaitGroup
+
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		// Start consumers first (they'll block until items are available)
+		for i := 0; i < numConsumers; i++ {
+			wg.Add(1)
+			go func(consumerID int) {
+				defer wg.Done()
+
+				for {
+					runnable, err := q.Dequeue(ctx)
+					if err != nil {
+						if errors.Is(err, context.DeadlineExceeded) {
+							return
+						}
+						require.NoError(t, err)
+						return
+					}
+
+					// Execute the runnable which will update our tracking
+					runnable(ctx)
+
+					// Check if we've processed all expected items
+					mu.Lock()
+					totalProcessed := 0
+					for _, count := range processedItems {
+						totalProcessed += count
+					}
+					done := totalProcessed >= totalItems
+					mu.Unlock()
+
+					if done {
+						return
+					}
+				}
+			}(i)
+		}
+
+		// Start producers
+		for i := 0; i < numProducers; i++ {
+			wg.Add(1)
+			go func(producerID int) {
+				defer wg.Done()
+				tenantID := fmt.Sprintf("tenant-%d", producerID%3) // Use 3 different tenants
+
+				for j := 0; j < itemsPerProducer; j++ {
+					itemID := fmt.Sprintf("p%d-item%d", producerID, j)
+
+					err := q.Enqueue(ctx, tenantID, func(ctx context.Context) {
+						mu.Lock()
+						processedItems[itemID] = 1
+						mu.Unlock()
+					})
+
+					if err != nil {
+						// Context might have been canceled if test is slow
+						if errors.Is(err, context.DeadlineExceeded) ||
+							errors.Is(err, ErrQueueClosed) {
+							return
+						}
+						require.NoError(t, err)
+					}
+
+					// Small sleep to reduce contention
+					time.Sleep(time.Millisecond)
+				}
+			}(i)
+		}
+
+		// Wait for all goroutines to finish
+		wg.Wait()
+
+		// Check that all items were processed
+		mu.Lock()
+		require.Equal(t, totalItems, len(processedItems),
+			"All enqueued items should have been processed")
+		mu.Unlock()
+
+		// Verify queue is now empty
+		require.Equal(t, 0, q.Len(), "Queue should be empty after processing all items")
+		require.Equal(t, 0, q.ActiveTenantsLen(), "No active tenants should remain after processing")
+		// Stop the queue
+		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), q))
+	})
+
+	t.Run("SlowDequeuerHandling", func(t *testing.T) {
+		t.Parallel()
+
+		q := NewQueue(QueueOptionsWithDefaults(&QueueOptions{MaxSizePerTenant: 5}))
+		require.NoError(t, q.StartAsync(context.Background()), "Queue should start")
+		require.NoError(t, q.AwaitRunning(context.Background()), "Queue should be running")
+
+		ctx := context.Background()
+		tenantA := "tenant-a"
+		tenantB := "tenant-b"
+		tenantC := "tenant-c"
+
+		// Create channels to track execution order
+		completionOrder := make(chan string, 10)
+		var wg sync.WaitGroup
+
+		// Enqueue a slow item for tenant A
+		wg.Add(1)
+		err := q.Enqueue(ctx, tenantA, func(ctx context.Context) {
+			defer wg.Done()
+			time.Sleep(300 * time.Millisecond) // Simulate slow processing
+			completionOrder <- "A-slow"
+		})
+		require.NoError(t, err)
+
+		// Enqueue regular items for other tenants
+		for i := 0; i < 2; i++ {
+			wg.Add(1)
+			err := q.Enqueue(ctx, tenantB, func(ctx context.Context) {
+				defer wg.Done()
+				completionOrder <- fmt.Sprintf("B-%d", i)
+			})
+			require.NoError(t, err)
+
+			wg.Add(1)
+			err = q.Enqueue(ctx, tenantC, func(ctx context.Context) {
+				defer wg.Done()
+				completionOrder <- fmt.Sprintf("C-%d", i)
+			})
+			require.NoError(t, err)
+		}
+
+		// Enqueue another item for tenant A
+		wg.Add(1)
+		err = q.Enqueue(ctx, tenantA, func(ctx context.Context) {
+			defer wg.Done()
+			completionOrder <- "A-fast"
+		})
+		require.NoError(t, err)
+
+		// Start multiple dequeuer goroutines
+		for i := 0; i < 3; i++ {
+			go func() {
+				for j := 0; j < 2; j++ { // Each will dequeue 2 items
+					dequeueCtx, cancel := context.WithTimeout(ctx, 2*time.Second)
+					runnable, err := q.Dequeue(dequeueCtx)
+					cancel()
+					if err != nil {
+						return
+					}
+					runnable(ctx)
+				}
+			}()
+		}
+
+		// Wait for all processing to complete
+		wg.Wait()
+		close(completionOrder)
+
+		// Collect completion order
+		execOrder := make([]string, 0, len(completionOrder))
+		for item := range completionOrder {
+			execOrder = append(execOrder, item)
+		}
+
+		// Verify that despite A-slow being dequeued first, other tenants' work completed while it was running
+		slowAPos := -1
+		fastAPos := -1
+		for i, item := range execOrder {
+			if item == "A-slow" {
+				slowAPos = i
+			}
+			if item == "A-fast" {
+				fastAPos = i
+			}
+		}
+
+		// The slow A task was started first but should finish later
+		require.True(t, slowAPos > 0, "A-slow should be in the execution order")
+		require.True(t, fastAPos > 0, "A-fast should be in the execution order")
+
+		// Verify some items from other tenants completed between the two A items
+		// This confirms round-robin fairness even with slow consumers
+		firstTenantBFoundPos := -1
+		for i, item := range execOrder {
+			if strings.HasPrefix(item, "B-") || strings.HasPrefix(item, "C-") {
+				firstTenantBFoundPos = i
+				break
+			}
+		}
+
+		// Make sure at least one other tenant item completed
+		require.True(t, firstTenantBFoundPos >= 0,
+			"At least one B or C item should be in execution order")
+
+		// Verify length is now 0
+		require.Equal(t, 0, q.Len())
+		require.Equal(t, 0, q.ActiveTenantsLen())
+
+		// Stop the queue
+		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), q))
+	})
+
+	t.Run("ActiveTenantsLength", func(t *testing.T) {
+		t.Parallel()
+
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, q.StartAsync(context.Background()), "Queue should start")
+		require.NoError(t, q.AwaitRunning(context.Background()), "Queue should be running")
+
+		// Enqueue items for different tenants
+		err := q.Enqueue(context.Background(), "tenant1", func(ctx context.Context) {})
+		require.NoError(t, err)
+		err = q.Enqueue(context.Background(), "tenant2", func(ctx context.Context) {})
+		require.NoError(t, err)
+
+		// Check active tenants
+		activeTenants := q.ActiveTenantsLen()
+		require.Equal(t, activeTenants, 2)
+
+		// Stop the queue
+		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), q))
+	})
+
+	t.Run("QueueLength", func(t *testing.T) {
+		t.Parallel()
+
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, q.StartAsync(context.Background()), "Queue should start")
+		require.NoError(t, q.AwaitRunning(context.Background()), "Queue should be running")
+
+		// Enqueue items
+		err := q.Enqueue(context.Background(), "tenant1", func(ctx context.Context) {})
+		require.NoError(t, err)
+		err = q.Enqueue(context.Background(), "tenant1", func(ctx context.Context) {})
+		require.NoError(t, err)
+
+		// Check queue length
+		queueLen := q.Len()
+		require.Equal(t, queueLen, 2)
+
+		// Stop the queue
+		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), q))
+	})
+
+	t.Run("GracefulShutdown", func(t *testing.T) {
+		t.Parallel()
+
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, q.StartAsync(context.Background()), "Queue should start")
+		require.NoError(t, q.AwaitRunning(context.Background()), "Queue should be running")
+
+		processed := make(chan struct{})
+
+		// Enqueue an item that signals when processed
+		err := q.Enqueue(context.Background(), "tenant1", func(ctx context.Context) {
+			close(processed)
+		})
+		require.NoError(t, err)
+
+		// Start a goroutine to dequeue and run the item
+		var wg sync.WaitGroup
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second)
+			defer cancel()
+			runnable, err := q.Dequeue(ctx)
+			require.NoError(t, err)
+			require.NotNil(t, runnable)
+			runnable(ctx)
+		}()
+
+		// Wait for the item to be processed
+		select {
+		case <-processed:
+		case <-time.After(2 * time.Second):
+			t.Fatal("Timed out waiting for item to be processed before shutdown")
+		}
+
+		// Now gracefully stop the queue
+		ctx, cancel := context.WithTimeout(context.Background(), 2*time.Second)
+		defer cancel()
+		require.NoError(t, services.StopAndAwaitTerminated(ctx, q))
+		wg.Wait()
+
+		// Check that the queue is closed
+		err = q.Enqueue(context.Background(), "tenant1", func(ctx context.Context) {})
+		require.ErrorIs(t, err, ErrQueueClosed)
+	})
+}

pkg/util/scheduler/scheduler.go

@@ -0,0 +1,178 @@
+package scheduler
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"sync"
+	"time"
+
+	"github.com/grafana/dskit/backoff"
+	"github.com/grafana/dskit/services"
+	"github.com/grafana/grafana/pkg/infra/log"
+)
+
+const (
+	// DefaultMaxBackoff is the default maximum backoff duration for workers.
+	DefaultMaxBackoff = 1 * time.Second
+	// DefaultMinBackoff is the default minimum backoff duration for workers.
+	DefaultMinBackoff = 100 * time.Millisecond
+	// DefaultNumWorkers is the default number of workers in the scheduler.
+	DefaultNumWorkers = 4
+	// DefaultMaxRetries is the default maximum number of retries for dequeue operations.
+	DefaultMaxRetries = 5
+)
+
+type WorkQueue interface {
+	services.Service
+
+	Dequeue(ctx context.Context) (runnable func(ctx context.Context), err error)
+}
+
+// Worker processes items from the QoS request queue
+type Worker struct {
+	id         int
+	queue      WorkQueue
+	wg         *sync.WaitGroup
+	maxBackoff time.Duration
+	maxRetries int
+	logger     log.Logger
+}
+
+func (w *Worker) run(ctx context.Context) {
+	defer w.wg.Done()
+	w.logger.Debug("worker started", "id", w.id)
+
+	for ctx.Err() == nil {
+		err := w.dequeueWithRetries(ctx)
+		if err != nil {
+			break
+		}
+	}
+	w.logger.Debug("worker stopped", "id", w.id)
+}
+
+func (w *Worker) dequeueWithRetries(ctx context.Context) error {
+	boff := backoff.New(ctx, backoff.Config{
+		MinBackoff: DefaultMinBackoff,
+		MaxBackoff: w.maxBackoff,
+		MaxRetries: w.maxRetries,
+	})
+
+	for boff.Ongoing() {
+		runnable, err := w.queue.Dequeue(ctx)
+		if err == nil {
+			runnable(ctx)
+			break
+		}
+
+		if errors.Is(err, ErrQueueClosed) {
+			w.logger.Error("queue closed, stopping worker", "id", w.id)
+			return fmt.Errorf("worker %d: queue closed", w.id)
+		}
+
+		w.logger.Error("retrying dequeue", "id", w.id, "error", err, "attempt", boff.NumRetries())
+		boff.Wait()
+	}
+	if err := boff.ErrCause(); err != nil {
+		w.logger.Error("failed to dequeue after retries", "id", w.id, "error", err)
+	}
+	return nil
+}
+
+// Scheduler manages a pool of Workers consuming from a Queue.
+type Scheduler struct {
+	services.Service
+
+	logger     log.Logger
+	queue      WorkQueue
+	wg         sync.WaitGroup
+	maxBackoff time.Duration
+
+	workers    []*Worker
+	numWorkers int
+}
+
+// Config holds configuration for the Scheduler.
+type Config struct {
+	NumWorkers int
+	MaxBackoff time.Duration
+	MaxRetries int
+	Logger     log.Logger
+}
+
+func (c *Config) validate() error {
+	if c.NumWorkers <= 0 {
+		c.NumWorkers = DefaultNumWorkers
+	}
+	if c.MaxBackoff <= 0 {
+		c.MaxBackoff = DefaultMaxBackoff
+	}
+	if c.MaxRetries <= 0 {
+		c.MaxRetries = DefaultMaxRetries
+	}
+	if c.Logger == nil {
+		c.Logger = log.New("scheduler")
+	}
+	return nil
+}
+
+// NewScheduler creates a new scheduler instance.
+func NewScheduler(queue WorkQueue, config *Config) (*Scheduler, error) {
+	if queue == nil {
+		return nil, errors.New("queue cannot be nil")
+	}
+	if err := config.validate(); err != nil {
+		return nil, fmt.Errorf("invalid config: %w", err)
+	}
+
+	s := &Scheduler{
+		logger:     config.Logger,
+		queue:      queue,
+		numWorkers: config.NumWorkers,
+		maxBackoff: config.MaxBackoff,
+		workers:    make([]*Worker, 0, config.NumWorkers),
+	}
+
+	s.Service = services.NewIdleService(s.starting, s.stopping)
+	return s, nil
+}
+
+// starting is called by the services.Service lifecycle to start the scheduler.
+func (s *Scheduler) starting(ctx context.Context) error {
+	s.logger.Info("scheduler starting", "numWorkers", s.numWorkers)
+
+	queueCtx, cancel := context.WithTimeout(ctx, 5*time.Second)
+	defer cancel()
+	if err := s.queue.AwaitRunning(queueCtx); err != nil {
+		return fmt.Errorf("queue is not running: %w", err)
+	}
+
+	s.workers = make([]*Worker, 0, s.numWorkers)
+	s.wg.Add(s.numWorkers)
+
+	for i := 0; i < s.numWorkers; i++ {
+		worker := &Worker{
+			id:         i,
+			queue:      s.queue,
+			wg:         &s.wg,
+			maxBackoff: s.maxBackoff,
+			logger:     s.logger,
+		}
+		s.workers = append(s.workers, worker)
+		go worker.run(ctx)
+	}
+
+	s.logger.Info("scheduler started")
+	return nil
+}
+
+// stopping is called by the services.Service lifecycle to stop the scheduler.
+func (s *Scheduler) stopping(_ error) error {
+	s.logger.Info("scheduler stopping")
+
+	s.wg.Wait()
+
+	s.logger.Info("scheduler stopped")
+	return nil
+}

pkg/util/scheduler/scheduler_bench_test.go

@@ -0,0 +1,289 @@
+package scheduler
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/grafana/dskit/services"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/stretchr/testify/require"
+)
+
+const defaultMaxSizePerTenant = 10000
+const defaultMaxBackoff = 1 * time.Second
+
+func benchScheduler(b *testing.B, numWorkers, numTenants, itemsPerTenant int) {
+	tenantIDs := make([]string, numTenants)
+	for i := range tenantIDs {
+		tenantIDs[i] = fmt.Sprintf("tenant-%d", i)
+	}
+
+	b.ResetTimer()
+	var processed atomic.Int64
+
+	for n := 0; n < b.N; n++ {
+		q := NewQueue(QueueOptionsWithDefaults(&QueueOptions{
+			MaxSizePerTenant: defaultMaxSizePerTenant,
+		}))
+		require.NoError(b, services.StartAndAwaitRunning(context.Background(), q))
+		scheduler, err := NewScheduler(q, &Config{
+			NumWorkers: numWorkers,
+			MaxBackoff: defaultMaxBackoff,
+		})
+		require.NoError(b, err)
+
+		require.NoError(b, scheduler.StartAsync(context.Background()))
+		require.NoError(b, scheduler.AwaitRunning(context.Background()))
+
+		var wg sync.WaitGroup
+		totalItems := numTenants * itemsPerTenant
+		wg.Add(totalItems)
+
+		for i := 0; i < numTenants; i++ {
+			tenantID := tenantIDs[i]
+			for j := 0; j < itemsPerTenant; j++ {
+				require.NoError(b, q.Enqueue(context.Background(), tenantID, func(_ context.Context) {
+					processed.Add(1)
+					wg.Done()
+				}))
+			}
+		}
+
+		done := make(chan struct{})
+		go func() {
+			wg.Wait()
+			close(done)
+		}()
+
+		select {
+		case <-done:
+		case <-time.After(30 * time.Second):
+			b.Fatalf("Timed out: Enqueued=%d, Processed=%d", totalItems, processed.Load())
+		}
+
+		scheduler.StopAsync()
+		require.NoError(b, scheduler.AwaitTerminated(context.Background()))
+		require.Equal(b, services.Terminated, scheduler.State())
+	}
+
+	b.ReportMetric(float64(processed.Load())/b.Elapsed().Seconds(), "items/sec")
+}
+
+func BenchmarkScheduler_1Worker_10Tenants(b *testing.B) {
+	benchScheduler(b, 1, 10, 1000)
+}
+
+func BenchmarkScheduler_2Workers_10Tenants(b *testing.B) {
+	benchScheduler(b, 2, 10, 1000)
+}
+
+func BenchmarkScheduler_4Workers_10Tenants(b *testing.B) {
+	benchScheduler(b, 4, 10, 1000)
+}
+
+func BenchmarkScheduler_8Workers_10Tenants(b *testing.B) {
+	benchScheduler(b, 8, 10, 1000)
+}
+
+func BenchmarkScheduler_16Workers_10Tenants(b *testing.B) {
+	benchScheduler(b, 16, 10, 1000)
+}
+
+func BenchmarkScheduler_4Workers_10Tenants_100ItemsPerTenant(b *testing.B) {
+	benchScheduler(b, 4, 10, 100)
+}
+
+func BenchmarkScheduler_4Workers_10Tenants_1000ItemsPerTenant(b *testing.B) {
+	benchScheduler(b, 4, 10, 1000)
+}
+
+func BenchmarkScheduler_4Workers_100Tenant_100ItemsPerTenant(b *testing.B) {
+	benchScheduler(b, 4, 100, 100)
+}
+
+func BenchmarkScheduler_4Workers_100Tenant_1000ItemsPerTenant(b *testing.B) {
+	benchScheduler(b, 4, 100, 1000)
+}
+
+func BenchmarkScheduler_4Workers_1000Tenant_100ItemsPerTenant(b *testing.B) {
+	benchScheduler(b, 4, 1000, 100)
+}
+
+func BenchmarkScheduler_4Workers_1000Tenant_1000ItemsPerTenant(b *testing.B) {
+	benchScheduler(b, 4, 1000, 1000)
+}
+
+func BenchmarkScheduler_4Workers_10000Tenant_10ItemsPerTenant(b *testing.B) {
+	benchScheduler(b, 4, 10000, 10)
+}
+
+func BenchmarkScheduler_4Workers_10000Tenant_100ItemsPerTenant(b *testing.B) {
+	benchScheduler(b, 4, 10000, 100)
+}
+
+// Benchmark comparing round-robin fairness among tenants
+func BenchmarkSchedulerFairness(b *testing.B) {
+	q := NewQueue(QueueOptionsWithDefaults(&QueueOptions{MaxSizePerTenant: 10000}))
+	require.NoError(b, services.StartAndAwaitRunning(context.Background(), q))
+	scheduler, err := NewScheduler(q, &Config{
+		NumWorkers: 4,
+		MaxBackoff: 100 * time.Millisecond,
+		Logger:     log.NewNopLogger(),
+	})
+	require.NoError(b, err)
+
+	require.NoError(b, scheduler.StartAsync(context.Background()))
+	require.NoError(b, scheduler.AwaitRunning(context.Background()))
+	defer func() {
+		scheduler.StopAsync()
+		require.NoError(b, scheduler.AwaitTerminated(context.Background()))
+	}()
+
+	const numTenants = 10
+	const itemsPerTenant = 1000
+
+	tenantIDs := make([]string, numTenants)
+	for i := range tenantIDs {
+		tenantIDs[i] = fmt.Sprintf("tenant-%d", i)
+	}
+	processedPerTenant := make([]atomic.Int64, numTenants)
+
+	b.ResetTimer()
+
+	for n := 0; n < b.N; n++ {
+		for i := range processedPerTenant {
+			processedPerTenant[i].Store(0)
+		}
+		var wg sync.WaitGroup
+		totalItems := numTenants * itemsPerTenant
+		wg.Add(totalItems)
+
+		for i := 0; i < numTenants; i++ {
+			tenantID := tenantIDs[i]
+			tenantIdx := i
+			for j := 0; j < itemsPerTenant; j++ {
+				require.NoError(b, q.Enqueue(context.Background(), tenantID, func(_ context.Context) {
+					processedPerTenant[tenantIdx].Add(1)
+					wg.Done()
+				}))
+			}
+		}
+
+		done := make(chan struct{})
+		go func() {
+			wg.Wait()
+			close(done)
+		}()
+
+		select {
+		case <-done:
+		case <-time.After(30 * time.Second):
+			b.Fatalf("Timed out waiting for items to be processed")
+		}
+
+		min, max, total := int64(itemsPerTenant+1), int64(0), int64(0)
+		for i := 0; i < numTenants; i++ {
+			count := processedPerTenant[i].Load()
+			total += count
+			if count < min {
+				min = count
+			}
+			if count > max {
+				max = count
+			}
+		}
+		fairnessRatio := float64(min) / float64(max)
+		b.ReportMetric(fairnessRatio, "fairness")
+		b.ReportMetric(float64(total)/b.Elapsed().Seconds(), "items/sec")
+	}
+
+	// Stop the scheduler and verify it's terminated
+	scheduler.StopAsync()
+	require.NoError(b, scheduler.AwaitTerminated(context.Background()))
+}
+
+// Add a new benchmark with alternating tenant enqueuing pattern
+func BenchmarkSchedulerFairnessAlternating(b *testing.B) {
+	q := NewQueue(QueueOptionsWithDefaults(nil))
+	require.NoError(b, services.StartAndAwaitRunning(context.Background(), q))
+	scheduler, err := NewScheduler(q, &Config{
+		NumWorkers: 4,
+		MaxBackoff: 100 * time.Millisecond,
+		Logger:     log.NewNopLogger(),
+	})
+	require.NoError(b, err)
+
+	require.NoError(b, scheduler.StartAsync(context.Background()))
+	require.NoError(b, scheduler.AwaitRunning(context.Background()))
+	defer func() {
+		scheduler.StopAsync()
+		require.NoError(b, scheduler.AwaitTerminated(context.Background()))
+	}()
+
+	const numTenants = 1000
+	const itemsPerTenant = 1000
+
+	tenantIDs := make([]string, numTenants)
+	for i := 0; i < numTenants; i++ {
+		tenantIDs[i] = fmt.Sprintf("tenant-%d", i)
+	}
+	processedPerTenant := make([]atomic.Int64, numTenants)
+
+	b.ResetTimer()
+
+	for n := 0; n < b.N; n++ {
+		for i := 0; i < numTenants; i++ {
+			processedPerTenant[i].Store(0)
+		}
+		var wg sync.WaitGroup
+		totalItems := numTenants * itemsPerTenant
+		wg.Add(totalItems)
+
+		// Enqueue in a round-robin pattern: 1 item per tenant per round
+		for j := 0; j < itemsPerTenant; j++ {
+			for i := 0; i < numTenants; i++ {
+				tenantID := tenantIDs[i]
+				tenantIdx := i
+				require.NoError(b, q.Enqueue(context.Background(), tenantID, func(_ context.Context) {
+					processedPerTenant[tenantIdx].Add(1)
+					wg.Done()
+				}))
+			}
+		}
+
+		done := make(chan struct{})
+		go func() {
+			wg.Wait()
+			close(done)
+		}()
+
+		select {
+		case <-done:
+		case <-time.After(30 * time.Second):
+			b.Fatalf("Timed out waiting for items to be processed")
+		}
+
+		min, max, total := int64(itemsPerTenant+1), int64(0), int64(0)
+		for i := 0; i < numTenants; i++ {
+			count := processedPerTenant[i].Load()
+			total += count
+			if count < min {
+				min = count
+			}
+			if count > max {
+				max = count
+			}
+		}
+		fairnessRatio := float64(min) / float64(max)
+		b.ReportMetric(fairnessRatio, "fairness")
+		b.ReportMetric(float64(total)/b.Elapsed().Seconds(), "items/sec")
+	}
+
+	// Stop the scheduler and verify it's terminated
+	scheduler.StopAsync()
+	require.NoError(b, scheduler.AwaitTerminated(context.Background()))
+}

pkg/util/scheduler/scheduler_test.go

@@ -0,0 +1,241 @@
+package scheduler
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/grafana/dskit/services"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/stretchr/testify/require"
+)
+
+func TestScheduler(t *testing.T) {
+	t.Parallel()
+
+	t.Run("ConfigValidation", func(t *testing.T) {
+		t.Parallel()
+
+		t.Run("ValidConfig", func(t *testing.T) {
+			cfg := Config{
+				NumWorkers: 5,
+				MaxBackoff: 1 * time.Second,
+				Logger:     log.New("qos.test"),
+			}
+			err := cfg.validate()
+			require.NoError(t, err)
+			require.NotNil(t, cfg.Logger)
+		})
+
+		t.Run("ZeroWorkersGetDefault", func(t *testing.T) {
+			cfg := Config{
+				NumWorkers: 0,
+				MaxBackoff: 1 * time.Second,
+			}
+			err := cfg.validate()
+			require.NoError(t, err)
+			require.Equal(t, cfg.NumWorkers, DefaultNumWorkers)
+			require.NotNil(t, cfg.Logger, "Logger should not be nil")
+		})
+
+		t.Run("NilLoggerGetsDefault", func(t *testing.T) {
+			cfg := Config{
+				NumWorkers: 1,
+				MaxBackoff: 1 * time.Second,
+				Logger:     nil,
+			}
+			err := cfg.validate()
+			require.NoError(t, err)
+			require.NotNil(t, cfg.Logger)
+		})
+
+		t.Run("ZeroTimeoutGetsDefault", func(t *testing.T) {
+			cfg := Config{
+				NumWorkers: 1,
+				MaxBackoff: 0,
+				Logger:     log.New("qos.test"),
+			}
+			err := cfg.validate()
+			require.NoError(t, err)
+			require.Equal(t, cfg.MaxBackoff, DefaultMaxBackoff)
+		})
+
+		t.Run("ZeroRetriesGetsDefault", func(t *testing.T) {
+			cfg := Config{
+				NumWorkers: 1,
+				MaxBackoff: 1 * time.Second,
+				MaxRetries: 0,
+				Logger:     log.New("qos.test"),
+			}
+			err := cfg.validate()
+			require.NoError(t, err)
+			require.Equal(t, cfg.MaxRetries, DefaultMaxRetries)
+		})
+	})
+
+	t.Run("NewScheduler", func(t *testing.T) {
+		t.Parallel()
+
+		t.Run("ValidParameters", func(t *testing.T) {
+			q := NewQueue(QueueOptionsWithDefaults(nil))
+			require.NoError(t, services.StartAndAwaitRunning(context.Background(), q))
+
+			cfg := Config{
+				NumWorkers: 2,
+				MaxBackoff: 1 * time.Second,
+				Logger:     log.New("qos.test"),
+			}
+			scheduler, err := NewScheduler(q, &cfg)
+			require.NoError(t, err)
+			require.NotNil(t, scheduler)
+			require.NoError(t, services.StartAndAwaitRunning(context.Background(), scheduler))
+			require.Equal(t, q, scheduler.queue)
+			require.Equal(t, cfg.NumWorkers, scheduler.numWorkers)
+			require.True(t, scheduler.State() == services.Running)
+			require.NoError(t, services.StopAndAwaitTerminated(context.Background(), scheduler))
+		})
+
+		t.Run("NilQueue", func(t *testing.T) {
+			cfg := Config{
+				NumWorkers: 2,
+				MaxBackoff: 1 * time.Second,
+			}
+			scheduler, err := NewScheduler(nil, &cfg)
+			require.Error(t, err)
+			require.Nil(t, scheduler)
+		})
+	})
+
+	t.Run("Lifecycle", func(t *testing.T) {
+		t.Parallel()
+
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, services.StartAndAwaitRunning(context.Background(), q))
+
+		scheduler, err := NewScheduler(q, &Config{
+			NumWorkers: 3,
+			MaxBackoff: 100 * time.Millisecond,
+			Logger:     log.New("qos.test"),
+		})
+		require.NoError(t, err)
+		require.True(t, scheduler.State() == services.New)
+		require.NoError(t, services.StartAndAwaitRunning(context.Background(), scheduler))
+		require.True(t, scheduler.State() == services.Running)
+		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), scheduler))
+		require.True(t, scheduler.State() == services.Terminated)
+	})
+
+	t.Run("ProcessItems", func(t *testing.T) {
+		t.Parallel()
+
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, services.StartAndAwaitRunning(context.Background(), q))
+
+		const itemCount = 10
+		var processed sync.Map
+		var wg sync.WaitGroup
+		wg.Add(itemCount)
+
+		scheduler, err := NewScheduler(q, &Config{
+			NumWorkers: 2,
+			MaxBackoff: 100 * time.Millisecond,
+			Logger:     log.New("qos.test"),
+		})
+		require.NoError(t, err)
+		require.NoError(t, services.StartAndAwaitRunning(context.Background(), scheduler))
+
+		for i := 0; i < itemCount; i++ {
+			itemID := i
+			require.NoError(t, q.Enqueue(context.Background(), "tenant-1", func(_ context.Context) {
+				processed.Store(itemID, true)
+				wg.Done()
+			}))
+		}
+
+		done := make(chan struct{})
+		go func() {
+			wg.Wait()
+			close(done)
+		}()
+
+		select {
+		case <-done:
+		case <-time.After(5 * time.Second):
+			t.Fatal("Timed out waiting for all items to be processed")
+		}
+
+		count := 0
+		processed.Range(func(_, _ any) bool {
+			count++
+			return true
+		})
+		require.Equal(t, itemCount, count, "Not all items were processed")
+
+		require.NoError(t, services.StopAndAwaitTerminated(context.Background(), scheduler))
+	})
+
+	t.Run("GracefulShutdown", func(t *testing.T) {
+		t.Parallel()
+
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+		require.NoError(t, services.StartAndAwaitRunning(context.Background(), q))
+
+		var processed atomic.Int32
+		taskStarted := make(chan struct{})
+		taskFinished := make(chan struct{})
+
+		scheduler, err := NewScheduler(q, &Config{
+			NumWorkers: 1,
+			MaxBackoff: 100 * time.Millisecond,
+			Logger:     log.New("qos.test"),
+		})
+		require.NoError(t, err)
+		require.NoError(t, services.StartAndAwaitRunning(context.Background(), scheduler))
+
+		for i := 0; i < 5; i++ {
+			require.NoError(t, q.Enqueue(context.Background(), "tenant-1", func(_ context.Context) {
+				processed.Add(1)
+			}))
+		}
+
+		require.NoError(t, q.Enqueue(context.Background(), "tenant-1", func(_ context.Context) {
+			close(taskStarted)
+			time.Sleep(1 * time.Second)
+			processed.Add(1)
+			close(taskFinished)
+		}))
+
+		select {
+		case <-taskStarted:
+		case <-time.After(2 * time.Second):
+			t.Fatal("Timed out waiting for long-running task to start")
+		}
+
+		scheduler.StopAsync()
+
+		select {
+		case <-taskFinished:
+		case <-time.After(2 * time.Second):
+			t.Fatal("Timed out waiting for long-running task to finish")
+		}
+
+		require.Equal(t, int32(6), processed.Load(), "Not all items were processed")
+		require.NoError(t, scheduler.AwaitTerminated(context.Background()))
+	})
+
+	t.Run("WithQueueClosed", func(t *testing.T) {
+		t.Parallel()
+
+		q := NewQueue(QueueOptionsWithDefaults(nil))
+
+		scheduler, err := NewScheduler(q, &Config{
+			NumWorkers: 2,
+			MaxBackoff: 100 * time.Millisecond,
+			Logger:     log.New("qos.test"),
+		})
+		require.NoError(t, err)
+		require.ErrorContains(t, services.StartAndAwaitRunning(context.Background(), scheduler), "queue is not running")
+	})
+}