feat(task_queue): implement queue whitelist for non-org metrics and update task queue worker initialization

Author: leo-marble

cmd/worker.go

@@ -7,6 +7,7 @@ import (
 	"net/http"
 	"os"
 	"os/signal"
+	"slices"
 	"syscall"
 	"time"
 
@@ -156,8 +157,13 @@ func RunTaskQueue(apiVersion string) error {
 		return err
 	}
 
-	// Add the global queue
-	maps.Copy(queues, usecases.QueuesGlobal())
+	// For non-org queues
+	nonOrgQueues := make(map[string]river.QueueConfig)
+	maps.Copy(nonOrgQueues, usecases.QueueMetrics())
+	queueWhitelist := slices.Collect(maps.Keys(nonOrgQueues))
+
+	// Add the metrics queue
+	maps.Copy(queues, nonOrgQueues)
 
 	// Periodics always contain the per-org tasks retrieved above. Add other, non-organization-scoped periodics below
 	periodics := append(
@@ -233,7 +239,7 @@ func RunTaskQueue(apiVersion string) error {
 	logger.InfoContext(ctx, "starting worker", slog.String("version", apiVersion))
 
 	// Asynchronously keep the task queue workers up to date with the orgs in the database
-	taskQueueWorker := uc.NewTaskQueueWorker(riverClient)
+	taskQueueWorker := uc.NewTaskQueueWorker(riverClient, queueWhitelist)
 	go taskQueueWorker.RefreshQueuesFromOrgIds(ctx)
 
 	// Start the cron jobs using the old entrypoint.

usecases/scheduled_execution/metrics_collection_job.go

@@ -17,7 +17,7 @@ func NewMetricsCollectionPeriodicJob() *river.PeriodicJob {
 		river.PeriodicInterval(METRICS_COLLECTION_WORKER_INTERVAL),
 		func() (river.JobArgs, *river.InsertOpts) {
 			return models.MetricsCollectionArgs{}, &river.InsertOpts{
-				Queue: "global",
+				Queue: "metrics",
 				UniqueOpts: river.UniqueOpts{
 					ByQueue:  true,
 					ByPeriod: METRICS_COLLECTION_WORKER_INTERVAL,

usecases/task_queue.go

@@ -3,6 +3,7 @@ package usecases
 import (
 	"context"
 	"fmt"
+	"slices"
 	"sync"
 	"time"
 
@@ -19,26 +20,27 @@ import (
 )
 
 const (
-	numberWorkersPerQueue       = 5
-	globalQueueName             = "global"
-	numberWorkersForGlobalQueue = 1
+	numberWorkersPerQueue = 5
 )
 
 type TaskQueueWorker struct {
 	executorFactory executor_factory.ExecutorFactory
 	orgRepository   repositories.OrganizationRepository
+	queueWhitelist  []string
 	riverClient     *river.Client[pgx.Tx]
 	mu              *sync.Mutex
 }
 
 func NewTaskQueueWorker(
 	executorFactory executor_factory.ExecutorFactory,
 	orgRepository repositories.OrganizationRepository,
+	queueWhitelist []string,
 	riverClient *river.Client[pgx.Tx],
 ) *TaskQueueWorker {
 	return &TaskQueueWorker{
 		executorFactory: executorFactory,
 		orgRepository:   orgRepository,
+		queueWhitelist:  queueWhitelist,
 		riverClient:     riverClient,
 		mu:              &sync.Mutex{},
 	}
@@ -138,8 +140,8 @@ func (w *TaskQueueWorker) removeQueuesFromMissingOrgs(ctx context.Context,
 			continue
 		}
 
-		// Ignore global queue
-		if q.Name == globalQueueName {
+		// Ignore whitelisted queues
+		if slices.Contains(w.queueWhitelist, q.Name) {
 			continue
 		}
 
@@ -186,10 +188,10 @@ func QueuesFromOrgs(ctx context.Context, orgsRepo repositories.OrganizationRepos
 	return queues, periodics, nil
 }
 
-func QueuesGlobal() map[string]river.QueueConfig {
+func QueueMetrics() map[string]river.QueueConfig {
 	queues := make(map[string]river.QueueConfig, 1)
-	queues[globalQueueName] = river.QueueConfig{
-		MaxWorkers: numberWorkersForGlobalQueue,
+	queues["metrics"] = river.QueueConfig{
+		MaxWorkers: 1,
 	}
 	return queues
 }

usecases/usecases.go

@@ -336,10 +336,11 @@ func (usecases *Usecases) NewLicenseUsecase() PublicLicenseUseCase {
 	)
 }
 
-func (usecases *Usecases) NewTaskQueueWorker(riverClient *river.Client[pgx.Tx]) *TaskQueueWorker {
+func (usecases *Usecases) NewTaskQueueWorker(riverClient *river.Client[pgx.Tx], queueWhitelist []string) *TaskQueueWorker {
 	return NewTaskQueueWorker(
 		usecases.NewExecutorFactory(),
 		&usecases.Repositories.MarbleDbRepository,
+		queueWhitelist,
 		riverClient,
 	)
 }