Only consider the top N orgs when dequeuing, to help mitigate large spikes in queues (like around the hour and half hour marks)

Author: ericallam

apps/webapp/app/env.server.ts

@@ -233,6 +233,7 @@ const EnvironmentSchema = z.object({
   MARQS_AVAILABLE_CAPACITY_BIAS: z.coerce.number().default(0.3),
   MARQS_QUEUE_AGE_RANDOMIZATION_BIAS: z.coerce.number().default(0.25),
   MARQS_REUSE_SNAPSHOT_COUNT: z.coerce.number().int().default(0),
+  MARQS_MAXIMUM_ORG_COUNT: z.coerce.number().int().optional(),
 
   PROD_TASK_HEARTBEAT_INTERVAL_MS: z.coerce.number().int().optional(),
 

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

@@ -45,6 +45,7 @@ export type FairDequeuingStrategyOptions = {
    */
   biases?: FairDequeuingStrategyBiases;
   reuseSnapshotCount?: number;
+  maximumOrgCount?: number;
 };
 
 type FairQueueConcurrency = {
@@ -342,14 +343,26 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
 
       const now = Date.now();
 
-      const queues = await this.#allChildQueuesByScore(parentQueue, consumerId, now);
+      let queues = await this.#allChildQueuesByScore(parentQueue, consumerId, now);
 
       span.setAttribute("parent_queue_count", queues.length);
 
       if (queues.length === 0) {
         return emptyFairQueueSnapshot;
       }
 
+      // Apply org selection if maximumOrgCount is specified
+      let selectedOrgIds: Set<string>;
+      if (this.options.maximumOrgCount && this.options.maximumOrgCount > 0) {
+        selectedOrgIds = this.#selectTopOrgs(queues, this.options.maximumOrgCount);
+        // Filter queues to only include selected orgs
+        queues = queues.filter((queue) => selectedOrgIds.has(queue.org));
+
+        span.setAttribute("selected_org_count", selectedOrgIds.size);
+      }
+
+      span.setAttribute("selected_queue_count", queues.length);
+
       const orgIds = new Set<string>();
       const envIds = new Set<string>();
       const envIdToOrgId = new Map<string, string>();
@@ -449,6 +462,52 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
     });
   }
 
