More v1 run engine fixes (#1644)

* Remove the socket count from shared queue consumer cause it can take up to 5s to complete

* Add support for re-using queue snapshots across dequeues

* 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

Diff for: apps/webapp/app/env.server.ts

@@ -232,6 +232,8 @@ 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),
+  MARQS_MAXIMUM_ORG_COUNT: z.coerce.number().int().optional(),
 
   PROD_TASK_HEARTBEAT_INTERVAL_MS: z.coerce.number().int().optional(),
 

Diff for: apps/webapp/app/v3/marqs/fairDequeuingStrategy.server.ts

@@ -44,6 +44,8 @@ export type FairDequeuingStrategyOptions = {
    * If not provided, no biasing will be applied (completely random shuffling)
    */
   biases?: FairDequeuingStrategyBiases;
+  reuseSnapshotCount?: number;
+  maximumOrgCount?: number;
 };
 
 type FairQueueConcurrency = {
@@ -90,6 +92,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,16 +316,53 @@ 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);
+      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>();
@@ -341,10 +384,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 +394,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 +421,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 +431,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,10 +448,66 @@ 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;
     });
   }
 
+  #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);

Diff for: apps/webapp/app/v3/marqs/index.server.ts

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

Diff for: apps/webapp/app/v3/marqs/sharedQueueConsumer.server.ts

@@ -1122,15 +1122,6 @@ export class SharedQueueConsumer {
         "emitResumeAfterDependencyWithAck",
         async (span) => {
           try {
-            const sockets = await this.#startActiveSpan("getCoordinatorSockets", async (span) => {
-              const sockets = await socketIo.coordinatorNamespace.fetchSockets();
-
-              span.setAttribute("socket_count", sockets.length);
-
-              return sockets;
-            });
-
-            span.setAttribute("socket_count", sockets.length);
             span.setAttribute("attempt_id", resumableAttempt.id);
             span.setAttribute(
               "timeout_in_ms",