feat: queuing of calls to python process

.changeset/gentle-clouds-bet.md

@@ -0,0 +1,5 @@
+---
+"apollo": patch
+---
+
+Queue requests to the python process

.env.example

@@ -45,6 +45,16 @@
 # See services/name_rules.py.
 APOLLO_UNICODE_STEP_NAMES=false
 
+# Max concurrent service calls bridge.ts runs at once, per Apollo process
+# (python or TypeScript - not python-specific). Calls past this cap queue in
+# FIFO order rather than running immediately. Defaults to 10.
+# See platform/src/util/python-queue.ts.
+# APOLLO_CONCURRENCY=10
+
+# How long a queued call waits for a free slot before giving up (ms).
+# Defaults to 5 minutes.
+# APOLLO_QUEUE_TIMEOUT_MS=300000
+
 ANTHROPIC_API_KEY=sk-YOUR-API-KEY-HERE
 
 OPENAI_API_KEY=sk-YOUR-API-KEY-HERE

platform/src/bridge.ts

@@ -6,13 +6,15 @@ import { getInternalToken } from "./auth/internal-token";
 import {
   emptyResult,
   malformedResult,
+  queueTimeout,
   subprocessCancelled,
   subprocessFailed,
   subprocessKilled,
   subprocessSpawnFailed,
 } from "./util/errors";
+import { enqueue, QueueTimeoutError } from "./util/python-queue";
 import pkg from "../../package.json";
-import { USE_FORK_SERVER, runForked } from "./bridge.fork";
+import { USE_FORK_SERVER, runForked as runWithFork } from "./bridge.fork";
 import { LOG_LINE } from "./util/log-line";
 
 /**
@@ -38,6 +40,10 @@ export const run = async (
   const inputPath = tmpfile.replace("{}", "input");
   const outputPath = tmpfile.replace("{}", "output");
 
+  // Cleanup shared by every path below that never spawns.
+  const removeTmpFiles = () =>
+    Promise.all([rm(inputPath).catch(() => {}), rm(outputPath).catch(() => {})]);
+
   // Outside the promise, deliberately. The Promise constructor only catches a
   // synchronous throw from its executor, so an await that rejects in there -
   // a full disk, a read-only tmp - leaves the promise pending for ever and
@@ -52,17 +58,48 @@ export const run = async (
 
     await Bun.write(outputPath, "");
   } catch (error) {
-    // Removed rather than left behind by a half-finished setup, for the same
-    // reason it is 0600 above.
-    await rm(inputPath).catch(() => {});
-    await rm(outputPath).catch(() => {});
+    await removeTmpFiles();
     throw subprocessSpawnFailed(scriptName, error);
   }
 
-  if (USE_FORK_SERVER) {
-    return runForked(scriptName, port, inputPath, outputPath, onLog, onEvent, signal);
+  // One cap for both paths below (see util/python-queue.ts): forking is cheap
+  // to start, but the work it does is just as heavy as a fresh interpreter.
+  const runFn = USE_FORK_SERVER ? runWithFork : runWithSpawn;
+
+  let started = false;
+  const { promise, cancel } = enqueue(scriptName, () => {
+    started = true;
+    return runFn(scriptName, port, inputPath, outputPath, onLog, onEvent, signal);
+  });
+
+  if (signal) {
+    if (signal.aborted) cancel();
+    else signal.addEventListener("abort", cancel, { once: true });
+  }
+
+  try {
+    return await promise;
+  } catch (error) {
+    if (!started) {
+      await removeTmpFiles();
+      if (error instanceof QueueTimeoutError) {
+        throw queueTimeout(scriptName);
+      }
+      throw subprocessCancelled(scriptName, "QUEUED");
+    }
+    throw error;
   }
+};
 
+const runWithSpawn = (
+  scriptName: string,
+  port: number,
+  inputPath: string,
+  outputPath: string,
+  onLog?: (str: string) => void,
+  onEvent?: (type: string, payload: any) => void,
+  signal?: AbortSignal
+): Promise<JSON | null> => {
   return new Promise<JSON | null>((resolve, reject) => {
     const proc = spawn(
       "poetry",

platform/src/util/errors.ts

@@ -94,6 +94,15 @@ export function subprocessCancelled(
   );
 }
 
+export function queueTimeout(service: string): ApolloThrowable {
+  return new ApolloThrowable(
+    503,
+    "QUEUE_TIMEOUT",
+    `Service "${service}" timed out waiting for a free slot`,
+    { service }
+  );
+}
+
 /** Something outside Apollo killed the service - most often the OOM killer.
  *  Distinct from a cancellation, which is us, and from a non-zero exit, which
  *  is the service deciding to stop. */