+  #selectTopOrgs(queues: FairQueue[], maximumOrgCount: number): Set<string> {
+    // Group queues by org
+    const queuesByOrg = queues.reduce((acc, queue) => {
+      if (!acc[queue.org]) {
+        acc[queue.org] = [];
+      }
+      acc[queue.org].push(queue);
+      return acc;
+    }, {} as Record<string, FairQueue[]>);
+
+    // Calculate average age for each org
+    const orgAverageAges = Object.entries(queuesByOrg).map(([orgId, orgQueues]) => {
+      const averageAge = orgQueues.reduce((sum, q) => sum + q.age, 0) / orgQueues.length;
+      return { orgId, averageAge };
+    });
+
+    // Perform weighted shuffle based on average ages
+    const maxAge = Math.max(...orgAverageAges.map((o) => o.averageAge));
+    const weightedOrgs = orgAverageAges.map((org) => ({
+      orgId: org.orgId,
+      weight: org.averageAge / maxAge, // Normalize weights
+    }));
+
+    // Select top N orgs using weighted shuffle
+    const selectedOrgs = new Set<string>();
+    let remainingOrgs = [...weightedOrgs];
+    let totalWeight = remainingOrgs.reduce((sum, org) => sum + org.weight, 0);
+
+    while (selectedOrgs.size < maximumOrgCount && remainingOrgs.length > 0) {
+      let random = this._rng() * totalWeight;
+      let index = 0;
+
+      while (random > 0 && index < remainingOrgs.length) {
+        random -= remainingOrgs[index].weight;
+        index++;
+      }
+      index = Math.max(0, index - 1);
+
+      selectedOrgs.add(remainingOrgs[index].orgId);
+      totalWeight -= remainingOrgs[index].weight;
+      remainingOrgs.splice(index, 1);
+    }
+
+    return selectedOrgs;
+  }
+
   async #getOrgConcurrency(orgId: string): Promise<FairQueueConcurrency> {
     return await startSpan(this.options.tracer, "getOrgConcurrency", async (span) => {
       span.setAttribute("org_id", orgId);

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

@@ -1626,6 +1626,7 @@ function getMarQSClient() {
             queueAgeRandomization: env.MARQS_QUEUE_AGE_RANDOMIZATION_BIAS,
           },
           reuseSnapshotCount: env.MARQS_REUSE_SNAPSHOT_COUNT,
+          maximumOrgCount: env.MARQS_MAXIMUM_ORG_COUNT,
         }),
         envQueuePriorityStrategy: new FairDequeuingStrategy({
           tracer: tracer,

apps/webapp/test/fairDequeuingStrategy.test.ts

@@ -779,4 +779,135 @@ describe("FairDequeuingStrategy", () => {
     expect(mixed["queue-1"][0]).toBeGreaterThan(mixed["queue-3"][0]); // Older preferred
     expect(mixed["queue-3"][0]).toBeGreaterThan(0); // But newer still gets chances
   });
+
+  redisTest(
+    "should respect maximumOrgCount and select orgs based on queue ages",
+    async ({ redis }) => {
+      const keyProducer = createKeyProducer("test");
+      const strategy = new FairDequeuingStrategy({
+        tracer,
+        redis,
+        keys: keyProducer,
+        defaultOrgConcurrency: 10,
+        defaultEnvConcurrency: 5,
+        parentQueueLimit: 100,
+        checkForDisabledOrgs: true,
+        seed: "test-seed-max-orgs",
+        maximumOrgCount: 2, // Only select top 2 orgs
+      });
+
+      const now = Date.now();
+
+      // Setup 4 orgs with different queue age profiles
+      const orgSetups = [
+        {
+          orgId: "org-1",
+          queues: [
+            { age: 1000 }, // Average age: 1000
+          ],
+        },
+        {
+          orgId: "org-2",
+          queues: [
+            { age: 5000 }, // Average age: 5000
+            { age: 5000 },
+          ],
+        },
+        {
+          orgId: "org-3",
+          queues: [
+            { age: 2000 }, // Average age: 2000
+            { age: 2000 },
+          ],
+        },
+        {
+          orgId: "org-4",
+          queues: [
+            { age: 500 }, // Average age: 500
+            { age: 500 },
+          ],
+        },
+      ];
+
+      // Setup queues and concurrency for each org
+      for (const setup of orgSetups) {
+        await setupConcurrency({
+          redis,
+          keyProducer,
+          org: { id: setup.orgId, currentConcurrency: 0, limit: 10 },
+          env: { id: "env-1", 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.orgId}-${i}`,
+            orgId: setup.orgId,
+            envId: "env-1",
+          });
+        }
+      }
+
+      // Run multiple iterations to verify consistent behavior
+      const iterations = 100;
+      const selectedOrgCounts: Record<string, number> = {};
+
+      for (let i = 0; i < iterations; i++) {
+        const result = await strategy.distributeFairQueuesFromParentQueue(
+          "parent-queue",
+          `consumer-${i}`
+        );
+
+        // Track which orgs were included in the result
+        const selectedOrgs = new Set(result.map((queueId) => keyProducer.orgIdFromQueue(queueId)));
+
+        // Verify we never get more than maximumOrgCount orgs
+        expect(selectedOrgs.size).toBeLessThanOrEqual(2);
+
+        for (const orgId of selectedOrgs) {
+          selectedOrgCounts[orgId] = (selectedOrgCounts[orgId] || 0) + 1;
+        }
+      }
+
+      console.log("Organization selection counts:", selectedOrgCounts);
+
+      // org-2 should be selected most often (highest average age)
+      expect(selectedOrgCounts["org-2"]).toBeGreaterThan(selectedOrgCounts["org-4"] || 0);
+
+      // org-4 should be selected least often (lowest average age)
+      const org4Count = selectedOrgCounts["org-4"] || 0;
+      expect(org4Count).toBeLessThan(selectedOrgCounts["org-2"]);
+
+      // Verify that orgs with higher average queue age are selected more frequently
+      const sortedOrgs = Object.entries(selectedOrgCounts).sort((a, b) => b[1] - a[1]);
+      console.log("Sorted organization frequencies:", sortedOrgs);
+
+      // The top 2 most frequently selected orgs should be org-2 and org-3
+      // as they have the highest average queue ages
+      const topTwoOrgs = new Set([sortedOrgs[0][0], sortedOrgs[1][0]]);
+      expect(topTwoOrgs).toContain("org-2"); // Highest average age
+      expect(topTwoOrgs).toContain("org-3"); // Second highest average age
+
+      // Calculate selection percentages
+      const totalSelections = Object.values(selectedOrgCounts).reduce((a, b) => a + b, 0);
+      const selectionPercentages = Object.entries(selectedOrgCounts).reduce(
+        (acc, [orgId, count]) => {
+          acc[orgId] = (count / totalSelections) * 100;
+          return acc;
+        },
+        {} as Record<string, number>
+      );
+
+      console.log("Organization selection percentages:", selectionPercentages);
+
+      // Verify that org-2 (highest average age) gets selected in at least 40% of iterations
+      expect(selectionPercentages["org-2"]).toBeGreaterThan(40);
+
+      // Verify that org-4 (lowest average age) gets selected in less than 20% of iterations
+      expect(selectionPercentages["org-4"] || 0).toBeLessThan(20);
+    }
+  );
 });