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/README.md b/README.md index 179bdddd..38c5a60f 100644 --- a/README.md +++ b/README.md @@ -65,7 +65,12 @@ bun dev To see an index of the available language services, head to `localhost:3000`. -The chat services search an index of the OpenFn documentation, which every environment populates itself. Two backends are supported: Pinecone (the default) and Postgres with the `pgvector` extension (opt-in, see [Database](#database)). Run `bun py embed_docsite` to build yours; see [embed_docsite](services/embed_docsite/README.md) for the write-target options and [search_docsite](services/search_docsite/README.md) for backend selection. +The chat services search an index of the OpenFn documentation, which every +environment populates itself. Two backends are supported: Pinecone (the default) +and Postgres with the `pgvector` extension (opt-in, see [Database](#database)). +Run `bun py embed_docsite` to build yours; see +[embed_docsite](services/embed_docsite/README.md) for the write-target options +and [search_docsite](services/search_docsite/README.md) for backend selection. ## Python Setup @@ -155,10 +160,10 @@ Controls which characters a workflow step name may contain. Apollo sanitises step names on the way out and Lightning validates them on the way in, so the two rules have to agree. The default is `false`. -| Value | Rule | -| --- | --- | -| `false` (default) | ASCII only. Letters, digits, spaces, hyphens, underscores. Accents are folded (`Café` becomes `Cafe`) and anything else is dropped. This is the rule Lightning enforces today. | -| `true` | Anything except control characters. Letters and marks from any script, all punctuation and symbols, emoji, `/`, `:`, `>`, `&`, quotes and apostrophes. `Vérifier l'état` and `患者確認` survive exactly as typed. | +| Value | Rule | +| ----------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | +| `false` (default) | ASCII only. Letters, digits, spaces, hyphens, underscores. Accents are folded (`Café` becomes `Cafe`) and anything else is dropped. This is the rule Lightning enforces today. | +| `true` | Anything except control characters. Letters and marks from any script, all punctuation and symbols, emoji, `/`, `:`, `>`, `&`, quotes and apostrophes. `Vérifier l'état` and `患者確認` survive exactly as typed. | Leave it off until Lightning ships Unicode step names (Lightning#4577), then turn it on. Turning it on first means Apollo emits names Lightning rejects. @@ -170,18 +175,18 @@ Lightning is the worse of the two failures: Lightning rejecting a name is loud and recoverable, whereas Apollo quietly renaming a valid name is the silent vandalism this flag exists to prevent. -The rejected control set is the same in both modes: C0 (`U+0000`-`U+001F`, -NUL included), DEL (`U+007F`), C1 (`U+0080`-`U+009F`), the noncharacters -`U+FFFE` and `U+FFFF`, the surrogates `U+D800`-`U+DFFF`, and the line and -paragraph separators `U+2028` and `U+2029`. A NUL byte in a name crashes the Postgres insert on -Lightning's side. Names are NFC-normalised in both modes so that Apollo and -Lightning agree on how to spell an accent, which is what step lookup matches -on, and capped at 100 graphemes because that is what Ecto's `validate_length` -counts. +The rejected control set is the same in both modes: C0 (`U+0000`-`U+001F`, NUL +included), DEL (`U+007F`), C1 (`U+0080`-`U+009F`), the noncharacters `U+FFFE` +and `U+FFFF`, the surrogates `U+D800`-`U+DFFF`, and the line and paragraph +separators `U+2028` and `U+2029`. A NUL byte in a name crashes the Postgres +insert on Lightning's side. Names are NFC-normalised in both modes so that +Apollo and Lightning agree on how to spell an accent, which is what step lookup +matches on, and capped at 100 graphemes because that is what Ecto's +`validate_length` counts. The rule lives in `services/name_rules.py`, and everything that states or -enforces it is derived from there: the sanitiser, the workflow-generation -prompt (`describe_rule_for_prompt`), the acceptance-test judges +enforces it is derived from there: the sanitiser, the workflow-generation prompt +(`describe_rule_for_prompt`), the acceptance-test judges (`describe_rule_for_judge`, substituted into the rubric markdown by `judges.load_judge`), and the `assert_no_special_chars` test assertion. Change the rule in that one file and all four follow. @@ -195,15 +200,15 @@ an emoji ZWJ run at the joiner unless a pictograph follows) and the whole point is to agree with Elixir, not with the spec. `tools/unicode_parity` is the harness that checks it. Run `python3 edges.py`, -then `elixir probe.exs`, then `python3 check.py` with the Elixir version Lightning runs; `--tables` -prints the literals to paste back into `name_rules`. It checks five things: -every codepoint's break class, the `Extended_Pictographic` set, the trim set, -what a GB11 emoji run may be separated from its joiner by, and cluster -boundaries over a generated corpus. Normalisation is not among them: -`normalize_nfc` is the standard library's, so there is no table of ours to -check against Elixir. - -`Extended_Pictographic` needs its own check because it is *not* a break class, +then `elixir probe.exs`, then `python3 check.py` with the Elixir version +Lightning runs; `--tables` prints the literals to paste back into `name_rules`. +It checks five things: every codepoint's break class, the +`Extended_Pictographic` set, the trim set, what a GB11 emoji run may be +separated from its joiner by, and cluster boundaries over a generated corpus. +Normalisation is not among them: `normalize_nfc` is the standard library's, so +there is no table of ours to check against Elixir. + +`Extended_Pictographic` needs its own check because it is _not_ a break class, so a per-codepoint sweep cannot see it — an over-broad set there silently changes clustering either side of a ZWJ and nothing else notices. That is exactly how a hand-written table with 531 wrong codepoints survived two rounds @@ -250,6 +255,29 @@ service: curl http://localhost:3000/services/status ``` +### Capacity testing + +`tools/capacity-test.js` fires parallel requests at the `echo` service of a +running server and logs each request's number, duration and response code, +followed by the total time and the count and percentage of requests that +succeeded (a 2xx response). It sends 100 requests by default; pass a number as +the first argument to change that: + +```bash +bun run test:capacity +bun run test:capacity 500 +``` + +It reads `APOLLO_ENDPOINT` (default `http://localhost:3000`) and sends +`AI_ASSISTANT_API_KEY` as the `api_key` in each payload. Requests are not +throttled, and log lines appear in completion order rather than request order. + +To restrict memory usage locally, launch with Docker: + +``` +docker run --memory=512m --memory-swap=512m -p 3000:3000 --network host --env-file .env apollo +``` + ## Troubleshooting If you get errors like `poetry: command not found` (error code 127), and poetry @@ -305,9 +333,8 @@ them is safe: `load_adaptor_docs` runs, so applying it by hand is optional. - `lightning_clients` - created and kept current by the migration runner (`platform/src/db/migrate.ts`, migrations under - [`migrations/clients/`](migrations/clients/)). It is applied - automatically at Apollo startup when a database URL is set; no manual `psql` - step is needed. + [`migrations/clients/`](migrations/clients/)). It is applied automatically at + Apollo startup when a database URL is set; no manual `psql` step is needed. First, make sure you've configured your desired `POSTGRES_URL` in your `.env` file. diff --git a/package.json b/package.json index d0a5841a..55b1a580 100644 --- a/package.json +++ b/package.json @@ -8,6 +8,7 @@ "dev": "bun --watch platform/src/index.ts", "build": "bun build platform/src/index.ts", "test": "bun test platform/test ", + "test:capacity": "bun tools/capacity-test.js", "client": "bun platform/src/auth/client/cli.ts", "migrate": "bun platform/src/db/migrate.ts", "py": "poetry run python services/entry.py " 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; + } + }); + }); +}); diff --git a/tools/capacity-test.js b/tools/capacity-test.js new file mode 100644 index 00000000..e7e64bb0 --- /dev/null +++ b/tools/capacity-test.js @@ -0,0 +1,47 @@ +// Calls the echo service repeatedly against a running server and logs timings +// Usage: APOLLO_ENDPOINT=http://localhost:3000 bun tools/capacity-test.js [iterations] +const endpoint = process.env.APOLLO_ENDPOINT ?? "http://localhost:3000"; +const api_key = process.env.AI_ASSISTANT_API_KEY; +const iterations = process.argv[2] ? Number(process.argv[2]) : 100; + +if (!Number.isInteger(iterations) || iterations < 1) { + console.error(`invalid iteration count: ${process.argv[2]}`); + process.exit(1); +} + +if (!endpoint) { + console.error("APOLLO_ENDPOINT is not set"); + process.exit(1); +} + +const url = `${endpoint.replace(/\/$/, "")}/services/echo`; + +const start = performance.now(); + +let succeeded = 0; + +const call = async (i) => { + const t = performance.now(); + let status; + try { + const res = await fetch(url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ n: i, api_key }), + }); + await res.text(); + status = res.status; + if (res.ok) succeeded++; + } catch (e) { + status = `error (${e.message})`; + } + const ms = (performance.now() - t).toFixed(0); + console.log(`#${i} ${ms}ms ${status}`); +}; + +await Promise.all(Array.from({ length: iterations }, (_, i) => call(i + 1))); + +const total = ((performance.now() - start) / 1000).toFixed(2); +const pct = ((succeeded / iterations) * 100).toFixed(1); +console.log(`completed ${iterations} requests in ${total}s`); +console.log(`succeeded: ${succeeded}/${iterations} (${pct}%)`);