Add ability to bias env selection by concurrency limit and capacity

Author: ericallam

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

@@ -8,6 +8,20 @@ import seedrandom from "seedrandom";
 import { Tracer } from "@opentelemetry/api";
 import { startSpan } from "../tracing.server";
 
+export type FairDequeuingStrategyBiases = {
+  /**
+   * How much to bias towards environments with higher concurrency limits
+   * 0 = no bias, 1 = full bias based on limit differences
+   */
+  concurrencyLimitBias: number;
+
+  /**
+   * How much to bias towards environments with more available capacity
+   * 0 = no bias, 1 = full bias based on available capacity
+   */
+  availableCapacityBias: number;
+};
+
 export type FairDequeuingStrategyOptions = {
   redis: Redis;
   keys: MarQSKeyProducer;
@@ -17,6 +31,11 @@ export type FairDequeuingStrategyOptions = {
   checkForDisabledOrgs: boolean;
   tracer: Tracer;
   seed?: string;
+  /**
+   * Configure biasing for environment shuffling
+   * If not provided, no biasing will be applied (completely random shuffling)
+   */
+  biases?: FairDequeuingStrategyBiases;
 };
 
 type FairQueueConcurrency = {
@@ -33,13 +52,23 @@ type FairQueueSnapshot = {
   queues: Array<FairQueue>;
 };
 
+type WeightedEnv = {
+  envId: string;
+  weight: number;
+};
+
 const emptyFairQueueSnapshot: FairQueueSnapshot = {
   id: "empty",
   orgs: {},
   envs: {},
   queues: [],
 };
 
+const defaultBiases: FairDequeuingStrategyBiases = {
+  concurrencyLimitBias: 0,
+  availableCapacityBias: 0,
+};
+
 export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
   private _cache: UnkeyCache<{
     concurrencyLimit: number;
@@ -107,30 +136,85 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
     );
   }
 
-  // Now we need to:
-  // 1. Shuffle the environments
-  // 2. Sort the queues by their environment order in the shuffled list
-  // 3. Keep the queues sorted by their age inside their "environment" slice of the final array
   #shuffleQueuesByEnv(snapshot: FairQueueSnapshot): Array<string> {
     const envs = Object.keys(snapshot.envs);
+    const biases = this.options.biases ?? defaultBiases;
+
+    if (biases.concurrencyLimitBias === 0 && biases.availableCapacityBias === 0) {
+      const shuffledEnvs = this.#shuffle(envs);
+      return this.#orderQueuesByEnvs(shuffledEnvs, snapshot);
+    }
 
-    const shuffledEnvs = this.#shuffle(envs);
+    // Find the maximum concurrency limit for normalization
+    const maxLimit = Math.max(...envs.map((envId) => snapshot.envs[envId].concurrency.limit));
 
+    // Calculate weights for each environment
+    const weightedEnvs: WeightedEnv[] = envs.map((envId) => {
+      const env = snapshot.envs[envId];
+
+      // Start with base weight of 1
+      let weight = 1;
+
+      // Add normalized concurrency limit bias if configured
+      if (biases.concurrencyLimitBias > 0) {
+        const normalizedLimit = env.concurrency.limit / maxLimit;
+        // Square or cube the bias to make it more pronounced at higher values
+        weight *= 1 + Math.pow(normalizedLimit * biases.concurrencyLimitBias, 2);
+      }
+
+      // Add available capacity bias if configured
+      if (biases.availableCapacityBias > 0) {
+        const usedCapacityPercentage = env.concurrency.current / env.concurrency.limit;
+        const availableCapacityBonus = 1 - usedCapacityPercentage;
+        // Square or cube the bias to make it more pronounced at higher values
+        weight *= 1 + Math.pow(availableCapacityBonus * biases.availableCapacityBias, 2);
+      }
+
+      return { envId, weight };
+    });
+
+    const shuffledEnvs = this.#weightedShuffle(weightedEnvs);
+    return this.#orderQueuesByEnvs(shuffledEnvs, snapshot);
+  }
+
+  #weightedShuffle(weightedItems: WeightedEnv[]): string[] {
+    const totalWeight = weightedItems.reduce((sum, item) => sum + item.weight, 0);
+    const result: string[] = [];
+    const items = [...weightedItems];
+
+    while (items.length > 0) {
+      let random = this._rng() * totalWeight;
+      let index = 0;
+
+      // Find item based on weighted random selection
+      while (random > 0 && index < items.length) {
+        random -= items[index].weight;
+        index++;
+      }
+      index = Math.max(0, index - 1);
+
+      // Add selected item to result and remove from items
+      result.push(items[index].envId);
+      items.splice(index, 1);
+    }
+
+    return result;
+  }
+
+  // Helper method to maintain DRY principle
+  #orderQueuesByEnvs(envs: string[], snapshot: FairQueueSnapshot): Array<string> {
     const queuesByEnv = snapshot.queues.reduce((acc, queue) => {
       if (!acc[queue.env]) {
         acc[queue.env] = [];
       }
-
       acc[queue.env].push(queue);
-
       return acc;
     }, {} as Record<string, Array<FairQueue>>);
 
