diff --git a/.changeset/mighty-buckets-wish.md b/.changeset/mighty-buckets-wish.md new file mode 100644 index 0000000000..3664b50da6 --- /dev/null +++ b/.changeset/mighty-buckets-wish.md @@ -0,0 +1,16 @@ +--- +"graphile-build-pg": patch +"postgraphile": patch +"grafserv": patch +"grafast": patch +--- + +Address a slow memory leak in GraphQL subscriptions that can cause OOM errors +when a single subscription hits many hundreds of thousands of events. The cause +was racing an abortPromise with each event, causing the abort promise to build +up a large list of callbacks in its internal Promise mechanics when the event +always won the race. The fix was to move from using a promise for abort to an +AbortController, where the abort task can be released after each event +successfully resolves, avoiding memory buildup. Also applied this fix to similar +patterns elsewhere in the codebase, albeit places much less sensitive to this +issue. diff --git a/grafast/grafast/src/engine/distributor.ts b/grafast/grafast/src/engine/distributor.ts index 6448ae540a..93d82f8336 100644 --- a/grafast/grafast/src/engine/distributor.ts +++ b/grafast/grafast/src/engine/distributor.ts @@ -1,12 +1,8 @@ import * as assert from "../assert.ts"; import { isDev, noop } from "../dev.ts"; import type { Maybe } from "../interfaces.ts"; -import { - type PromiseWithResolve, - promiseWithResolve, -} from "../promiseWithResolve.ts"; import type { Step } from "../step.ts"; -import { arrayOfLength, isPromiseLike, sleep } from "../utils.ts"; +import { arrayOfLength, isPromiseLike } from "../utils.ts"; const DEFAULT_DISTRIBUTOR_BUFFER_SIZE = 1001; const DEFAULT_DISTRIBUTOR_BUFFER_SIZE_INCREMENT = 1001; @@ -27,10 +23,12 @@ export function isDistributor( } // Save on garbage collection by just using this promise for everything -const DONE_PROMISE: Promise> = Promise.resolve({ - done: true, - value: undefined, -}); +const DONE_PROMISE: Promise> = Promise.resolve( + Object.freeze({ + done: true, + value: undefined, + }), +); /** * Creates a "distributor" for the sourceIterable such that the dependent steps @@ -99,13 +97,19 @@ export function distributor( */ const buffer: Array>> = []; - // Easy way to resolve a promise for slowing down the fastest consumer - let wmi: PromiseWithResolve | null = null; - function lowWaterMarkIncreased(): PromiseLike { - if (wmi === null) { - wmi = promiseWithResolve(); - } - return wmi.promise; + // Consumers waiting for the low-water mark to advance. + let lowWaterMarkWaiters: Set<() => void> | null = null; + function waitForLowWaterMarkOrPause(): Promise { + return new Promise((resolve) => { + const waiters = (lowWaterMarkWaiters ??= new Set()); + const done = () => { + clearTimeout(timeout); + waiters.delete(done); + resolve(); + }; + const timeout = setTimeout(done, pauseDuration); + waiters.add(done); + }); } /** @@ -185,11 +189,14 @@ export function distributor( } // Announce that the lowWaterMark advanced - if (advanced && wmi !== null) { + if (advanced && lowWaterMarkWaiters !== null) { // Avoid race condition - const deferred = wmi; - wmi = null; - deferred.resolve(); + const waiters = lowWaterMarkWaiters; + lowWaterMarkWaiters = null; + + for (const resolve of waiters) { + resolve(); + } } } } @@ -286,10 +293,7 @@ export function distributor( // Whoa there! Getting a little ahead of ourselves! Wait for the slowest // consumer to advance (or for it to time out), before resolving. // const oldLowWaterMark = lowWaterMark; - const next = Promise.race([ - lowWaterMarkIncreased(), - sleep(pauseDuration), - ]).then( + const next = waitForLowWaterMarkOrPause().then( // const advanced = lowWaterMark > oldLowWaterMark; // TODO: should we wait a little longer if we did advance so we're // not creating a new timer for each and every low watermark diff --git a/grafast/grafast/src/index.ts b/grafast/grafast/src/index.ts index 11aac9113c..f3cf52a159 100644 --- a/grafast/grafast/src/index.ts +++ b/grafast/grafast/src/index.ts @@ -268,6 +268,7 @@ import type { ObjectTypeSpec, } from "./utils.ts"; import { + abortable, arrayOfLength, arraysMatch, asyncIteratorWithCleanup, @@ -414,6 +415,7 @@ export { $$idempotent, $$inhibit, $$verbatim, + abortable, access, AccessStep, applyInput, @@ -586,6 +588,7 @@ exportAsMany("grafast", { __TrackedValueStep, __ValueStep, inspect, + abortable, access, get, AccessStep, diff --git a/grafast/grafast/src/prepare.ts b/grafast/grafast/src/prepare.ts index ef571ed146..e94b91388b 100644 --- a/grafast/grafast/src/prepare.ts +++ b/grafast/grafast/src/prepare.ts @@ -49,9 +49,9 @@ import type { StreamMaybeMoreableArray, StreamMoreableArray, } from "./interfaces.ts"; -import { promiseWithResolve } from "./promiseWithResolve.ts"; import { timeSource } from "./timeSource.ts"; import { + abortable, arrayOfLength, asyncIteratorWithCleanup, isPromiseLike, @@ -325,7 +325,7 @@ function executePreemptive( onError: ErrorBehavior, outputDataAsString: boolean, executionTimeout: number | null, - abortSignal: AbortSignal, + requestAbortSignal: AbortSignal, ): PromiseOrDirect< ExecutionResult | AsyncGenerator > { @@ -368,7 +368,7 @@ function executePreemptive( stopTime, // toSerialize: [], eventEmitter: args[$$eventEmitter], - abortSignal, + abortSignal: requestAbortSignal, }; const bucketPromise = executeBucket(rootBucket, requestContext); @@ -480,13 +480,21 @@ function executePreemptive( // `releaseUnusedIterators(rootBucket, rootBucketIndex, null)` here. const arr = bucketRootValue as StreamMoreableArray; const stream = arr[$$streamMore]; - // Do the async iterable - let stopped = false; - const { promise: abortPromise, resolve: resolveAbort } = - promiseWithResolve(); + const iteratorAbortController = new AbortController(); + const abortIteratorWhenRequestAborts = () => + iteratorAbortController.abort(); + requestAbortSignal.addEventListener( + "abort", + abortIteratorWhenRequestAborts, + { once: true }, + ); + const iteratorAbortSignal = iteratorAbortController.signal; const iterator = newIterator((e) => { - stopped = true; - resolveAbort(); + iteratorAbortController.abort(); + requestAbortSignal.removeEventListener( + "abort", + abortIteratorWhenRequestAborts, + ); if (e != null) { try { const result = stream.throw?.(e); @@ -511,38 +519,61 @@ function executePreemptive( let i = 0; // eslint-disable-next-line no-constant-condition while (true) { - const next = await Promise.race([abortPromise, stream.next()]); - if (stopped || !next) { + const rawNext = stream.next(); + const next = isPromiseLike(rawNext) + ? await abortable(iteratorAbortSignal, undefined, rawNext) + : rawNext; + if (next?.done) { + // Stream already exited break; } - if (!next) { - iterator.throw(new Error("Invalid iteration")).then(null, noop); - break; - } - const { done, value } = next; - if (done) { - break; - } - const payload = await Promise.race([ - abortPromise, - executeStreamPayload(value, i), - ]); - if (payload === undefined) { + if (next === undefined || iteratorAbortSignal.aborted) { + const result = stream.return?.(); + if (isPromiseLike(result)) { + result.then(null, noop); + } break; } - if (isAsyncIterable(payload)) { - // FIXME: do we need to avoid 'for await' because it can cause the - // stream to exit late if we're waiting on a promise and the stream - // exits in the interrim? We're assuming that no promises will be - // sufficiently long-lived for this to be an issue right now. - // TODO: should probably tie all this into an AbortController/signal too - for await (const entry of payload) { - iterator.push(entry); + try { + const rawPayload = executeStreamPayload(next.value, i); + const payload = isPromiseLike(rawPayload) + ? await abortable(iteratorAbortSignal, undefined, rawPayload) + : rawPayload; + if (payload === undefined) { + break; + } + if (isAsyncIterable(payload)) { + const payloadIterator = payload[Symbol.asyncIterator](); + while (true) { + const next = await abortable( + iteratorAbortSignal, + undefined, + payloadIterator.next(), + ); + if (next?.done) { + // Iterator already exited + break; + } + if (next === undefined || iteratorAbortSignal.aborted) { + const result = payloadIterator.return?.(undefined); + if (isPromiseLike(result)) { + result.then(null, noop); + } + break; + } + iterator.push(next.value); + } + } else { + iterator.push(payload); + } + i++; + } catch (error) { + const result = iterator.return?.(); + if (isPromiseLike(result)) { + result.then(null, noop); } - } else { - iterator.push(payload); + throw error; } - i++; } })() .then( diff --git a/grafast/grafast/src/utils.ts b/grafast/grafast/src/utils.ts index 834ef60237..790adc18e7 100644 --- a/grafast/grafast/src/utils.ts +++ b/grafast/grafast/src/utils.ts @@ -1458,10 +1458,12 @@ export const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); // Save on garbage collection by just using this promise for everything -const DONE_PROMISE: Promise> = Promise.resolve({ - done: true, - value: undefined, -}); +const DONE_PROMISE: Promise> = Promise.resolve( + Object.freeze({ + done: true, + value: undefined, + }), +); /** * Returns a new version of `iterable` that calls `callback()` on termination, @@ -1549,3 +1551,32 @@ export function markSyncAndSafe< } return fn; } + +/** + * If `signal` is aborted, returns `valueForAbort`. Otherwise returns + * equivalent to `promiseOrValue` except if the signal is aborted first it will + * resolve with `valueForAbort` immediately. + */ +export function abortable( + signal: AbortSignal, + valueForAbort: F, + promise: PromiseLike, +): Promise { + if (signal.aborted) { + return Promise.resolve(valueForAbort); + } + return new Promise((resolve, reject) => { + const resolveWithoutArgs = () => resolve(valueForAbort); + signal.addEventListener("abort", resolveWithoutArgs, { once: true }); + return promise.then( + (val) => { + signal.removeEventListener("abort", resolveWithoutArgs); + resolve(val); + }, + (e) => { + signal.removeEventListener("abort", resolveWithoutArgs); + reject(e); + }, + ); + }); +} diff --git a/grafast/grafserv/src/core/base.ts b/grafast/grafserv/src/core/base.ts index 1be1e77832..3db65f9ead 100644 --- a/grafast/grafserv/src/core/base.ts +++ b/grafast/grafserv/src/core/base.ts @@ -572,7 +572,7 @@ function defaultMakeGetExecutionConfig(): ( let latestSchema: GraphQLSchema; let latestSchemaOrPromise: PromiseOrDirect; let latestParseAndValidate: ReturnType; - let schemaPrepare: Promise | null = null; + let schemaPrepareWaiters: Set> | null = null; return function getExecutionConfig(this: GrafservBase) { // Get up to date schema, in case we're in watch mode @@ -581,16 +581,40 @@ function defaultMakeGetExecutionConfig(): ( if (schemaOrPromise !== latestSchemaOrPromise) { latestSchemaOrPromise = schemaOrPromise; if ("then" in schemaOrPromise) { - schemaPrepare = (async () => { - latestSchema = await schemaOrPromise; - latestSchemaOrPromise = schemaOrPromise; - latestParseAndValidate = makeParseAndValidateFunction( - latestSchema, - resolvedPreset, - dynamicOptions, - ); - schemaPrepare = null; - return true; + const prepareWaiters = new Set>(); + schemaPrepareWaiters = prepareWaiters; + const releaseWaiters = (error?: Error) => { + // Make sure new waiters have their own batch + if (schemaPrepareWaiters === prepareWaiters) { + schemaPrepareWaiters = null; + } + if (error) { + for (const waiter of prepareWaiters) { + waiter.reject(error); + } + } else { + for (const waiter of prepareWaiters) { + waiter.resolve(true); + } + } + prepareWaiters.clear(); + }; + + // Kick off an async task that waits for the schema to be ready, + // completes setup, then informs all waiters of the result. + void (async () => { + try { + latestSchema = await schemaOrPromise; + latestSchemaOrPromise = schemaOrPromise; + latestParseAndValidate = makeParseAndValidateFunction( + latestSchema, + resolvedPreset, + dynamicOptions, + ); + releaseWaiters(); + } catch (error) { + releaseWaiters(error); + } })(); } else { if (latestSchema === schemaOrPromise) { @@ -605,10 +629,16 @@ function defaultMakeGetExecutionConfig(): ( } } } - if (schemaPrepare !== null) { + if (schemaPrepareWaiters !== null) { + const prepareWaiters = schemaPrepareWaiters; const sleeper = sleep(dynamicOptions.schemaWaitTime); - const schemaReadyPromise = Promise.race([schemaPrepare, sleeper.promise]); - return schemaReadyPromise.then((schemaReady) => { + const waiter = Promise.withResolvers(); + prepareWaiters.add(waiter); + sleeper.promise.then(() => { + prepareWaiters.delete(waiter); + waiter.resolve(false); + }, waiter.reject); + return waiter.promise.then((schemaReady) => { sleeper.release(); if (schemaReady !== true) { // Handle missing schema diff --git a/graphile-build/graphile-build-pg/src/plugins/PgIntrospectionPlugin.ts b/graphile-build/graphile-build-pg/src/plugins/PgIntrospectionPlugin.ts index ab7f92feb0..7192b9b009 100644 --- a/graphile-build/graphile-build-pg/src/plugins/PgIntrospectionPlugin.ts +++ b/graphile-build/graphile-build-pg/src/plugins/PgIntrospectionPlugin.ts @@ -5,7 +5,7 @@ import { withSuperuserPgClientFromPgService, } from "@dataplan/pg"; import type { PromiseOrDirect, Step } from "grafast"; -import { constant, context, noop, object, promiseWithResolve } from "grafast"; +import { abortable, constant, context, noop, object } from "grafast"; import type { GatherPluginContext } from "graphile-build"; import { EXPORTABLE, gatherConfig } from "graphile-build"; import type { @@ -35,7 +35,8 @@ import { import { version } from "../version.ts"; import { watchFixtures } from "../watchFixtures.ts"; -const sleep = (ms: number) => new Promise((resolve) => setTimeout(resolve, ms)); +const sleep = (ms: number) => + new Promise((resolve) => setTimeout(resolve, ms)); /** Someone else created */ const CLASH_CODES = ["23505", "42P06", "42P07", "42710"]; @@ -685,9 +686,11 @@ export const PgIntrospectionPlugin: GraphileConfig.Plugin = { let eventStream = await pgService.pgSubscriber.subscribe("postgraphile_watch"); const $$stop = Symbol("stop"); - const { resolve, promise: abort } = - promiseWithResolve(); - unlistens.push(() => resolve($$stop)); + + const controller = new AbortController(); + const signal = controller.signal; + unlistens.push(() => void controller.abort()); + const regather = () => { // Delete the introspection results info.cache.introspectionResultsPromise = null; @@ -697,7 +700,7 @@ export const PgIntrospectionPlugin: GraphileConfig.Plugin = { callback(); }; const waitNext = () => { - const next = Promise.race([abort, eventStream.next()]); + const next = abortable(signal, $$stop, eventStream.next()); next.then( (event) => { if (event === $$stop) { @@ -739,7 +742,7 @@ export const PgIntrospectionPlugin: GraphileConfig.Plugin = { console.error( `postgraphile_watch subscription failed (${e}); waiting ${delay.toFixed(0)}ms then re-establishing`, ); - const result = await Promise.race([sleep(delay), abort]); + const result = await abortable(signal, $$stop, sleep(delay)); if (result === $$stop) { return; } diff --git a/postgraphile/postgraphile/graphile.config.ts b/postgraphile/postgraphile/graphile.config.ts index 1ec76f7d7c..04e7edf1fd 100644 --- a/postgraphile/postgraphile/graphile.config.ts +++ b/postgraphile/postgraphile/graphile.config.ts @@ -434,6 +434,11 @@ const preset: GraphileConfig.Preset = { sub(topic: String!): Int gql(max: Int! = 10): Int slow: String + fast: FastSubscriptionPayload + } + type FastSubscriptionPayload { + datetime: String + count: Int } `, objects: { @@ -499,6 +504,27 @@ const preset: GraphileConfig.Preset = { [sleep], ), }, + fast: { + resolve: EXPORTABLE( + () => + function resolve(e) { + return e; + }, + [], + ), + subscribe: EXPORTABLE( + (sleep) => + async function* subscribe() { + let count = 0; + while (true) { + ++count; + yield { datetime: new Date().toISOString(), count }; + await sleep(5); + } + }, + [sleep], + ), + }, }, }, },