platform/src/util/python-queue.ts

@@ -0,0 +1,99 @@
+// FIFO queue capping how many service calls bridge.ts runs at once.
+// Calls past the cap wait here instead of running immediately.
+
+import { captureMessage } from "./sentry";
+
+const DEFAULT_CONCURRENCY = 10;
+
+export const maxConcurrency = (): number =>
+  Number(process.env.APOLLO_CONCURRENCY) || DEFAULT_CONCURRENCY;
+
+const DEFAULT_QUEUE_TIMEOUT_MS = 5 * 60 * 1000;
+
+export const queueTimeoutMs = (): number =>
+  Number(process.env.APOLLO_QUEUE_TIMEOUT_MS) || DEFAULT_QUEUE_TIMEOUT_MS;
+
+let active = 0;
+const pending: Array<() => void> = [];
+
+export const activeCount = (): number => active;
+export const queueDepth = (): number => pending.length;
+
+// Distinct from cancel()'s rejection, so a caller can tell "the server gave
+// up" apart from "the client went away".
+export class QueueTimeoutError extends Error {
+  constructor() {
+    super("timed out waiting in the queue");
+    this.name = "QueueTimeoutError";
+  }
+}
+
+export type Enqueued<T> = {
+  promise: Promise<T>;
+  cancel: () => void;
+};
+
+// cancel() drops a still-queued task without ever running it; calling it
+// once the task has started is a no-op.
+export const enqueue = <T>(
+  scriptName: string,
+  task: () => Promise<T>
+): Enqueued<T> => {
+  let removeFromQueue: ((reason: Error) => void) | null = null;
+  let timer: ReturnType<typeof setTimeout> | undefined;
+
+  const promise = new Promise<T>((resolve, reject) => {
+    const run = () => {
+      if (timer) clearTimeout(timer);
+      active++;
+      task().then(
+        (value) => {
+          active--;
+          pending.shift()?.();
+          resolve(value);
+        },
+        (error) => {
+          active--;
+          pending.shift()?.();
+          reject(error);
+        }
+      );
+    };
+
+    if (active < maxConcurrency()) {
+      run();
+      return;
+    }
+
+    console.log(
+      `queueing ${scriptName}: ${active}/${maxConcurrency()} running, ${
+        pending.length + 1
+      } waiting`
+    );
+    captureMessage(`service call queued: ${scriptName}`, {
+      scriptName,
+      active,
+      queueDepth: pending.length + 1,
+    });
+
+    pending.push(run);
+    timer = setTimeout(
+      () => removeFromQueue?.(new QueueTimeoutError()),
+      queueTimeoutMs()
+    );
+    timer.unref?.();
+
+    removeFromQueue = (reason: Error) => {
+      const i = pending.indexOf(run);
+      if (i === -1) return; // already started - too late to cancel
+      pending.splice(i, 1);
+      if (timer) clearTimeout(timer);
+      reject(reason);
+    };
+  });
+
+  return {
+    promise,
+    cancel: () => removeFromQueue?.(new Error("queue wait cancelled")),
+  };
+};

platform/src/util/sentry.ts

@@ -46,3 +46,13 @@ export const captureException = (
   if (!enabled) return;
   Sentry.captureException(err, extras ? { extra: extras } : undefined);
 };
+
+// Same no-op-when-disabled contract as captureException, for events that
+// aren't an error.
+export const captureMessage = (
+  message: string,
+  extras?: Record<string, unknown>
+): void => {
+  if (!enabled) return;
+  Sentry.captureMessage(message, extras ? { extra: extras } : undefined);
+};

platform/test/bridge-queue-fork.test.ts

