Fix memory leak in subscriptions

.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.

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<TData = any>(
 }
 
 // Save on garbage collection by just using this promise for everything
-const DONE_PROMISE: Promise<IteratorReturnResult<void>> = Promise.resolve({
-  done: true,
-  value: undefined,
-});
+const DONE_PROMISE: Promise<IteratorReturnResult<void>> = 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<TData>(
    */
   const buffer: Array<Promise<IteratorResult<TData, void>>> = [];
 
-  // Easy way to resolve a promise for slowing down the fastest consumer
-  let wmi: PromiseWithResolve<void> | null = null;
-  function lowWaterMarkIncreased(): PromiseLike<void> {
-    if (wmi === null) {
-      wmi = promiseWithResolve<void>();
-    }
-    return wmi.promise;
+  // Consumers waiting for the low-water mark to advance.
+  let lowWaterMarkWaiters: Set<() => void> | null = null;
+  function waitForLowWaterMarkOrPause(): Promise<void> {
+    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<TData>(
       }
 
       // 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<TData>(
           // 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

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,

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<AsyncExecutionResult, void, void>
 > {
@@ -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<void>();
+      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(

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<IteratorReturnResult<void>> = Promise.resolve({
-  done: true,
-  value: undefined,
-});
+const DONE_PROMISE: Promise<IteratorReturnResult<void>> = 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<T, F>(
+  signal: AbortSignal,
+  valueForAbort: F,
+  promise: PromiseLike<T>,
+): Promise<T | F> {
+  if (signal.aborted) {
+    return Promise.resolve(valueForAbort);
+  }
+  return new Promise<T | F>((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);
+      },
+    );
+  });
+}