v3: WIP maximum queues per env implementation

Author: ericallam

apps/webapp/app/v3/marqs/fairDequeuingStrategy.server.ts

@@ -43,6 +43,11 @@ export type FairDequeuingStrategyOptions = {
   biases?: FairDequeuingStrategyBiases;
   reuseSnapshotCount?: number;
   maximumEnvCount?: number;
+  /**
+   * Maximum number of queues to process per environment
+   * If not provided, all queues in an environment will be processed
+   */
+  maximumQueuePerEnvCount?: number;
 };
 
 type FairQueueConcurrency = {
@@ -216,8 +221,6 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
     return result;
   }
 
-  // Helper method to maintain DRY principle
-  // Update return type
   #orderQueuesByEnvs(envs: string[], snapshot: FairQueueSnapshot): Array<EnvQueues> {
     const queuesByEnv = snapshot.queues.reduce((acc, queue) => {
       if (!acc[queue.env]) {
@@ -231,11 +234,17 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
       if (queuesByEnv[envId]) {
         // Get ordered queues for this env
         const orderedQueues = this.#weightedRandomQueueOrder(queuesByEnv[envId]);
+
+        // Apply queue limit if maximumQueuePerEnvCount is set
+        const limitedQueues = this.options.maximumQueuePerEnvCount
+          ? orderedQueues.slice(0, this.options.maximumQueuePerEnvCount)
+          : orderedQueues;
+
         // Only add the env if it has queues
-        if (orderedQueues.length > 0) {
+        if (limitedQueues.length > 0) {
           acc.push({
             envId,
-            queues: orderedQueues.map((queue) => queue.id),
+            queues: limitedQueues.map((queue) => queue.id),
           });
         }
       }
@@ -512,6 +521,10 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
 
       span.setAttribute("queue_count", result.length);
 
+      if (result.length === this.options.parentQueueLimit) {
+        span.setAttribute("parent_queue_limit_reached", true);
+      }
+
       return result;
     });
   }

apps/webapp/test/fairDequeuingStrategy.test.ts

@@ -1032,6 +1032,93 @@ describe("FairDequeuingStrategy", () => {
       expect(selectionPercentages["env-4"] || 0).toBeLessThan(20);
     }
   );
+
+  redisTest(
+    "should respect maximumQueuePerEnvCount when distributing queues",
+    async ({ redisOptions }) => {
+      const redis = createRedisClient(redisOptions);
+
+      const keyProducer = createKeyProducer("test");
+      const strategy = new FairDequeuingStrategy({
+        tracer,
+        redis,
+        keys: keyProducer,
+        defaultEnvConcurrency: 5,
+        parentQueueLimit: 100,
+        seed: "test-seed-max-queues",
+        maximumQueuePerEnvCount: 2, // Only take 2 queues per env
+      });
+
+      const now = Date.now();
+
+      // Setup two environments with different numbers of queues
+      const envSetups = [
+        {
+          envId: "env-1",
+          queues: [
+            { age: 5000 }, // Oldest
+            { age: 4000 },
+            { age: 3000 }, // This should be excluded due to maximumQueuePerEnvCount
+          ],
+        },
+        {
+          envId: "env-2",
+          queues: [
+            { age: 2000 },
+            { age: 1000 }, // Newest
+          ],
+        },
+      ];
+
+      // Setup queues and concurrency for each env
+      for (const setup of envSetups) {
+        await setupConcurrency({
+          redis,
+          keyProducer,
+          env: { id: setup.envId, currentConcurrency: 0, limit: 5 },
+        });
+
+        for (let i = 0; i < setup.queues.length; i++) {
+          await setupQueue({
+            redis,
+            keyProducer,
+            parentQueue: "parent-queue",
+            score: now - setup.queues[i].age,
+            queueId: `queue-${setup.envId}-${i}`,
+            orgId: `org-${setup.envId}`,
+            envId: setup.envId,
+          });
+        }
+      }
+
+      const result = await strategy.distributeFairQueuesFromParentQueue(
+        "parent-queue",
+        "consumer-1"
+      );
+
+      // Verify that each environment has at most 2 queues
+      for (const envQueues of result) {
+        expect(envQueues.queues.length).toBeLessThanOrEqual(2);
+      }
+
+      // Get queues for env-1 (which had 3 queues originally)
+      const env1Queues = result.find((eq) => eq.envId === "env-1")?.queues ?? [];
+
+      // Should have exactly 2 queues
+      expect(env1Queues.length).toBe(2);
+
+      // The queues should be the two oldest ones (queue-env-1-0 and queue-env-1-1)
+      expect(env1Queues).toContain(keyProducer.queueKey("org-env-1", "env-1", "queue-env-1-0"));
+      expect(env1Queues).toContain(keyProducer.queueKey("org-env-1", "env-1", "queue-env-1-1"));
+      expect(env1Queues).not.toContain(keyProducer.queueKey("org-env-1", "env-1", "queue-env-1-2"));
+
+      // Get queues for env-2 (which had 2 queues originally)
+      const env2Queues = result.find((eq) => eq.envId === "env-2")?.queues ?? [];
+
+      // Should still have both queues since it was within the limit
+      expect(env2Queues.length).toBe(2);
+    }
+  );
 });
 
 // Helper function to flatten results for counting