@@ -0,0 +1,118 @@
+// `bridge.fork.ts` reads APOLLO_FORK_SERVER into a top-level `const` at
+// import time, so it can only be toggled by an env var set in the shell
+// *before* bun starts - not from inside a test. That also means this file
+// can't share a `bun test` process with bridge-queue.test.ts or
+// bridge-cancel.test.ts: whichever file bun imports "../src/bridge" from
+// first freezes USE_FORK_SERVER for every file in that run. Run this one on
+// its own:
+//
+//   OBJC_DISABLE_INITIALIZE_FORK_SAFETY=YES APOLLO_FORK_SERVER=true \
+//     bun test platform/test/bridge-queue-fork.test.ts
+//
+// (the OBJC_DISABLE_... var is macOS-only; see .env.example)
+//
+// No pgrep-based process counting here, unlike bridge-queue.test.ts: a
+// forked child never execs, so `ps`/`pgrep -f` only ever sees the master's
+// original "entry.fork.py" command line, not the service name - there is no
+// OS-level signal that distinguishes "a request is running" from "the master
+// is idle". activeCount()/queueDepth() are the only reliable signal here.
+
+import { afterEach, describe, expect, it } from "bun:test";
+
+import { run } from "../src/bridge";
+import { USE_FORK_SERVER } from "../src/bridge.fork";
+import { ApolloThrowable } from "../src/util/errors";
+import { activeCount, queueDepth } from "../src/util/python-queue";
+
+const PORT = 9873;
+
+const started = () => {
+  let resolve!: () => void;
+  const promise = new Promise<void>((r) => {
+    resolve = r;
+  });
+  return {
+    promise,
+    onEvent: (type: string) => {
+      if (type === "probe_started") resolve();
+    },
+  };
+};
+
+afterEach(() => {
+  delete process.env.APOLLO_CONCURRENCY;
+});
+
+// Skips cleanly (not a failure) when run as part of the default `bun test
+// platform/test` sweep, where APOLLO_FORK_SERVER is unset - see the file
+// header for why this suite needs its own process to run for real.
+describe.skipIf(!USE_FORK_SERVER)("queueing python runs (fork path)", () => {
+  it("holds a second forked run queued until the first (over the cap) finishes", async () => {
+    process.env.APOLLO_CONCURRENCY = "1";
+
+    const firstProbe = started();
+    const first = run(
+      "_cancel_probe",
+      PORT,
+      { sleep_for: 1 },
+      undefined,
+      firstProbe.onEvent
+    );
+    await firstProbe.promise; // the fork really did start and is doing work
+    expect(activeCount()).toBe(1);
+
+    const second = run("_cancel_probe", PORT, { sleep_for: 0 });
+    await Bun.sleep(200);
+    // Still held by `first` - `second` is in line, not a second fork.
+    expect(activeCount()).toBe(1);
+    expect(queueDepth()).toBe(1);
+
+    await first;
+    await second; // released straight to `second`, per releaseSlot()'s contract
+
+    expect(activeCount()).toBe(0);
+    expect(queueDepth()).toBe(0);
+  }, 30_000);
+
+  it("rejects a queued forked run on abort without ever forking it", async () => {
+    process.env.APOLLO_CONCURRENCY = "1";
+
+    const holderProbe = started();
+    const holder = run(
+      "_cancel_probe",
+      PORT,
+      { sleep_for: 2 },
+      undefined,
+      holderProbe.onEvent
+    );
+    await holderProbe.promise;
+    expect(activeCount()).toBe(1);
+
+    const abort = new AbortController();
+    const queued = run(
+      "_cancel_probe",
+      PORT,
+      { sleep_for: 0 },
+      undefined,
+      undefined,
+      abort.signal
+    );
+    // run() awaits writing its tmp files before it ever reaches
+    // acquireSlot() - give that a tick before checking it actually queued.
+    await Bun.sleep(50);
+    expect(queueDepth()).toBe(1);
+
+    abort.abort();
+    const error = (await queued.catch((e) => e)) as ApolloThrowable;
+    expect(error).toBeInstanceOf(ApolloThrowable);
+    expect(error.type).toBe("SUBPROCESS_CANCELLED");
+
+    // The cancelled wait never became a fork - the slot is still only
+    // `holder`'s.
+    expect(queueDepth()).toBe(0);
+    expect(activeCount()).toBe(1);
+
+    await holder;
+    expect(activeCount()).toBe(0);
+  }, 30_000);
+});