Add support for re-using queue snapshots across dequeues

Author: ericallam

apps/webapp/app/env.server.ts

@@ -232,6 +232,7 @@ const EnvironmentSchema = z.object({
   MARQS_CONCURRENCY_LIMIT_BIAS: z.coerce.number().default(0.75),
   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),
 
   PROD_TASK_HEARTBEAT_INTERVAL_MS: z.coerce.number().int().optional(),
 

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

@@ -44,6 +44,7 @@ export type FairDequeuingStrategyOptions = {
    * If not provided, no biasing will be applied (completely random shuffling)
    */
   biases?: FairDequeuingStrategyBiases;
+  reuseSnapshotCount?: number;
 };
 
 type FairQueueConcurrency = {
@@ -90,6 +91,10 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
   }>;
 
   private _rng: seedrandom.PRNG;
+  private _reusedSnapshotForConsumer: Map<
+    string,
+    { snapshot: FairQueueSnapshot; reuseCount: number }
+  > = new Map();
 
   constructor(private options: FairDequeuingStrategyOptions) {
     const ctx = new DefaultStatefulContext();
@@ -310,6 +315,31 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
       span.setAttribute("consumer_id", consumerId);
       span.setAttribute("parent_queue", parentQueue);
 
+      if (
+        typeof this.options.reuseSnapshotCount === "number" &&
+        this.options.reuseSnapshotCount > 0
+      ) {
+        const key = `${parentQueue}:${consumerId}`;
+        const reusedSnapshot = this._reusedSnapshotForConsumer.get(key);
+
+        if (reusedSnapshot) {
+          if (reusedSnapshot.reuseCount < this.options.reuseSnapshotCount) {
+            span.setAttribute("reused_snapshot", true);
+
+            this._reusedSnapshotForConsumer.set(key, {
+              snapshot: reusedSnapshot.snapshot,
+              reuseCount: reusedSnapshot.reuseCount + 1,
+            });
+
+            return reusedSnapshot.snapshot;
+          } else {
+            this._reusedSnapshotForConsumer.delete(key);
+          }
+        }
+      }
+
+      span.setAttribute("reused_snapshot", false);
+
       const now = Date.now();
 
       const queues = await this.#allChildQueuesByScore(parentQueue, consumerId, now);
@@ -341,10 +371,6 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
         (org) => org.concurrency.current >= org.concurrency.limit
       );
 
-      span.setAttributes({
-        ...flattenAttributes(orgsAtFullConcurrency, "orgs_at_full_concurrency"),
-      });
-
       const orgIdsAtFullConcurrency = new Set(orgsAtFullConcurrency.map((org) => org.id));
 
       const orgsSnapshot = orgs.reduce((acc, org) => {
@@ -355,6 +381,12 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
         return acc;
       }, {} as Record<string, { concurrency: FairQueueConcurrency }>);
 
+      span.setAttributes({
+        org_count: orgs.length,
+        orgs_at_full_concurrency_count: orgsAtFullConcurrency.length,
+        orgs_snapshot_count: Object.keys(orgsSnapshot).length,
+      });
+
       if (Object.keys(orgsSnapshot).length === 0) {
         return emptyFairQueueSnapshot;
       }
@@ -376,10 +408,6 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
         (env) => env.concurrency.current >= env.concurrency.limit
       );
 
-      span.setAttributes({
-        ...flattenAttributes(envsAtFullConcurrency, "envs_at_full_concurrency"),
-      });
-
       const envIdsAtFullConcurrency = new Set(envsAtFullConcurrency.map((env) => env.id));
 
       const envsSnapshot = envs.reduce((acc, env) => {
@@ -390,6 +418,11 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
         return acc;
       }, {} as Record<string, { concurrency: FairQueueConcurrency }>);
 
+      span.setAttributes({
+        env_count: envs.length,
+        envs_at_full_concurrency_count: envsAtFullConcurrency.length,
+      });
+
       const queuesSnapshot = queues.filter(
         (queue) =>
           !orgIdsAtFullConcurrency.has(queue.org) && !envIdsAtFullConcurrency.has(queue.env)
@@ -402,6 +435,16 @@ export class FairDequeuingStrategy implements MarQSFairDequeueStrategy {
         queues: queuesSnapshot,
       };
 
+      if (
+        typeof this.options.reuseSnapshotCount === "number" &&
+        this.options.reuseSnapshotCount > 0
+      ) {
+        this._reusedSnapshotForConsumer.set(`${parentQueue}:${consumerId}`, {
+          snapshot,
+          reuseCount: 0,
+        });
+      }
+
       return snapshot;
     });
   }

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