-    const queues = shuffledEnvs.reduce((acc, envId) => {
+    const queues = envs.reduce((acc, envId) => {
       if (queuesByEnv[envId]) {
         acc.push(...queuesByEnv[envId].sort((a, b) => b.age - a.age));
       }
-
       return acc;
     }, [] as Array<FairQueue>);
 

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

@@ -69,9 +69,6 @@ export class MarQS {
   public keys: MarQSKeyProducer;
   #rebalanceWorkers: Array<AsyncWorker> = [];
 
-  private _staleQueues: Set<string> = new Set();
-  private _staleQueueHits: Map<string, number> = new Map();
-
   constructor(private readonly options: MarQSOptions) {
     this.redis = options.redis;
 
@@ -372,6 +369,8 @@ export class MarQS {
                 [SemanticAttributes.PARENT_QUEUE]: message.parentQueue,
                 age_in_seconds: ageOfMessageInMs / 1000,
                 attempted_queues: queues.indexOf(messageQueue) + 1, // How many queues we tried before success
+                message_timestamp: message.timestamp,
+                message_age: Date.now() - message.timestamp,
               });
 
               await this.options.subscriber?.messageDequeued(message);

apps/webapp/test/fairDequeuingStrategy.test.ts

@@ -463,4 +463,134 @@ describe("FairDequeuingStrategy", () => {
       }
     }
   );
+
+  redisTest(
+    "should bias shuffling based on concurrency limits and available capacity",
+    async ({ redis }) => {
+      const keyProducer = createKeyProducer("test");
+      const now = Date.now();
+
+      // Setup three environments with different concurrency settings
+      const envSetups = [
+        {
+          envId: "env-1",
+          limit: 100,
+          current: 20, // Lots of available capacity
+          queueCount: 3,
+        },
+        {
+          envId: "env-2",
+          limit: 50,
+          current: 40, // Less available capacity
+          queueCount: 3,
+        },
+        {
+          envId: "env-3",
+          limit: 10,
+          current: 5, // Some available capacity
+          queueCount: 3,
+        },
+      ];
+
+      // Setup queues and concurrency for each environment
+      for (const setup of envSetups) {
+        await setupConcurrency({
+          redis,
+          keyProducer,
+          org: { id: "org-1", currentConcurrency: 0, limit: 200 },
+          env: {
+            id: setup.envId,
+            currentConcurrency: setup.current,
+            limit: setup.limit,
+          },
+        });
+
+        for (let i = 0; i < setup.queueCount; i++) {
+          await setupQueue({
+            redis,
+            keyProducer,
+            parentQueue: "parent-queue",
+            score: now - 1000 * (i + 1),
+            queueId: `queue-${i}`,
+            orgId: "org-1",
+            envId: setup.envId,
+          });
+        }
+      }
+
+      // Create multiple strategies with different seeds
+      const numStrategies = 5;
+      const strategies = Array.from(
+        { length: numStrategies },
+        (_, i) =>
+          new FairDequeuingStrategy({
+            tracer,
+            redis,
+            keys: keyProducer,
+            defaultOrgConcurrency: 10,
+            defaultEnvConcurrency: 5,
+            parentQueueLimit: 100,
+            checkForDisabledOrgs: true,
+            seed: `test-seed-${i}`,
+            biases: {
+              concurrencyLimitBias: 0.8,
+              availableCapacityBias: 0.5,
+            },
+          })
+      );
+
+      // Run iterations across all strategies
+      const iterationsPerStrategy = 100;
+      const allResults: Record<string, number>[] = [];
+
+      for (const strategy of strategies) {
+        const firstPositionCounts: Record<string, number> = {};
+
+        for (let i = 0; i < iterationsPerStrategy; i++) {
+          const result = await strategy.distributeFairQueuesFromParentQueue(
+            "parent-queue",
+            `consumer-${i % 3}`
+          );
+
+          expect(result.length).toBeGreaterThan(0);
+
+          const firstEnv = keyProducer.envIdFromQueue(result[0]);
+          firstPositionCounts[firstEnv] = (firstPositionCounts[firstEnv] || 0) + 1;
+        }
+
+        allResults.push(firstPositionCounts);
+      }
+
+      // Calculate average distributions across all strategies
+      const avgDistribution: Record<string, number> = {};
+      const envIds = ["env-1", "env-2", "env-3"];
+
+      for (const envId of envIds) {
+        const sum = allResults.reduce((acc, result) => acc + (result[envId] || 0), 0);
+        avgDistribution[envId] = sum / numStrategies;
+      }
+
+      // Log individual strategy results and the average
+      console.log("\nResults by strategy:");
+      allResults.forEach((result, i) => {
+        console.log(`Strategy ${i + 1}:`, result);
+      });
+
+      console.log("\nAverage distribution:", avgDistribution);
+
+      // Calculate percentages from average distribution
+      const totalCount = Object.values(avgDistribution).reduce((sum, count) => sum + count, 0);
+      const highLimitPercentage = (avgDistribution["env-1"] / totalCount) * 100;
+      const lowLimitPercentage = (avgDistribution["env-3"] / totalCount) * 100;
+
+      console.log("\nPercentages:");
+      console.log("High limit percentage:", highLimitPercentage);
+      console.log("Low limit percentage:", lowLimitPercentage);
+
+      // Verify distribution across all strategies
+      expect(highLimitPercentage).toBeLessThan(60);
+      expect(lowLimitPercentage).toBeGreaterThan(10);
+      expect(highLimitPercentage).toBeGreaterThan(lowLimitPercentage);
+    }
+  );
 });

references/v3-catalog/trigger.config.ts

@@ -13,7 +13,8 @@ export { handleError } from "./src/handleError.js";
 
 export default defineConfig({
   runtime: "node",
-  project: "yubjwjsfkxnylobaqvqz",
+  // project: "yubjwjsfkxnylobaqvqz",
+  project: "proj_oveexhsnehnarnusktpc",
   machine: "medium-1x",
   instrumentations: [new OpenAIInstrumentation()],
   additionalFiles: ["wrangler/wrangler.toml"],