From 4ba378603dfb9896745701cd11b56d606dbfa657 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Mustafa=20Sencer=20=C3=96zcan?= <32759850+mustafasencer@users.noreply.github.com> Date: Tue, 27 May 2025 21:15:03 +0200 Subject: [PATCH] feat(unified-storage): round robin queue and scheduler package (#105544) --- pkg/util/scheduler/queue.go | 363 +++++++++++++ pkg/util/scheduler/queue_test.go | 601 +++++++++++++++++++++ pkg/util/scheduler/scheduler.go | 178 ++++++ pkg/util/scheduler/scheduler_bench_test.go | 289 ++++++++++ pkg/util/scheduler/scheduler_test.go | 241 +++++++++ 5 files changed, 1672 insertions(+) create mode 100644 pkg/util/scheduler/queue.go create mode 100644 pkg/util/scheduler/queue_test.go create mode 100644 pkg/util/scheduler/scheduler.go create mode 100644 pkg/util/scheduler/scheduler_bench_test.go create mode 100644 pkg/util/scheduler/scheduler_test.go diff --git a/pkg/util/scheduler/queue.go b/pkg/util/scheduler/queue.go new file mode 100644 index 00000000000..eceb0bcf38a --- /dev/null +++ b/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 +} diff --git a/pkg/util/scheduler/queue_test.go b/pkg/util/scheduler/queue_test.go new file mode 100644 index 00000000000..e703501747e --- /dev/null +++ b/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) + }) +} diff --git a/pkg/util/scheduler/scheduler.go b/pkg/util/scheduler/scheduler.go new file mode 100644 index 00000000000..6f781c2536e --- /dev/null +++ b/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 +} diff --git a/pkg/util/scheduler/scheduler_bench_test.go b/pkg/util/scheduler/scheduler_bench_test.go new file mode 100644 index 00000000000..64d79a2bd5d --- /dev/null +++ b/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())) +} diff --git a/pkg/util/scheduler/scheduler_test.go b/pkg/util/scheduler/scheduler_test.go new file mode 100644 index 00000000000..f0fa02d6e4c --- /dev/null +++ b/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") + }) +}