@@ -1625,6 +1625,7 @@ function getMarQSClient() {
             availableCapacityBias: env.MARQS_AVAILABLE_CAPACITY_BIAS,
             queueAgeRandomization: env.MARQS_QUEUE_AGE_RANDOMIZATION_BIAS,
           },
+          reuseSnapshotCount: env.MARQS_REUSE_SNAPSHOT_COUNT,
         }),
         envQueuePriorityStrategy: new FairDequeuingStrategy({
           tracer: tracer,

apps/webapp/test/fairDequeuingStrategy.test.ts

@@ -203,6 +203,96 @@ describe("FairDequeuingStrategy", () => {
     expect(result).toEqual([queue1, queue2]);
   });
 
+  redisTest("should reuse snapshots across calls for the same consumer", async ({ redis }) => {
+    const keyProducer = createKeyProducer("test");
+    const strategy = new FairDequeuingStrategy({
+      tracer,
+      redis,
+      keys: keyProducer,
+      defaultOrgConcurrency: 10,
+      defaultEnvConcurrency: 5,
+      parentQueueLimit: 10,
+      checkForDisabledOrgs: true,
+      seed: "test-seed-reuse-1",
+      reuseSnapshotCount: 1,
+    });
+
+    const now = Date.now();
+
+    await setupQueue({
+      redis,
+      keyProducer,
+      parentQueue: "parent-queue",
+      score: now - 3000,
+      queueId: "queue-1",
+      orgId: "org-1",
+      envId: "env-1",
+    });
+
+    await setupQueue({
+      redis,
+      keyProducer,
+      parentQueue: "parent-queue",
+      score: now - 2000,
+      queueId: "queue-2",
+      orgId: "org-2",
+      envId: "env-2",
+    });
+
+    await setupQueue({
+      redis,
+      keyProducer,
+      parentQueue: "parent-queue",
+      score: now - 1000,
+      queueId: "queue-3",
+      orgId: "org-3",
+      envId: "env-3",
+    });
+
+    const startDistribute1 = performance.now();
+
+    const result = await strategy.distributeFairQueuesFromParentQueue("parent-queue", "consumer-1");
+
+    const distribute1Duration = performance.now() - startDistribute1;
+
+    console.log("First distribution took", distribute1Duration, "ms");
+
+    expect(result).toHaveLength(3);
+    // Should only get the two oldest queues
+    const queue1 = keyProducer.queueKey("org-1", "env-1", "queue-1");
+    const queue2 = keyProducer.queueKey("org-2", "env-2", "queue-2");
+    const queue3 = keyProducer.queueKey("org-3", "env-3", "queue-3");
+    expect(result).toEqual([queue2, queue1, queue3]);
+
+    const startDistribute2 = performance.now();
+
+    const result2 = await strategy.distributeFairQueuesFromParentQueue(
+      "parent-queue",
+      "consumer-1"
+    );
+
+    const distribute2Duration = performance.now() - startDistribute2;
+
+    console.log("Second distribution took", distribute2Duration, "ms");
+
+    // Make sure the second call is more than 10 times faster than the first
+    expect(distribute2Duration).toBeLessThan(distribute1Duration / 10);
+
+    const startDistribute3 = performance.now();
+
+    const result3 = await strategy.distributeFairQueuesFromParentQueue(
+      "parent-queue",
+      "consumer-1"
+    );
+
+    const distribute3Duration = performance.now() - startDistribute3;
+
+    console.log("Third distribution took", distribute3Duration, "ms");
+
+    // Make sure the third call is more than 4 times the second
+    expect(distribute3Duration).toBeGreaterThan(distribute2Duration * 4);
+  });
+
   redisTest("should fairly distribute queues across environments over time", async ({ redis }) => {
     const keyProducer = createKeyProducer("test");
     const strategy = new FairDequeuingStrategy({