diff --git a/.changeset/gentle-clouds-bet.md b/.changeset/gentle-clouds-bet.md new file mode 100644 index 00000000..94e57df9 --- /dev/null +++ b/.changeset/gentle-clouds-bet.md @@ -0,0 +1,5 @@ +--- +"apollo": patch +--- + +Queue requests to the python process diff --git a/.env.example b/.env.example index 5e945f07..078a1d3c 100644 --- a/.env.example +++ b/.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 diff --git a/platform/src/bridge.ts b/platform/src/bridge.ts index bd79ad6e..127253e8 100644 --- a/platform/src/bridge.ts +++ b/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 => { return new Promise((resolve, reject) => { const proc = spawn( "poetry", diff --git a/platform/src/util/errors.ts b/platform/src/util/errors.ts index 3bf1fb0a..ce59b67e 100644 --- a/platform/src/util/errors.ts +++ b/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. */ diff --git a/platform/src/util/python-queue.ts b/platform/src/util/python-queue.ts new file mode 100644 index 00000000..2d91ab48 --- /dev/null +++ b/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 = { + promise: Promise; + 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 = ( + scriptName: string, + task: () => Promise +): Enqueued => { + let removeFromQueue: ((reason: Error) => void) | null = null; + let timer: ReturnType | undefined; + + const promise = new Promise((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")), + }; +}; diff --git a/platform/src/util/sentry.ts b/platform/src/util/sentry.ts index 23634609..91d6f8d3 100644 --- a/platform/src/util/sentry.ts +++ b/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 +): void => { + if (!enabled) return; + Sentry.captureMessage(message, extras ? { extra: extras } : undefined); +}; diff --git a/platform/test/bridge-queue-fork.test.ts b/platform/test/bridge-queue-fork.test.ts new file mode 100644 index 00000000..364c1c91 --- /dev/null +++ b/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((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); +}); diff --git a/platform/test/bridge-queue.test.ts b/platform/test/bridge-queue.test.ts new file mode 100644 index 00000000..f6175969 --- /dev/null +++ b/platform/test/bridge-queue.test.ts @@ -0,0 +1,134 @@ +import { afterEach, describe, expect, it } from "bun:test"; + +import { run } from "../src/bridge"; +import { ApolloThrowable } from "../src/util/errors"; + +const PORT = 9872; + +// Mirrors bridge-cancel.test.ts's probe helpers: the spawned command is +// `poetry run python ...`, so matching the process list alone hits the +// poetry wrapper seconds before the interpreter exists. +const started = () => { + let resolve!: () => void; + const promise = new Promise((r) => { + resolve = r; + }); + return { + promise, + onEvent: (type: string) => { + if (type === "probe_started") resolve(); + }, + }; +}; + +const probeProcessCount = async () => { + const proc = Bun.spawn(["pgrep", "-f", "_cancel_probe"], { + stdout: "pipe", + stderr: "pipe", + }); + const out = await new Response(proc.stdout).text(); + return out.trim().split("\n").filter(Boolean).length; +}; + +const waitForProbesToClear = async (limitMs = 15_000) => { + const deadline = Date.now() + limitMs; + while (Date.now() < deadline) { + if ((await probeProcessCount()) === 0) return true; + await Bun.sleep(200); + } + return false; +}; + +afterEach(() => { + delete process.env.APOLLO_CONCURRENCY; +}); + +describe("queueing python runs", () => { + it("holds a second 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: 2 }, + undefined, + firstProbe.onEvent + ); + await firstProbe.promise; + + // The cap is 1 and it's held by `first` - a second run must queue rather + // than spawn a second interpreter on top of it. + const second = run("_cancel_probe", PORT, { sleep_for: 0 }); + await Bun.sleep(300); + expect(await probeProcessCount()).toBe(1); + + await first; + // Freed now that `first` is done - `second` should get its slot and run. + await second; + + expect(await waitForProbesToClear()).toBe(true); + }, 30_000); + + it("rejects a queued run the moment its caller aborts, without ever spawning it", async () => { + process.env.APOLLO_CONCURRENCY = "1"; + + const holderProbe = started(); + const holder = run( + "_cancel_probe", + PORT, + { sleep_for: 3 }, + undefined, + holderProbe.onEvent + ); + await holderProbe.promise; + expect(await probeProcessCount()).toBe(1); + + const abort = new AbortController(); + const queued = run( + "_cancel_probe", + PORT, + { sleep_for: 0 }, + undefined, + undefined, + abort.signal + ); + + abort.abort(); + const error = (await queued.catch((e) => e)) as ApolloThrowable; + expect(error).toBeInstanceOf(ApolloThrowable); + expect(error.type).toBe("SUBPROCESS_CANCELLED"); + + // The cancelled queue entry never became a process - `holder` is still + // the only one on the machine. + expect(await probeProcessCount()).toBe(1); + + await holder; // let it finish naturally (sleep_for: 3) + expect(await waitForProbesToClear()).toBe(true); + }, 30_000); + + it("runs up to the cap concurrently and queues the rest, with a cap above 1", async () => { + process.env.APOLLO_CONCURRENCY = "2"; + + let startedCount = 0; + let resolveTwoStarted!: () => void; + const twoStarted = new Promise((r) => (resolveTwoStarted = r)); + + const probes = [started(), started(), started()]; + const runs = probes.map((probe) => + run("_cancel_probe", PORT, { sleep_for: 1 }, undefined, (type) => { + probe.onEvent(type); + if (type === "probe_started" && ++startedCount === 2) { + resolveTwoStarted(); + } + }) + ); + + await twoStarted; + await Bun.sleep(200); + expect(await probeProcessCount()).toBe(2); + + await Promise.all(runs); + expect(await waitForProbesToClear()).toBe(true); + }, 30_000); +}); diff --git a/platform/test/util/python-queue.test.ts b/platform/test/util/python-queue.test.ts new file mode 100644 index 00000000..19a53afd --- /dev/null +++ b/platform/test/util/python-queue.test.ts @@ -0,0 +1,268 @@ +import { describe, expect, it, spyOn } from "bun:test"; + +import { + activeCount, + enqueue, + maxConcurrency, + queueDepth, + queueTimeoutMs, + QueueTimeoutError, +} from "../../src/util/python-queue"; +import * as sentry from "../../src/util/sentry"; + +// A task that stays running until `release()` is called, so a test can hold +// a slot open for as long as it likes. +const holdable = () => { + let release!: () => void; + const task = () => + new Promise((resolve) => { + release = resolve; + }); + return { task, release: () => release() }; +}; + +// Each test drives the module-level counters up and lets everything it +// enqueues finish (or cancels it), so the singleton returns to baseline and +// tests don't leak into each other - including bridge.ts's own use of the +// same module elsewhere in the suite. +describe("python-queue", () => { + it("runs a task immediately while under the cap", async () => { + process.env.APOLLO_CONCURRENCY = "2"; + try { + const before = activeCount(); + const holder = holdable(); + const { promise } = enqueue("test", holder.task); + // Granted synchronously - active bumps before anything is awaited. + expect(activeCount()).toBe(before + 1); + holder.release(); + await promise; + expect(activeCount()).toBe(before); + } finally { + delete process.env.APOLLO_CONCURRENCY; + } + }); + + it("queues past the cap and runs in FIFO order", async () => { + process.env.APOLLO_CONCURRENCY = "1"; + try { + const first = holdable(); + enqueue("first", first.task); + expect(activeCount()).toBe(1); + + const order: string[] = []; + const second = enqueue("second", async () => { + order.push("second"); + }); + const third = enqueue("third", async () => { + order.push("third"); + }); + + // Neither can have run yet - the only slot is held by "first". + await Bun.sleep(10); + expect(order).toEqual([]); + expect(queueDepth()).toBe(2); + + first.release(); // frees "first"'s slot - "second" then "third" cascade + // through on the same slot, since neither holds it open. + await third.promise; + expect(order).toEqual(["second", "third"]); + expect(queueDepth()).toBe(0); + + expect(activeCount()).toBe(0); + } finally { + delete process.env.APOLLO_CONCURRENCY; + } + }); + + it("cancels a queued task without ever running it, and without consuming a slot", async () => { + process.env.APOLLO_CONCURRENCY = "1"; + try { + const holder = holdable(); + enqueue("holder", holder.task); + expect(activeCount()).toBe(1); + + let ran = false; + const { promise, cancel } = enqueue("waiter", async () => { + ran = true; + }); + expect(queueDepth()).toBe(1); + + cancel(); + await expect(promise).rejects.toBeTruthy(); + + // The cancel dropped it from the line - it must not sit there forever + // waiting for a slot that will now never come looking for it. + expect(queueDepth()).toBe(0); + // Still 1: the held slot was never touched by the cancelled waiter. + expect(activeCount()).toBe(1); + expect(ran).toBe(false); + + holder.release(); + await Bun.sleep(5); + expect(activeCount()).toBe(0); + } finally { + delete process.env.APOLLO_CONCURRENCY; + } + }); + + it("treats cancelling an already-started task as a no-op", async () => { + process.env.APOLLO_CONCURRENCY = "1"; + try { + const { promise, cancel } = enqueue("immediate", async () => "done"); + cancel(); // already running - nothing left to cancel + await expect(promise).resolves.toBe("done"); + expect(activeCount()).toBe(0); + } finally { + delete process.env.APOLLO_CONCURRENCY; + } + }); + + it("drops a cancelled middle waiter without disturbing the others", async () => { + process.env.APOLLO_CONCURRENCY = "1"; + try { + const holder = holdable(); + enqueue("holder", holder.task); + + const order: string[] = []; + const first = enqueue("first", async () => { + order.push("first"); + }); + const middle = enqueue("middle", async () => { + order.push("middle"); + }); + const last = enqueue("last", async () => { + order.push("last"); + }); + expect(queueDepth()).toBe(3); + + middle.cancel(); + await expect(middle.promise).rejects.toBeTruthy(); + expect(queueDepth()).toBe(2); // only the cancelled one left the line + + holder.release(); // "holder" done - "first" is still at the front + await first.promise; + await last.promise; // "middle" is gone, so "last" is next + + expect(order).toEqual(["first", "last"]); + expect(activeCount()).toBe(0); + } finally { + delete process.env.APOLLO_CONCURRENCY; + } + }); + + it("never exceeds the cap under concurrent load, and admits everyone eventually", async () => { + process.env.APOLLO_CONCURRENCY = "2"; + try { + let observedPeak = 0; + let completed = 0; + const total = 6; + + const runs = Array.from({ length: total }, (_, i) => + enqueue(`req-${i}`, async () => { + observedPeak = Math.max(observedPeak, activeCount()); + await Bun.sleep(5); // hold the slot briefly so others really do queue + completed++; + }) + ); + + await Promise.all(runs.map((r) => r.promise)); + + expect(completed).toBe(total); + expect(observedPeak).toBeLessThanOrEqual(2); + expect(observedPeak).toBeGreaterThan(0); + expect(activeCount()).toBe(0); + expect(queueDepth()).toBe(0); + } finally { + delete process.env.APOLLO_CONCURRENCY; + } + }); + + it("rejects a queued task with QueueTimeoutError once it waits past the timeout, without ever running it or consuming a slot", async () => { + process.env.APOLLO_CONCURRENCY = "1"; + process.env.APOLLO_QUEUE_TIMEOUT_MS = "20"; + try { + const holder = holdable(); + enqueue("holder", holder.task); + + let ran = false; + const { promise } = enqueue("waiter", async () => { + ran = true; + }); + + await expect(promise).rejects.toBeInstanceOf(QueueTimeoutError); + expect(queueDepth()).toBe(0); + expect(activeCount()).toBe(1); // holder's slot is untouched + expect(ran).toBe(false); + + holder.release(); + await Bun.sleep(5); + expect(activeCount()).toBe(0); + } finally { + delete process.env.APOLLO_CONCURRENCY; + delete process.env.APOLLO_QUEUE_TIMEOUT_MS; + } + }, 10_000); + + it("reports a Sentry message when a task is queued", async () => { + process.env.APOLLO_CONCURRENCY = "1"; + const capture = spyOn(sentry, "captureMessage").mockImplementation( + () => {} + ); + try { + const holder = holdable(); + enqueue("holder", holder.task); + + const { promise } = enqueue("waiter", async () => {}); + expect(capture).toHaveBeenCalledTimes(1); + expect(String(capture.mock.calls[0]?.[0])).toContain("waiter"); + + holder.release(); + await promise; + } finally { + capture.mockRestore(); + delete process.env.APOLLO_CONCURRENCY; + } + }); + + describe("maxConcurrency", () => { + it("falls back to the default when unset, empty, non-numeric, or zero", () => { + for (const value of [undefined, "", "not-a-number", "0"]) { + if (value === undefined) delete process.env.APOLLO_CONCURRENCY; + else process.env.APOLLO_CONCURRENCY = value; + + expect(maxConcurrency()).toBe(10); + } + delete process.env.APOLLO_CONCURRENCY; + }); + + it("uses a configured positive integer", () => { + process.env.APOLLO_CONCURRENCY = "3"; + try { + expect(maxConcurrency()).toBe(3); + } finally { + delete process.env.APOLLO_CONCURRENCY; + } + }); + }); + + describe("queueTimeoutMs", () => { + it("falls back to the default when unset, empty, non-numeric, or zero", () => { + for (const value of [undefined, "", "not-a-number", "0"]) { + if (value === undefined) delete process.env.APOLLO_QUEUE_TIMEOUT_MS; + else process.env.APOLLO_QUEUE_TIMEOUT_MS = value; + + expect(queueTimeoutMs()).toBe(5 * 60 * 1000); + } + delete process.env.APOLLO_QUEUE_TIMEOUT_MS; + }); + + it("uses a configured value", () => { + process.env.APOLLO_QUEUE_TIMEOUT_MS = "1234"; + try { + expect(queueTimeoutMs()).toBe(1234); + } finally { + delete process.env.APOLLO_QUEUE_TIMEOUT_MS; + } + }); + }); +});